[Mlir-commits] [mlir] [mlir] Extend `tile_using_for` verifier to fix a crash (PR #98366)

[Felix Schneider]

https://github.com/ubfx updated https://github.com/llvm/llvm-project/pull/98366

>From 79a4667ae685539a95f8a4abce10955dee922bc5 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jul 2024 20:48:37 +0200
Subject: [PATCH 1/2] [mlir] Extend `tile_using_for` verifier to fix a crash

This patch adds a check for the correct number of `loops` results of the
`transform.structured.tile_using_for` Op to the verifier, fixing a crash.

Fix https://github.com/llvm/llvm-project/issues/98008
---
 mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 4eb334f8bbbfa..e53426083527a 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2890,6 +2890,12 @@ LogicalResult transform::TileUsingForOp::verify() {
     return emitOpError("expected same number of sizes (")
            << getMixedSizes().size() << ") and scalable sizes ()"
            << getScalableSizes().size() << ")";
+  auto staticSizes = getStaticSizes();
+  unsigned numExpectedLoops = staticSizes.size() - llvm::count(staticSizes, 0);
+  if (getLoops().size() != numExpectedLoops)
+    return emitOpError("expected number of loops to tile (")
+           << numExpectedLoops << ") to match number of `loops` results ("
+           << getLoops().size() << ")";
   return success();
 }
 

>From 9f5001d77aaa3108339cfb8dc104cec8819d396e Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jul 2024 20:58:50 +0200
Subject: [PATCH 2/2] Add test

---
 .../Dialect/Linalg/transform-op-tile.mlir     | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/test/Dialect/Linalg/transform-op-tile.mlir b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
index d244670f73754..955ea6b0ebbbd 100644
--- a/mlir/test/Dialect/Linalg/transform-op-tile.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
@@ -253,3 +253,23 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    // expected-error @below {{op expected number of loops to tile (3) to match number of `loops` results (1)}}
+    %1, %loops = transform.structured.tile_using_for %0 tile_sizes [4, 4, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+func.func @tile_linalg_matmul(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32> {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  return %0 : tensor<128x128xf32>
+}