[Mlir-commits] [mlir] [mlir][ArmSME] fix f64 scalable matmul crashes in `VectorLegalizationPass` (PR #207947)

Federico Bruzzone llvmlistbot at llvm.org
Tue Jul 7 07:46:15 PDT 2026


https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/207947

>From 5463401c92343bcf98437f9aca344c09b496358c Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 7 Jul 2026 11:20:48 +0200
Subject: [PATCH 1/3] [mlir][ArmSME] f64 scalable matmul crashes in
 VectorLegalizationPass

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 .../ArmSME/Transforms/VectorLegalization.cpp  |  4 ++-
 .../Dialect/ArmSME/vector-legalization.mlir   | 33 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
index 6e2654ab468f9..7a6c510920e22 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
@@ -943,7 +943,9 @@ struct LowerColumnTransferReadToLoops
     auto upperBound = createVscaleMultiple(numRows);
     auto step = arith::ConstantIndexOp::create(rewriter, loc, 1);
     Value init = arith::ConstantOp::create(
-        rewriter, loc, newResType, DenseElementsAttr::get(newResType, 0.0f));
+        rewriter, loc, newResType,
+        DenseElementsAttr::get(newResType,
+                               rewriter.getZeroAttr(resType.getElementType())));
 
     scf::ForOp loadLoop;
     {
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 6cdf576272ebc..3490f67cf68ff 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -650,6 +650,39 @@ func.func @xfer_read_scalable_column(%a: index, %b: index, %pad: f32, %src: memr
 
 // -----
 
+/// Same as above but for f64, which exercises a different zero-init constant
+/// (see https://github.com/iree-org/iree/issues/24689).
+
+// CHECK-LABEL:   func.func @xfer_read_scalable_column_f64(
+// CHECK-SAME:      %[[IDX_0:[a-zA-Z0-9]+]]: index,
+// CHECK-SAME:      %[[IDX_1:[a-zA-Z0-9]+]]: index,
+// CHECK-SAME:      %[[PAD:.*]]: f64,
+// CHECK-SAME:      %[[SRC:.*]]: memref<?x?xf64>) -> vector<[4]x1xf64> {
+func.func @xfer_read_scalable_column_f64(%a: index, %b: index, %pad: f64, %src: memref<?x?xf64>) -> (vector<[4]x1xf64>) {
+  // CHECK:           %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf64>
+  // CHECK:           %[[STEP:.*]] = arith.constant 1 : index
+  // CHECK:           %[[C4:.*]] = arith.constant 4 : index
+  // CHECK:           %[[LB:.*]] = arith.constant 0 : index
+  // CHECK:           %[[VSCALE:.*]] = vector.vscale
+  // CHECK:           %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+
+  // <scf.for>
+  // CHECK:           %[[SCF:.*]] = scf.for %[[IND_VAR:.*]] = %[[LB]] to %[[C4_VSCALE]] step %[[STEP]] iter_args(%[[SCF_RES:.*]] = %[[INIT]]) -> (vector<[4]xf64>) {
+  // CHECK:             %[[IDX_0_UPDATED:.*]] = arith.addi %[[IND_VAR]], %[[IDX_0]] : index
+  // CHECK:             %[[VAL_10:.*]] = memref.load %[[SRC]][%[[IDX_0_UPDATED]], %[[IDX_1]]] : memref<?x?xf64>
+  // CHECK:             %[[RES_UPDATED:.*]] = vector.insert %[[VAL_10]], %[[SCF_RES]] [%[[IND_VAR]]] : f64 into vector<[4]xf64>
+  // CHECK:             scf.yield %[[RES_UPDATED]] : vector<[4]xf64>
+  // CHECK:           }
+
+  // <shape-cast>
+  // CHECK:           %[[SC:.*]] = vector.shape_cast %[[SCF]] : vector<[4]xf64> to vector<[4]x1xf64>
+  // CHECK:           return %[[SC]]
+  %read = vector.transfer_read %src[%a, %b], %pad : memref<?x?xf64>, vector<[4]x1xf64>
+  return %read : vector<[4]x1xf64>
+}
+
+// -----
+
 // CHECK-LABEL:   func.func @negative_xfer_read_scalable_column_x2
 func.func @negative_xfer_read_scalable_column_x2(%a: index, %b: index, %pad: f32, %src: memref<?x?xf32>) -> (vector<[4]x2xf32>) {
   // CHECK-NOT: scf.for

>From f85aab3c8d460f844675c31eed12015b3979b4b4 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 7 Jul 2026 11:37:02 +0200
Subject: [PATCH 2/3] Update mlir/test/Dialect/ArmSME/vector-legalization.mlir

Co-authored-by: Artem Gindinson <gindinson at roofline.ai>
---
 mlir/test/Dialect/ArmSME/vector-legalization.mlir | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 3490f67cf68ff..a75590770799e 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -651,7 +651,8 @@ func.func @xfer_read_scalable_column(%a: index, %b: index, %pad: f32, %src: memr
 // -----
 
 /// Same as above but for f64, which exercises a different zero-init constant
-/// (see https://github.com/iree-org/iree/issues/24689).
+/// Same as xfer_read_scalable_column but for f64, which exercises a different
+/// zero-init constant (see https://github.com/iree-org/iree/issues/24689).
 
 // CHECK-LABEL:   func.func @xfer_read_scalable_column_f64(
 // CHECK-SAME:      %[[IDX_0:[a-zA-Z0-9]+]]: index,

>From a95e020ac4c1cea33bbef45b1826ad558e68788f Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 7 Jul 2026 16:44:09 +0200
Subject: [PATCH 3/3] Address comments

Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
 mlir/test/Dialect/ArmSME/vector-legalization.mlir | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index a75590770799e..a31417f106a96 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -650,10 +650,6 @@ func.func @xfer_read_scalable_column(%a: index, %b: index, %pad: f32, %src: memr
 
 // -----
 
-/// Same as above but for f64, which exercises a different zero-init constant
-/// Same as xfer_read_scalable_column but for f64, which exercises a different
-/// zero-init constant (see https://github.com/iree-org/iree/issues/24689).
-
 // CHECK-LABEL:   func.func @xfer_read_scalable_column_f64(
 // CHECK-SAME:      %[[IDX_0:[a-zA-Z0-9]+]]: index,
 // CHECK-SAME:      %[[IDX_1:[a-zA-Z0-9]+]]: index,



More information about the Mlir-commits mailing list