[Mlir-commits] [mlir] [MLIR][Linalg] Bail out if the tiles provided are more than the number (PR #66007)

[lorenzo chelini]

https://github.com/chelini created https://github.com/llvm/llvm-project/pull/66007:

of loops

Currently, the compiler crashes if the number of tiles provided exceeds the number of loops.

>From 4e8be7bf666c35f30abc780f5204cebabe00e9d0 Mon Sep 17 00:00:00 2001
From: Lorenzo Chelini <l.chelini at icloud.com>
Date: Mon, 11 Sep 2023 22:41:57 +0200
Subject: [PATCH] [MLIR][Linalg] Bail out if the tiles provided are more than
 the number of loops

Currently, the compiler crashes if the number of tiles provided exceeds
the number of loops.
---
 .../Linalg/TransformOps/LinalgTransformOps.cpp  |  6 ++++++
 mlir/test/Dialect/Linalg/transform-op-tile.mlir | 17 +++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 1b2283c054c7d34..6539641030f905b 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2533,6 +2533,12 @@ transform::TileOp::apply(transform::TransformRewriter &rewriter,
       diag.attachNote(op->getLoc()) << "target op";
       return diag;
     }
+    if (tileSizes.size() > tilingInterface.getLoopIteratorTypes().size()) {
+      DiagnosedSilenceableFailure diag = emitSilenceableError()
+                                         << "too many tiles for";
+      diag.attachNote(op->getLoc()) << "target op";
+      return diag;
+    }
 
     scf::SCFTilingOptions tilingOptions;
     if (!tileSizes.empty()) {
diff --git a/mlir/test/Dialect/Linalg/transform-op-tile.mlir b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
index d4629dcb29c3efc..1ed2ff3732fc324 100644
--- a/mlir/test/Dialect/Linalg/transform-op-tile.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
@@ -220,3 +220,20 @@ transform.sequence failures(propagate) {
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1, %loops:3 = transform.structured.tile %0 [4, 4, [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
 }
+
+// -----
+
+func.func @matmul(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
+                    %arg2: tensor<128x128xf32>) ->  tensor<128x128xf32> {
+  // expected-note @below {{target op}}
+  %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>) -> tensor<128x128xf32>
+  return %0 : tensor<128x128xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  // expected-error @below {{too many tiles for}}
+  %1, %loops = transform.structured.tile %0 [1, 0, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+}