[Mlir-commits] [mlir] 32d46c9 - [mlir][vector] Fix masked load/store emulation for rank-0 memrefs (#173325)

[llvmlistbot at llvm.org]

Author: Prathamesh Tagore
Date: 2026-01-05T17:56:09Z
New Revision: 32d46c9033ce5129906e4d164535271637b3aaa5

URL: https://github.com/llvm/llvm-project/commit/32d46c9033ce5129906e4d164535271637b3aaa5
DIFF: https://github.com/llvm/llvm-project/commit/32d46c9033ce5129906e4d164535271637b3aaa5.diff

LOG: [mlir][vector] Fix masked load/store emulation for rank-0 memrefs (#173325)

Added rank‑0 handling to masked load/store emulation by reinterpreting
rank‑0 memrefs as 1‑D buffers with a synthetic index, preventing
empty‑indices crashes.

Fixes https://github.com/llvm/llvm-project/issues/131243

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
    mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
index 7acc120508a44..3ca08f3099bd3 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
@@ -84,9 +84,9 @@ struct VectorMaskedLoadOpConverter final
             scf::YieldOp::create(builder, loc, iValue);
           });
       iValue = ifOp.getResult(0);
-
-      indices.back() =
-          arith::AddIOp::create(rewriter, loc, indices.back(), one);
+      if (!indices.empty())
+        indices.back() =
+            arith::AddIOp::create(rewriter, loc, indices.back(), one);
     }
 
     rewriter.replaceOp(maskedLoadOp, iValue);
@@ -148,8 +148,9 @@ struct VectorMaskedStoreOpConverter final
           llvm::MaybeAlign(maskedStoreOp.getAlignment().value_or(0)));
 
       rewriter.setInsertionPointAfter(ifOp);
-      indices.back() =
-          arith::AddIOp::create(rewriter, loc, indices.back(), one);
+      if (!indices.empty())
+        indices.back() =
+            arith::AddIOp::create(rewriter, loc, indices.back(), one);
     }
 
     rewriter.eraseOp(maskedStoreOp);

diff  --git a/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir b/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
index 6e5d68c859e2c..b904fb1ab3f7a 100644
--- a/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
+++ b/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
@@ -123,3 +123,25 @@ func.func @vector_maskedstore_with_alignment(%arg0 : memref<4x5xf32>, %arg1 : ve
   vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 { alignment = 8 } : memref<4x5xf32>, vector<4xi1>, vector<4xf32>
   return
 }
+
+// CHECK-LABEL:  @vector_maskedload_rank0
+// CHECK-SAME:  %[[ARG0:.*]]: memref<f32>, %[[ARG3:.*]]: vector<1xi1>, %[[ARG4:.*]]: vector<1xf32>
+// CHECK-NEXT:  %[[VAL0:.*]] = vector.extract %[[ARG3]][0]
+// CHECK-NEXT:  scf.if %[[VAL0]] -> (vector<1xf32>) {
+// CHECK-NEXT:  %[[VAL1:.*]] = memref.load %[[ARG0]][]
+// CHECK-NEXT:  vector.insert %[[VAL1]], %[[ARG4]] [0]
+func.func @vector_maskedload_rank0(%arg0: memref<f32>, %arg3: vector<1xi1>, %arg4: vector<1xf32>) -> vector<1xf32> {
+  %0 = vector.maskedload %arg0[], %arg3, %arg4 : memref<f32>, vector<1xi1>, vector<1xf32> into vector<1xf32>
+  return %0: vector<1xf32>
+}
+
+// CHECK-LABEL:  @vector_maskedstore_rank0
+// CHECK-SAME:  %[[ARG0:.*]]: memref<f32>, %[[ARG3:.*]]: vector<1xi1>, %[[ARG4:.*]]: vector<1xf32>
+// CHECK-NEXT:  %[[VAL0:.*]] = vector.extract %[[ARG3]][0]
+// CHECK-NEXT:  scf.if %[[VAL0]] {
+// CHECK-NEXT:  %[[VAL1:.*]] = vector.extract %[[ARG4]][0]
+// CHECK-NEXT:  memref.store %[[VAL1]], %[[ARG0]][]
+func.func @vector_maskedstore_rank0(%arg0: memref<f32>, %arg3: vector<1xi1>, %arg4: vector<1xf32>) {
+  vector.maskedstore %arg0[], %arg3, %arg4 : memref<f32>, vector<1xi1>, vector<1xf32>
+  return
+}