[mlir][vector] Propagate alignment when emulating masked{load,stores}. #155648

amd-eochoalo · 2025-08-27T16:21:38Z

Propagate alignment from vector.maskedload and vector.maskedstore to memref.load and memref.store during VectorEmulateMaskedLoadStore pass.

github-actions · 2025-08-27T16:24:53Z

✅ With the latest revision this PR passed the C/C++ code formatter.

llvmbot · 2025-08-27T17:13:36Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Erick Ochoa Lopez (amd-eochoalo)

Changes

Improvements that may follow in another NFC pull request

Add constructors for llvm::MaybeAlign for the Memref dialect just like in the Vector dialect.
Use the IntValidAlign from [mlir][vector] Improve vector dialect documentation and constraint predicates #155259

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

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp (+10-3)
(modified) mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir (+30)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
index cb3e8dc67a1ae..a00a1352ce40d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp
@@ -64,6 +64,9 @@ struct VectorMaskedLoadOpConverter final
     Value mask = maskedLoadOp.getMask();
     Value base = maskedLoadOp.getBase();
     Value iValue = maskedLoadOp.getPassThru();
+    bool nontemporal = false;
+    auto alignment = maskedLoadOp.getAlignment();
+    uint64_t align = alignment.has_value() ? alignment.value() : 0;
     auto indices = llvm::to_vector_of<Value>(maskedLoadOp.getIndices());
     Value one = arith::ConstantOp::create(rewriter, loc, indexType,
                                           IntegerAttr::get(indexType, 1));
@@ -73,8 +76,8 @@ struct VectorMaskedLoadOpConverter final
       auto ifOp = scf::IfOp::create(
           rewriter, loc, maskBit,
           [&](OpBuilder &builder, Location loc) {
-            auto loadedValue =
-                memref::LoadOp::create(builder, loc, base, indices);
+            auto loadedValue = memref::LoadOp::create(
+                builder, loc, base, indices, nontemporal, align);
             auto combinedValue =
                 vector::InsertOp::create(builder, loc, loadedValue, iValue, i);
             scf::YieldOp::create(builder, loc, combinedValue.getResult());
@@ -132,6 +135,9 @@ struct VectorMaskedStoreOpConverter final
     Value mask = maskedStoreOp.getMask();
     Value base = maskedStoreOp.getBase();
     Value value = maskedStoreOp.getValueToStore();
+    bool nontemporal = false;
+    auto alignment = maskedStoreOp.getAlignment();
+    uint64_t align = alignment.has_value() ? alignment.value() : 0;
     auto indices = llvm::to_vector_of<Value>(maskedStoreOp.getIndices());
     Value one = arith::ConstantOp::create(rewriter, loc, indexType,
                                           IntegerAttr::get(indexType, 1));
@@ -141,7 +147,8 @@ struct VectorMaskedStoreOpConverter final
       auto ifOp = scf::IfOp::create(rewriter, loc, maskBit, /*else=*/false);
       rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
       auto extractedValue = vector::ExtractOp::create(rewriter, loc, value, i);
-      memref::StoreOp::create(rewriter, loc, extractedValue, base, indices);
+      memref::StoreOp::create(rewriter, loc, extractedValue, base, indices,
+                              nontemporal, align);
 
       rewriter.setInsertionPointAfter(ifOp);
       indices.back() =
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 3867f075af8e4..e74eb08339684 100644
--- a/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
+++ b/mlir/test/Dialect/Vector/vector-emulate-masked-load-store.mlir
@@ -54,6 +54,22 @@ func.func @vector_maskedload(%arg0 : memref<4x5xf32>) -> vector<4xf32> {
   return %0: vector<4xf32>
 }
 
+// CHECK-LABEL:  @vector_maskedload_with_alignment
+//       CHECK:       memref.load
+//       CHECK-SAME:  {alignment = 8 : i64}
+//       CHECK:       memref.load
+//       CHECK-SAME:  {alignment = 8 : i64}
+func.func @vector_maskedload_with_alignment(%arg0 : memref<4x5xf32>) -> vector<4xf32> {
+  %idx_0 = arith.constant 0 : index
+  %idx_1 = arith.constant 1 : index
+  %idx_4 = arith.constant 4 : index
+  %mask = vector.create_mask %idx_1 : vector<4xi1>
+  %s = arith.constant 0.0 : f32
+  %pass_thru = vector.splat %s : vector<4xf32>
+  %0 = vector.maskedload %arg0[%idx_0, %idx_4], %mask, %pass_thru {alignment = 8}: memref<4x5xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+  return %0: vector<4xf32>
+}
+
 // CHECK-LABEL:  @vector_maskedstore
 //  CHECK-SAME:  (%[[ARG0:.*]]: memref<4x5xf32>, %[[ARG1:.*]]: vector<4xf32>) {
 //   CHECK-DAG:  %[[C7:.*]] = arith.constant 7 : index
@@ -93,3 +109,17 @@ func.func @vector_maskedstore(%arg0 : memref<4x5xf32>, %arg1 : vector<4xf32>) {
   vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 : memref<4x5xf32>, vector<4xi1>, vector<4xf32>
   return
 }
+
+// CHECK-LABEL:  @vector_maskedstore_with_alignment
+//       CHECK:       memref.store
+//       CHECK-SAME:  {alignment = 8 : i64}
+//       CHECK:       memref.store
+//       CHECK-SAME:  {alignment = 8 : i64}
+func.func @vector_maskedstore_with_alignment(%arg0 : memref<4x5xf32>, %arg1 : vector<4xf32>) {
+  %idx_0 = arith.constant 0 : index
+  %idx_1 = arith.constant 1 : index
+  %idx_4 = arith.constant 4 : index
+  %mask = vector.create_mask %idx_1 : vector<4xi1>
+  vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 { alignment = 8 } : memref<4x5xf32>, vector<4xi1>, vector<4xf32>
+  return
+}

banach-space

Thanks, just a minor suggestion.

mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp

…-alignment

banach-space · 2025-09-08T07:09:00Z

This has landed without an explicit approval. Please note that we do require explicit approvals in LLVM:

https://llvm.org/docs/GitHub.html#approvals

This change LGTM, so this is also a note to myself to be more explicit next time. Thanks for contributing!

Groverkss · 2025-09-08T07:23:48Z

This has landed without an explicit approval. Please note that we do require explicit approvals in LLVM:

https://llvm.org/docs/GitHub.html#approvals

This change LGTM, so this is also a note to myself to be more explicit next time. Thanks for contributing!

I don't know if we have to be super strict about how that rule is enforced. LLVM allows post-review commit for small, trivial patches (like this one). A passing by comment about code style doesn't really count as a review that needs a follow up IMO. Waiting for a small trivial change like this does have developer cost.

amd-eochoalo · 2025-09-08T14:18:12Z

@banach-space sorry about this. I must have believed that you had approved after your review, but I was clearly wrong. Not my intention @banach-space , happy to make changes if you think they are necessary.

banach-space · 2025-09-08T14:38:18Z

Not my intention @banach-space , happy to make changes if you think they are necessary.

No worries - this was also a note to myself:

This change LGTM, so this is also a note to myself to be more explicit next time.

As in, all good, approved :)

[mlir][vector] Propagate alignment when emulating masked{load,stores}.

52d6a7b

amd-eochoalo force-pushed the eochoa/2025-08-27/vector-emulate-masked-load-alignment branch from d2dc9bd to 52d6a7b Compare August 27, 2025 16:27

amd-eochoalo self-assigned this Aug 27, 2025

amd-eochoalo marked this pull request as ready for review August 27, 2025 17:13

amd-eochoalo requested review from banach-space, dcaballe, hanhanW and nicolasvasilache as code owners August 27, 2025 17:13

llvmbot added mlir:vectorops mlir mlir:vector labels Aug 27, 2025

amd-eochoalo added 2 commits August 28, 2025 05:21

Use explicit type instead of auto

6432c87

Style

7983260

banach-space reviewed Sep 5, 2025

View reviewed changes

mlir/lib/Dialect/Vector/Transforms/VectorEmulateMaskedLoadStore.cpp Outdated Show resolved Hide resolved

amd-eochoalo added 2 commits September 5, 2025 06:38

Propagate constant

0f5578e

Merge branch 'main' into eochoa/2025-08-27/vector-emulate-masked-load…

4988e8f

…-alignment

amd-eochoalo enabled auto-merge (squash) September 5, 2025 13:44

amd-eochoalo merged commit c5e6d56 into llvm:main Sep 5, 2025
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][vector] Propagate alignment when emulating masked{load,stores}. #155648

[mlir][vector] Propagate alignment when emulating masked{load,stores}. #155648

Uh oh!

amd-eochoalo commented Aug 27, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Aug 27, 2025 •

edited

Loading

Uh oh!

llvmbot commented Aug 27, 2025 •

edited

Loading

Uh oh!

banach-space left a comment

Uh oh!

Uh oh!

Uh oh!

banach-space commented Sep 8, 2025

Uh oh!

Groverkss commented Sep 8, 2025 •

edited

Loading

Uh oh!

amd-eochoalo commented Sep 8, 2025

Uh oh!

banach-space commented Sep 8, 2025

Uh oh!

Uh oh!

[mlir][vector] Propagate alignment when emulating masked{load,stores}. #155648

[mlir][vector] Propagate alignment when emulating masked{load,stores}. #155648

Uh oh!

Conversation

amd-eochoalo commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

banach-space commented Sep 8, 2025

Uh oh!

Groverkss commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

amd-eochoalo commented Sep 8, 2025

Uh oh!

banach-space commented Sep 8, 2025

Uh oh!

Uh oh!

amd-eochoalo commented Aug 27, 2025 •

edited

Loading

github-actions bot commented Aug 27, 2025 •

edited

Loading

llvmbot commented Aug 27, 2025 •

edited

Loading

Groverkss commented Sep 8, 2025 •

edited

Loading