[Mlir-commits] [mlir] [MLIR][Affine] Preserve alignment attribute in affine-pipeline-data-transfer (PR #189220)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/189220

>From 4d61725fed6f340e890c64e2d19d47a22aed1073 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Mar 2026 12:55:02 -0700
Subject: [PATCH] [MLIR][Affine] Preserve alignment attribute in
 affine-pipeline-data-transfer

The `doubleBuffer` function in `PipelineDataTransfer.cpp` creates a new
doubled-up memref alloc but discards the alignment attribute of the
original alloc. When the source memref was allocated with
`{alignment = N}`, the double buffer should carry that same alignment
to avoid silent ABI mismatches and to maintain the programmer's intent.

Fix by propagating the alignment attribute from the original AllocOp
(if its defining op is an AllocOp) to the new double-buffer AllocOp.

Fixes #146015

Assisted-by: Claude Code
---
 .../Transforms/PipelineDataTransfer.cpp       |  9 +++-
 .../Affine/pipeline-data-transfer.mlir        | 48 +++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index d84cb4f0cde5f..660418480be02 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -99,9 +99,14 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
           forOp.getLoc(), oldMemRef, dim.index()));
   }
 
+  // Propagate alignment from the original alloc op, if any.
+  IntegerAttr alignment;
+  if (auto oldAllocOp = oldMemRef.getDefiningOp<memref::AllocOp>())
+    alignment = oldAllocOp.getAlignmentAttr();
+
   // Create and place the alloc right before the 'affine.for' operation.
-  Value newMemRef = memref::AllocOp::create(bOuter, forOp.getLoc(),
-                                            newMemRefType, allocOperands);
+  Value newMemRef = memref::AllocOp::create(
+      bOuter, forOp.getLoc(), newMemRefType, allocOperands, alignment);
 
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);
diff --git a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
index 35507c37be79b..120cde8639bae 100644
--- a/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
+++ b/mlir/test/Dialect/Affine/pipeline-data-transfer.mlir
@@ -396,3 +396,51 @@ func.func @same_memref_source_and_tag(%arg0: index, %arg1: index) {
   return
 }
 // CHECK: affine.for
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/146015.
+// The double-buffer alloc created by pipeline-data-transfer should preserve
+// the alignment attribute from the original alloc.
+// CHECK-LABEL: func @preserve_alignment
+func.func @preserve_alignment() {
+  %A = memref.alloc() : memref<256 x f32>
+  // CHECK: memref.alloc() {alignment = 1024 : i64} : memref<2x32xf32, 1>
+  %Ah = memref.alloc() {alignment = 1024} : memref<32 x f32, 1>
+  %tag = memref.alloc() : memref<1 x f32>
+  %zero = arith.constant 0 : index
+  %num_elts = arith.constant 32 : index
+
+  affine.for %i = 0 to 8 {
+    affine.dma_start %A[%i], %Ah[%i], %tag[%zero], %num_elts : memref<256 x f32>, memref<32 x f32, 1>, memref<1 x f32>
+    affine.dma_wait %tag[%zero], %num_elts : memref<1 x f32>
+    %v = affine.load %Ah[%i] : memref<32 x f32, 1>
+  }
+  memref.dealloc %tag : memref<1 x f32>
+  memref.dealloc %Ah : memref<32 x f32, 1>
+  return
+}
+
+// -----
+
+// Negative test: alloc without alignment must NOT gain a spurious alignment
+// attribute on the double-buffer alloc.
+// CHECK-LABEL: func @no_alignment_not_propagated
+func.func @no_alignment_not_propagated() {
+  %A = memref.alloc() : memref<256 x f32>
+  // CHECK: memref.alloc() : memref<2x32xf32, 1>
+  // CHECK-NOT: {alignment
+  %Ah = memref.alloc() : memref<32 x f32, 1>
+  %tag = memref.alloc() : memref<1 x f32>
+  %zero = arith.constant 0 : index
+  %num_elts = arith.constant 32 : index
+
+  affine.for %i = 0 to 8 {
+    affine.dma_start %A[%i], %Ah[%i], %tag[%zero], %num_elts : memref<256 x f32>, memref<32 x f32, 1>, memref<1 x f32>
+    affine.dma_wait %tag[%zero], %num_elts : memref<1 x f32>
+    %v = affine.load %Ah[%i] : memref<32 x f32, 1>
+  }
+  memref.dealloc %tag : memref<1 x f32>
+  memref.dealloc %Ah : memref<32 x f32, 1>
+  return
+}