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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-sme

@llvm/pr-subscribers-mlir

Author: Federico Bruzzone (FedericoBruzzone)

<details>
<summary>Changes</summary>

This PR simply fix what we identified in iree-org/iree#<!-- -->24689.

@<!-- -->banach-space @<!-- -->egebeysel 


---
Full diff: https://github.com/llvm/llvm-project/pull/207947.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp (+3-1) 
- (modified) mlir/test/Dialect/ArmSME/vector-legalization.mlir (+33) 


``````````diff
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

``````````

</details>


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