[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
This commit is contained in:
Uday Bondhugula 2020-06-24 02:22:03 +05:30
parent 87340a2bf1
commit aec5344f48
2 changed files with 38 additions and 11 deletions

View File

@ -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;

View File

@ -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