[mlir][SCF] Adding custom builder to SCF::WhileOp.

[Moerafaat]

Summary

This is a similar builder to the one for SCF::IfOp which allows users to pass region builders to it. Refer to the builders for IfOp.

[tpopp]

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

988

ForOp uses a BodyBuilderFn to represent this type. Let's use that here also, and move the definition of BodyBuilderFn inside of a let extraClassDeclaration = ... inside of the SCF_Dialect definition.

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

2692

Here and below, I don't actually know the signature of this method, so can you modify them to beforeRegion, /*parameterNameHere=*/{}, resultTypes?

[Moerafaat]

[mlir][SCF] Updating one of the usages of SCF::WhileOp to use the new builder.

[Moerafaat]

[mlir][SCF] Resolving comments

[Moerafaat]

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

988

Done

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

2692

Done

[tpopp]

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

988

Can you use BodyBuilderFn here now?

[tpopp]

It looks like you will also need to run clang-format on this. I use git clang-format --style=LLVM HEAD to format my patches before landing (then you can use arc diff to upload the final patch this time and I can submit it on your behalf).

[Moerafaat]

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

988

I tried this before and it fails in SCFOps.h.inc saying "unknown type name 'BodyBuilderFn'". I checked ForOp and it also does it the way I am currently doing. My guess is the extra declarations reflect in the .cpp file which are not visible to the header file. What do you think?

[Moerafaat]

Fixing clang formatting

[tpopp]

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

988

Right, it looks like extraClassDeclarations come after the generated methods. I think it would be more ideal for extraClassDeclarations to come before, but that would only work if extraClassDeclarations didn't have definitions also, which it often does, I believe, so let's leave this as it is.

We could also get around this by putting this in the extraClassDeclarations of the SCF dialect, I believe, but that is not necesarilly a good choice, as there's no guarantee that all SCF ops have the same body builder function.

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

^bb0(%arg2: f32):
       // "After" region.
       // Forwards the values back to "before" region unmodified.
       scf.yield %arg2 : f32
     }
     ```
 
     Note that the types of region arguments need not to match with each other.
     The op expects the operand types to match with argument types of the
-    "before" region"; the result types to match with the trailing operand types
+    "before" region; the result types to match with the trailing operand types
     of the terminator of the "before" region, and with the argument types of the
     "after" region. The following scheme can be used to share the results of
     some operations executed in the "before" region with the "after" region,
     avoiding the need to recompute them.
 
     ```mlir
     %res = scf.while (%arg1 = %init1) : (f32) -> i64 {
       // One can perform some computations, e.g., necessary to evaluate the
let description = [{
     assignment ::= ssa-value `=` ssa-value
     ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$inits);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$before, SizedRegion<1>:$after);
 
+  let builders = [
+    OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>":$beforeBuilder,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>":$afterBuilder)>
+  ];
+
   let extraClassDeclaration = [{
+    using BodyBuilderFn =
+        function_ref<void(OpBuilder &, Location, ValueRange)>;
+
     OperandRange getSuccessorEntryOperands(Optional<unsigned> index);
     ConditionOp getConditionOp();
     YieldOp getYieldOp();
     Block::BlockArgListType getBeforeArguments();
     Block::BlockArgListType getAfterArguments();
   }];
 
   let hasCanonicalizer = 1;

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

auto zero =
       rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(elementTy, 0));
   auto leadingZeros = rewriter.create<arith::ConstantOp>(
       loc, IntegerAttr::get(elementTy, bitWidth));
 
   SmallVector<Value> operands = {operand, leadingZeros, zero};
   SmallVector<Type> types = {elementTy, elementTy, elementTy};
   SmallVector<Location> locations = {loc, loc, loc};
 
-  auto whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
-  Block *before =
-      rewriter.createBlock(&whileOp.getBefore(), {}, types, locations);
-  Block *after =
-      rewriter.createBlock(&whileOp.getAfter(), {}, types, locations);
-
+  auto whileOp = rewriter.create<scf::WhileOp>(
+      loc, types, operands,
+      [&](OpBuilder &beforeBuilder, Location beforeLoc, ValueRange args) {
         // The conditional block of the while loop.
-  {
-    rewriter.setInsertionPointToStart(&whileOp.getBefore().front());
-    Value input = before->getArgument(0);
-    Value zero = before->getArgument(2);
+        Value input = args[0];
+        Value zero = args[2];
 
-    Value inputNotZero = rewriter.create<arith::CmpIOp>(
+        Value inputNotZero = beforeBuilder.create<arith::CmpIOp>(
             loc, arith::CmpIPredicate::ne, input, zero);
-    rewriter.create<scf::ConditionOp>(loc, inputNotZero,
-                                      before->getArguments());
-  }
-
+        beforeBuilder.create<scf::ConditionOp>(loc, inputNotZero, args);
+      },
+      [&](OpBuilder &afterBuilder, Location afterLoc, ValueRange args) {
         // The body of the while loop: shift right until reaching a value of 0.
-  {
-    rewriter.setInsertionPointToStart(&whileOp.getAfter().front());
-    Value input = after->getArgument(0);
-    Value leadingZeros = after->getArgument(1);
-
-    auto one =
-        rewriter.create<arith::ConstantOp>(loc, IntegerAttr::get(elementTy, 1));
-    auto shifted = rewriter.create<arith::ShRUIOp>(loc, resultTy, input, one);
-    auto leadingZerosMinusOne =
-        rewriter.create<arith::SubIOp>(loc, resultTy, leadingZeros, one);
-
-    rewriter.create<scf::YieldOp>(
-        loc,
-        ValueRange({shifted, leadingZerosMinusOne, after->getArgument(2)}));
-  }
+        Value input = args[0];
+        Value leadingZeros = args[1];
+
+        auto one = afterBuilder.create<arith::ConstantOp>(
+            loc, IntegerAttr::get(elementTy, 1));
+        auto shifted =
+            afterBuilder.create<arith::ShRUIOp>(loc, resultTy, input, one);
+        auto leadingZerosMinusOne = afterBuilder.create<arith::SubIOp>(
+            loc, resultTy, leadingZeros, one);
+
+        afterBuilder.create<scf::YieldOp>(
+            loc, ValueRange({shifted, leadingZerosMinusOne, args[2]}));
+      });
 
   rewriter.setInsertionPointAfter(whileOp);
   rewriter.replaceOp(op, whileOp->getResult(1));
   return success();
 }
 
 void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) {
   patterns.add(convertCtlzOp);
 }
 
 void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }

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

return emitOpError() << "needs to have type " << reduceType
                          << " (the type of the enclosing ReduceOp)";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // WhileOp
 //===----------------------------------------------------------------------===//
 
+void WhileOp::build(::mlir::OpBuilder &odsBuilder,
+                    ::mlir::OperationState &odsState, TypeRange resultTypes,
+                    ValueRange operands, BodyBuilderFn beforeBuilder,
+                    BodyBuilderFn afterBuilder) {
+  assert(beforeBuilder && "the builder callback for 'before' must be present");
+  assert(afterBuilder && "the builder callback for 'after' must be present");
+
+  odsState.addOperands(operands);
+  odsState.addTypes(resultTypes);
+
+  OpBuilder::InsertionGuard guard(odsBuilder);
+
+  SmallVector<Location, 4> blockArgLocs;
+  for (Value operand : operands) {
+    blockArgLocs.push_back(operand.getLoc());
+  }
+
+  Region *beforeRegion = odsState.addRegion();
+  Block *beforeBlock = odsBuilder.createBlock(beforeRegion, /*insertPt=*/{},
+                                              resultTypes, blockArgLocs);
+  beforeBuilder(odsBuilder, odsState.location, beforeBlock->getArguments());
+
+  Region *afterRegion = odsState.addRegion();
+  Block *afterBlock = odsBuilder.createBlock(afterRegion, /*insertPt=*/{},
+                                             resultTypes, blockArgLocs);
+  afterBuilder(odsBuilder, odsState.location, afterBlock->getArguments());
+}
+
 OperandRange WhileOp::getSuccessorEntryOperands(Optional<unsigned> index) {
   assert(index && *index == 0 &&
          "WhileOp is expected to branch only to the first region");
 
   return getInits();
 }
 
 ConditionOp WhileOp::getConditionOp() {