[mlir][TilingInterface] Fix a crash in PartialTilingInterface for some inputs

[vmurali]

Summary

Previously, mlir-opt crashes for the following example:

mlir-opt --test-transform-dialect-interpreter Test.mlir

Test.mlir:

func.func @reduction_bug(%arg0: tensor<32x32xi32>, %arg1: tensor<32xi32>, %out: tensor<32xi32>) -> tensor<32xi32> {
  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
                                          affine_map<(d0, d1, d2) -> (d1)>,
                                          affine_map<(d0, d1, d2) -> (d0)>],
   iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<32x32xi32>, tensor<32xi32>) outs(%out : tensor<32xi32>) {
    ^bb0(%a: i32, %b: i32, %c: i32):
      %r2 = arith.addi %c, %a : i32
      linalg.yield %r2 : i32
    } -> tensor<32xi32>
  return %red : tensor<32xi32>
}

transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation):
  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
  %1, %2, %3 = transform.structured.tile_reduction_using_scf %0 { tile_sizes = [0, 0, 8] }
}

It was because the result tensor's rank was supposed to be exactly one less than the number of loops. This differential makes it fail gracefully.

[ThomasRaoux]

Let's find a way to add a negative test for this :)

[vmurali]

In D137730#3934924, @ThomasRaoux wrote:

Let's find a way to add a negative test for this :)

I tried transform.sequence failures(suppress). It still fails with the same error. Any idea how to make transform ignore errors (by bailing out without doing any transformation)?

[vmurali]

In D137730#3934924, @ThomasRaoux wrote:

Let's find a way to add a negative test for this :)

I tried transform.sequence failures(suppress). It still fails with the same error. Any idea how to make transform ignore errors (by bailing out without doing any transformation)? @nicolasvasilache

[nicolasvasilache]

suppress can only suppress silenceable failures, which can be ignored and leave the IR in a valid state.
Definite failures cannot be suppressed as they leave the IR in an invalid state.

What you may want is to refactor the transform to safelyreturn a silenceable failure in your use case.

[vmurali]

In D137730#3944846, @nicolasvasilache wrote:

suppress can only suppress silenceable failures, which can be ignored and leave the IR in a valid state.
Definite failures cannot be suppressed as they leave the IR in an invalid state.

What you may want is to refactor the transform to safelyreturn a silenceable failure in your use case.

I want to suppress a match failure (see changes in this differential). The transform pipeline should technically see the match failure and fail to perform any subsequent transformations (like how a normal pass pipeline would). I just want it to not bail out in the very end, and instead return the state the IR was in before the pipeline was applied.

More generally, if I have to apply different pipelines in some priority order (the first one to succeed will change the IR, and the rest of the pipelines will not be applied), how do I do that? (In my case, the second pipeline is a no-op transformation.)

[nicolasvasilache]

In D137730#3945094, @vmurali wrote:

In D137730#3944846, @nicolasvasilache wrote:

suppress can only suppress silenceable failures, which can be ignored and leave the IR in a valid state.
Definite failures cannot be suppressed as they leave the IR in an invalid state.

What you may want is to refactor the transform to safelyreturn a silenceable failure in your use case.

I want to suppress a match failure (see changes in this differential).

mlir::scf::tileReductionUsingScf returns a failure to its caller which turns it into a definite error that you cannot recover from.
As I wrote earlier, What you may want is to refactor the transform to safely return a silenceable failure in your use case..
This means making mlir::scf::tileReductionUsingScf return "more information" than a binary success/fail and us that in the transform op.
The part that is annoying is that the whole MLIR infra is very binary in its return results (i.e. there is no difference between a match failure and an execution failure, both are failures and there is no type guarantee that IR is in a correct state)

The transform pipeline should technically see the match failure and fail to perform any subsequent transformations (like how a normal pass pipeline would).

There is no transform pipeline, just a list of transforms that are applied unconditionally.
This is very much unlike any "normal pass pipeline".
You may want to use transform.alternatives which will clone and try its branches until the first one succeeds.

I just want it to not bail out in the very end, and instead return the state the IR was in before the pipeline was applied.
More generally, if I have to apply different pipelines in some priority order (the first one to succeed will change the IR, and the rest of the pipelines will not be applied), how do I do that? (In my case, the second pipeline is a no-op transformation.)

transform.alternatives is likely what you are looking for ,then.

[nicolasvasilache]

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

216

you may want to add a comment that this is supposed to fail gracefully and not change the IR

222

double comments here and below

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

if (!identity.has_value())
       return op->emitOpError(
           "Failed to get an identity value for the reduction operation.");
 
     // Calculate the new shape, we insert the new dimension based on the index
     // of the reduction dimension.
     SmallVector<int64_t> newOutputShape;
     ArrayRef<int64_t> oldShape =
         linalgOp.getShape(linalgOp.getDpsInitOperand(0));
+    assert(sizes.size() == oldShape.size() + 1 &&
+           "result tensor should have rank exactly one dimension smaller than "
+           "the number of loops.");
     SmallVector<Value> dynamicDims;
     for (int64_t idx : llvm::seq<int64_t>(0, oldShape.size() + 1)) {
       if (idx == insertSplitDimension) {
         dispatchIndexOpFoldResults(sizes[idx], dynamicDims, newOutputShape,
                                    ShapedType::kDynamicStrideOrOffset);
         continue;
       }
       int64_t oldIdx = idx < insertSplitDimension ? idx : idx - 1;

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

mlir::scf::tileReductionUsingScf(PatternRewriter &b,
   int reductionDim;
   for (auto &[idx, iteratorType] :
        llvm::enumerate(tilingInterfaceOp.getLoopIteratorTypes())) {
     if (iteratorType == utils::IteratorType::reduction) {
       reductionDim = idx;
       break;
     }
   }
+  {
+    auto origResultTensor = cast<DestinationStyleOpInterface>(op.getOperation())
+                                .getDpsInitOperand(0);
+    size_t origResultSize = 0;
+    if (auto shapedType =
+            origResultTensor->get().getType().dyn_cast<ShapedType>())
+      origResultSize = shapedType.getShape().size();
+    if (iterationDomain.size() != origResultSize + 1) {
+      return b.notifyMatchFailure(
+          op, "only support result tensor whose rank is exactly one dimension "
+              "smaller than the number of loops.");
+    }
+  }
   // 1. create the inital tensor value.
   FailureOr<Operation *> identityTensor =
       op.generateInitialTensorForPartialReduction(b, loc, tileSize,
                                                   reductionDim);
   if (failed(identityTensor))
     return b.notifyMatchFailure(op,
                                 "cannot create a tensor of identity value.");
   // 2. Create the nested loops.

mlir/test/Dialect/Linalg/transform-tile-reduction.mlir

 //     CHECK:       tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, 0, %[[IV]]] [%[[D0]], %[[D2]], 1] [1, 1, 1] : tensor<?x?xf32> into tensor<?x?x5xf32>
 //     CHECK:     }
 //     CHECK:   }
 //     CHECK:   %[[R:.*]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[L]] : tensor<?x?x5xf32>) outs(%[[ARG2]] : tensor<?x?xf32>) {
 //     CHECK:     arith.addf
 //     CHECK:     linalg.yield
 //     CHECK:   } -> tensor<?x?xf32>
 //     CHECK:   return %[[R]] : tensor<?x?xf32>
+
+// -----
+
+func.func @reduction_bug(%arg0: tensor<32x32xi32>, %arg1: tensor<32x32xi32>, %out: tensor<32xi32>) -> tensor<32xi32> {
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
+                                          affine_map<(d0, d1, d2) -> (d0, d1)>,
+                                          affine_map<(d0, d1, d2) -> (d0)>],
+   iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<32x32xi32>, tensor<32x32xi32>) outs(%out : tensor<32xi32>) {
+    ^bb0(%a: i32, %b: i32, %c: i32):
+      %r1 = arith.muli %a, %b: i32
+      %r2 = arith.addi %c, %r1 : i32
+      linalg.yield %r2 : i32
+    } -> tensor<32xi32>
+  return %red : tensor<32xi32>
+}
+
+transform.sequence failures(suppress) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  %1, %2, %3 = transform.structured.tile_reduction_using_scf %0 { tile_sizes = [0, 0, 8] }
+}
+    
+// //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>  
+// //   CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>  
+// //   CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0)>  
+// // CHECK-LABEL:  func @reduction_bug
+// //       CHECK: %[[RED:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "reduction"]}  
+// //  CHECK-SAME:   ins(%[[I1]], %[[I2]] : tensor<32x32xi32>, tensor<32x32xi32>) outs(%[[F]] : tensor<32xi32>) { 
+// //       CHECK:   arith.muli
+// //       CHECK:   arith.addi
+// //       CHECK:   linalg.yield 
+// //       CHECK: } -> tensor<32xi32> 
+// //       CHECK: return %[[RED]] : tensor<32xi32>