[mlir][memref] Extend multi-buffering transform

[kile]

Summary

Extend multi-buffering to simplify the affine map created if any of its operands are constants. This avoids downstream problems where more complex affine.apply operations cannot be expanded.
Transfer attributes from the old allocation to the new allocation.

[ThomasRaoux]

Please include context by using -U9999999 when creating patch:
https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

[kile]

Added additional context to patch

[ThomasRaoux]

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp

118

We should bail out before if the alloc doesn't dominate the loop as the transformation doesn't make sense.

[kile]

Bail when allocations are inside of the loop

[ThomasRaoux]

This looks good to me once the style comments are addressed

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp

108

nit: please remove brace for single line if:
https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

129

nit: please spell out the types here:
https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

[kile]

Style Updates

[kile]

@ThomasRaoux thank you very much for the review! Would you mind being the one to actually commit the change? Unfortunately I don't have commit access.

[ThomasRaoux]

In D134894#3831280, @kile wrote:

@ThomasRaoux thank you very much for the review! Would you mind being the one to actually commit the change? Unfortunately I don't have commit access.

Not a problem. Could you send me the name and email to use for the commit? (for reference: https://llvm.org/docs/Phabricator.html#committing-someone-s-change-from-phabricator)

[kile]

@ThomasRaoux sure thing: email - [e-mail], name: Kirsten Lee

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp

FailureOr<memref::AllocOp> mlir::memref::multiBuffer(memref::AllocOp allocOp,
   }
   if (!candidateLoop)
     return failure();
   llvm::Optional<Value> inductionVar = candidateLoop.getSingleInductionVar();
   llvm::Optional<OpFoldResult> lowerBound = candidateLoop.getSingleLowerBound();
   llvm::Optional<OpFoldResult> singleStep = candidateLoop.getSingleStep();
   if (!inductionVar || !lowerBound || !singleStep)
     return failure();
+
+  if (!dom.dominates(allocOp.getOperation(), candidateLoop))
+    return failure();
+
   OpBuilder builder(candidateLoop);
-  Value stepValue =
-      getOrCreateValue(*singleStep, builder, candidateLoop->getLoc());
-  Value lowerBoundValue =
-      getOrCreateValue(*lowerBound, builder, candidateLoop->getLoc());
   SmallVector<int64_t, 4> newShape(1, multiplier);
   ArrayRef<int64_t> oldShape = allocOp.getType().getShape();
   newShape.append(oldShape.begin(), oldShape.end());
   auto newMemref = MemRefType::get(newShape, allocOp.getType().getElementType(),
                                    MemRefLayoutAttrInterface(),
                                    allocOp.getType().getMemorySpace());
   builder.setInsertionPoint(allocOp);
   Location loc = allocOp->getLoc();
-  auto newAlloc = builder.create<memref::AllocOp>(loc, newMemref);
+  auto newAlloc = builder.create<memref::AllocOp>(loc, newMemref, ValueRange{},
+                                                  allocOp->getAttrs());
   builder.setInsertionPoint(&candidateLoop.getLoopBody().front(),
                             candidateLoop.getLoopBody().front().begin());
+
+  SmallVector<Value> operands = {*inductionVar};
   AffineExpr induc = getAffineDimExpr(0, allocOp.getContext());
-  AffineExpr init = getAffineDimExpr(1, allocOp.getContext());
-  AffineExpr step = getAffineDimExpr(2, allocOp.getContext());
+  unsigned dimCount = 1;
+  auto getAffineExpr = [&](OpFoldResult e) -> AffineExpr {
+    if (Optional<int64_t> constValue = getConstantIntValue(e)) {
+      return getAffineConstantExpr(*constValue, allocOp.getContext());
+    } else {
+      auto value = getOrCreateValue(e, builder, candidateLoop->getLoc());
+      operands.push_back(value);
+      return getAffineDimExpr(dimCount++, allocOp.getContext());
+    }
+  };
+  auto init = getAffineExpr(*lowerBound);
+  auto step = getAffineExpr(*singleStep);
+
   AffineExpr expr = ((induc - init).floorDiv(step)) % multiplier;
-  auto map = AffineMap::get(3, 0, expr);
-  std::array<Value, 3> operands = {*inductionVar, lowerBoundValue, stepValue};
+  auto map = AffineMap::get(dimCount, 0, expr);
   Value bufferIndex = builder.create<AffineApplyOp>(loc, map, operands);
   SmallVector<OpFoldResult> offsets, sizes, strides;
   offsets.push_back(bufferIndex);
   offsets.append(oldShape.size(), builder.getIndexAttr(0));
   strides.assign(oldShape.size() + 1, builder.getIndexAttr(1));
   sizes.push_back(builder.getIndexAttr(1));
   for (int64_t size : oldShape)
     sizes.push_back(builder.getIndexAttr(size));

mlir/test/Dialect/MemRef/multibuffer.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -test-multi-buffering=multiplier=5 -cse -split-input-file | FileCheck %s
 
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (((d0 - 1) floordiv 3) mod 5)>
 
 // CHECK-LABEL: func @multi_buffer
 func.func @multi_buffer(%a: memref<1024x1024xf32>) {
-// CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
+// CHECK-DAG: %[[A:.*]] = memref.alloc() {someAttribute} : memref<5x4x128xf32>
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
-  %0 = memref.alloc() : memref<4x128xf32>
+  %0 = memref.alloc() {someAttribute} : memref<4x128xf32>
   %c1024 = arith.constant 1024 : index
   %c1 = arith.constant 1 : index
   %c3 = arith.constant 3 : index
 // CHECK: scf.for %[[IV:.*]] = %[[C1]]
   scf.for %arg2 = %c1 to %c1024 step %c3 {
-// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]])
 // CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
     memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
 // CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
 // CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, strided{{.*}}>) -> ()
     "some_use"(%0) : (memref<4x128xf32>) -> ()
 // CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, strided{{.*}}>) -> ()
    "some_use"(%0) : (memref<4x128xf32>) -> ()
   }
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @multi_buffer_affine
 func.func @multi_buffer_affine(%a: memref<1024x1024xf32>) {
 // CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
-// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
   %0 = memref.alloc() : memref<4x128xf32>
   %c1024 = arith.constant 1024 : index
   %c1 = arith.constant 1 : index
   %c3 = arith.constant 3 : index
 // CHECK: affine.for %[[IV:.*]] = 1
   affine.for %arg2 = 1 to 1024 step 3 {
-// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]])
 // CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
     memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
 // CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
 // CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, strided{{.*}}>) -> ()
     "some_use"(%0) : (memref<4x128xf32>) -> ()
 // CHECK: "some_use"(%[[SV]]) : (memref<4x128xf32, strided{{.*}}>) -> ()
    "some_use"(%0) : (memref<4x128xf32>) -> ()
   }
   return
 }
 
 // -----
 
-// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 5)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (((d0 - 1) floordiv 3) mod 5)>
 
 // CHECK-LABEL: func @multi_buffer_subview_use
 func.func @multi_buffer_subview_use(%a: memref<1024x1024xf32>) {
 // CHECK-DAG: %[[A:.*]] = memref.alloc() : memref<5x4x128xf32>
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
   %0 = memref.alloc() : memref<4x128xf32>
   %c1024 = arith.constant 1024 : index
   %c1 = arith.constant 1 : index
   %c3 = arith.constant 3 : index
 // CHECK: scf.for %[[IV:.*]] = %[[C1]]
   scf.for %arg2 = %c1 to %c1024 step %c3 {
-// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]], %[[C1]], %[[C3]])
+// CHECK: %[[I:.*]] = affine.apply #[[$MAP1]](%[[IV]])
 // CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0, 0] [1, 4, 128] [1, 1, 1] : memref<5x4x128xf32> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    %1 = memref.subview %a[%arg2, 0] [4, 128] [1, 1] :
     memref<1024x1024xf32> to memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>>
 // CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4x128xf32, #{{.*}}> to memref<4x128xf32, strided<[128, 1], offset: ?>>
    memref.copy %1, %0 : memref<4x128xf32, affine_map<(d0, d1)[s0] -> (d0 * 1024 + s0 + d1)>> to memref<4x128xf32>
 // CHECK: %[[SV1:.*]] = memref.subview %[[SV]][0, 1] [4, 127] [1, 1] : memref<4x128xf32, strided<[128, 1], offset: ?>> to memref<4x127xf32, strided<[128, 1], offset: ?>>
    %s = memref.subview %0[0, 1] [4, 127] [1, 1] :
       memref<4x128xf32> to memref<4x127xf32, affine_map<(d0, d1) -> (d0 * 128 + d1 + 1)>>

mlir/test/Dialect/MemRef/transform-ops.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -verify-diagnostics -allow-unregistered-dialect | FileCheck %s
 
-// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (((d0 - d1) floordiv d2) mod 2)>
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> ((d0 floordiv 4) mod 2)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 
 // CHECK-LABEL: func @multi_buffer
 func.func @multi_buffer(%in: memref<16xf32>) {
   // CHECK: %[[A:.*]] = memref.alloc() : memref<2x4xf32>
   // expected-remark @below {{transformed}}
   %tmp = memref.alloc() : memref<4xf32>
 
   // CHECK: %[[C0:.*]] = arith.constant 0 : index
   // CHECK: %[[C4:.*]] = arith.constant 4 : index
   %c0 = arith.constant 0 : index
   %c4 = arith.constant 4 : index
   %c16 = arith.constant 16 : index
 
   // CHECK: scf.for %[[IV:.*]] = %[[C0]]
   scf.for %i0 = %c0 to %c16 step %c4 {
-    // CHECK: %[[I:.*]] = affine.apply #[[$MAP0]](%[[IV]], %[[C0]], %[[C4]])
+    // CHECK: %[[I:.*]] = affine.apply #[[$MAP0]](%[[IV]])
     // CHECK: %[[SV:.*]] = memref.subview %[[A]][%[[I]], 0] [1, 4] [1, 1] : memref<2x4xf32> to memref<4xf32, strided<[1], offset: ?>>
     %1 = memref.subview %in[%i0] [4] [1] : memref<16xf32> to memref<4xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
     // CHECK: memref.copy %{{.*}}, %[[SV]] : memref<4xf32, #[[$MAP1]]> to memref<4xf32, strided<[1], offset: ?>>
     memref.copy %1, %tmp :  memref<4xf32, affine_map<(d0)[s0] -> (d0 + s0)>> to memref<4xf32>
 
     "some_use"(%tmp) : (memref<4xf32>) ->()
   }
   return