[MLIR] Fix affine loop fusion private memref alloc

[bondhugula]

Summary

Drop stale code that provided the wrong operands to alloc.

Reported-by: rjnw on discourse

[andydavis1]

mlir/lib/Transforms/LoopFusion.cpp

915–916

Should you assert that 'newShape' is always constant?

[dcaballe]

Thanks for fixing this! LGTM

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

915–916

It's coming right out of getConstantBoundingSizeAndShape above and there is an assert on the numElements return value being constant; so I didn't add an assert here.

mlir/lib/Transforms/LoopFusion.cpp

static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   unsigned newMemSpace;
   if (bufSize <= localBufSizeThreshold && fastMemorySpace.hasValue()) {
     newMemSpace = fastMemorySpace.getValue();
   } else {
     newMemSpace = oldMemRefType.getMemorySpace();
   }
   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;
   remapExprs.reserve(rank);
   unsigned zeroOffsetCount = 0;
   for (unsigned i = 0; i < rank; i++) {
     if (auto constExpr = offsets[i].dyn_cast<AffineConstantExpr>())
       if (constExpr.getValue() == 0)

mlir/test/Transforms/loop-fusion.mlir

 // CHECK-NOT:  affine.for
 // CHECK:        addf
 // CHECK-NOT:  affine.for
 // CHECK:        subf
 // CHECK-NOT:  affine.for
 // 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