[MLIR] Add InferTypeOpInterface to scf.if op

[pifon2a]

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

1483

can we use getTerminator here?

[frgossen]

Address comments

[pifon2a]

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

1480

this should be either a llvm::cast, otherwise you need to check whether the dynamic cast actually happened.

[jpienaar]

Ah seems I never hit submit yesterday. Same comment as Alex on one and general question on verification part.

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

1476

Is the op invalid in that case? So here we are duplicating a bit of verification? Remember for builder methods invalid input is UB while for verification there is an order in which these run and so we should have checked that there is at least 1 region and that it satisfies SizedRegion<1>, given that do we need to verify these?

1484

If you use result of dyn_cast without checking if null, then cast is preferred (you could move the the trait that says it should have yield before type inference one to have this verified before you get here).

[frgossen]

Thanks!

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

1476

For the builders, I agree. inferReturnTypes is used when the regions and blocks are fully set up.
For the verifiers, do we want to rely on the order in which they are called? That seems fragile to me and not worth avoiding the cheap checks here.
Addressed this in https://reviews.llvm.org/D142155

1480

Addressed in https://reviews.llvm.org/D142155

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

 #ifndef MLIR_DIALECT_SCF_SCF_H
 #define MLIR_DIALECT_SCF_SCF_H
 
 #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/ParallelCombiningOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 namespace scf {
 void buildTerminatedBody(OpBuilder &builder, Location loc);

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

 
 #ifndef MLIR_DIALECT_SCF_SCFOPS
 #define MLIR_DIALECT_SCF_SCFOPS
 
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 def SCF_Dialect : Dialect {
   let name = "scf";
   let cppNamespace = "::mlir::scf";
   let dependentDialects = ["arith::ArithDialect"];
let extraClassDeclaration = [{
     ::mlir::OpResult getParentResult(int64_t idx);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // IfOp
 //===----------------------------------------------------------------------===//
 
-def IfOp : SCF_Op<"if",
-      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
-                                 ["getNumRegionInvocations",
-                                  "getRegionInvocationBounds"]>,
-       SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveMemoryEffects,
-       NoRegionArguments]> {
+def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
+    "getNumRegionInvocations", "getRegionInvocationBounds"]>,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
   let summary = "if-then-else operation";
   let description = [{
     The `scf.if` operation represents an if-then-else construct for
     conditionally executing two regions of code. The operand to an if operation
     is a boolean value. For example:
 
     ```mlir
     scf.if %b  {

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

while (ifOp) {
     // Check next enclosing IfOp.
     ifOp = ifOp->getParentOfType<IfOp>();
   }
 
   // Could not find a common IfOp among a's and b's ancestors.
   return false;
 }
 
+LogicalResult
+IfOp::inferReturnTypes(MLIRContext *ctx, std::optional<Location> loc,
+                       ValueRange operands, DictionaryAttr attrs,
+                       RegionRange regions,
+                       SmallVectorImpl<Type> &inferredReturnTypes) {
+  if (regions.empty())
+    return failure();
+  Region *r = regions.front();
+  assert(!r->empty());
+  Block &b = r->front();
+  auto yieldOp = llvm::dyn_cast<YieldOp>(b.getTerminator());
+  TypeRange types = yieldOp.getOperandTypes();
+  inferredReturnTypes.insert(inferredReturnTypes.end(), types.begin(),
+                             types.end());
+  return success();
+}
+
 void IfOp::build(OpBuilder &builder, OperationState &result, Value cond,
                  bool withElseRegion) {
   build(builder, result, /*resultTypes=*/std::nullopt, cond, withElseRegion);
 }
 
 void IfOp::build(OpBuilder &builder, OperationState &result,
                  TypeRange resultTypes, Value cond, bool withElseRegion) {
   auto addTerminator = [&](OpBuilder &nested, Location loc) {
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);
 
   // Build then region.
   OpBuilder::InsertionGuard guard(builder);
   Region *thenRegion = result.addRegion();
-  Block *thenBlock = builder.createBlock(thenRegion);
+  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;
+  if (elseBuilder) {
     builder.createBlock(elseRegion);
     elseBuilder(builder, result.location);
   }
 
+  // Infer result types.
+  SmallVector<Type> inferredReturnTypes;
+  MLIRContext *ctx = builder.getContext();
+  auto attrDict = DictionaryAttr::get(ctx, result.attributes);
+  if (succeeded(inferReturnTypes(ctx, std::nullopt, result.operands, attrDict,
+                                 result.regions, inferredReturnTypes))) {
+    result.addTypes(inferredReturnTypes);
+  }
+}
+
 LogicalResult IfOp::verify() {
   if (getNumResults() != 0 && getElseRegion().empty())
     return emitOpError("must have an else block if defining values");
   return success();
 }
 
 ParseResult IfOp::parse(OpAsmParser &parser, OperationState &result) {
   // Create the regions for 'then'.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

td_library(
     name = "SCFTdFiles",
     srcs = [
         "include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td",
         "include/mlir/Dialect/SCF/IR/SCFOps.td",
     ],
     includes = ["include"],
     deps = [
         ":ControlFlowInterfacesTdFiles",
+        ":InferTypeOpInterfaceTdFiles",
         ":LoopLikeInterfaceTdFiles",
         ":ParallelCombiningOpInterfaceTdFiles",
         ":SideEffectInterfacesTdFiles",
         ":ViewLikeInterfaceTdFiles",
     ],
 )
 
 gentbl_cc_library(
cc_library(
     deps = [
         ":ArithDialect",
         ":ArithUtils",
         ":BufferizationDialect",
         ":ControlFlowDialect",
         ":ControlFlowInterfaces",
         ":FuncDialect",
         ":IR",
+        ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":ParallelCombiningOpInterface",
         ":Pass",
         ":SCFDeviceMappingInterfacesIncGen",
         ":SCFIncGen",
         ":SCFPassIncGen",
         ":Support",