[Mlir-commits] [mlir] [mlir] emit better errors in transform.structured.interchange (PR #67315)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/67315

The implementation doesn't emit any diagnostics as it is shared with the pattern-based implementation. Check preconditions early and emit diagnostics from the transform op instead. Without this change, the op would produce a definite failure and no error message.

>From 1d97c434dd3670639da8b2265b4979a4e71df254 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Mon, 25 Sep 2023 11:37:14 +0000
Subject: [PATCH] [mlir] emit better errors in transform.structured.interchange

The implementation doesn't emit any diagnostics as it is shared with the
pattern-based implementation. Check preconditions early and emit
diagnostics from the transform op instead. Without this change, the op
would produce a definite failure and no error message.
---
 .../TransformOps/LinalgTransformOps.cpp       | 11 +++++++++-
 .../Linalg/transform-op-interchange.mlir      | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 8bfd731fcc81213..66bbaed3797f3ba 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1037,12 +1037,21 @@ transform::InterchangeOp::applyToOne(transform::TransformRewriter &rewriter,
     results.push_back(target);
     return DiagnosedSilenceableFailure::success();
   }
+
+  unsigned numLoops = cast<LinalgOp>(target.getOperation()).getNumLoops();
+  if (interchangeVector.size() != numLoops) {
+    return emitSilenceableError()
+           << getIteratorInterchangeAttrName() << " has length ("
+           << interchangeVector.size()
+           << ") different from the number of loops in the target operation ("
+           << numLoops << ")";
+  }
   FailureOr<GenericOp> res =
       interchangeGenericOp(rewriter, target,
                            SmallVector<unsigned>(interchangeVector.begin(),
                                                  interchangeVector.end()));
   if (failed(res))
-    return DiagnosedSilenceableFailure::definiteFailure();
+    return emitDefiniteFailure() << "failed to apply";
   results.push_back(res->getOperation());
   return DiagnosedSilenceableFailure::success();
 }
diff --git a/mlir/test/Dialect/Linalg/transform-op-interchange.mlir b/mlir/test/Dialect/Linalg/transform-op-interchange.mlir
index 7966b22a257ab5c..3efb76afdfea76e 100644
--- a/mlir/test/Dialect/Linalg/transform-op-interchange.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-interchange.mlir
@@ -38,3 +38,25 @@ transform.sequence failures(propagate) {
   // expected-error @below {{transform applied to the wrong op kind}}
   transform.structured.interchange %0 iterator_interchange = [1, 0] : (!transform.any_op) -> !transform.any_op
 }
+
+// -----
+
+func.func @too_many_iters(%0: tensor<?x?xf32>, %1: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %r = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%0: tensor<?x?xf32>) outs(%1: tensor<?x?xf32>) {
+  ^bb0(%2: f32, %3: f32):
+    %4 = arith.mulf %2, %2 : f32
+    linalg.yield %4 : f32
+  } -> tensor<?x?xf32>
+  return %r : tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  // expected-error @below {{"iterator_interchange" has length (3) different from the number of loops in the target operation (2)}}
+  transform.structured.interchange %0 iterator_interchange = [2,1,0] : (!transform.any_op) -> !transform.any_op
+  transform.yield
+}