[Mlir-commits] [mlir] aec5344 - [MLIR] Fix affine loop fusion private memref alloc

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2020-06-24T22:19:29+05:30
New Revision: aec5344f48a23217a20ece5219f9e1e2f149c552

URL: https://github.com/llvm/llvm-project/commit/aec5344f48a23217a20ece5219f9e1e2f149c552
DIFF: https://github.com/llvm/llvm-project/commit/aec5344f48a23217a20ece5219f9e1e2f149c552.diff

LOG: [MLIR] Fix affine loop fusion private memref alloc

Drop stale code that provided the wrong operands to alloc.

Reported-by: rjnw on discourse

Differential Revision: https://reviews.llvm.org/D82409

Added: 
    

Modified: 
    mlir/lib/Transforms/LoopFusion.cpp
    mlir/test/Transforms/loop-fusion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index 0ae6657babcc..4ca27b479494 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -911,22 +911,14 @@ static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   }
   auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
                                        {}, newMemSpace);
-  // Gather alloc operands for the dynamic dimensions of the memref.
-  SmallVector<Value, 4> allocOperands;
-  unsigned dynamicDimCount = 0;
-  for (auto dimSize : oldMemRefType.getShape()) {
-    if (dimSize == -1)
-      allocOperands.push_back(
-          top.create<DimOp>(forOp.getLoc(), oldMemRef, dynamicDimCount++));
-  }
 
-  // Create new private memref for fused loop 'forOp'.
+  // Create new private memref for fused loop 'forOp'. 'newShape' is always
+  // a constant shape.
   // TODO(andydavis) Create/move alloc ops for private memrefs closer to their
   // consumer loop nests to reduce their live range. Currently they are added
   // at the beginning of the function, because loop nests can be reordered
   // during the fusion pass.
-  Value newMemRef =
-      top.create<AllocOp>(forOp.getLoc(), newMemRefType, allocOperands);
+  Value newMemRef = top.create<AllocOp>(forOp.getLoc(), newMemRefType);
 
   // Build an AffineMap to remap access functions based on lower bound offsets.
   SmallVector<AffineExpr, 4> remapExprs;

diff  --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index 0e79e0b67ba2..a153b96bf362 100644
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -2535,3 +2535,38 @@ func @multi_outgoing_edges(%in0 : memref<32xf32>,
 // CHECK:        mulf
 // CHECK-NOT:  affine.for
 // CHECK:        divf
+
+// -----
+
+// Test fusion when dynamically shaped memrefs are used with constant trip count loops.
+
+// CHECK-LABEL: func @calc
+func @calc(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %len: index) {
+  %c1 = constant 1 : index
+  %1 = alloc(%len) : memref<?xf32>
+  affine.for %arg4 = 1 to 10 {
+    %7 = affine.load %arg0[%arg4] : memref<?xf32>
+    %8 = affine.load %arg1[%arg4] : memref<?xf32>
+    %9 = addf %7, %8 : f32
+    affine.store %9, %1[%arg4] : memref<?xf32>
+  }
+  affine.for %arg4 = 1 to 10 {
+    %7 = affine.load %1[%arg4] : memref<?xf32>
+    %8 = affine.load %arg1[%arg4] : memref<?xf32>
+    %9 = mulf %7, %8 : f32
+    affine.store %9, %arg2[%arg4] : memref<?xf32>
+  }
+  return
+}
+// CHECK:       alloc() : memref<1xf32>
+// CHECK:       affine.for %arg{{.*}} = 1 to 10 {
+// CHECK-NEXT:    affine.load %arg{{.*}}
+// CHECK-NEXT:    affine.load %arg{{.*}}
+// CHECK-NEXT:    addf
+// CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
+// CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
+// CHECK-NEXT:    affine.load %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
+// CHECK-NEXT:    mulf
+// CHECK-NEXT:    affine.store %{{.*}}, %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return