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

[Mehdi Amini]

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

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

>From d36d00cc1f8750cc841386c2a6951e7ec417ffc7 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
+}