[mlir] SCF: provide function_ref builders for IfOp

[ftynse]

Summary

Now that OpBuilder is available in build functions, it becomes possible to
populate the "then" and "else" regions directly when building the "if"
operation. This is desirable in more structured forms of builders, especially
in when conditionals are mixed with loops. Provide new build APIs taking
callbacks for body constructors, similarly to scf::ForOp, and replace more
clunky edsc::BlockBuilder uses with these. The original APIs remain available
and go through the new implementation.

[nicolasvasilache]

Nice, thanks!

mlir/include/mlir/Dialect/SCF/EDSC/Builders.h

scf::ValueVector loopNestBuilder(ValueRange lbs, ValueRange ubs,
                                  ValueRange steps,
                                  function_ref<void(ValueRange)> fun = nullptr);
 scf::ValueVector loopNestBuilder(Value lb, Value ub, Value step,
                                  function_ref<void(Value)> fun = nullptr);
 scf::ValueVector loopNestBuilder(
     Value lb, Value ub, Value step, ValueRange iterArgInitValues,
     function_ref<scf::ValueVector(Value, ValueRange)> fun = nullptr);
 
+/// Adapters for building if conditions using the builder and the location
+/// stored in ScopedContext. 'thenBody' is mandatory, 'elseBody' can be omitted
+/// if the condition should not have an 'else' part.
+ValueRange
+conditionBuilder(TypeRange results, Value condition,
+                 function_ref<scf::ValueVector()> thenBody,
+                 function_ref<scf::ValueVector()> elseBody = nullptr);
+ValueRange conditionBuilder(Value condition, function_ref<void()> thenBody,
+                            function_ref<void()> elseBody = nullptr);
+
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_EDSC_BUILDERS_H_

mlir/include/mlir/Dialect/SCF/SCF.h

 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace scf {
 
+void buildTerminatedBody(OpBuilder &builder, Location loc);
+
 #include "mlir/Dialect/SCF/SCFOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/SCFOps.h.inc"
 
 // Insert `loop.yield` at the end of the only region's only block if it
 // does not have a terminator already.  If a new `loop.yield` is inserted,
 // the location is specified by `loc`. If the region is empty, insert a new

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

def IfOp : SCF_Op<"if",
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
               "Value cond, bool withElseRegion">,
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "TypeRange resultTypes, Value cond, bool withElseRegion">
+              "TypeRange resultTypes, Value cond, bool withElseRegion">,
+    OpBuilder<
+        "OpBuilder &builder, OperationState &result, TypeRange resultTypes, "
+        "Value cond, "
+        "function_ref<void(OpBuilder &, Location)> thenBuilder "
+        "    = buildTerminatedBody, "
+        "function_ref<void(OpBuilder &, Location)> elseBuilder = nullptr">,
+    OpBuilder<
+        "OpBuilder &builder, OperationState &result, Value cond, "
+        "function_ref<void(OpBuilder &, Location)> thenBuilder "
+        "    = buildTerminatedBody, "
+        "function_ref<void(OpBuilder &, Location)> elseBuilder = nullptr">
   ];
 
   let extraClassDeclaration = [{
     OpBuilder getThenBodyBuilder() {
       Block* body = getBody(0);
       return results().empty() ? OpBuilder::atBlockTerminator(body)
                                : OpBuilder::atBlockEnd(body);
     }

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

Value inBoundsCondition = emitInBoundsCondition(
         majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
 
     if (inBoundsCondition) {
       // 2. If the condition is not null, we need an IfOp, which may yield
       // if `options.unroll` is true.
       SmallVector<Type, 1> resultType;
       if (options.unroll)
         resultType.push_back(vectorType);
-      auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
-          ScopedContext::getLocation(), resultType, inBoundsCondition,
-          /*withElseRegion=*/true);
 
-      // 3.a. If in-bounds, progressively lower to a 1-D transfer read.
-      BlockBuilder(&ifOp.thenRegion().front(), Append())([&] {
+      // 3. If in-bounds, progressively lower to a 1-D transfer read, otherwise
+      // splat a 1-D vector.
+      ValueRange ifResults = conditionBuilder(
+          resultType, inBoundsCondition,
+          [&]() -> scf::ValueVector {
             Value vector = load1DVector(majorIvsPlusOffsets);
-        // 3.a.i. If `options.unroll` is true, insert the 1-D vector in the
+            // 3.a. If `options.unroll` is true, insert the 1-D vector in the
             // aggregate. We must yield and merge with the `else` branch.
             if (options.unroll) {
               vector = vector_insert(vector, result, majorIvs);
-          (loop_yield(vector));
-          return;
+              return {vector};
             }
-        // 3.a.ii. Otherwise, just go through the temporary `alloc`.
+            // 3.b. Otherwise, just go through the temporary `alloc`.
             std_store(vector, alloc, majorIvs);
-      });
-
-      // 3.b. If not in-bounds, splat a 1-D vector.
-      BlockBuilder(&ifOp.elseRegion().front(), Append())([&] {
+            return {};
+          },
+          [&]() -> scf::ValueVector {
             Value vector = std_splat(minorVectorType, xferOp.padding());
-        // 3.a.i. If `options.unroll` is true, insert the 1-D vector in the
+            // 3.c. If `options.unroll` is true, insert the 1-D vector in the
             // aggregate. We must yield and merge with the `then` branch.
             if (options.unroll) {
               vector = vector_insert(vector, result, majorIvs);
-          (loop_yield(vector));
-          return;
+              return {vector};
             }
-        // 3.b.ii. Otherwise, just go through the temporary `alloc`.
+            // 3.d. Otherwise, just go through the temporary `alloc`.
             std_store(vector, alloc, majorIvs);
+            return {};
           });
+
       if (!resultType.empty())
-        result = *ifOp.results().begin();
+        result = *ifResults.begin();
     } else {
       // 4. Guaranteed in-bounds, progressively lower to a 1-D transfer read.
       Value loaded1D = load1DVector(majorIvsPlusOffsets);
       // 5.a. If `options.unroll` is true, insert the 1-D vector in the
       // aggregate.
       if (options.unroll)
         result = vector_insert(loaded1D, result, majorIvs);
       // 5.b. Otherwise, just go through the temporary `alloc`.
emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
     // context.
     SmallVector<Value, 4> majorIvsPlusOffsets;
     Value inBoundsCondition = emitInBoundsCondition(
         majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
 
     if (inBoundsCondition) {
       // 2.a. If the condition is not null, we need an IfOp, to write
       // conditionally. Progressively lower to a 1-D transfer write.
-      auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
-          ScopedContext::getLocation(), TypeRange{}, inBoundsCondition,
-          /*withElseRegion=*/false);
-      BlockBuilder(&ifOp.thenRegion().front(),
-                   Append())([&] { emitTransferWrite(majorIvsPlusOffsets); });
+      conditionBuilder(inBoundsCondition,
+                       [&] { emitTransferWrite(majorIvsPlusOffsets); });
     } else {
       // 2.b. Guaranteed in-bounds. Progressively lower to a 1-D transfer write.
       emitTransferWrite(majorIvsPlusOffsets);
     }
   });
 
   rewriter.eraseOp(op);
 

mlir/lib/Dialect/SCF/EDSC/Builders.cpp

return mlir::scf::buildLoopNest(
         assert(ivs.size() == 1 && "expected one induction variable");
         ScopedContext context(builder, loc);
         if (fun)
           return fun(ivs[0], args);
         return scf::ValueVector(iterArgInitValues.begin(),
                                 iterArgInitValues.end());
       });
 }
+
+static std::function<void(OpBuilder &, Location)>
+wrapIfBody(function_ref<scf::ValueVector()> body, TypeRange expectedTypes) {
+  (void)expectedTypes;
+  return [=](OpBuilder &builder, Location loc) {
+    ScopedContext context(builder, loc);
+    scf::ValueVector returned = body();
+    assert(ValueRange(returned).getTypes() == expectedTypes &&
+           "'if' body builder returned values of unexpected type");
+    builder.create<scf::YieldOp>(loc, returned);
+  };
+}
+
+ValueRange
+mlir::edsc::conditionBuilder(TypeRange results, Value condition,
+                             function_ref<scf::ValueVector()> thenBody,
+                             function_ref<scf::ValueVector()> elseBody) {
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  assert(thenBody && "thenBody is mandatory");
+
+  auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
+      ScopedContext::getLocation(), results, condition,
+      wrapIfBody(thenBody, results), wrapIfBody(elseBody, results));
+  return ifOp.getResults();
+}
+
+static std::function<void(OpBuilder &, Location)>
+wrapZeroResultIfBody(function_ref<void()> body) {
+  return [=](OpBuilder &builder, Location loc) {
+    ScopedContext context(builder, loc);
+    body();
+    builder.create<scf::YieldOp>(loc);
+  };
+}
+
+ValueRange mlir::edsc::conditionBuilder(Value condition,
+                                        function_ref<void()> thenBody,
+                                        function_ref<void()> elseBody) {
+  assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
+  assert(thenBody && "thenBody is mandatory");
+
+  ScopedContext::getBuilderRef().create<scf::IfOp>(
+      ScopedContext::getLocation(), condition, wrapZeroResultIfBody(thenBody),
+      elseBody ? llvm::function_ref<void(OpBuilder &, Location)>(
+                     wrapZeroResultIfBody(elseBody))
+               : llvm::function_ref<void(OpBuilder &, Location)>(nullptr));
+  return {};
+}

mlir/lib/Dialect/SCF/SCF.cpp

 SCFDialect::SCFDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"
       >();
 }
 
+/// Default callback for IfOp builders. Inserts a yield without arguments.
+void mlir::scf::buildTerminatedBody(OpBuilder &builder, Location loc) {
+  builder.create<scf::YieldOp>(loc);
+}
+
 //===----------------------------------------------------------------------===//
 // ForOp
 //===----------------------------------------------------------------------===//
 
 void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
                   Value ub, Value step, ValueRange iterArgs,
                   BodyBuilderFn bodyBuilder) {
   result.addOperands({lb, ub, step});

 
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {
   build(builder, result, /*resultTypes=*/llvm::None, cond, withElseRegion);
 }
 
 void IfOp::build(OpBuilder &builder, OperationState &result,
                  TypeRange resultTypes, Value cond, bool withElseRegion) {
+  auto addTerminator = [&](OpBuilder &nested, Location loc) {
+    if (resultTypes.empty())
+      IfOp::ensureTerminator(*nested.getInsertionBlock()->getParent(), nested,
+                             loc);
+  };
+
+  build(builder, result, resultTypes, cond, addTerminator,
+        withElseRegion ? addTerminator
+                       : function_ref<void(OpBuilder &, Location)>());
+}
+
+void IfOp::build(OpBuilder &builder, OperationState &result,
+                 TypeRange resultTypes, Value cond,
+                 function_ref<void(OpBuilder &, Location)> thenBuilder,
+                 function_ref<void(OpBuilder &, Location)> elseBuilder) {
+  assert(thenBuilder && "the builder callback for 'then' must be present");
+
   result.addOperands(cond);
   result.addTypes(resultTypes);
 
+  OpBuilder::InsertionGuard guard(builder);
   Region *thenRegion = result.addRegion();
-  thenRegion->push_back(new Block());
-  if (resultTypes.empty())
-    IfOp::ensureTerminator(*thenRegion, builder, result.location);
+  builder.createBlock(thenRegion);
+  thenBuilder(builder, result.location);
 
   Region *elseRegion = result.addRegion();
-  if (withElseRegion) {
-    elseRegion->push_back(new Block());
-    if (resultTypes.empty())
-      IfOp::ensureTerminator(*elseRegion, builder, result.location);
-  }
+  if (!elseBuilder)
+    return;
+
+  builder.createBlock(elseRegion);
+  elseBuilder(builder, result.location);
+
+void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
+                 function_ref<void(OpBuilder &, Location)> thenBuilder,
+                 function_ref<void(OpBuilder &, Location)> elseBuilder) {
+  build(builder, result, TypeRange(), cond, thenBuilder, elseBuilder);
 }
 
 static LogicalResult verify(IfOp op) {
   // Verify that the entry of each child region does not have arguments.
   for (auto &region : op.getOperation()->getRegions()) {
     if (region.empty())
       continue;