[Mlir-commits] [mlir] [MLIR][LLVM] Propagate alignment attribute from memref to LLVM (PR #151380)

[Erick Ochoa Lopez]

https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/151380

Propagate alignment attribute from operations in the memref dialect to the LLVM dialect.

>From 24971e2dfe87d146322297a35cb0318974fd440f Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 30 Jul 2025 15:04:35 -0400
Subject: [PATCH 1/2] Propagate alignment from memref.load to llvm.load.

---
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp    |  4 ++--
 .../test/Conversion/MemRefToLLVM/memref-to-llvm.mlir | 12 ++++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index dc2035b0700d0..6e67df478f082 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -895,8 +895,8 @@ struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
                                          adaptor.getMemref(),
                                          adaptor.getIndices(), kNoWrapFlags);
     rewriter.replaceOpWithNewOp<LLVM::LoadOp>(
-        loadOp, typeConverter->convertType(type.getElementType()), dataPtr, 0,
-        false, loadOp.getNontemporal());
+        loadOp, typeConverter->convertType(type.getElementType()), dataPtr,
+        loadOp.getAlignment().value_or(0), false, loadOp.getNontemporal());
     return success();
   }
 };
diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index 51d56389dac9e..65f75e467de69 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -753,6 +753,18 @@ func.func @load_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>) {
 
 // -----
 
+// CHECK-LABEL: func @load_with_alignment(
+// CHECK-INTERFACE-LABEL: func @load_with_alignment(
+func.func @load_with_alignment(%arg0 : memref<32xf32>) {
+  %1 = arith.constant 7 : index
+  // CHECK: llvm.load %{{.*}} {alignment = 32 : i64} : !llvm.ptr -> f32
+  // CHECK-INTERFACE: llvm.load
+  %2 = memref.load %arg0[%1] {alignment = 32} : memref<32xf32>
+  func.return
+}
+
+// -----
+
 // CHECK-LABEL: func @store_non_temporal(
 // CHECK-INTERFACE-LABEL: func @store_non_temporal(
 func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {

>From a1df1b8dd627a5fa9b0074c3cf3c228308f0ac2d Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 30 Jul 2025 15:16:17 -0400
Subject: [PATCH 2/2] Pass on alignment from memref.store to llvm.store.

---
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp    |  3 ++-
 .../test/Conversion/MemRefToLLVM/memref-to-llvm.mlir | 12 ++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 6e67df478f082..d6bdd342c9045 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -918,7 +918,8 @@ struct StoreOpLowering : public LoadStoreOpLowering<memref::StoreOp> {
         getStridedElementPtr(rewriter, op.getLoc(), type, adaptor.getMemref(),
                              adaptor.getIndices(), kNoWrapFlags);
     rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, adaptor.getValue(), dataPtr,
-                                               0, false, op.getNontemporal());
+                                               op.getAlignment().value_or(0),
+                                               false, op.getNontemporal());
     return success();
   }
 };
diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
index 65f75e467de69..959cc21acdae2 100644
--- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir
@@ -778,6 +778,18 @@ func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>,
 
 // -----
 
+// CHECK-LABEL: func @store_with_alignment(
+// CHECK-INTERFACE-LABEL: func @store_with_alignment(
+func.func @store_with_alignment(%arg0 : memref<32xf32>, %arg1 : f32) {
+  %1 = arith.constant 7 : index
+  // CHECK: llvm.store %{{.*}}, %{{.*}} {alignment = 32 : i64} : f32, !llvm.ptr
+  // CHECK-INTERFACE: llvm.store
+  memref.store %arg1, %arg0[%1] {alignment = 32} : memref<32xf32>
+  func.return
+}
+
+// -----
+
 // Ensure unconvertable memory space not cause a crash
 
 // CHECK-LABEL: @alloca_unconvertable_memory_space