[MLIR][LLVM] Propagate alignment attribute from memref to LLVM (#151380)

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

Possible improvements: maybe the alignment attribute in LLVM's store and
load operations should be confined/constrained to i64? I believe that
way one can avoid typing the value in the attribute dictionary. I.e.,
from `{ alignment = 32 : i64 }` to `{ alignment = 32}`

Author: amd-eochoalo

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();
   }
 };
@@ -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();
   }
 };

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

@@ -753,6 +753,17 @@ 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>, %arg1 : index) {
+  // CHECK: llvm.load %{{.*}} {alignment = 32 : i64} : !llvm.ptr -> f32
+  // CHECK-INTERFACE: llvm.load
+  %1 = memref.load %arg0[%arg1] {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)>>) {
@@ -766,6 +777,17 @@ 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, %arg2 : index) {
+  // CHECK: llvm.store %{{.*}}, %{{.*}} {alignment = 32 : i64} : f32, !llvm.ptr
+  // CHECK-INTERFACE: llvm.store
+  memref.store %arg1, %arg0[%arg2] {alignment = 32} : memref<32xf32>
+  func.return
+}
+
+// -----
+
 // Ensure unconvertable memory space not cause a crash
 
 // CHECK-LABEL: @alloca_unconvertable_memory_space