[mlir][scf] Canonicalize scf.for last tensor iteration result.

[nicolasvasilache]

Summary

Canonicalize the iter_args of an scf::ForOp that involve a tensor_load and
for which only the last loop iteration is actually visible outside of the
loop. The canonicalization looks for a pattern such as:

%t0 = ... : tensor_type
%0 = scf.for ... iter_args(%bb0 : %t0) -> (tensor_type) {
  ...
  // %m is either tensor_to_memref(%bb00) or defined above the loop
  %m... : memref_type
  ... // uses of %m with potential inplace updates
  %new_tensor = tensor_load %m : memref_type
  ...
  scf.yield %new_tensor : tensor_type
}

%bb0 may have either 0 or 1 use. If it has 1 use it must be exactly a
%m = tensor_to_memref %bb0 op that feeds into the yielded tensor_load
op.

If no aliasing write of %new_tensor occurs between tensor_load and yield
then the value %0 visible outside of the loop is the last tensor_load
produced in the loop.

For now, we approximate the absence of aliasing by only supporting the case
when the tensor_load is the operation immediately preceding the yield.

The canonicalization rewrites the pattern as:

// %m is either a tensor_to_memref or defined above
%m... : memref_type
scf.for ... { // no iter_args
  ... // uses of %m with potential inplace updates
}
%0 = tensor_load %m : memref_type

This revision also exhibits the incorrect behavior of the folding:
tensor_load(tensor_to_memref(x)) -> x which ignores aliasing.

The test is altered to pass despite the erroneous folding, but should be revisited once the folding is hardened.

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

605–607

This should never happen, I'd assert instead.

641

Code motion is not undoable by "root update", only modifications to the op state are, you'll have to recreate the op at the new position.

647

Same

659–662

I'm confused by this comment. You want to replace a subset of forOp results with values produced by hoisted ops. I don't see how side-effects are relevant here. Furthermore, the documentation above and the test show that iter_args are removed, but I don't see the corresponding bbargs removed from the loop, it is just cloned. I expect something weird to happen if the bbargs that correspond to hoisted tensor updates are not trailing in the list.

[nicolasvasilache]

mlir/lib/Dialect/SCF/SCF.cpp

641

thanks!

659–662

I'll rephrase but basically rewriter only has a replaceOp method that unconditionally deletes the op.
So I I replace forOp it just gets dropped, even it has side effects.

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

659–662

Can't you use replaceOpWithIf for this? Or provide an easier entry point to that function for what you want to do here? It isn't supported in dialect conversion ATM, but should work for canonicalization patterns.

[silvas]

mlir/test/Dialect/SCF/canonicalize.mlir

375

is this needed after the recent patch you landed?

[nicolasvasilache]

mlir/lib/Dialect/SCF/SCF.cpp

659–662

Actually, I can just do nothin and let other canonicalizations kick in when appropriate...

mlir/test/Dialect/SCF/canonicalize.mlir

375

I reverted it locally so when I update it'll go away.

[nicolasvasilache]

Address review.

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

605–607

Why assert at all? Isn't this already verified by the op?

[nicolasvasilache]

Dropp redundant assert.

mlir/lib/Dialect/SCF/SCF.cpp

605–607

fair enough.

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

600

This pattern looks like something that would be better handled by LICM+a simpler canonicalization that forwarded yield results that are from outside of the loop.

[silvas]

LGTM conceptually. Please wait for River/Alex LGTM for the rewrite mechanics since they already started looking at it.

mlir/lib/Dialect/SCF/SCF.cpp

584

it is a bit strange to talk about aliasing of a tensor. can you reword?

mlir/test/Dialect/SCF/canonicalize.mlir

345

add a test case where the safety check fires and the pattern fails to apply?

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

600

Wait, nvm. Missed the iteration handling here.

[nicolasvasilache]

mlir/lib/Dialect/SCF/SCF.cpp

659–662

Ugh no .. that is wrong, I am stupidly missing replacements .. sorry

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

581

Where are you checking that it feeds into the tensor_load?

[nicolasvasilache]

Address review.

mlir/lib/Dialect/SCF/SCF.cpp

659–662

@rriddle ok replaceOpWithIf wins.

It would be great to find a naming scheme that conveys the fact that replaceOp replace unconditionally whereas replaceOpWithIf does not replace ever and it is the user's responsibility to eraseOp if allUsesReplaced is set to true (if they so wish).

I had completely overlooked the fact that not all replaceOpXXX do a blanket eraseOp.

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

605–607

I've seen numerous asserts (and wrote some) on op state that are enforced by the verifier. I do this as part of "assert liberally" policy in the LLVM style guide, often in situations where I suspect the op may evolve and having an assertion will make it easier to find the places that need changing.

[nicolasvasilache]

Fix comment.

[nicolasvasilache]

mlir/lib/Dialect/SCF/SCF.cpp

605–607

I'm def. +1 on defensive asserts.

After considering your comments here I found it quite unlikely that scf::ForOp will evolve to multi-blocks in the future.
Given the current state of the codebase I'd prefer to see a new type of op.

I'd still lean towards defensive style in general, I can revive the assert if there is a stronger feeling, I don't have a strong preference here.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

605–607

I'm +1 on assert as well: these are both debugging and documentation help.

It duplicates the verifier sometimes, but that should be OK: if it wasn't in the verifier then it couldn't be an assert in the first place :)

mlir/lib/Dialect/SCF/SCF.cpp

if ((lbValue + stepValue).sge(ubValue)) {
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }
 
     return failure();
   }
 };
+
+/// Canonicalize the iter_args of an scf::ForOp that involve a tensor_load and
+/// for which only the last loop iteration is actually visible outside of the
+/// loop. The canonicalization looks for a pattern such as:
+/// ```
+///    %t0 = ... : tensor_type
+///    %0 = scf.for ... iter_args(%bb0 : %t0) -> (tensor_type) {
+///      ...
+///      // %m is either tensor_to_memref(%bb00) or defined above the loop
+///      %m... : memref_type
+///      ... // uses of %m with potential inplace updates
+///      %new_tensor = tensor_load %m : memref_type
+///      ...
+///      scf.yield %new_tensor : tensor_type
+///    }
+/// ```
+///
+/// `%bb0` may have either 0 or 1 use. If it has 1 use it must be exactly a
+/// `%m = tensor_to_memref %bb0` op that feeds into the yielded `tensor_load`
+/// op.
+///
+/// If no aliasing write to the memref `%m`, from which `%new_tensor`is loaded,
+/// occurs between tensor_load and yield then the value %0 visible outside of
+/// the loop is the last `tensor_load` produced in the loop.
+///
+/// For now, we approximate the absence of aliasing by only supporting the case
+/// when the tensor_load is the operation immediately preceding the yield.
+///
+/// The canonicalization rewrites the pattern as:
+/// ```
+///    // %m is either a tensor_to_memref or defined above
+///    %m... : memref_type
+///    scf.for ... iter_args(%bb0 : %t0) -> (tensor_type) {
+///      ... // uses of %m with potential inplace updates
+///      scf.yield %bb0: tensor_type
+///    }
+///    %0 = tensor_load %m : memref_type
+/// ```
+///
+/// A later bbArg canonicalization will further rewrite as:
+/// ```
+///    // %m is either a tensor_to_memref or defined above
+///    %m... : memref_type
+///    scf.for ... { // no iter_args
+///      ... // uses of %m with potential inplace updates
+///    }
+///    %0 = tensor_load %m : memref_type
+/// ```
+struct LastTensorLoadCanonicalization : public OpRewritePattern<ForOp> {
+  using OpRewritePattern<ForOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForOp forOp,
+                                PatternRewriter &rewriter) const override {
+    assert(std::next(forOp.region().begin()) == forOp.region().end() &&
+           "unexpected multiple blocks");
+
+    Location loc = forOp.getLoc();
+    DenseMap<Value, Value> replacements;
+    for (BlockArgument bbArg : forOp.getRegionIterArgs()) {
+      unsigned idx = bbArg.getArgNumber() - /*numIv=*/1;
+      auto yieldOp = cast<scf::YieldOp>(forOp.region().front().getTerminator());
+      Value yieldVal = yieldOp->getOperand(idx);
+      auto tensorLoadOp = yieldVal.getDefiningOp<TensorLoadOp>();
+      bool isTensor = bbArg.getType().isa<TensorType>();
+
+      TensorToMemrefOp tensorToMemRefOp;
+      // Either bbArg has no use or it has a single tensor_to_memref use.
+      if (bbArg.hasOneUse())
+        tensorToMemRefOp =
+            dyn_cast<TensorToMemrefOp>(*bbArg.getUsers().begin());
+      if (!isTensor || !tensorLoadOp ||
+          (!bbArg.use_empty() && !tensorToMemRefOp))
+        continue;
+      // If tensorToMemRefOp is present, it must feed into the `tensorLoadOp`.
+      if (tensorToMemRefOp && tensorLoadOp.memref() != tensorToMemRefOp)
+        continue;
+      // TODO: Any aliasing write of tensorLoadOp.memref() nested under `forOp`
+      // must be before `tensorLoadOp` in the block so that the lastWrite
+      // property is not subject to additional side-effects.
+      // For now, we only support the case when tensorLoadOp appears immediately
+      // before the terminator.
+      if (tensorLoadOp->getNextNode() != yieldOp)
+        continue;
+
+      // Clone the optional tensorToMemRefOp before forOp.
+      if (tensorToMemRefOp) {
+        rewriter.setInsertionPoint(forOp);
+        rewriter.replaceOpWithNewOp<TensorToMemrefOp>(
+            tensorToMemRefOp, tensorToMemRefOp.memref().getType(),
+            tensorToMemRefOp.tensor());
+      }
+
+      // Clone the tensorLoad after forOp.
+      rewriter.setInsertionPointAfter(forOp);
+      Value newTensorLoad =
+          rewriter.create<TensorLoadOp>(loc, tensorLoadOp.memref());
+      Value forOpResult = forOp.getResult(bbArg.getArgNumber() - /*iv=*/1);
+      replacements.insert(std::make_pair(forOpResult, newTensorLoad));
+
+      // Make the terminator just yield the bbArg, the old tensorLoadOp + the
+      // old bbArg (that is now directly yielded) will canonicalize away.
+      rewriter.startRootUpdate(yieldOp);
+      yieldOp.setOperand(idx, bbArg);
+      rewriter.finalizeRootUpdate(yieldOp);
+    }
+    if (replacements.empty())
+      return failure();
+
+    // We want to replace a subset of the results of `forOp`. rewriter.replaceOp
+    // replaces the whole op and erase it unconditionally. This is wrong for
+    // `forOp` as it generally contains ops with side effects.
+    // Instead, use `rewriter.replaceOpWithIf`.
+    SmallVector<Value> newResults;
+    newResults.reserve(forOp.getNumResults());
+    for (Value v : forOp.getResults()) {
+      auto it = replacements.find(v);
+      newResults.push_back((it != replacements.end()) ? it->second : v);
+    }
+    unsigned idx = 0;
+    rewriter.replaceOpWithIf(forOp, newResults, [&](OpOperand &op) {
+      return op.get() != newResults[idx++];
+    });
+    return success();
+  }
+};
 } // namespace
 
 void ForOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
-  results.insert<ForOpIterArgsFolder, SimplifyTrivialLoops>(context);
+  results.insert<ForOpIterArgsFolder, SimplifyTrivialLoops,
+                 LastTensorLoadCanonicalization>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {

mlir/test/Dialect/SCF/canonicalize.mlir

-// RUN: mlir-opt %s -pass-pipeline='func(canonicalize)' | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='func(canonicalize)' -split-input-file | FileCheck %s
+
+
+// -----
 
 func @single_iteration(%A: memref<?x?x?xi32>) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c6 = constant 6 : index
   %c7 = constant 7 : index

 //  CHECK-NEXT:     %[[a:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     %[[b:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     %[[r1:.*]] = scf.for {{.*}} iter_args(%arg4 = %[[a]]) -> (i32) {
 //  CHECK-NEXT:       %[[c:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:       scf.yield %[[c]] : i32
 //  CHECK-NEXT:     }
 //  CHECK-NEXT:     return %[[a]], %[[r1]], %[[b]] : i32, i32, i32
 
+// -----
+
 // CHECK-LABEL: @replace_true_if
 func @replace_true_if() {
   %true = constant true
   // CHECK-NOT: scf.if
   // CHECK: "test.op"
   scf.if %true {
     "test.op"() : () -> ()
     scf.yield
   }
   return
 }
 
+// -----
+
 // CHECK-LABEL: @remove_false_if
 func @remove_false_if() {
   %false = constant false
   // CHECK-NOT: scf.if
   // CHECK-NOT: "test.op"
   scf.if %false {
     "test.op"() : () -> ()
     scf.yield
   }
   return
 }
 
+// -----
+
 // CHECK-LABEL: @replace_true_if_with_values
 func @replace_true_if_with_values() {
   %true = constant true
   // CHECK-NOT: scf.if
   // CHECK: %[[VAL:.*]] = "test.op"
   %0 = scf.if %true -> (i32) {
     %1 = "test.op"() : () -> i32
     scf.yield %1 : i32
   } else {
     %2 = "test.other_op"() : () -> i32
     scf.yield %2 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @replace_false_if_with_values
 func @replace_false_if_with_values() {
   %false = constant false
   // CHECK-NOT: scf.if
   // CHECK: %[[VAL:.*]] = "test.other_op"
   %0 = scf.if %false -> (i32) {
     %1 = "test.op"() : () -> i32
     scf.yield %1 : i32
   } else {
     %2 = "test.other_op"() : () -> i32
     scf.yield %2 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @remove_zero_iteration_loop
 func @remove_zero_iteration_loop() {
   %c42 = constant 42 : index
   %c1 = constant 1 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   %0 = scf.for %i = %c42 to %c1 step %c1 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @remove_zero_iteration_loop_vals
 func @remove_zero_iteration_loop_vals(%arg0: index) {
   %c2 = constant 2 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK-NOT: test.op
   %0 = scf.for %i = %arg0 to %arg0 step %c2 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @replace_single_iteration_loop_1
 func @replace_single_iteration_loop_1() {
   // CHECK: %[[LB:.*]] = constant 42
   %c42 = constant 42 : index
   %c43 = constant 43 : index
   %c1 = constant 1 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
   %0 = scf.for %i = %c42 to %c43 step %c1 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @replace_single_iteration_loop_2
 func @replace_single_iteration_loop_2() {
   // CHECK: %[[LB:.*]] = constant 5
   %c5 = constant 5 : index
   %c6 = constant 6 : index
   %c11 = constant 11 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
   %0 = scf.for %i = %c5 to %c11 step %c6 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
 
 // CHECK-LABEL: @replace_single_iteration_loop_non_unit_step
 func @replace_single_iteration_loop_non_unit_step() {
   // CHECK: %[[LB:.*]] = constant 42
   %c42 = constant 42 : index
   %c47 = constant 47 : index
   %c5 = constant 5 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
   %0 = scf.for %i = %c42 to %c47 step %c5 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// -----
+
 // CHECK-LABEL: @remove_empty_parallel_loop
 func @remove_empty_parallel_loop(%lb: index, %ub: index, %s: index) {
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> f32
   // CHECK-NOT: scf.parallel
   // CHECK-NOT: test.produce
   // CHECK-NOT: test.transform
   %0 = scf.parallel (%i, %j, %k) = (%lb, %ub, %lb) to (%ub, %ub, %ub) step (%s, %s, %s) init(%init) -> f32 {
     %1 = "test.produce"() : () -> f32
     scf.reduce(%1) : f32 {
     ^bb0(%lhs: f32, %rhs: f32):
       %2 = "test.transform"(%lhs, %rhs) : (f32, f32) -> f32
       scf.reduce.return %2 : f32
     }
     scf.yield
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (f32) -> ()
   return
 }
+
+// -----
+func private @process(%0 : memref<128x128xf32>)
+func private @process_tensor(%0 : tensor<128x128xf32>) -> memref<128x128xf32>
+
+// CHECK-LABEL: last_value
+//  CHECK-SAME:   %[[T0:[0-9a-z]*]]: tensor<128x128xf32>
+//  CHECK-SAME:   %[[T1:[0-9a-z]*]]: tensor<128x128xf32>
+//  CHECK-SAME:   %[[T2:[0-9a-z]*]]: tensor<128x128xf32>
+//  CHECK-SAME:   %[[M0:[0-9a-z]*]]: memref<128x128xf32>
+func @last_value(%t0: tensor<128x128xf32>, %t1: tensor<128x128xf32>,
+                 %t2: tensor<128x128xf32>, %m0: memref<128x128xf32>,
+                 %lb : index, %ub : index, %step : index)
+  -> (tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>)
+{
+  // CHECK-NEXT: %[[M1:.*]] = tensor_to_memref %[[T1]] : memref<128x128xf32>
+  // CHECK-NEXT: %[[FOR_RES:.*]] = scf.for {{.*}} iter_args(%[[BBARG_T2:.*]] = %[[T2]]) -> (tensor<128x128xf32>) {
+  %0:3 = scf.for %arg0 = %lb to %ub step %step iter_args(%arg1 = %t0, %arg2 = %t1, %arg3 = %t2)
+    -> (tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>)
+  {
+    %m1 = tensor_to_memref %arg2 : memref<128x128xf32>
+
+    // CHECK-NEXT:   call @process(%[[M0]]) : (memref<128x128xf32>) -> ()
+    call @process(%m0) : (memref<128x128xf32>) -> ()
+
+    // CHECK-NEXT:   call @process(%[[M1]]) : (memref<128x128xf32>) -> ()
+    call @process(%m1) : (memref<128x128xf32>) -> ()
+
+    // This does not hoist (fails the bbArg has at most a single check).
+    // CHECK-NEXT:   %[[T:.*]] = call @process_tensor(%[[BBARG_T2]]) : (tensor<128x128xf32>) -> memref<128x128xf32>
+    // CHECK-NEXT:   %[[YIELD_T:.*]] = tensor_load %[[T:.*]]
+    %m2 = call @process_tensor(%arg3): (tensor<128x128xf32>) -> memref<128x128xf32>
+    %3 = tensor_load %m2 : memref<128x128xf32>
+
+    // All this stuff goes away, incrementally
+    %1 = tensor_load %m0 : memref<128x128xf32>
+    %2 = tensor_load %m1 : memref<128x128xf32>
+
+    // CHECK-NEXT:   scf.yield %[[YIELD_T]] : tensor<128x128xf32>
+    scf.yield %1, %2, %3 : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
+
+  // CHECK-NEXT: }
+  }
+
+  // CHECK-NEXT: %[[R0:.*]] = tensor_load %[[M0]] : memref<128x128xf32>
+  // CHECK-NEXT: %[[R1:.*]] = tensor_load %[[M1]] : memref<128x128xf32>
+  // CHECK-NEXT: return %[[R0]], %[[R1]], %[[FOR_RES]] : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
+  return %0#0, %0#1, %0#2 : tensor<128x128xf32>, tensor<128x128xf32>, tensor<128x128xf32>
+}