[mlir][llvm] Add `align` attribute to `llvm.intr.masked.{expandload,compressstore}` #153063

amd-eochoalo · 2025-08-11T18:14:18Z

Add requiresArgsAndResultsAttr to LLVM_OneResultIntrOp
Add args_attrs to llvm.intr.masked.{expandload,compressstore}

The LLVM intrinsics llvm.intr.masked.expandload and llvm.intr.masked.compressstore both allow an optional align parameter attribute to be set which defaults to one.

Inlining the documentation below for llvm.intr.masked.expandload 's
and llvm.intr.masked.compressstore's arguments respectively

The align parameter attribute can be provided for the first argument. The pointer alignment defaults to 1.

The align parameter attribute can be provided for the second argument. The pointer alignment defaults to 1.

llvmbot · 2025-08-11T18:25:04Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Erick Ochoa Lopez (amd-eochoalo)

Changes

The LLVM intrinsics llvm.intr.masked.expandload and llvm.intr.masked.compressstore both allow an optional align parameter attribute to be set which defaults to one.

Inlining the documentation below for llvm.intr.masked.expandload 's
and llvm.intr.masked.compressstore's arguments respectively

> The align parameter attribute can be provided for the first argument. The pointer alignment defaults to 1.

> The align parameter attribute can be provided for the second argument. The pointer alignment defaults to 1.

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

4 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td (+38-2)
(modified) mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir (+4-4)
(modified) mlir/test/Target/LLVMIR/Import/intrinsic.ll (+11-2)
(modified) mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir (+11)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index 76b08e664ee76..e2ad3a112ada2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -1070,13 +1070,49 @@ def LLVM_masked_scatter : LLVM_ZeroResultIntrOp<"masked.scatter"> {
 
 /// Create a call to Masked Expand Load intrinsic.
 def LLVM_masked_expandload : LLVM_IntrOp<"masked.expandload", [0], [], [], 1> {
-  let arguments = (ins LLVM_AnyPointer, LLVM_VectorOf<I1>, LLVM_AnyVector);
+  let arguments = (ins LLVM_AnyPointer:$base,
+                       LLVM_VectorOf<I1>:$mask,
+                       LLVM_AnyVector:$passthru,
+                       DefaultValuedAttr<I32Attr, "1">:$align);
+
+  let assemblyFormat = "`(` $base`,` $mask`,` $passthru `)` attr-dict `:` functional-type(operands, results)";
+
+  string llvmBuilder = [{
+    $res = builder.CreateMaskedExpandLoad(
+      $_resultType, $base, $align != 1 ? llvm::Align($align) : llvm::MaybeAlign(), $mask, $passthru);
+  }];
+  string mlirBuilder = [{
+    auto *intrinInst = dyn_cast<llvm::IntrinsicInst>(inst);
+    llvm::MaybeAlign alignment = intrinInst->getParamAlign(0);
+    IntegerAttr attr = $_builder.getI32IntegerAttr(alignment.has_value() ? alignment->value() : 1);
+    $res = LLVM::masked_expandload::create($_builder, $_location,
+      $_resultType, $base, $mask, $passthru, attr);
+  }];
 }
 
 /// Create a call to Masked Compress Store intrinsic.
 def LLVM_masked_compressstore
     : LLVM_IntrOp<"masked.compressstore", [], [0], [], 0> {
-  let arguments = (ins LLVM_AnyVector, LLVM_AnyPointer, LLVM_VectorOf<I1>);
+  let arguments = (ins LLVM_AnyVector:$vector,
+                       LLVM_AnyPointer:$base,
+                       LLVM_VectorOf<I1>:$mask,
+                       DefaultValuedAttr<I32Attr, "1">:$align);
+
+  let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
+
+  let assemblyFormat = "`(` $vector `,` $base `,` $mask `)` attr-dict `:` functional-type(operands, results)";
+
+  string llvmBuilder = [{
+    builder.CreateMaskedCompressStore(
+      $vector, $base, $align != 1 ? llvm::Align($align) : llvm::MaybeAlign(), $mask);
+  }];
+  string mlirBuilder = [{
+    auto *intrinInst = dyn_cast<llvm::IntrinsicInst>(inst);
+    llvm::MaybeAlign alignment = intrinInst->getParamAlign(1);
+    IntegerAttr attr = $_builder.getI32IntegerAttr(alignment.has_value() ? alignment->value() : 1);
+    $_op = LLVM::masked_compressstore::create($_builder, $_location,
+      $vector, $base, $mask, attr);
+  }];
 }
 
 //
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
index 5a424a8ac0d5f..60dce615cfc05 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm-interface.mlir
@@ -2134,7 +2134,7 @@ func.func @expand_load_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2: vec
 // CHECK: %[[CO:.*]] = arith.constant 0 : index
 // CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
 // CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
-// CHECK: %[[E:.*]] = "llvm.intr.masked.expandload"(%[[P]], %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<11xi1>, vector<11xf32>) -> vector<11xf32>
+// CHECK: %[[E:.*]] = llvm.intr.masked.expandload(%[[P]], %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<11xi1>, vector<11xf32>) -> vector<11xf32>
 // CHECK: return %[[E]] : vector<11xf32>
 
 // -----
@@ -2145,7 +2145,7 @@ func.func @expand_load_op_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>, %a
   return %0 : vector<11xindex>
 }
 // CHECK-LABEL: func @expand_load_op_index
-// CHECK: %{{.*}} = "llvm.intr.masked.expandload"(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<11xi1>, vector<11xi64>) -> vector<11xi64>
+// CHECK: %{{.*}} = llvm.intr.masked.expandload(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<11xi1>, vector<11xi64>) -> vector<11xi64>
 
 // -----
 
@@ -2163,7 +2163,7 @@ func.func @compress_store_op(%arg0: memref<?xf32>, %arg1: vector<11xi1>, %arg2:
 // CHECK: %[[CO:.*]] = arith.constant 0 : index
 // CHECK: %[[C:.*]] = builtin.unrealized_conversion_cast %[[CO]] : index to i64
 // CHECK: %[[P:.*]] = llvm.getelementptr %{{.*}}[%[[C]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
-// CHECK: "llvm.intr.masked.compressstore"(%{{.*}}, %[[P]], %{{.*}}) : (vector<11xf32>, !llvm.ptr, vector<11xi1>) -> ()
+// CHECK: llvm.intr.masked.compressstore(%{{.*}}, %[[P]], %{{.*}}) : (vector<11xf32>, !llvm.ptr, vector<11xi1>) -> ()
 
 // -----
 
@@ -2173,7 +2173,7 @@ func.func @compress_store_op_index(%arg0: memref<?xindex>, %arg1: vector<11xi1>,
   return
 }
 // CHECK-LABEL: func @compress_store_op_index
-// CHECK: "llvm.intr.masked.compressstore"(%{{.*}}, %{{.*}}, %{{.*}}) : (vector<11xi64>, !llvm.ptr, vector<11xi1>) -> ()
+// CHECK: llvm.intr.masked.compressstore(%{{.*}}, %{{.*}}, %{{.*}}) : (vector<11xi64>, !llvm.ptr, vector<11xi1>) -> ()
 
 // -----
 
diff --git a/mlir/test/Target/LLVMIR/Import/intrinsic.ll b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
index 9f882ad6f22e8..63bfdfce59521 100644
--- a/mlir/test/Target/LLVMIR/Import/intrinsic.ll
+++ b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
@@ -538,13 +538,22 @@ define void @masked_gather_scatter_intrinsics(<7 x ptr> %vec, <7 x i1> %mask) {
 
 ; CHECK-LABEL:  llvm.func @masked_expand_compress_intrinsics
 define void @masked_expand_compress_intrinsics(ptr %0, <7 x i1> %1, <7 x float> %2) {
-  ; CHECK: %[[val1:.+]] = "llvm.intr.masked.expandload"(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
+  ; CHECK: %[[val1:.+]] = llvm.intr.masked.expandload(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
   %4 = call <7 x float> @llvm.masked.expandload.v7f32(ptr %0, <7 x i1> %1, <7 x float> %2)
-  ; CHECK: "llvm.intr.masked.compressstore"(%[[val1]], %{{.*}}, %{{.*}}) : (vector<7xf32>, !llvm.ptr, vector<7xi1>) -> ()
+  ; CHECK: llvm.intr.masked.compressstore(%[[val1]], %{{.*}}, %{{.*}}) : (vector<7xf32>, !llvm.ptr, vector<7xi1>) -> ()
   call void @llvm.masked.compressstore.v7f32(<7 x float> %4, ptr %0, <7 x i1> %1)
   ret void
 }
 
+; CHECK-LABEL:  llvm.func @masked_expand_compress_intrinsics_with_alignment
+define void @masked_expand_compress_intrinsics_with_alignment(ptr %0, <7 x i1> %1, <7 x float> %2) {
+  ; CHECK: %[[val1:.+]] = llvm.intr.masked.expandload(%{{.*}}, %{{.*}}, %{{.*}}) {align = 8 : i32} : (!llvm.ptr, vector<7xi1>, vector<7xf32>) -> vector<7xf32>
+  %4 = call <7 x float> @llvm.masked.expandload.v7f32(ptr align 8 %0, <7 x i1> %1, <7 x float> %2)
+  ; CHECK: llvm.intr.masked.compressstore(%[[val1]], %{{.*}}, %{{.*}}) {align = 8 : i32} : (vector<7xf32>, !llvm.ptr, vector<7xi1>) -> ()
+  call void @llvm.masked.compressstore.v7f32(<7 x float> %4, ptr align 8 %0, <7 x i1> %1)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @annotate_intrinsics
 define void @annotate_intrinsics(ptr %var, ptr %ptr, i16 %int, ptr %annotation, ptr %fileName, i32 %line, ptr %args) {
   ; CHECK: "llvm.intr.var.annotation"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i32, !llvm.ptr) -> ()
diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index 2b420ed246fb2..11590391493ce 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -577,6 +577,17 @@ llvm.func @masked_expand_compress_intrinsics(%ptr: !llvm.ptr, %mask: vector<7xi1
   llvm.return
 }
 
+// CHECK-LABEL: @masked_expand_compress_intrinsics_with_alignment
+llvm.func @masked_expand_compress_intrinsics_with_alignment(%ptr: !llvm.ptr, %mask: vector<7xi1>, %passthru: vector<7xf32>) {
+  // CHECK: call <7 x float> @llvm.masked.expandload.v7f32(ptr align 8 %{{.*}}, <7 x i1> %{{.*}}, <7 x float> %{{.*}})
+  %0 = "llvm.intr.masked.expandload"(%ptr, %mask, %passthru) { align = 8 : i32 }
+    : (!llvm.ptr, vector<7xi1>, vector<7xf32>) -> (vector<7xf32>)
+  // CHECK: call void @llvm.masked.compressstore.v7f32(<7 x float> %{{.*}}, ptr align 8 %{{.*}}, <7 x i1> %{{.*}})
+  "llvm.intr.masked.compressstore"(%0, %ptr, %mask) { align = 8 : i32 }
+    : (vector<7xf32>, !llvm.ptr, vector<7xi1>) -> ()
+  llvm.return
+}
+
 // CHECK-LABEL: @annotate_intrinsics
 llvm.func @annotate_intrinsics(%var: !llvm.ptr, %int: i16, %ptr: !llvm.ptr, %annotation: !llvm.ptr, %fileName: !llvm.ptr, %line: i32, %attr: !llvm.ptr) {
   // CHECK: call void @llvm.var.annotation.p0.p0(ptr %{{.*}}, ptr %{{.*}}, ptr %{{.*}}, i32 %{{.*}}, ptr %{{.*}})

gysit

Thanks for improving the intrinsics.

Note that there is now a more generic mechanism for importing such LLVM parameter attributes. You can set the requiresArgAndResultAttrs flag on the LLVM_IntrOp base class and add the result and argument dictionaries to the operation argument list. Once this is done import and export should work for all kinds of parameter attributes.

Have a look at the memcpy inline intrinsic for an example:

def LLVM_MemcpyInlineOp :
    ...
    /*requiresArgAndResultAttrs=*/1
    ...
    > {
  dag args = (ins ...);
  // Append the arguments defined by LLVM_IntrOpBase.
  let arguments = !con(args, baseArgs);
  ...

Let me know if you have questions or if it doesn't work.

amd-eochoalo · 2025-08-11T19:41:47Z

@gysit thank you! I will give it a try.

* masked.compressstore * masked.expandload

amd-eochoalo · 2025-08-12T19:03:55Z

@gysit I need some clarification. What you are asking is for the align parameter attribute found in LLVM IR to be represented as an argument attribute in MLIR's LLVM dialect. E.g.,

"llvm.intr.masked.compressstore"(%arg2, %3, %arg1) {arg_attrs = [{}, {align = 1 : i32}, {}]} : (vector<11xf32>, !llvm.ptr, vector<11xi1>) -> ()

Would this be correct? Happy to continue working on this, but I noticed that for example llvm.intr.masked.store 's operation in the LLVM dialect does not follow this approach and instead has an attribute in the attr-dict, which is similar to how I coded it here.

gysit · 2025-08-12T19:47:01Z

Yes that looks good!

The masked store is modeled differently in LLVM (compared to the intrinsics you are updating in this PR):
https://llvm.org/docs/LangRef.html#llvm-masked-store-intrinsics
It has a separate alignment argument instead of an parameter attribute on the pointer.

We should ideally always follow the LLVM IR design. That means if the alignment is a parameter argument, then let's import it to an argument attribute. If the alignment is its own argument, then let's import it as separate argument.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

This commit reverts and fixed up the following commits: * 420351b * 6e4966c * 7467a42 * bddaec9

krzysz00

This version looks pretty good to me

gysit

Thanks for the changes!

LGTM modulo nits

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

amd-eochoalo marked this pull request as ready for review August 11, 2025 18:24

amd-eochoalo requested review from banach-space, dcaballe and nicolasvasilache as code owners August 11, 2025 18:24

llvmbot added mlir:llvm mlir labels Aug 11, 2025

gysit reviewed Aug 11, 2025

View reviewed changes

amd-eochoalo added 4 commits August 11, 2025 18:57

Adds names to parameters of llvm intrinsics:

bddaec9

* masked.compressstore * masked.expandload

Add alignment to llvm.intr.masked.compressstore

7467a42

Add align attribute to llvm.intr.masked.expandload

6e4966c

Add target test

420351b

amd-eochoalo requested a review from gysit August 13, 2025 13:46

kuhar requested a review from krzysz00 August 13, 2025 15:51

gysit reviewed Aug 13, 2025

View reviewed changes

amd-eochoalo commented Aug 13, 2025

View reviewed changes

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td Outdated Show resolved Hide resolved

amd-eochoalo mentioned this pull request Aug 13, 2025

[mlir][vector] Propagate alignment from vector to llvm dialects. #153482

Merged

amd-eochoalo added 4 commits August 13, 2025 14:09

Revert previous changes

a007d04

This commit reverts and fixed up the following commits: * 420351b * 6e4966c * 7467a42 * bddaec9

[mlir][llvm] Add align to masked.compressstore.

ef4b273

Add requiresArgAndResultAttrs to LLVM_OneResultIntrOp

fc9ab38

Add alignment to intr.masked.expandload.

dd50b3d

krzysz00 approved these changes Aug 13, 2025

View reviewed changes

amd-eochoalo added 2 commits August 13, 2025 19:10

Fix placement of attribute alignment.

876962b

Fix typos

990dea1

amd-eochoalo added 7 commits August 13, 2025 19:22

Use base class's builder

37ce068

Use function to get attribute name.

06d938f

Move masked_expandload::build out of tblgen

aa6c08d

Move masked_compressstore::build out of tblgen

483c1c7

Outline the generation of callArgs into a static helper function

be27a32

Use the base class's llvmBuilder.

e2afeb6

Fix return type.

c78134a

gysit approved these changes Aug 14, 2025

View reviewed changes

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp Outdated Show resolved Hide resolved

amd-eochoalo added 3 commits August 14, 2025 13:51

Style: address review comments.

954a26b

Merge branch 'main' into eochoa/2025-08-11/intrinsics-v2

dfac231

Merge branch 'main' into eochoa/2025-08-11/intrinsics-v2

401b9da

amd-eochoalo merged commit 61caab7 into llvm:main Aug 15, 2025
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][llvm] Add `align` attribute to `llvm.intr.masked.{expandload,compressstore}` #153063

[mlir][llvm] Add `align` attribute to `llvm.intr.masked.{expandload,compressstore}` #153063

Uh oh!

amd-eochoalo commented Aug 11, 2025 •

edited

Loading

Uh oh!

llvmbot commented Aug 11, 2025 •

edited

Loading

Uh oh!

gysit left a comment

Uh oh!

amd-eochoalo commented Aug 11, 2025

Uh oh!

amd-eochoalo commented Aug 12, 2025 •

edited

Loading

Uh oh!

gysit commented Aug 12, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

krzysz00 left a comment

Uh oh!

gysit left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[mlir][llvm] Add align attribute to llvm.intr.masked.{expandload,compressstore} #153063

[mlir][llvm] Add align attribute to llvm.intr.masked.{expandload,compressstore} #153063

Uh oh!

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

gysit left a comment

Choose a reason for hiding this comment

Uh oh!

amd-eochoalo commented Aug 11, 2025

Uh oh!

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

Uh oh!

Uh oh!

gysit commented Aug 12, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

krzysz00 left a comment

Choose a reason for hiding this comment

Uh oh!

gysit left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[mlir][llvm] Add `align` attribute to `llvm.intr.masked.{expandload,compressstore}` #153063

[mlir][llvm] Add `align` attribute to `llvm.intr.masked.{expandload,compressstore}` #153063

amd-eochoalo commented Aug 11, 2025 •

edited

Loading

llvmbot commented Aug 11, 2025 •

edited

Loading

amd-eochoalo commented Aug 12, 2025 •

edited

Loading