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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Yuxi Sun (sherylll)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/146318.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp (+36-3) 


``````````diff
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);

``````````

</details>


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