[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 9fcf0bb6e27dfdc3937452d4ca6543f24dd5f775 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     | 22 +++++++++++++++++++
 2 files changed, 28 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..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
+}