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()) + 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 +}