Address comments (round 1).

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

@@ -31,7 +31,7 @@ namespace tensor {
 FailureOr<TilingResult> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
-/// Method to swap an `tensor.insert_slice`s with its consumer when the
+/// Method to swap `tensor.insert_slice`s with their consumers when the
 /// consumer implements the `TilingInterface`. The size of `sliceOps` and
 /// `consumerOperands` is expected to be the same. Every entry in
 /// `consumerOperands` represents a use of the the corresponding

mlir/include/mlir/Interfaces/TilingInterface.td

@@ -202,7 +202,7 @@ def TilingInterface : OpInterface<"TilingInterface"> {
       InterfaceMethod<
         /*desc=*/[{
           Method to generate the tiled implementation of an operation that uses
-          exactly tiles of the given operands.
+          the exact tiles of the given operands.
 
           This method is required to allow operations to be "tiled and fused"
           with an (already tiled) producer. Given tiles of the producer, this

mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp

@@ -22,8 +22,11 @@
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/Support/Debug.h"
 #include <optional>
 
+#define DEBUG_TYPE "linalg-tiling-interface-impl"
+
 using namespace mlir;
 using namespace mlir::linalg;
 
@@ -170,9 +173,11 @@ struct LinalgOpTilingInterface
           OpFoldResult seenOffset = it->second;
           OpFoldResult seenSize = mappedSizes.lookup(position);
           if (seenOffset != offset || seenSize != size) {
-            return linalgOp->emitOpError(
-                "inconsistent iteration space mapping from offsets/sizes of "
-                "operands/results");
+            LLVM_DEBUG({
+              llvm::dbgs() << "inconsistent iteration space mapping from "
+                              "offsets/sizes of operands/results";
+            });
+            return failure();
           }
         } else {
           mappedOffsets[position] = offset;
@@ -874,8 +879,11 @@ struct PackOpTiling
       ArrayRef<SmallVector<OpFoldResult>> allSizes,
       SmallVectorImpl<OpFoldResult> &resultOffsets,
       SmallVectorImpl<OpFoldResult> &resultSizes) const {
-    if (operandNumbers.size() != 1 || operandNumbers[0] != 0)
-      return op->emitOpError("unsupporeted operands for consumer fusion");
+    if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
+      LLVM_DEBUG(
+          { llvm::dbgs() << "unsupported operands for consumer fusion"; });
+      return failure();
+    }
 
     ArrayRef<OpFoldResult> offsets(allOffsets[0]);
     ArrayRef<OpFoldResult> sizes(allSizes[0]);
@@ -943,8 +951,11 @@ struct PackOpTiling
       Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
       ArrayRef<SmallVector<OpFoldResult>> allOffsets,
       ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
-    if (operandNumbers.size() != 1 || operandNumbers[0] != 0)
-      return op->emitOpError("unhandled operands for consumer fusion");
+    if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
+      LLVM_DEBUG(
+          { llvm ::dbgs() << "unhandled operands for consumer fusion"; });
+      return failure();
+    }
 
     ArrayRef<OpFoldResult> offsets(allOffsets[0]);
     ArrayRef<OpFoldResult> sizes(allSizes[0]);
@@ -1228,7 +1239,8 @@ struct UnPackOpTiling
       SmallVectorImpl<OpFoldResult> &resultOffsets,
       SmallVectorImpl<OpFoldResult> &resultSizes) const {
     if (operandNumbers.size() != 1) {
-      return op->emitOpError("unable to handle multiple operands");
+      LLVM_DEBUG({ llvm::dbgs() << "unable to handle multiple operands"; });
+      return failure();
     }
     auto unPackOp = cast<UnPackOp>(op);
     unsigned operandNumber = operandNumbers[0];
@@ -1293,7 +1305,8 @@ struct UnPackOpTiling
       ArrayRef<SmallVector<OpFoldResult>> allOffsets,
       ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
     if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
-      return op->emitOpError("unhandled operands for consumer fusion");
+      LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
+      return failure();
     }
     auto unPackOp = cast<UnPackOp>(op);
     ArrayRef<OpFoldResult> offsets(allOffsets[0]);

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -2060,15 +2060,16 @@ static FailureOr<SmallVector<OpOperand *>> getUntiledConsumerOperandsFromSlices(
                 [&](auto op) {
                   return getUntiledConsumerFromSlice(rewriter, op, loops);
                 })
-            .Default([](Operation *op) {
-              return op->emitOpError("unhandled slice type");
+            .Default([&](Operation *op) {
+              return rewriter.notifyMatchFailure(op, "unhandled slice type");
             });
     if (failed(fusedOperand)) {
       return failure();
     }
     if (!fusedOperands.empty() &&
         fusedOperand.value()->getOwner() != fusedOperands.front()->getOwner()) {
-      return fusedOperands.front()->getOwner()->emitOpError(
+      return rewriter.notifyMatchFailure(
+          fusedOperand.value()->getOwner(),
           "all candidate slices must be to the same consumer");
     }
     fusedOperands.push_back(fusedOperand.value());
@@ -2125,23 +2126,23 @@ mlir::scf::tileAndFuseConsumerOfSlices(
     RewriterBase &rewriter, ArrayRef<Operation *> candidateSlices,
     MutableArrayRef<LoopLikeOpInterface> loops) {
   if (candidateSlices.empty()) {
-    return emitError(rewriter.getUnknownLoc(),
-                     "no candidate slices provided for consumer fusion");
+    return rewriter.notifyMatchFailure(
+        rewriter.getUnknownLoc(),
+        "no candidate slices provided for consumer fusion");
   }
   // Return if `loops` is empty, return an error for now. Caller is expected
   // to handle this case.
   if (loops.empty()) {
-    return candidateSlices.front()->emitOpError(
+    return rewriter.notifyMatchFailure(
+        candidateSlices.front(),
         "cannot call tile and fuse consumer with an empty loop nest");
   }
 
-  if (!(llvm::all_of(
-            candidateSlices,
-            [](Operation *op) { return isa<tensor::InsertSliceOp>(op); }) ||
-        llvm::all_of(candidateSlices, [](Operation *op) {
-          return isa<tensor::ParallelInsertSliceOp>(op);
-        }))) {
-    return candidateSlices.front()->emitOpError(
+  if (!(llvm::all_of(candidateSlices, llvm::IsaPred<tensor::InsertSliceOp>) ||
+        llvm::all_of(candidateSlices,
+                     llvm::IsaPred<tensor::ParallelInsertSliceOp>))) {
+    return rewriter.notifyMatchFailure(
+        candidateSlices.front(),
         "candidates slices need to be all `tensor.extract_slice`s or "
         "`tensor.parallel_insert_slice`s");
   }
@@ -2261,8 +2262,14 @@ mlir::scf::tileAndFuseConsumerOfSlices(
     for (auto candidateSliceOp : clonedInsertSlices) {
       SmallVector<OpFoldResult> offsets = candidateSliceOp.getMixedOffsets();
       SmallVector<OpFoldResult> sizes = candidateSliceOp.getMixedSizes();
+      SmallVector<OpFoldResult> strides = candidateSliceOp.getMixedStrides();
 
       // 9. Check all insert stride is 1.
+      if (!llvm::all_of(strides, isOneInteger)) {
+        return rewriter.notifyMatchFailure(
+            candidateSliceOp, "containingOp's result yield with stride");
+      }
+
       allOffsets.emplace_back(std::move(offsets));
       allSizes.emplace_back(std::move(sizes));
     }

mlir/lib/Dialect/Tensor/Transforms/SwapExtractSliceWithProducerPatterns.cpp

@@ -17,6 +17,9 @@
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "tensor-swap-slices"
 
 using namespace mlir;
 
@@ -43,21 +46,28 @@ FailureOr<TilingResult> tensor::replaceInsertSlicesWithTiledConsumer(
     OpBuilder &builder, ArrayRef<tensor::InsertSliceOp> sliceOps,
     ArrayRef<OpOperand *> consumerOperands) {
   if (sliceOps.empty()) {
-    return emitError(builder.getUnknownLoc(),
-                     "expected candidate slices list to be non-empty");
+    LLVM_DEBUG(
+        { llvm::dbgs() << "expected candidate slices list to be non-empty"; });
+    return failure();
   }
   if (sliceOps.size() != consumerOperands.size()) {
-    return sliceOps.front()->emitOpError(
-        "expected as many operands as the number of slices passed");
+    LLVM_DEBUG({
+      llvm::dbgs()
+          << "expected as many operands as the number of slices passed";
+    });
+    return failure();
   }
   auto consumerOp =
       dyn_cast<TilingInterface>(consumerOperands.front()->getOwner());
   if (!consumerOp)
     return failure();
   for (auto opOperand : consumerOperands.drop_front()) {
     if (opOperand->getOwner() != consumerOp) {
-      return consumerOp->emitOpError(
-          "expected all consumer operands to be from the same operation");
+      LLVM_DEBUG({
+        llvm::dbgs()
+            << "expected all consumer operands to be from the same operation";
+      });
+      return failure();
     }
   }
 

mlir/test/Interfaces/TilingInterface/tile-and-fuse-consumer.mlir

@@ -688,7 +688,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// Check that when the given operand tiles are incosistent, tiling fails.
+// Check that when the given operand tiles are inconsistent, tiling fails.
 
 func.func @multi_slice_fusion2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>, %arg3 : index) -> tensor<?xf32> {
   %c0 = arith.constant 0 : index
@@ -881,14 +881,14 @@ func.func @multi_slice_fusion_invalid(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<
 	      linalg.yield %0: f32
     } -> tensor<?x?xf32>
     scf.forall.in_parallel {
+      // expected-error @below {{failed to fuse consumer of slice}}
       tensor.parallel_insert_slice %generic0 into %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]
           : tensor<?x?xf32> into tensor<?x?xf32>
       tensor.parallel_insert_slice %generic1 into %init1[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]
           : tensor<?x?xf32> into tensor<?x?xf32>
     }
   }
   %empty = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>
-  // expected-error @below {{inconsistent iteration space mapping from offsets/sizes of operands/results}}
   %result = linalg.generic {
       indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
       iterator_types = ["parallel", "parallel"]}

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp

@@ -21,6 +21,9 @@
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "test-tiling-interface"
 
 #define GET_OP_CLASSES
 #include "TestTilingInterfaceTransformOps.h.inc"
@@ -182,7 +185,7 @@ static LogicalResult applyFuseConsumer(
         scf::tileAndFuseConsumerOfSlices(rewriter, slices, loops);
 
     if (failed(fuseConsumerResults))
-      return failure();
+      return slices.front()->emitOpError("failed to fuse consumer of slice");
 
     // Report back the relevant handles to the transform op.
     for (OpOperand *origConsumerOperand :