[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 dc644957f12c39e798d561d43684a5e23ab1bc3b 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 patterns by
special-casing logic.
---
 .../VectorEmulateMaskedLoadStore.cpp           | 11 ++++++-----
 .../vector-emulate-masked-load-store.mlir      | 18 ++++++++++++++++++
 2 files changed, 24 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..bf5a5d22e3446 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,21 @@ 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_rank0
+// CHECK-SAME: (%[[ARG0:.*]]: memref<f32, strided<[], offset: ?>>
+// CHECK: memref.load %[[ARG0]][]
+func.func @vector_masked_load_rank0(%arg0: memref<f32, strided<[], offset: ?>>, %arg3: vector<1xi1>,
+                                   %arg4: vector<1xf32>) -> vector<1xf32> {
+  %0 = vector.maskedload %arg0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32> into vector<1xf32>
+  return %0: vector<1xf32>
+}
+
+// CHECK-LABEL:  @vector_masked_store_rank0
+// CHECK-SAME: (%[[ARG0:.*]]: memref<f32, strided<[], offset: ?>>
+// CHECK: memref.store %{{.*}}, %[[ARG0]][]
+func.func @vector_masked_store_rank0(%arg0: memref<f32, strided<[], offset: ?>>, %arg3: vector<1xi1>,
+                                   %arg4: vector<1xf32>) {
+  vector.maskedstore %arg0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<1xi1>, vector<1xf32>
+  return
+}