[MLIR] Add return type inference to scf.if builder

[frgossen]

We cannot generally infer the return type as there are use cases relying on the old behaviour.
In particular, cases which expect an empty yield op to be generated just to replace it afterwards.

[jpienaar]

Could you perhaps expand on that exception or point to where this is assumed? Are these where during a pass an invalid state is created or is this across passes?

[frgossen]

In D141928#4059151, @jpienaar wrote:

Could you perhaps expand on that exception or point to where this is assumed? Are these where during a pass an invalid state is created or is this across passes?

It's around 100 test failures internally. The cases I looked at were pass-internal.

[jpienaar]

SG, could you add a TODO here documenting this so that we know when we look at it again?

[frgossen]

Add TODO

mlir/include/mlir/Dialect/SCF/IR/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<(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,

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

build(builder, result, resultTypes, cond, 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);
 
+  // Build then region.
   OpBuilder::InsertionGuard guard(builder);
   Region *thenRegion = result.addRegion();
   builder.createBlock(thenRegion);
   thenBuilder(builder, result.location);
 
+  // Build else region.
   Region *elseRegion = result.addRegion();
   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);
+  assert(thenBuilder && "the builder callback for 'then' must be present");
+  result.addOperands(cond);
+
+  // Build then region.
+  OpBuilder::InsertionGuard guard(builder);
+  Region *thenRegion = result.addRegion();
+  Block *thenBlock = builder.createBlock(thenRegion);
+  thenBuilder(builder, result.location);
+
+  // Infer types if there are any.
+  if (auto yieldOp = llvm::dyn_cast<YieldOp>(thenBlock->getTerminator()))
+    result.addTypes(yieldOp.getOperandTypes());
+
+  // Build else region.
+  Region *elseRegion = result.addRegion();
+  if (!elseBuilder)
+    return;
+  builder.createBlock(elseRegion);
+  elseBuilder(builder, result.location);
 }
 
 LogicalResult IfOp::verify() {
   if (getNumResults() != 0 && getElseRegion().empty())
     return emitOpError("must have an else block if defining values");
   return success();
 }