[mlir] Introduce callback-based builders to SCF Parallel and Reduce ops

[ftynse]

Summary

Similarly to scf::ForOp, introduce additional function_ref arguments to
::build functions of SCF ParallelOp and ReduceOp. The provided functions
will be called to construct the body of the respective operations while
constructing the operation itself. Exercise them in LoopUtils.

[pifon2a]

thanks, Alex!

mlir/lib/Dialect/SCF/SCF.cpp

541

nit: will smth like that work?

function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)> wrapper;
if (bodyBuilderFn) {
  wrapper = [&bodyBuilderFn](OpBuilder &nestedBuilder, Location nestedLoc,
                                  ValueRange ivs, ValueRange) {
    bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
  };
}
 build(builder, result, lowerBounds, upperBounds, steps, ValueRange(), wrapper);

[ftynse]

Address review

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

541

This way, the wrapper will refer to a temporary scoped within the if block, so it will become dangling after that.

I'll factor out the lambda into a function-level variable to help readability.

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

let arguments = (ins Variadic<Index>:$lowerBound,
                        Variadic<AnyType>:$initVals);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "ValueRange lowerBounds, ValueRange upperBounds, "
-              "ValueRange steps, ValueRange initVals = {}">,
+              "ValueRange steps, ValueRange initVals, "
+              "function_ref<void (OpBuilder &, Location, "
+                                 "ValueRange, ValueRange)>"
+              "  bodyBuilderFn = nullptr">,
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "ValueRange lowerBounds, ValueRange upperBounds, "
+              "ValueRange steps, "
+              "function_ref<void (OpBuilder &, Location, ValueRange)>"
+              "  bodyBuilderFn = nullptr">,
   ];
 
   let extraClassDeclaration = [{
     ValueRange getInductionVars() {
       return getBody()->getArguments();
     }
     unsigned getNumLoops() { return step().size(); }
     unsigned getNumReductions() { return initVals().size(); }
scf.reduce(%operand) : f32 {
         scf.reduce.return %res : f32
     }
     ```
   }];
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "Value operand">
+              "Value operand, "
+              "function_ref<void (OpBuilder &, Location, Value, Value)>"
+              "  bodyBuilderFn = nullptr">
   ];
 
   let arguments = (ins AnyType:$operand);
   let regions = (region SizedRegion<1>:$reductionOperator);
 }
 
 def ReduceReturnOp :
     SCF_Op<"reduce.return", [HasParent<"ReduceOp">, NoSideEffect,

mlir/lib/Dialect/SCF/SCF.cpp

void IfOp::getSuccessorRegions(Optional<unsigned> index,
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
-void ParallelOp::build(OpBuilder &builder, OperationState &result,
-                       ValueRange lbs, ValueRange ubs, ValueRange steps,
-                       ValueRange initVals) {
-  result.addOperands(lbs);
-  result.addOperands(ubs);
+void ParallelOp::build(
+    OpBuilder &builder, OperationState &result, ValueRange lowerBounds,
+    ValueRange upperBounds, ValueRange steps, ValueRange initVals,
+    function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>
+        bodyBuilderFn) {
+  result.addOperands(lowerBounds);
+  result.addOperands(upperBounds);
   result.addOperands(steps);
   result.addOperands(initVals);
   result.addAttribute(
       ParallelOp::getOperandSegmentSizeAttr(),
-      builder.getI32VectorAttr({static_cast<int32_t>(lbs.size()),
-                                static_cast<int32_t>(ubs.size()),
+      builder.getI32VectorAttr({static_cast<int32_t>(lowerBounds.size()),
+                                static_cast<int32_t>(upperBounds.size()),
                                 static_cast<int32_t>(steps.size()),
                                 static_cast<int32_t>(initVals.size())}));
+  result.addTypes(initVals.getTypes());
+
+  OpBuilder::InsertionGuard guard(builder);
+  unsigned numIVs = steps.size();
+  SmallVector<Type, 8> argTypes(numIVs, builder.getIndexType());
   Region *bodyRegion = result.addRegion();
+  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes);
+
+  if (bodyBuilderFn) {
+    builder.setInsertionPointToStart(bodyBlock);
+    bodyBuilderFn(builder, result.location,
+                  bodyBlock->getArguments().take_front(numIVs),
+                  bodyBlock->getArguments().drop_front(numIVs));
+  }
   ParallelOp::ensureTerminator(*bodyRegion, builder, result.location);
-  for (size_t i = 0, e = steps.size(); i < e; ++i)
-    bodyRegion->front().addArgument(builder.getIndexType());
-  for (Value init : initVals)
-    result.addTypes(init.getType());
+}
+
+void ParallelOp::build(
+    OpBuilder &builder, OperationState &result, ValueRange lowerBounds,
+    ValueRange upperBounds, ValueRange steps,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+  // Only pass a non-null wrapper if bodyBuilderFn is non-null itself. Make sure
+  // we don't capture a reference to a temporary by constructing the lambda at
+  // function level.
+  auto wrappedBuilderFn = [&bodyBuilderFn](OpBuilder &nestedBuilder,
+                                           Location nestedLoc, ValueRange ivs,
+                                           ValueRange) {
+    bodyBuilderFn(nestedBuilder, nestedLoc, ivs);
+  };
+  function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)> wrapper;
+  if (bodyBuilderFn)
+    wrapper = wrappedBuilderFn;
+
+  build(builder, result, lowerBounds, upperBounds, steps, ValueRange(),
+        wrapper);
 }
 
 static LogicalResult verify(ParallelOp op) {
   // Check that there is at least one value in lowerBound, upperBound and step.
   // It is sufficient to test only step, because it is ensured already that the
   // number of elements in lowerBound, upperBound and step are the same.
   Operation::operand_range stepValues = op.step();
   if (stepValues.empty())
ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast<ParallelOp>(containingOp);
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
-void ReduceOp::build(OpBuilder &builder, OperationState &result,
-                     Value operand) {
+void ReduceOp::build(
+    OpBuilder &builder, OperationState &result, Value operand,
+    function_ref<void(OpBuilder &, Location, Value, Value)> bodyBuilderFn) {
   auto type = operand.getType();
   result.addOperands(operand);
-  Region *bodyRegion = result.addRegion();
 
-  Block *b = new Block();
-  b->addArguments(ArrayRef<Type>{type, type});
-  bodyRegion->getBlocks().insert(bodyRegion->end(), b);
+  OpBuilder::InsertionGuard guard(builder);
+  Region *bodyRegion = result.addRegion();
+  Block *body = builder.createBlock(bodyRegion, {}, ArrayRef<Type>{type, type});
+  if (bodyBuilderFn)
+    bodyBuilderFn(builder, result.location, body->getArgument(0),
+                  body->getArgument(1));
 }
 
 static LogicalResult verify(ReduceOp op) {
   // The region of a ReduceOp has two arguments of the same type as its operand.
   auto type = op.operand().getType();
   Block &block = op.reductionOperator().front();
   if (block.empty())
     return op.emitOpError("the block inside reduce should not be empty");

mlir/lib/Transforms/Utils/LoopUtils.cpp

void mlir::collapseParallelLoops(
   }
 
   // Create new ParallelLoop with conversions to the original induction values.
   // The loop below uses divisions to get the relevant range of values in the
   // new induction value that represent each range of the original induction
   // value. The remainders then determine based on that range, which iteration
   // of the original induction value this represents. This is a normalized value
   // that is un-normalized already by the previous logic.
-  auto newPloop = outsideBuilder.create<scf::ParallelOp>(loc, lowerBounds,
-                                                         upperBounds, steps);
-  OpBuilder insideBuilder(newPloop.region());
+  auto newPloop = outsideBuilder.create<scf::ParallelOp>(
+      loc, lowerBounds, upperBounds, steps,
+      [&](OpBuilder &insideBuilder, Location, ValueRange ploopIVs) {
         for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
-    Value previous = newPloop.getBody()->getArgument(i);
+          Value previous = ploopIVs[i];
           unsigned numberCombinedDimensions = combinedDimensions[i].size();
           // Iterate over all except the last induction value.
           for (unsigned j = 0, e = numberCombinedDimensions - 1; j < e; ++j) {
             unsigned idx = combinedDimensions[i][j];
 
             // Determine the current induction value's current loop iteration
-      Value iv = insideBuilder.create<SignedRemIOp>(loc, previous,
-                                                    normalizedUpperBounds[idx]);
+            Value iv = insideBuilder.create<SignedRemIOp>(
+                loc, previous, normalizedUpperBounds[idx]);
             replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx), iv,
                                        loops.region());
 
-      // Remove the effect of the current induction value to prepare for the
-      // next value.
+            // Remove the effect of the current induction value to prepare for
+            // the next value.
             previous = insideBuilder.create<SignedDivIOp>(
                 loc, previous, normalizedUpperBounds[idx + 1]);
           }
 
           // The final induction value is just the remaining value.
           unsigned idx = combinedDimensions[i][numberCombinedDimensions - 1];
-    replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx), previous,
-                               loops.region());
+          replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx),
+                                     previous, loops.region());
         }
+      });
 
   // Replace the old loop with the new loop.
   loops.getBody()->back().erase();
   newPloop.getBody()->getOperations().splice(
       Block::iterator(newPloop.getBody()->back()),
       loops.getBody()->getOperations());
   loops.erase();
 }