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

[Prathamesh Tagore]

https://github.com/meshtag updated https://github.com/llvm/llvm-project/pull/173325

>From 9ad35d3c4b5cfd28b6b03562024c66d43be1aead Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Tue, 23 Dec 2025 02:30:55 +0100
Subject: [PATCH] [mlir][vector] Fix masked load/store emulation for rank-0
 memrefs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../Transforms/VectorEmulateMaskedLoadStore.cpp     | 11 ++++++-----
 .../Vector/vector-emulate-masked-load-store.mlir    | 13 +++++++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)

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..a07fd6997cdb9 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,16 @@ 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_masked_load_store_rank0
+// CHECK: %[[SUBVIEW:.*]] = memref.subview
+// CHECK: memref.load %[[SUBVIEW]][]
+// CHECK: memref.store %{{.*}}, %[[SUBVIEW]][]
+func.func @vector_masked_load_store_rank0(%arg0: memref<12x32xf32>, %arg1: index,
+                                   %arg2: index, %arg3: vector<1xi1>,
+                                   %arg4: vector<1xf32>) -> vector<1xf32> {
+  %subview = memref.subview %arg0[%arg1, %arg2] [1, 1] [1, 1] : memref<12x32xf32> to memref<f32, strided<[], offset: ?>>
+  %0 = vector.maskedload %subview[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32> into vector<1xf32>
+  vector.maskedstore %subview[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32>
+  return %0: vector<1xf32>
+}