Tosa Changes Integration

mlir/cmake/modules/AddMLIR.cmake

@@ -732,7 +732,8 @@ function(mlir_target_link_libraries target type)
   endif()
 
   if (MLIR_LINK_MLIR_DYLIB)
-    target_link_libraries(${target} ${type} MLIR)
+    # AMD: Do not link shared, as this casues linking errors
+    target_link_libraries(${target} ${type} ${ARGN})
   else()
     target_link_libraries(${target} ${type} ${ARGN})
   endif()

mlir/cmake/modules/AddMLIRPython.cmake

@@ -23,6 +23,11 @@
 #     grouping. Source groupings form a DAG.
 #   SOURCES: List of specific source files relative to ROOT_DIR to include.
 #   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
+
+if (POLICY CMP0175)
+  cmake_policy(SET CMP0175 OLD)
+endif()
+
 function(declare_mlir_python_sources name)
   cmake_parse_arguments(ARG
     ""

mlir/include/mlir/Analysis/AffineExprBounds.h

@@ -0,0 +1,91 @@
+//===- AffineExprBounds.h - Compute bounds of affine expressions *- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines an analysis of affine expressions to compute their
+// ranges (lower/upper bounds) in a given context.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H
+#define MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H
+
+#include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/Support/LogicalResult.h"
+
+using namespace mlir;
+
+/// This visitor computes the bounds of affine expressions, using as context the
+/// bounds of the dimensions of the expression.
+///
+/// Example:
+/// Given bounds 0 <= d0 <= 99 and 0 <= d1 <= 199, we can compute the bounds
+/// of the following expression:
+/// lb(2 * d0 + 3 * d1) = 0
+/// ub(2 * d0 + 3 * d1) = 795
+///
+///  * The bounds given in the context are inclusive, and the bounds returned
+///  are also inclusive.
+///  * If bounds are not available for a dimension, std::nullopt can be used
+///  instead. The bounds of an expression that involves it will be std::nullopt.
+///  * Limitations:
+///    - Parametric expressions (using symbols) are not supported.
+///    - Unsigned FloorDiv is currently not supported.
+class AffineExprBoundsVisitor
+    : public AffineExprVisitor<AffineExprBoundsVisitor, LogicalResult> {
+public:
+  /// Initialize the context (bounds) with APInt. All bounds must have the same
+  /// signedness and bit width.
+  AffineExprBoundsVisitor(ArrayRef<std::optional<APInt>> constLowerBounds,
+                          ArrayRef<std::optional<APInt>> constUpperBounds,
+                          bool boundsSigned, uint64_t bitWidth,
+                          MLIRContext *context);
+
+  /// Initialize the context (bounds) with 64-bit signed integers. This allows
+  /// to directly map index-type values such as Linalg op bounds, which are
+  /// represented as int64_t.
+  AffineExprBoundsVisitor(ArrayRef<std::optional<int64_t>> constLowerBounds,
+                          ArrayRef<std::optional<int64_t>> constUpperBounds,
+                          MLIRContext *context);
+
+  /// Get the upper bound of \p expr using the context bounds.
+  std::optional<APInt> getUpperBound(AffineExpr expr);
+  std::optional<int64_t> getIndexUpperBound(AffineExpr expr);
+
+  /// Get the lower bound of \p expr using the context bounds.
+  std::optional<APInt> getLowerBound(AffineExpr expr);
+  std::optional<int64_t> getIndexLowerBound(AffineExpr expr);
+
+  // These methods are directly called by the AffineExprVisitor base class.
+  LogicalResult visitMulExpr(AffineBinaryOpExpr expr);
+  LogicalResult visitAddExpr(AffineBinaryOpExpr expr);
+  LogicalResult visitDimExpr(AffineDimExpr expr);
+  LogicalResult visitSymbolExpr(AffineSymbolExpr expr);
+  LogicalResult visitConstantExpr(AffineConstantExpr expr);
+  LogicalResult visitCeilDivExpr(AffineBinaryOpExpr expr);
+  LogicalResult visitFloorDivExpr(AffineBinaryOpExpr expr);
+  LogicalResult visitModExpr(AffineBinaryOpExpr expr);
+
+private:
+  bool boundsSigned;
+  uint64_t bitWidth;
+  void inferBinOpRange(
+      AffineBinaryOpExpr expr,
+      const std::function<ConstantIntRanges(ArrayRef<ConstantIntRanges>)>
+          &opInference);
+
+  /// Bounds that have been computed for subexpressions are memoized and reused.
+  llvm::DenseMap<AffineExpr, APInt> lb;
+  llvm::DenseMap<AffineExpr, APInt> ub;
+};
+
+#endif // MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

@@ -860,11 +860,6 @@ def LinalgStructuredInterface
     /// `createFlatListOfOperandDims`.
     SmallVector<Range, 4> createLoopRanges(OpBuilder &b, Location loc);
 
-    /// Compute the static loop sizes necessary to vectorize the computation.
-    /// This is done by applying `getShapesToLoopsMap` to
-    /// `createFlatListOfOperandStaticDims`.
-    SmallVector<int64_t, 4> computeStaticLoopSizes();
-
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
     LogicalResult reifyResultShapes(OpBuilder &b,

mlir/include/mlir/Dialect/PDL/IR/Builtins.h

@@ -43,8 +43,9 @@ enum class UnaryOpKind {
 LogicalResult addEntryToDictionaryAttr(PatternRewriter &rewriter,
                                        PDLResultList &results,
                                        ArrayRef<PDLValue> args);
-Attribute addElemToArrayAttr(PatternRewriter &rewriter, Attribute attr,
-                             Attribute element);
+LogicalResult addElemToArrayAttr(PatternRewriter &rewriter,
+                                 PDLResultList &results,
+                                 ArrayRef<PDLValue> args);
 LogicalResult mul(PatternRewriter &rewriter, PDLResultList &results,
                   llvm::ArrayRef<PDLValue> args);
 LogicalResult div(PatternRewriter &rewriter, PDLResultList &results,

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1724,6 +1724,7 @@ def Tosa_TileOp : Tosa_InferShapedTypeOp<"tile"> {
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
@@ -1877,6 +1878,15 @@ def Tosa_CastOp: Tosa_Op<"cast", [Pure,
     | signed 16 to float       | int16   | float   |
     | float 32 to float 64     | float32 | float64 |
     | float 64 to float 32     | float64 | float32 |
+
+    AMD extensions:
+    | signed to unsigned       | signed  | unsigned|
+    | unsigned to signed       | unsigned| signed  |
+    | unsigned to float        | unsigned| float  |
+      - unsigned to signed integer and signed to unsigned integer:
+        wrap on overflow
+      - unsigned to float:
+        uses llvm's float to int conversion with TOSA rounding mode
   }];
 
   let arguments = (ins

mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h

@@ -84,6 +84,18 @@ LogicalResult EqualizeRanks(PatternRewriter &rewriter, Location loc,
 LogicalResult EqualizeRanks(ImplicitLocOpBuilder &builder, Value &input1,
                             Value &input2);
 
+Value getTosaConstShape(ImplicitLocOpBuilder &builder,
+                        llvm::ArrayRef<int64_t> shape);
+
+Value getTosaConstShape(PatternRewriter &rewriter, Location loc,
+                        llvm::ArrayRef<int64_t> shape);
+
+// Get accumulator type for TOSA convolution ops
+LogicalResult getConvOpsAccType(PatternRewriter &rewriter,
+                                RankedTensorType inputTy,
+                                RankedTensorType weightTy,
+                                RankedTensorType outputTy, TypeAttr &accType);
+
 namespace {
 
 // Creates a TOSA operation and performs shape inference on the individual
@@ -217,7 +229,8 @@ TosaOp CreateOpAndInferShape(PatternRewriter &rewriter, Location loc,
 }
 
 // Apply an int32_t permutation to some input, that should be of the same
-// size as perms. Perms should contain some permutation of 0 - perms.size() - 1.
+// size as perms. Perms should contain some permutation of 0 - perms.size()
+// - 1.
 template <typename T>
 SmallVector<T> applyTOSAPermutation(ArrayRef<T> input,
                                     ArrayRef<int32_t> perms) {

mlir/include/mlir/IR/AffineExpr.h

@@ -110,6 +110,11 @@ class AffineExpr {
   /// floordiv, ceildiv, and mod is only allowed w.r.t constants.
   bool isPureAffine() const;
 
+  /// Returns true if this expression is monotonicically increasing with respect
+  /// to the AffineDimExprs, i.e. increasing the value of any AffineDimExpr will
+  /// never decrease the value of the result.
+  bool isMonotonicallyIncreasing() const;
+
   /// Returns the greatest known integral divisor of this affine expression. The
   /// result is always positive.
   int64_t getLargestKnownDivisor() const;