[Mlir-commits] [mlir] [mlir][vector][memref] Add `alignment` attribute to memory access ops (PR #144344)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 17 13:30:36 PDT 2025
https://github.com/tyb0807 updated https://github.com/llvm/llvm-project/pull/144344
>From 2585e3224410cba09e9812c27dcb305f61cf5946 Mon Sep 17 00:00:00 2001
From: tyb0807 <sontuan.vu119 at gmail.com>
Date: Mon, 16 Jun 2025 14:45:40 +0200
Subject: [PATCH] [mlir][memref][vectpr] Add alignment attribute to memory
access ops
---
.../mlir/Dialect/MemRef/IR/MemRefOps.td | 67 ++++++++++++++++++-
.../mlir/Dialect/Vector/IR/VectorOps.td | 55 ++++++++++++++-
.../Dialect/MemRef/load-store-alignment.mlir | 27 ++++++++
.../Dialect/Vector/load-store-alignment.mlir | 27 ++++++++
4 files changed, 170 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Dialect/MemRef/load-store-alignment.mlir
create mode 100644 mlir/test/Dialect/Vector/load-store-alignment.mlir
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 77e3074661abf..487aecb075ccf 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1217,6 +1217,11 @@ def LoadOp : MemRef_Op<"load",
be reused in the cache. For details, refer to the
[https://llvm.org/docs/LangRef.html#load-instruction](LLVM load instruction).
+ An optional `alignment` attribute allows to specify the byte alignment of the
+ load operation. It must be a positive power of 2. The operation must access
+ memory at an address aligned to this boundary. Violations may lead to
+ architecture-specific faults or performance penalties.
+ A value of 0 indicates no specific alignment requirement.
Example:
```mlir
@@ -1227,7 +1232,39 @@ def LoadOp : MemRef_Op<"load",
let arguments = (ins Arg<AnyMemRef, "the reference to load from",
[MemRead]>:$memref,
Variadic<Index>:$indices,
- DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
+ DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [IntMinValue<0>]>:$alignment);
+
+ let builders = [
+ OpBuilder<(ins "Value":$memref,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, memref, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "Type":$resultType,
+ "Value":$memref,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, resultType, memref, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "TypeRange":$resultTypes,
+ "Value":$memref,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, resultTypes, memref, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>
+ ];
+
let results = (outs AnyType:$result);
let extraClassDeclaration = [{
@@ -1913,6 +1950,11 @@ def MemRef_StoreOp : MemRef_Op<"store",
be reused in the cache. For details, refer to the
[https://llvm.org/docs/LangRef.html#store-instruction](LLVM store instruction).
+ An optional `alignment` attribute allows to specify the byte alignment of the
+ store operation. It must be a positive power of 2. The operation must access
+ memory at an address aligned to this boundary. Violations may lead to
+ architecture-specific faults or performance penalties.
+ A value of 0 indicates no specific alignment requirement.
Example:
```mlir
@@ -1924,13 +1966,32 @@ def MemRef_StoreOp : MemRef_Op<"store",
Arg<AnyMemRef, "the reference to store to",
[MemWrite]>:$memref,
Variadic<Index>:$indices,
- DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
+ DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [IntMinValue<0>]>:$alignment);
let builders = [
+ OpBuilder<(ins "Value":$valueToStore,
+ "Value":$memref,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, valueToStore, memref, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>,
+ // OpBuilder<(ins "Value":$valueToStore,
+ // "Value":$memref,
+ // "ValueRange":$indices,
+ // "uint64_t":$alignment), [{
+ // return build($_builder, $_state, valueToStore, memref, indices, false,
+ // $_builder.getI64IntegerAttr(alignment));
+ // }]>,
OpBuilder<(ins "Value":$valueToStore, "Value":$memref), [{
$_state.addOperands(valueToStore);
$_state.addOperands(memref);
- }]>];
+ }]>
+ ];
let extraClassDeclaration = [{
Value getValueToStore() { return getOperand(0); }
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 8353314ed958b..a6117855810f2 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1734,12 +1734,42 @@ def Vector_LoadOp : Vector_Op<"load"> {
```mlir
%result = vector.load %memref[%c0] : memref<7xf32>, vector<8xf32>
```
+
+ An optional `alignment` attribute allows to specify the byte alignment of the
+ load operation. It must be a positive power of 2. The operation must access
+ memory at an address aligned to this boundary. Violations may lead to
+ architecture-specific faults or performance penalties.
+ A value of 0 indicates no specific alignment requirement.
}];
let arguments = (ins Arg<AnyMemRef, "the reference to load from",
[MemRead]>:$base,
Variadic<Index>:$indices,
- DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
+ DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [IntMinValue<0>]>:$alignment);
+
+ let builders = [
+ OpBuilder<(ins "VectorType":$resultType,
+ "Value":$base,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, resultType, base, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>,
+ OpBuilder<(ins "TypeRange":$resultTypes,
+ "Value":$base,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, resultTypes, base, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>
+ ];
+
let results = (outs AnyVectorOfAnyRank:$result);
let extraClassDeclaration = [{
@@ -1818,6 +1848,12 @@ def Vector_StoreOp : Vector_Op<"store"> {
```mlir
vector.store %valueToStore, %memref[%c0] : memref<7xf32>, vector<8xf32>
```
+
+ An optional `alignment` attribute allows to specify the byte alignment of the
+ store operation. It must be a positive power of 2. The operation must access
+ memory at an address aligned to this boundary. Violations may lead to
+ architecture-specific faults or performance penalties.
+ A value of 0 indicates no specific alignment requirement.
}];
let arguments = (ins
@@ -1825,8 +1861,21 @@ def Vector_StoreOp : Vector_Op<"store"> {
Arg<AnyMemRef, "the reference to store to",
[MemWrite]>:$base,
Variadic<Index>:$indices,
- DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal
- );
+ DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal,
+ ConfinedAttr<OptionalAttr<I64Attr>,
+ [IntMinValue<0>]>:$alignment);
+
+ let builders = [
+ OpBuilder<(ins "Value":$valueToStore,
+ "Value":$base,
+ "ValueRange":$indices,
+ CArg<"bool", "false">:$nontemporal,
+ CArg<"uint64_t", "0">:$alignment), [{
+ return build($_builder, $_state, valueToStore, base, indices, nontemporal,
+ alignment != 0 ? $_builder.getI64IntegerAttr(alignment) :
+ nullptr);
+ }]>
+ ];
let extraClassDeclaration = [{
MemRefType getMemRefType() {
diff --git a/mlir/test/Dialect/MemRef/load-store-alignment.mlir b/mlir/test/Dialect/MemRef/load-store-alignment.mlir
new file mode 100644
index 0000000000000..57947e894c4dd
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/load-store-alignment.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: func @test_load_store_alignment
+// CHECK: memref.load {{.*}} {alignment = 16 : i64}
+// CHECK: memref.store {{.*}} {alignment = 16 : i64}
+func.func @test_load_store_alignment(%memref: memref<4xi32>) {
+ %c0 = arith.constant 0 : index
+ %val = memref.load %memref[%c0] { alignment = 16 } : memref<4xi32>
+ memref.store %val, %memref[%c0] { alignment = 16 } : memref<4xi32>
+ return
+}
+
+// -----
+
+func.func @test_invalid_load_alignment(%memref: memref<4xi32>) {
+ // expected-error @+1 {{custom op 'memref.load' 'memref.load' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
+ %val = memref.load %memref[%c0] { alignment = -1 } : memref<4xi32>
+ return
+}
+
+// -----
+
+func.func @test_invalid_store_alignment(%memref: memref<4xi32>, %val: memref<4xi32>) {
+ // expected-error @+1 {{custom op 'memref.store' 'memref.store' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
+ memref.store %val, %memref[%c0] { alignment = -1 } : memref<4xi32>
+ return
+}
diff --git a/mlir/test/Dialect/Vector/load-store-alignment.mlir b/mlir/test/Dialect/Vector/load-store-alignment.mlir
new file mode 100644
index 0000000000000..dbc9be8c56eaf
--- /dev/null
+++ b/mlir/test/Dialect/Vector/load-store-alignment.mlir
@@ -0,0 +1,27 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: func @test_load_store_alignment
+// CHECK: vector.load {{.*}} {alignment = 16 : i64}
+// CHECK: vector.store {{.*}} {alignment = 16 : i64}
+func.func @test_load_store_alignment(%memref: memref<4xi32>) {
+ %c0 = arith.constant 0 : index
+ %val = vector.load %memref[%c0] { alignment = 16 } : memref<4xi32>, vector<4xi32>
+ vector.store %val, %memref[%c0] { alignment = 16 } : memref<4xi32>, vector<4xi32>
+ return
+}
+
+// -----
+
+func.func @test_invalid_load_alignment(%memref: memref<4xi32>) {
+ // expected-error @+1 {{custom op 'vector.load' 'vector.load' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
+ %val = vector.load %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
+ return
+}
+
+// -----
+
+func.func @test_invalid_store_alignment(%memref: memref<4xi32>, %val: vector<4xi32>) {
+ // expected-error @+1 {{custom op 'vector.store' 'vector.store' op attribute 'alignment' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
+ vector.store %val, %memref[%c0] { alignment = -1 } : memref<4xi32>, vector<4xi32>
+ return
+}
More information about the Mlir-commits
mailing list