[MLIR] Remove scf.if builder with explicit result types and callbacks

[frgossen]

Summary

Instead, use the builder and infer the return type based on the inner yield ops.
Also, fix uses that do not create the terminator as required for the callback builders.

Depends On D142049

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

let description = [{
     ```
   }];
   let arguments = (ins I1:$condition);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
 
   let skipDefaultBuilders = 1;
   let builders = [
+    OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond)>,
     OpBuilder<(ins "Value":$cond, "bool":$withElseRegion)>,
     OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond,
       "bool":$withElseRegion)>,
-    // TODO: Remove builder when it is no longer used to create invalid `if` ops
-    // (with a type mispatch between the op and it's inner `yield` op).
-    OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond,
-      CArg<"function_ref<void(OpBuilder &, Location)>",
-           "buildTerminatedBody">:$thenBuilder,
-      CArg<"function_ref<void(OpBuilder &, Location)>",
-           "nullptr">:$elseBuilder)>,
     OpBuilder<(ins "Value":$cond,
       CArg<"function_ref<void(OpBuilder &, Location)>",
            "buildTerminatedBody">:$thenBuilder,
       CArg<"function_ref<void(OpBuilder &, Location)>",
-           "nullptr">:$elseBuilder)>
+           "nullptr">:$elseBuilder)>,
   ];
 
   let extraClassDeclaration = [{
     OpBuilder getThenBodyBuilder(OpBuilder::Listener *listener = nullptr) {
       Block* body = getBody(0);
       return getResults().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                   : OpBuilder::atBlockEnd(body, listener);
     }

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

Value getBroadcastedDim(ImplicitLocOpBuilder lb, ValueRange extentTensors,
   for (auto tup : llvm::zip(extentTensors, rankDiffs)) {
     Value shape = std::get<0>(tup);
     Value rankDiff = std::get<1>(tup);
     Value outOfBounds = lb.create<arith::CmpIOp>(arith::CmpIPredicate::ult,
                                                  outputDimension, rankDiff);
     Type indexTy = lb.getIndexType();
     broadcastedDim =
         lb.create<IfOp>(
-              TypeRange{indexTy}, outOfBounds,
+              outOfBounds,
               [&](OpBuilder &b, Location loc) {
                 b.create<scf::YieldOp>(loc, broadcastedDim);
               },
               [&](OpBuilder &b, Location loc) {
                 // The broadcasting logic is:
                 // - if one extent (here we arbitrarily choose the
                 // extent from the greater-rank operand) is equal to 1,
                 // then take the extent from the other operand
auto reduceResult = lb.create<ForOp>(
         Value broadcastable = iterArgs[0];
         for (auto tup : llvm::zip(adaptor.getShapes(), rankDiffs)) {
           Value shape, rankDiff;
           std::tie(shape, rankDiff) = tup;
           Value outOfBounds = b.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ult, iv, rankDiff);
           broadcastable =
               b.create<IfOp>(
-                   loc, TypeRange{i1Ty}, outOfBounds,
+                   loc, outOfBounds,
                    [&](OpBuilder &b, Location loc) {
                      // Non existent dimensions are always broadcastable
                      b.create<scf::YieldOp>(loc, broadcastable);
                    },
                    [&](OpBuilder &b, Location loc) {
                      // Every value needs to be either 1, or the same non-1
                      // value to be broadcastable in this dim.
                      Value operandDimension =
Value firstRank =
       rewriter.create<tensor::DimOp>(loc, indexTy, firstShape, zero);
   Value result = nullptr;
   // Generate a linear sequence of compares, all with firstShape as lhs.
   for (Value shape : adaptor.getShapes().drop_front(1)) {
     Value rank = rewriter.create<tensor::DimOp>(loc, indexTy, shape, zero);
     Value eqRank = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                   firstRank, rank);
     auto same = rewriter.create<IfOp>(
-        loc, i1Ty, eqRank,
+        loc, eqRank,
         [&](OpBuilder &b, Location loc) {
           Value one = b.create<arith::ConstantIndexOp>(loc, 1);
           Value init =
               b.create<arith::ConstantOp>(loc, i1Ty, b.getBoolAttr(true));
           auto loop = b.create<scf::ForOp>(
               loc, zero, firstRank, one, ValueRange{init},
               [&](OpBuilder &b, Location nestedLoc, Value iv, ValueRange args) {
                 Value conj = args[0];

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

if (cond)
       cond = lb.create<arith::AndIOp>(cond, maskCond);
     else
       cond = maskCond;
   }
 
   // If the condition is non-empty, generate an SCF::IfOp.
   if (cond) {
     auto check = lb.create<scf::IfOp>(
-        resultTypes, cond,
+        cond,
         /*thenBuilder=*/
         [&](OpBuilder &b, Location loc) {
           maybeYieldValue(b, loc, hasRetVal, inBoundsCase(b, loc));
         },
         /*elseBuilder=*/
         [&](OpBuilder &b, Location loc) {
           if (outOfBoundsCase) {
             maybeYieldValue(b, loc, hasRetVal, outOfBoundsCase(b, loc));

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

auto asyncDispatch = [&](OpBuilder &nestedBuilder, Location loc) {
 
     // Wait for the completion of all parallel compute operations.
     b.create<AwaitAllOp>(group);
 
     b.create<scf::YieldOp>();
   };
 
   // Dispatch either single block compute function, or launch async dispatch.
-  b.create<scf::IfOp>(TypeRange(), isSingleBlock, syncDispatch, asyncDispatch);
+  b.create<scf::IfOp>(isSingleBlock, syncDispatch, asyncDispatch);
 }
 
 // Dispatch parallel compute functions by submitting all async compute tasks
 // from a simple for loop in the caller thread.
 static void
 doSequentialDispatch(ImplicitLocOpBuilder &b, PatternRewriter &rewriter,
                      ParallelComputeFunction &parallelComputeFunction,
                      scf::ParallelOp op, Value blockSize, Value blockCount,
auto dispatch = [&](OpBuilder &nestedBuilder, Location loc) {
     // size is larger than the number of iterations in the unrollable inner
     // loops, because otherwise it can reduce the available parallelism.
     if (numUnrollableLoops > 0) {
       Value numIters = b.create<arith::ConstantIndexOp>(
           numIterations[op.getNumLoops() - numUnrollableLoops]);
       Value useBlockAlignedComputeFn = b.create<arith::CmpIOp>(
           arith::CmpIPredicate::sge, blockSize, numIters);
 
-      b.create<scf::IfOp>(TypeRange(), useBlockAlignedComputeFn,
-                          dispatchBlockAligned, dispatchDefault);
+      b.create<scf::IfOp>(useBlockAlignedComputeFn, dispatchBlockAligned,
+                          dispatchDefault);
       b.create<scf::YieldOp>();
     } else {
       dispatchDefault(b, loc);
     }
   };
 
   // Replace the `scf.parallel` operation with the parallel compute function.
-  b.create<scf::IfOp>(TypeRange(), isZeroIterations, noOp, dispatch);
+  b.create<scf::IfOp>(isZeroIterations, noOp, dispatch);
 
   // Parallel operation was replaced with a block iteration loop.
   rewriter.eraseOp(op);
 
   return success();
 }
 
 void AsyncParallelForPass::runOnOperation() {

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

IfOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
   if (!yieldOp)
     return failure();
   TypeRange types = yieldOp.getOperandTypes();
   inferredReturnTypes.insert(inferredReturnTypes.end(), types.begin(),
                              types.end());
   return success();
 }
 
+void IfOp::build(OpBuilder &builder, OperationState &result,
+                 TypeRange resultTypes, Value cond) {
+  result.addTypes(resultTypes);
+  result.addOperands(cond);
+
+  // Build regions.
+  OpBuilder::InsertionGuard guard(builder);
+  result.addRegion();
+  result.addRegion();
+}
+
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {
-  build(builder, result, /*resultTypes=*/std::nullopt, cond, withElseRegion);
+  build(builder, result, TypeRange{}, 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);
+  result.addOperands(cond);
 
   // Build then region.
   OpBuilder::InsertionGuard guard(builder);
   Region *thenRegion = result.addRegion();
   builder.createBlock(thenRegion);
-  thenBuilder(builder, result.location);
+  if (resultTypes.empty())
+    IfOp::ensureTerminator(*thenRegion, builder, result.location);
 
   // Build else region.
   Region *elseRegion = result.addRegion();
-  if (!elseBuilder)
-    return;
+  if (withElseRegion) {
     builder.createBlock(elseRegion);
-  elseBuilder(builder, result.location);
+    if (resultTypes.empty())
+      IfOp::ensureTerminator(*elseRegion, builder, result.location);
+  }
 }
 
 void IfOp::build(OpBuilder &builder, OperationState &result, 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);
 
if (usedResults.size() == op.getNumResults())
       return failure();
 
     // Compute the result types of the replacement operation.
     SmallVector<Type, 4> newTypes;
     llvm::transform(usedResults, std::back_inserter(newTypes),
                     [](OpResult result) { return result.getType(); });
 
     // Create a replacement operation with empty then and else regions.
-    auto emptyBuilder = [](OpBuilder &, Location) {};
-    auto newOp = rewriter.create<IfOp>(op.getLoc(), newTypes, op.getCondition(),
-                                       emptyBuilder, emptyBuilder);
+    auto newOp =
+        rewriter.create<IfOp>(op.getLoc(), newTypes, op.getCondition());
+    rewriter.createBlock(&newOp.getThenRegion());
+    rewriter.createBlock(&newOp.getElseRegion());
 
     // Move the bodies and replace the terminators (note there is a then and
     // an else region since the operation returns results).
     transferBody(op.getBody(0), newOp.getBody(0), usedResults, rewriter);
     transferBody(op.getBody(1), newOp.getBody(1), usedResults, rewriter);
 
     // Replace the operation by the new one.
     SmallVector<Value, 4> repResults(op.getNumResults());
for (const auto &it :
           &op.getElseRegion() == falseVal.getParentRegion())
         nonHoistable.push_back(trueVal.getType());
     }
     // Early exit if there aren't any yielded values we can
     // hoist outside the if.
     if (nonHoistable.size() == op->getNumResults())
       return failure();
 
-    IfOp replacement = rewriter.create<IfOp>(op.getLoc(), nonHoistable, cond);
+    IfOp replacement = rewriter.create<IfOp>(op.getLoc(), nonHoistable, cond,
+                                             /*withElseRegion=*/false);
     if (replacement.thenBlock())
       rewriter.eraseBlock(replacement.thenBlock());
     replacement.getThenRegion().takeBody(op.getThenRegion());
     replacement.getElseRegion().takeBody(op.getElseRegion());
 
     SmallVector<Value> results(op->getNumResults());
     assert(thenYieldArgs.size() == results.size());
     assert(elseYieldArgs.size() == results.size());
for (const auto &tup : llvm::enumerate(thenYield)) {
       }
       elseYieldsToUpgradeToSelect.push_back(tup.index());
     }
 
     Location loc = op.getLoc();
     Value newCondition = rewriter.create<arith::AndIOp>(
         loc, op.getCondition(), nestedIf.getCondition());
     auto newIf = rewriter.create<IfOp>(loc, op.getResultTypes(), newCondition);
+    Block *newIfBlock = rewriter.createBlock(&newIf.getThenRegion());
 
     SmallVector<Value> results;
     llvm::append_range(results, newIf.getResults());
     rewriter.setInsertionPoint(newIf);
 
     for (auto idx : elseYieldsToUpgradeToSelect)
       results[idx] = rewriter.create<arith::SelectOp>(
           op.getLoc(), op.getCondition(), thenYield[idx], elseYield[idx]);
 
-    Block *newIfBlock = newIf.thenBlock();
-    if (newIfBlock)
-      rewriter.eraseOp(newIfBlock->getTerminator());
-    else
-      newIfBlock = rewriter.createBlock(&newIf.getThenRegion());
     rewriter.mergeBlocks(nestedIf.thenBlock(), newIfBlock);
     rewriter.setInsertionPointToEnd(newIf.thenBlock());
     rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenYield(), thenYield);
     if (!elseYield.empty()) {
       rewriter.createBlock(&newIf.getElseRegion());
       rewriter.setInsertionPointToEnd(newIf.elseBlock());
       rewriter.create<YieldOp>(loc, elseYield);
     }

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Operation *tensor::bubbleUpPadSlice(OpBuilder &b, tensor::PadOp padOp,
   // the original data source x is not used.
   if (hasZeroLen)
     return createGenerateOp();
 
   // If there are dynamic dimensions: Generate an scf.if check to avoid
   // creating SliceOps with result dimensions of size 0 at runtime.
   if (generateZeroSliceGuard && dynHasZeroLenCond) {
     auto result = b.create<scf::IfOp>(
-        loc, resultType, dynHasZeroLenCond,
+        loc, dynHasZeroLenCond,
         /*thenBuilder=*/
         [&](OpBuilder &b, Location loc) {
           b.create<scf::YieldOp>(loc, createGenerateOp()->getResult(0));
         },
         /*elseBuilder=*/
         [&](OpBuilder &b, Location loc) {
           b.create<scf::YieldOp>(loc, createPadOfExtractSlice()->getResult(0));
         });

mlir/lib/Dialect/Tensor/Transforms/SplitPaddingPatterns.cpp

LogicalResult matchAndRewrite(tensor::PadOp padOp,
     // padding. For the "else" branch, we retain the clone op.
     auto thenBuilder = [&padOp](OpBuilder &builder, Location loc) {
       builder.create<scf::YieldOp>(loc, padOp.getSource());
     };
     auto elseBuilder = [&padOp](OpBuilder &builder, Location loc) {
       Operation *newOp = builder.clone(*padOp);
       builder.create<scf::YieldOp>(loc, newOp->getResults());
     };
-    rewriter.replaceOpWithNewOp<scf::IfOp>(padOp, padOp.getType(), ifCond,
-                                           thenBuilder, elseBuilder);
+    rewriter.replaceOpWithNewOp<scf::IfOp>(padOp, ifCond, thenBuilder,
+                                           elseBuilder);
     return success();
   }
 };
 
 } // namespace
 
 void tensor::populateSplitPaddingPatterns(RewritePatternSet &patterns,
                                           PatternBenefit baseBenefit) {
   patterns.add<SplitPadding>(patterns.getContext(), baseBenefit);
 }

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

Value pos =
             : rewriter
                   .create<AffineApplyOp>(loc, sym0 % elementsPerLane, newPos)
                   .getResult();
     Value isInsertingLane = rewriter.create<arith::CmpIOp>(
         loc, arith::CmpIPredicate::eq, newWarpOp.getLaneid(), insertingLane);
     Value newResult =
         rewriter
             .create<scf::IfOp>(
-                loc, distrType, isInsertingLane,
+                loc, isInsertingLane,
                 /*thenBuilder=*/
                 [&](OpBuilder &builder, Location loc) {
                   Value newInsert = builder.create<vector::InsertElementOp>(
                       loc, newSource, distributedVec, pos);
                   builder.create<scf::YieldOp>(loc, newInsert);
                 },
                 /*elseBuilder=*/
                 [&](OpBuilder &builder, Location loc) {
if (distrSrcDim >= 0) {
         Value newInsert = builder.create<vector::InsertOp>(
             loc, distributedSrc, distributedDest, newPos);
         builder.create<scf::YieldOp>(loc, newInsert);
       };
       auto nonInsertingBuilder = [&](OpBuilder &builder, Location loc) {
         builder.create<scf::YieldOp>(loc, distributedDest);
       };
       newResult = rewriter
-                      .create<scf::IfOp>(loc, distrDestType, isInsertingLane,
+                      .create<scf::IfOp>(loc, isInsertingLane,
                                          /*thenBuilder=*/insertingBuilder,
                                          /*elseBuilder=*/nonInsertingBuilder)
                       .getResult(0);
     }
 
     newWarpOp->getResult(operandNumber).replaceAllUsesWith(newResult);
     return success();
   }

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

 static scf::IfOp
 createFullPartialLinalgCopy(RewriterBase &b, vector::TransferReadOp xferOp,
                             TypeRange returnTypes, Value inBoundsCond,
                             MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b.create<scf::IfOp>(
-      loc, returnTypes, inBoundsCond,
+      loc, inBoundsCond,
       [&](OpBuilder &b, Location loc) {
         Value res = memref;
         if (compatibleMemRefType != xferOp.getShapedType())
           res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
         scf::ValueVector viewAndIndices{res};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getIndices().begin(),
                               xferOp.getIndices().end());
         b.create<scf::YieldOp>(loc, viewAndIndices);

 static scf::IfOp createFullPartialVectorTransferRead(
     RewriterBase &b, vector::TransferReadOp xferOp, TypeRange returnTypes,
     Value inBoundsCond, MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   scf::IfOp fullPartialIfOp;
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b.create<scf::IfOp>(
-      loc, returnTypes, inBoundsCond,
+      loc, inBoundsCond,
       [&](OpBuilder &b, Location loc) {
         Value res = memref;
         if (compatibleMemRefType != xferOp.getShapedType())
           res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
         scf::ValueVector viewAndIndices{res};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getIndices().begin(),
                               xferOp.getIndices().end());
         b.create<scf::YieldOp>(loc, viewAndIndices);

 getLocationToWriteFullVec(RewriterBase &b, vector::TransferWriteOp xferOp,
                           TypeRange returnTypes, Value inBoundsCond,
                           MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b
       .create<scf::IfOp>(
-          loc, returnTypes, inBoundsCond,
+          loc, inBoundsCond,
           [&](OpBuilder &b, Location loc) {
             Value res = memref;
             if (compatibleMemRefType != xferOp.getShapedType())
               res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
             scf::ValueVector viewAndIndices{res};
             viewAndIndices.insert(viewAndIndices.end(),
                                   xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());