[Mlir-commits] [mlir] [MLIR][Linalg] Fix crash in packMatmulGreedily on non-matmul ops (PR #218947)

[Chibuoyim Ogbonna]

https://github.com/bruteforceboy created https://github.com/llvm/llvm-project/pull/218947

The following case currently causes a crash:
```
func.func @f(%a: tensor<1x1x1xf32>, %b: tensor<1x1x1xf32>, %c: tensor<1x1x1xf32>) -> tensor<1x1x1xf32> {
  %0 = linalg.add ins(%a, %b : tensor<1x1x1xf32>, tensor<1x1x1xf32>) outs(%c : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
  return %0 : tensor<1x1x1xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match interface{LinalgOp} in %arg0 : (!transform.any_op) -> !transform.any_op
    transform.structured.pack_greedily %0
        matmul_packed_sizes = [8, 8, 8] matmul_inner_dims_order = [0, 1, 2]
      : (!transform.any_op) -> !transform.any_op
    transform.yield
  }
}
```
with `bin/mlir-opt repro.mlir --transform-interpreter`

`transform.structured.pack_greedily` crashes on any op with 2 inputs, 1 init and 3+ loops that is not actually matmul like. The `linalg.add` here is rank 3, so it has 3 loops, passes the `numLoops <= 2` check, and reaches the `m/n/k` usage. The crash happens at:
```
  int64_t mPos = maybeDimensions->m.back(), nPos = maybeDimensions->n.back(),
          kPos = maybeDimensions->k.back();
```
when trying to get the last element of an *empty* vector.

This patch checks that none of the inferred `m/n/k` vectors is empty, and if any of them is, the transformation will not pack this op.

>From 13f71d57ae1660a45207b0e3bc26adc98dc1489f Mon Sep 17 00:00:00 2001
From: bruteforceboy <chibuoyim.faith.ogbonna at huawei.com>
Date: Wed, 26 Aug 2026 14:24:07 +0100
Subject: [PATCH] [MLIR][Linalg] Fix crash in packMatmulGreedily on non-matmul
 ops

---
 .../Dialect/Linalg/Transforms/Transforms.cpp  |  4 +-
 .../Linalg/transform-pack-greedily.mlir       | 86 +++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index aa021cbad24a7..537b54d02fc4d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -806,7 +806,9 @@ linalg::packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
   // 1. Infer dims that are important for matmul.
   FailureOr<ContractionDimensions> maybeDimensions =
       inferContractionDims(linalgOp);
-  if (failed(maybeDimensions)) {
+  // The inferred m/n/k may be empty when the op is not actually matmul-like.
+  if (failed(maybeDimensions) || maybeDimensions->m.empty() ||
+      maybeDimensions->n.empty() || maybeDimensions->k.empty()) {
     LDBG() << "couldn't infer matmul iterators in: " << linalgOp;
     return rewriter.notifyMatchFailure(linalgOp,
                                        "couldn't infer matmul iterators");
diff --git a/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir b/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
index 5812c4db88247..8f350744b62d7 100644
--- a/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
+++ b/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
@@ -369,6 +369,92 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+!A = tensor<1023x255x127xf32>
+
+// CHECK-LABEL: @elementwise_no_mnk(
+func.func @elementwise_no_mnk(%A : !A, %B : !A, %C : !A) -> !A {
+  //      CHECK: linalg.add
+  %0 = linalg.add ins(%A, %B : !A, !A) outs(%C : !A) -> !A
+  return %0 : !A
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+    %add = transform.structured.match ops{["linalg.add"]} in %module_op
+      : (!transform.any_op) -> !transform.op<"linalg.add">
+    transform.structured.pack_greedily %add
+        matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
+      : (!transform.op<"linalg.add">) -> !transform.any_op
+      transform.yield
+  }
+}
+
+// -----
+
+!A = tensor<42x1023x255xf32>
+!X = tensor<42x255xf32>
+!Y = tensor<42x1023xf32>
+
+// CHECK-LABEL: @batch_matvec_no_n(
+func.func @batch_matvec_no_n(%A : !A, %x : !X, %y : !Y) -> !Y {
+  //      CHECK: linalg.batch_matvec
+  %0 = linalg.batch_matvec ins(%A, %x : !A, !X) outs(%y : !Y) -> !Y
+  return %0 : !Y
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+    %batch_matvec = transform.structured.match ops{["linalg.batch_matvec"]} in %module_op
+      : (!transform.any_op) -> !transform.op<"linalg.batch_matvec">
+    transform.structured.pack_greedily %batch_matvec
+        matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
+      : (!transform.op<"linalg.batch_matvec">) -> !transform.any_op
+      transform.yield
+  }
+}
+
+// -----
+
+!A = tensor<1023x255xf32>
+!B = tensor<127xf32>
+!C = tensor<1023x127xf32>
+
+#mnr_accesses = [
+  affine_map<(m, n, r) -> (m, r)>,
+  affine_map<(m, n, r) -> (n)>,
+  affine_map<(m, n, r) -> (m, n)>
+]
+#mnr_trait = {
+  indexing_maps = #mnr_accesses,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+// CHECK-LABEL: @lhs_only_reduction_no_k(
+func.func @lhs_only_reduction_no_k(%A : !A, %B : !B, %C : !C) -> !C {
+  //  CHECK-NOT: linalg.pack
+  //      CHECK: linalg.generic
+  %0 = linalg.generic #mnr_trait ins(%A, %B : !A, !B) outs(%C : !C) {
+    ^bb0(%a: f32, %b: f32, %c: f32):
+      %d = arith.mulf %a, %b : f32
+      %e = arith.addf %c, %d : f32
+      linalg.yield %e : f32
+  } -> !C
+  return %0 : !C
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+    %generic = transform.structured.match ops{["linalg.generic"]} in %module_op
+      : (!transform.any_op) -> !transform.op<"linalg.generic">
+    transform.structured.pack_greedily %generic
+        matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
+      : (!transform.op<"linalg.generic">) -> !transform.any_op
+      transform.yield
+  }
+}
+
+// -----
+
 func.func @no_padding_on_packs(%A: tensor<32x32xf32>, %B: tensor<32x32xf32>, %C: tensor<32x32xf32>)
     -> tensor<32x32xf32> {
   %0 = linalg.matmul  ins(%A, %B: tensor<32x32xf32>, tensor<32x32xf32>)