[mlir][SCF] Fix incorrect API usage in RewritePatterns

[springerm]

Summary

Incorrect API usage was detected by D144552.

Depends On: D144549

[nicolasvasilache]

Fantastic, thank you @springerm !

mlir/include/mlir/IR/PatternMatch.h

public:
   }
   template <typename OperandType, typename ValueT>
   void replaceAllUsesWith(IRObjectWithUseList<OperandType> *from, ValueT &&to) {
     for (OperandType &operand : llvm::make_early_inc_range(from->getUses())) {
       Operation *op = operand.getOwner();
       updateRootInPlace(op, [&]() { operand.set(to); });
     }
   }
+  void replaceAllUsesWith(ValueRange from, ValueRange to) {
+    assert(from.size() == to.size() && "incorrect number of replacements");
+    for (auto it : llvm::zip(from, to))
+      replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
+  }
 
   /// Find uses of `from` and replace them with `to` if the `functor` returns
   /// true. It also marks every modified uses and notifies the rewriter that an
   /// in-place operation modification is about to happen.
   void
   replaceUsesWithIf(Value from, Value to,
                     llvm::unique_function<bool(OpOperand &) const> functor);
 

mlir/lib/Dialect/SCF/IR/SCF.cpp

LogicalResult matchAndRewrite(ForallOp op,
     SmallVector<OpFoldResult> mixedLowerBound(op.getMixedLowerBound());
     SmallVector<OpFoldResult> mixedUpperBound(op.getMixedUpperBound());
     SmallVector<OpFoldResult> mixedStep(op.getMixedStep());
     if (failed(foldDynamicIndexList(rewriter, mixedLowerBound)) &&
         failed(foldDynamicIndexList(rewriter, mixedUpperBound)) &&
         failed(foldDynamicIndexList(rewriter, mixedStep)))
       return failure();
 
+    rewriter.updateRootInPlace(op, [&]() {
       SmallVector<Value> dynamicLowerBound, dynamicUpperBound, dynamicStep;
       SmallVector<int64_t> staticLowerBound, staticUpperBound, staticStep;
       dispatchIndexOpFoldResults(mixedLowerBound, dynamicLowerBound,
                                  staticLowerBound);
       op.getDynamicLowerBoundMutable().assign(dynamicLowerBound);
       op.setStaticLowerBound(staticLowerBound);
 
       dispatchIndexOpFoldResults(mixedUpperBound, dynamicUpperBound,
                                  staticUpperBound);
       op.getDynamicUpperBoundMutable().assign(dynamicUpperBound);
       op.setStaticUpperBound(staticUpperBound);
 
       dispatchIndexOpFoldResults(mixedStep, dynamicStep, staticStep);
       op.getDynamicStepMutable().assign(dynamicStep);
       op.setStaticStep(staticStep);
+    });
     return success();
   }
 };
 
 } // namespace
 
 void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
for (auto yieldedAndBlockArgs :
          llvm::zip(term.getArgs(), op.getAfterArguments())) {
       if (std::get<0>(yieldedAndBlockArgs) == term.getCondition()) {
         if (!std::get<1>(yieldedAndBlockArgs).use_empty()) {
           if (!constantTrue)
             constantTrue = rewriter.create<arith::ConstantOp>(
                 op.getLoc(), term.getCondition().getType(),
                 rewriter.getBoolAttr(true));
 
-          std::get<1>(yieldedAndBlockArgs).replaceAllUsesWith(constantTrue);
+          rewriter.replaceAllUsesWith(std::get<1>(yieldedAndBlockArgs),
+                                      constantTrue);
           replaced = true;
         }
       }
     }
     return success(replaced);
   }
 };
 

mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp

LogicalResult matchAndRewrite(ForOp forOp,
     // Add induction variable incrementation
     rewriter.setInsertionPointToEnd(afterBlock);
     auto ivIncOp = rewriter.create<arith::AddIOp>(
         whileOp.getLoc(), afterBlock->getArgument(0), forOp.getStep());
 
     // Rewrite uses of the for-loop block arguments to the new while-loop
     // "after" arguments
     for (const auto &barg : enumerate(forOp.getBody(0)->getArguments()))
-      barg.value().replaceAllUsesWith(afterBlock->getArgument(barg.index()));
+      rewriter.replaceAllUsesWith(barg.value(),
+                                  afterBlock->getArgument(barg.index()));
 
     // Inline for-loop body operations into 'after' region.
     for (auto &arg : llvm::make_early_inc_range(*forOp.getBody()))
       arg.moveBefore(afterBlock, afterBlock->end());
 
     // Add incremented IV to yield operations
     for (auto yieldOp : afterBlock->getOps<scf::YieldOp>()) {
       SmallVector<Value> yieldOperands = yieldOp.getOperands();
       yieldOperands.insert(yieldOperands.begin(), ivIncOp.getResult());
-      yieldOp->setOperands(yieldOperands);
+      rewriter.updateRootInPlace(
+          yieldOp, [&]() { yieldOp->setOperands(yieldOperands); });
     }
 
     // We cannot do a direct replacement of the forOp since the while op returns
     // an extra value (the induction variable escapes the loop through being
     // carried in the set of iterargs). Instead, rewrite uses of the forOp
     // results.
     for (const auto &arg : llvm::enumerate(forOp.getResults()))
-      arg.value().replaceAllUsesWith(whileOp.getResult(arg.index() + 1));
+      rewriter.replaceAllUsesWith(arg.value(),
+                                  whileOp.getResult(arg.index() + 1));
 
     rewriter.eraseOp(forOp);
     return success();
   }
 };
 
 struct ForToWhileLoop : public impl::SCFForToWhileLoopBase<ForToWhileLoop> {
   void runOnOperation() override {

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

splitBound = b.createOrFold<AffineApplyOp>(loc, modMap,
                                              ValueRange{forOp.getLowerBound(),
                                                         forOp.getUpperBound(),
                                                         forOp.getStep()});
 
   // Create ForOp for partial iteration.
   b.setInsertionPointAfter(forOp);
   partialIteration = cast<ForOp>(b.clone(*forOp.getOperation()));
   partialIteration.getLowerBoundMutable().assign(splitBound);
-  forOp.replaceAllUsesWith(partialIteration->getResults());
+  b.replaceAllUsesWith(forOp.getResults(), partialIteration->getResults());
   partialIteration.getInitArgsMutable().assign(forOp->getResults());
 
   // Set new upper loop bound.
   b.updateRootInPlace(
       forOp, [&]() { forOp.getUpperBoundMutable().assign(splitBound); });
 
   return success();
 }
if (skipPartial) {
           return failure();
       }
     }
     // Apply loop peeling.
     scf::ForOp partialIteration;
     if (failed(peelAndCanonicalizeForLoop(rewriter, forOp, partialIteration)))
       return failure();
     // Apply label, so that the same loop is not rewritten a second time.
+    rewriter.updateRootInPlace(partialIteration, [&]() {
       partialIteration->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
+      partialIteration->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());
+    });
     rewriter.updateRootInPlace(forOp, [&]() {
       forOp->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
     });
-    partialIteration->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());
     return success();
   }
 
   /// If set to true, loops inside partial iterations of another peeled loop
   /// are not peeled. This reduces the size of the generated code. Partial
   /// iterations are not usually performance critical.
   /// Note: Takes into account the entire chain of parent operations, not just
   /// the direct parent.

mlir/test/Dialect/SCF/canonicalize.mlir

 }
 
 // CHECK-LABEL:   func @to_select_same_val
 // CHECK-DAG:       [[C0:%.*]] = arith.constant 0 : index
 // CHECK-DAG:       [[C1:%.*]] = arith.constant 1 : index
 // CHECK:           [[V0:%.*]] = arith.select {{.*}}, [[C0]], [[C1]]
 // CHECK:           return [[V0]], [[C1]] : index, index
 
+// -----
 
 func.func @to_select_with_body(%cond: i1) -> index {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %0 = scf.if %cond -> index {
     "test.op"() : () -> ()
     scf.yield %c0 : index
   } else {
     scf.yield %c1 : index
   }
   return %0 : index
 }
 
 // CHECK-LABEL:   func @to_select_with_body
 // CHECK-DAG:       [[C0:%.*]] = arith.constant 0 : index
 // CHECK-DAG:       [[C1:%.*]] = arith.constant 1 : index
 // CHECK:           [[V0:%.*]] = arith.select {{.*}}, [[C0]], [[C1]]
 // CHECK:           scf.if {{.*}} {
 // CHECK:             "test.op"() : () -> ()
 // CHECK:           }
 // CHECK:           return [[V0]] : index
+
 // -----
 
 func.func @to_select2(%cond: i1) -> (index, index) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   %0, %1 = scf.if %cond -> (index, index) {
func.func @for_yields_2(%lb : index, %ub : index, %step : index) -> i32 {
   }
   return %b : i32
 }
 
 // CHECK-LABEL:   func @for_yields_2
 //  CHECK-NEXT:     %[[R:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     return %[[R]] : i32
 
+// -----
+
+func.func private @make_i32() -> i32
+
 func.func @for_yields_3(%lb : index, %ub : index, %step : index) -> (i32, i32, i32) {
   %a = call @make_i32() : () -> (i32)
   %b = call @make_i32() : () -> (i32)
   %r:3 = scf.for %i = %lb to %ub step %step iter_args(%0 = %a, %1 = %a, %2 = %b) -> (i32, i32, i32) {
     %c = func.call @make_i32() : () -> (i32)
     scf.yield %0, %c, %2 : i32, i32, i32
   } {some_attr}
   return %r#0, %r#1, %r#2 : i32, i32, i32
%r:4 = scf.if %arg0 -> (i32, f32, i32, i8) {
     }
     scf.yield %0, %a#1, %a#0, %3 : i32, f32, i32, i8
   } else {
     scf.yield %0, %1, %2, %3 : i32, f32, i32, i8
   }
   return %r#0, %r#1, %r#2, %r#3 : i32, f32, i32, i8
 }
 
+// -----
+
 // CHECK-LABEL: @merge_yielding_nested_if_nv1
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 func.func @merge_yielding_nested_if_nv1(%arg0: i1, %arg1: i1) {
 // CHECK: %[[PRE0:.*]] = "test.op"() : () -> i32
 // CHECK: %[[PRE1:.*]] = "test.op1"() : () -> f32
 // CHECK: %[[COND:.*]] = arith.andi %[[ARG0]], %[[ARG1]]
 // CHECK: scf.if %[[COND]]
 // CHECK:   %[[IN0:.*]] = "test.inop"() : () -> i32
 // CHECK:   %[[IN1:.*]] = "test.inop1"() : () -> f32
 // CHECK: }
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (f32)
   scf.if %arg0 {
     %a:2 = scf.if %arg1 -> (i32, f32) {
       %i = "test.inop"() : () -> (i32)
       %i1 = "test.inop1"() : () -> (f32)
       scf.yield %i, %i1 : i32, f32
     } else {
       scf.yield %0, %1 : i32, f32
     }
   }
   return
 }
 
+// -----
+
 // CHECK-LABEL: @merge_yielding_nested_if_nv2
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 func.func @merge_yielding_nested_if_nv2(%arg0: i1, %arg1: i1) -> i32 {
 // CHECK: %[[PRE0:.*]] = "test.op"() : () -> i32
 // CHECK: %[[PRE1:.*]] = "test.op1"() : () -> i32
 // CHECK: %[[COND:.*]] = arith.andi %[[ARG0]], %[[ARG1]]
 // CHECK: %[[RES:.*]] = arith.select %[[ARG0]], %[[PRE0]], %[[PRE1]]
 // CHECK: scf.if %[[COND]]
 // CHECK:   "test.run"() : () -> ()
 // CHECK: }
 // CHECK: return %[[RES]]
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (i32)
   %r = scf.if %arg0 -> i32 {
     scf.if %arg1 {
       "test.run"() : () -> ()
     }
     scf.yield %0 : i32
   } else {
     scf.yield %1 : i32
   }
   return %r : i32
 }
 
+// -----
+
 // CHECK-LABEL: @merge_fail_yielding_nested_if
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 func.func @merge_fail_yielding_nested_if(%arg0: i1, %arg1: i1) -> (i32, f32, i32, i8) {
 // CHECK-NOT: andi
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (f32)
   %2 = "test.op2"() : () -> (i32)
   %3 = "test.op3"() : () -> (i8)

 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%[[arg:.*]]: i32):
 // CHECK-NEXT:           "test.use"(%[[arg]]) : (i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 // CHECK-NEXT:         return %[[res]] : i32
 
+// -----
+
 // CHECK-LABEL: @while_cmp_lhs
 func.func @while_cmp_lhs(%arg0 : i32) {
   %0 = scf.while () : () -> i32 {
     %val = "test.val"() : () -> i32
     %condition = arith.cmpi ne, %val, %arg0 : i32
     scf.condition(%condition) %val : i32
   } do {
   ^bb0(%val2: i32):

 // CHECK-NEXT:         %[[cmp:.+]] = arith.cmpi ne, %[[val]], %arg0 : i32
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%arg1: i32):
 // CHECK-NEXT:           "test.use"(%[[true]], %[[false]], %arg1) : (i1, i1, i32) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
+// -----
+
 // CHECK-LABEL: @while_cmp_rhs
 func.func @while_cmp_rhs(%arg0 : i32) {
   %0 = scf.while () : () -> i32 {
     %val = "test.val"() : () -> i32
     %condition = arith.cmpi ne, %arg0, %val : i32
     scf.condition(%condition) %val : i32
   } do {
   ^bb0(%val2: i32):

 // CHECK-NEXT:       scf.yield %[[tval0]], %[[tval]] : i32, i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[fval0:.+]] = "test.firstCodeFalse"() : () -> i32
 // CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[fval0]], %[[fval]] : i32, i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[res]]#0, %[[res]]#1 : i32, i32
 
+// -----
 
 // CHECK-LABEL: @combineIfs2
 func.func @combineIfs2(%arg0 : i1, %arg2: i64) -> i32 {
   scf.if %arg0 {
     "test.firstCodeTrue"() : () -> ()
     scf.yield
   }
   %res = scf.if %arg0 -> i32 {

 // CHECK-NEXT:       %[[tval:.+]] = "test.secondCodeTrue"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[tval]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[fval]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[res]] : i32
 
+// -----
 
 // CHECK-LABEL: @combineIfs3
 func.func @combineIfs3(%arg0 : i1, %arg2: i64) -> i32 {
   %res = scf.if %arg0 -> i32 {
     %v = "test.firstCodeTrue"() : () -> i32
     scf.yield %v : i32
   } else {
     %v2 = "test.firstCodeFalse"() : () -> i32

 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:       scf.yield %[[tval]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[fval:.+]] = "test.firstCodeFalse"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[fval]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[res]] : i32
 
+// -----
+
 // CHECK-LABEL: @combineIfs4
 func.func @combineIfs4(%arg0 : i1, %arg2: i64) {
   scf.if %arg0 {
     "test.firstCodeTrue"() : () -> ()
     scf.yield
   }
   scf.if %arg0 {
     "test.secondCodeTrue"() : () -> ()
     scf.yield
   }
   return
 }
 
 // CHECK-NEXT:     scf.if %arg0 {
 // CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:     }
 
+// -----
+
 // CHECK-LABEL: @combineIfsUsed
 // CHECK-SAME: %[[arg0:.+]]: i1
 func.func @combineIfsUsed(%arg0 : i1, %arg2: i64) -> (i32, i32) {
   %res = scf.if %arg0 -> i32 {
     %v = "test.firstCodeTrue"() : () -> i32
     scf.yield %v : i32
   } else {
     %v2 = "test.firstCodeFalse"() : () -> i32

 // CHECK-NEXT:       scf.yield %[[tval0]], %[[tval]] : i32, i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[fval0:.+]] = "test.firstCodeFalse"() : () -> i32
 // CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"(%[[fval0]]) : (i32) -> i32
 // CHECK-NEXT:       scf.yield %[[fval0]], %[[fval]] : i32, i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[res]]#0, %[[res]]#1 : i32, i32
 
+// -----
+
 // CHECK-LABEL: @combineIfsNot
 // CHECK-SAME: %[[arg0:.+]]: i1
 func.func @combineIfsNot(%arg0 : i1, %arg2: i64) {
   %true = arith.constant true
   %not = arith.xori %arg0, %true : i1
   scf.if %arg0 {
     "test.firstCodeTrue"() : () -> ()
     scf.yield
   }
   scf.if %not {
     "test.secondCodeTrue"() : () -> ()
     scf.yield
   }
   return
 }
 
 // CHECK-NEXT:     scf.if %[[arg0]] {
 // CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:     }
 
+// -----
+
 // CHECK-LABEL: @combineIfsNot2
 // CHECK-SAME: %[[arg0:.+]]: i1
 func.func @combineIfsNot2(%arg0 : i1, %arg2: i64) {
   %true = arith.constant true
   %not = arith.xori %arg0, %true : i1
   scf.if %not {
     "test.firstCodeTrue"() : () -> ()
     scf.yield
   }
   scf.if %arg0 {
     "test.secondCodeTrue"() : () -> ()
     scf.yield
   }
   return
 }
 
 // CHECK-NEXT:     scf.if %[[arg0]] {
 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
 // CHECK-NEXT:     }
+
 // -----
 
 // CHECK-LABEL: func @propagate_into_execute_region
 func.func @propagate_into_execute_region() {
   %cond = arith.constant 0 : i1
   affine.for %i = 0 to 100 {
     "test.foo"() : () -> ()
     %v = scf.execute_region -> i64 {

 }
 
 // CHECK-NEXT:     affine.for %arg0 = 0 to 100 {
 // CHECK-NEXT:       "test.foo"() : () -> ()
 // CHECK-NEXT:       %[[VAL:.*]] = "test.val"() : () -> i64
 // CHECK-NEXT:       "test.bar"(%[[VAL]]) : (i64) -> ()
 // CHECK-NEXT:     }
 
-
 // -----
 
 // CHECK-LABEL: func @func_execute_region_elim
 func.func @func_execute_region_elim() {
     "test.foo"() : () -> ()
     %v = scf.execute_region -> i64 {
       %c = "test.cmp"() : () -> i1
       cf.cond_br %c, ^bb2, ^bb3

 // CHECK:     cf.br ^[[bb3:.+]](%[[x]] : i64)
 // CHECK:   ^[[bb2]]:
 // CHECK:     %[[y:.+]] = "test.val2"
 // CHECK:     cf.br ^[[bb3]](%[[y:.+]] : i64)
 // CHECK:   ^[[bb3]](%[[z:.+]]: i64):
 // CHECK:     "test.bar"(%[[z]])
 // CHECK:     return
 
-
 // -----
 
 // CHECK-LABEL: func @func_execute_region_elim_multi_yield
 func.func @func_execute_region_elim_multi_yield() {
     "test.foo"() : () -> ()
     %v = scf.execute_region -> i64 {
       %c = "test.cmp"() : () -> i1
       cf.cond_br %c, ^bb2, ^bb3