[Mlir-commits] [mlir] [Affine] [PipelineDataTransfer] Small fixes (PR #146318)

[Yuxi Sun]

https://github.com/sherylll created https://github.com/llvm/llvm-project/pull/146318

Current implementation drops attributes from memref.alloc: https://github.com/llvm/llvm-project/issues/146015#issuecomment-3011613318

During double buffering:
1. preserve attributes (e.g. layout) of memref.alloc op
2. support strided layout

>From 9bb5795990980923d0a86ac4a0147f989edd43b0 Mon Sep 17 00:00:00 2001
From: sherylll <sherylsun14 at gmail.com>
Date: Mon, 30 Jun 2025 08:16:10 +0000
Subject: [PATCH] [Affine] pipeline-data-transfer pass: preserve attributes of
 ops and memref strided layout during double buffer

---
 .../Transforms/PipelineDataTransfer.cpp       | 39 +++++++++++++++++--
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index 92cb7075005a3..05f2f1e7057cc 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -85,7 +85,34 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
     SmallVector<int64_t, 4> newShape(1 + oldMemRefType.getRank());
     newShape[0] = 2;
     std::copy(oldShape.begin(), oldShape.end(), newShape.begin() + 1);
-    return MemRefType::Builder(oldMemRefType).setShape(newShape).setLayout({});
+
+    bool isDynamic = false;
+    for (int64_t dim : oldShape) {
+      if (dim == ShapedType::kDynamic) {
+        isDynamic = true;
+        break;
+      }
+    }
+    MemRefLayoutAttrInterface newLayout = {};
+    if (auto oldLayout = oldMemRefType.getLayout()) {
+      if (auto stridedLayout = dyn_cast<StridedLayoutAttr>(oldLayout)) {
+        // Calculate leading stride
+        ArrayRef<int64_t> oldStrides = stridedLayout.getStrides();
+        int64_t bufferStride =
+            isDynamic ? ShapedType::kDynamic : oldShape[0] * oldStrides[0];
+        SmallVector<int64_t> newStrides;
+        newStrides.push_back(bufferStride);
+        newStrides.append(oldStrides.begin(), oldStrides.end());
+
+        MLIRContext *context = oldMemRefType.getContext();
+        newLayout = StridedLayoutAttr::get(context, stridedLayout.getOffset(),
+                                           newStrides);
+      }
+    }
+
+    return MemRefType::Builder(oldMemRefType)
+        .setShape(newShape)
+        .setLayout(newLayout);
   };
 
   auto oldMemRefType = cast<MemRefType>(oldMemRef.getType());
@@ -102,8 +129,14 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   }
 
   // Create and place the alloc right before the 'affine.for' operation.
-  Value newMemRef = bOuter.create<memref::AllocOp>(
-      forOp.getLoc(), newMemRefType, allocOperands);
+  Value newMemRef;
+  if (auto *definingOp = oldMemRef.getDefiningOp()) {
+    newMemRef = bOuter.create<memref::AllocOp>(
+        forOp.getLoc(), newMemRefType, allocOperands, definingOp->getAttrs());
+  } else {
+    newMemRef = bOuter.create<memref::AllocOp>(forOp.getLoc(), newMemRefType,
+                                               allocOperands);
+  }
 
   // Create 'iv mod 2' value to index the leading dimension.
   auto d0 = bInner.getAffineDimExpr(0);