Merge remote-tracking branch 'origin/feature/fused-ops' into bump_to_f4943464

Author: jorickert

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H
 
 #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
+#include "mlir/Dialect/Affine/IR/AffineTraits.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -13,6 +13,7 @@
 #ifndef AFFINE_OPS
 #define AFFINE_OPS
 
+include "mlir/Dialect/Affine/IR/AffineTraits.td"
 include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"

mlir/include/mlir/Dialect/Affine/IR/AffineTraits.h

@@ -0,0 +1,27 @@
+//===- AffineTraits.h - MLIR Affine Traits --------------------------------===//
+//
+// 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 file defines traits brought in by the Affine dialect.
+//
+//===----------------------------------------------------------------------===//
+#ifndef AFFINE_TRAITS_H
+#define AFFINE_TRAITS_H
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir::OpTrait {
+
+template <typename ConcreteType>
+class AffineDim : public TraitBase<ConcreteType, AffineDim> {
+public:
+  static LogicalResult verifyTrait(Operation *op) { return success(); }
+};
+
+} // namespace mlir::OpTrait
+
+#endif // AFFINE_TRAITS_H

mlir/include/mlir/Dialect/Affine/IR/AffineTraits.td

@@ -0,0 +1,25 @@
+//===- AffineTraits.td - Affine dialect traits -------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines traits brought in by the MLIR Affine dialect.
+//
+//===----------------------------------------------------------------------===//
+#ifndef AFFINE_TRAITS
+#define AFFINE_TRAITS
+
+include "mlir/IR/OpBase.td"
+
+// Trait to declare that an op result is an affine dimension identifier.
+// Prevents the result from being seen as a symbol into AffineMaps
+// and IntegerSets.
+// This is a deviation from upstream to consider linalg.index as
+// a dimension rather than a symbol. See this PR:
+// https://github.com/Xilinx/llvm-project/pull/537
+def AffineDim : NativeOpTrait<"AffineDim">;
+
+#endif // AFFINE_TRAITS

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_LINALG_IR_LINALG_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Affine/IR/AffineTraits.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"

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

@@ -13,6 +13,7 @@
 #ifndef LINALG_OPS
 #define LINALG_OPS
 
+include "mlir/Dialect/Affine/IR/AffineTraits.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
@@ -46,7 +47,7 @@ def Linalg_YieldOp : Linalg_Op<"yield", [Pure, ReturnLike, Terminator]>,
   let hasVerifier = 1;
 }
 
-def Linalg_IndexOp : Linalg_Op<"index", [Pure]>,
+def Linalg_IndexOp : Linalg_Op<"index", [Pure, AffineDim]>,
     Arguments<(ins ConfinedAttr<I64Attr, [IntMinValue<0>]>:$dim)>,
     Results<(outs Index:$result)> {
   let summary = "linalg index operation";
@@ -100,7 +101,8 @@ def Linalg_SoftmaxOp : Linalg_Op<"softmax",
       ["getIterationDomain",
        "getLoopIteratorTypes",
        "getResultTilePosition",
-       "getTiledImplementation"]>]> {
+       "getTiledImplementation",
+       "generateResultTileValue"]>]> {
   let summary = "Softmax operator";
   let description = [{
     linalg.softmax computes a numerically stable version of softmax.

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_SCF_TRANSFORMS_TILEUSINGINTERFACE_H
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
@@ -194,6 +195,21 @@ struct SCFTileAndFuseOptions {
   /// before fusion. This will track deleted and newly inserted
   /// `tensor.extract_slice` ops and update the worklist.
   std::optional<FrozenRewritePatternSet> cleanupPatterns = std::nullopt;
+
+  /// A function to insert a tilable node into a list of nodes to be tiled.
+  /// This controls the order in which tiling and fusion happen.
+  using WorklistInsertFnTy = std::function<void(
+      tensor::ExtractSliceOp op, std::deque<tensor::ExtractSliceOp> &worklist)>;
+  /// By default, simply append the op at the end of the queue.
+  WorklistInsertFnTy worklistInsertFn =
+      [](tensor::ExtractSliceOp op,
+         std::deque<tensor::ExtractSliceOp> &worklist) {
+        worklist.push_back(op);
+      };
+  SCFTileAndFuseOptions &setWorklistInsertFn(WorklistInsertFnTy insertFn) {
+    worklistInsertFn = insertFn;
+    return *this;
+  }
 };
 
 /// Fuse the producer of the source of `candidateSliceOp` by computing the

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h

@@ -26,7 +26,6 @@ namespace tosa {
 
 // Expose Rewrite Functions that decompose TOSA Ops into further TOSA Ops.
 // The rewrites can be selectively added to a conversion pass.
-void populateTosaDecomposeConv2D(MLIRContext *ctx, RewritePatternSet &patterns);
 void populateTosaDecomposeTransposeConv(MLIRContext *ctx,
                                         RewritePatternSet &patterns);
 void populateTosaDecomposeDepthwise(MLIRContext *ctx,

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
@@ -312,6 +313,10 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
     return isa<AffineForOp, AffineParallelOp>(parentOp);
   }
 
+  // Remove me: linalg.index ops are valid affine dim identifiers
+  if (op->hasTrait<OpTrait::AffineDim>())
+    return true;
+
   // Affine apply operation is ok if all of its operands are ok.
   if (auto applyOp = dyn_cast<AffineApplyOp>(op))
     return applyOp.isValidDim(region);
@@ -439,6 +444,10 @@ bool mlir::affine::isValidSymbol(Value value, Region *region) {
     return false;
   }
 
+  // Remove me: linalg.index ops are not valid affine symbols
+  if (defOp->hasTrait<OpTrait::AffineDim>())
+    return false;
+
   // Constant operation is ok.
   Attribute operandCst;
   if (matchPattern(defOp, m_Constant(&operandCst)))

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -2648,39 +2648,56 @@ SmallVector<utils::IteratorType> SoftmaxOp::getLoopIteratorTypes() {
   return iteratorTypes;
 }
 
-FailureOr<TilingResult>
-SoftmaxOp::getTiledImplementation(OpBuilder &builder,
-                                  ArrayRef<OpFoldResult> offsets,
-                                  ArrayRef<OpFoldResult> sizes) {
-  int64_t rank = getInputOperandRank();
+static FailureOr<TilingResult>
+implementTiledSoftMax(SoftmaxOp &op, OpBuilder &builder,
+                      ArrayRef<OpFoldResult> offsets,
+                      ArrayRef<OpFoldResult> sizes) {
+  int64_t rank = op.getInputOperandRank();
   auto oneAttr = builder.getI64IntegerAttr(1);
   SmallVector<OpFoldResult> strides(rank, oneAttr);
   SmallVector<Value> tiledOperands;
   Operation *inputSlice =
-      getSlice(builder, getLoc(), getInput(), offsets, sizes, strides);
+      getSlice(builder, op.getLoc(), op.getInput(), offsets, sizes, strides);
   if (!inputSlice) {
-    return emitOpError("failed to compute input slice");
+    return op.emitOpError("failed to compute input slice");
   }
   tiledOperands.emplace_back(inputSlice->getResult(0));
   Operation *outputSlice =
-      getSlice(builder, getLoc(), getOutput(), offsets, sizes, strides);
+      getSlice(builder, op.getLoc(), op.getOutput(), offsets, sizes, strides);
   if (!outputSlice) {
-    return emitOpError("failed to compute output slice");
+    return op.emitOpError("failed to compute output slice");
   }
   tiledOperands.emplace_back(outputSlice->getResult(0));
 
   SmallVector<Type, 4> resultTypes;
-  if (hasPureTensorSemantics())
+  if (op.hasPureTensorSemantics())
     resultTypes.push_back(tiledOperands[1].getType());
   Operation *tiledOp =
-      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
+      mlir::clone(builder, op.getOperation(), resultTypes, tiledOperands);
 
   return TilingResult{
       {tiledOp},
       SmallVector<Value>(tiledOp->getResults()),
       llvm::to_vector(ArrayRef<Operation *>{inputSlice, outputSlice})};
 }
 
+FailureOr<TilingResult>
+SoftmaxOp::getTiledImplementation(OpBuilder &builder,
+                                  ArrayRef<OpFoldResult> offsets,
+                                  ArrayRef<OpFoldResult> sizes) {
+  return implementTiledSoftMax(*this, builder, offsets, sizes);
+}
+
+FailureOr<TilingResult>
+SoftmaxOp::generateResultTileValue(OpBuilder &builder, unsigned resultNumber,
+                                   ArrayRef<OpFoldResult> offsets,
+                                   ArrayRef<OpFoldResult> sizes) {
+  if (resultNumber != 0)
+    return failure();
+
+  return implementTiledSoftMax(*this, builder, offsets, sizes);
+}
+
 LogicalResult SoftmaxOp::getResultTilePosition(
     OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
     ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,