[mlir][scf] Update IfOp to have getInvocationBounds

[Mogball]

Summary

This allows scf.if to be used by Control-Flow sink.

Depends on D115088

[Mogball]

Fix the build.

[jpienaar]

mlir/test/Transforms/control-flow-sink.mlir

140

Could we make this change independent and independently tested from sinking?

[Mogball]

mlir/test/Transforms/control-flow-sink.mlir

140

What do you mean? This diff depends on the sinking patch.

[rriddle]

mlir/test/Transforms/control-flow-sink.mlir

138–140

This doesn't look like the right place for this. This file shouldn't contain test code for all possible dialects, the SCF/ folder should have a file testing control-flow-sink for its operations.

[Mogball]

Move tests to test/Dialect/SCF

[rriddle]

mlir/test/Dialect/SCF/control-flow-sink.mlir

4 ↗

(On Diff #392439)

Can you fix up all of the raw SSA names? We shouldn't be using the names directly from the output, we should use FileCheck captures instead.

[jpienaar]

mlir/test/Transforms/control-flow-sink.mlir

140

I meant in general, adding bound queries is not depend on sinking. It could be motivated, added and tested independently of sinking.

[jpienaar]

mlir/lib/Dialect/SCF/SCF.cpp

1206

Would BoolAttr work instead? (Given https://mlir.llvm.org/docs/Dialects/SCFDialect/#scfif-mlirscfifop it should be single bit integer here so would expect bool to work)

I see this is following the above, and almost exactly the same, could this be refactored into small helper?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1206

Do we even need this condition here?
If the condition is constant the if should be folded already.

[Mogball]

review comments

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

1204

As mentioned in another comment I made somewhere, this is assuming we never call this method speculatively. Not that this is bad, but it is a specific invariant of the API that should be documented (on the interface doc itself).

[Mogball]

mlir/lib/Dialect/SCF/SCF.cpp

1204

Ah, right. I forgot about that

[Mogball]

Add back the constant cond

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

let extraClassDeclaration = [{
     void getNumRegionInvocations(ArrayRef<Attribute> operands,
                                  SmallVectorImpl<int64_t> &countPerRegion);
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def IfOp : SCF_Op<"if",
-      [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+      [DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                                 ["getNumRegionInvocations",
+                                  "getRegionInvocationBounds"]>,
        SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveSideEffects,
        NoRegionArguments]> {
   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:
 
OpBuilder getElseBodyBuilder(OpBuilder::Listener *listener = nullptr) {
       Block* body = getBody(1);
       return results().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                : OpBuilder::atBlockEnd(body, listener);
     }
     Block* thenBlock();
     YieldOp thenYield();
     Block* elseBlock();
     YieldOp elseYield();
-
-    /// If the condition is a constant, returns 1 for the executed block and 0
-    /// for the other. Otherwise, returns `kUnknownNumRegionInvocations` for
-    /// both successors.
-    void getNumRegionInvocations(ArrayRef<Attribute> operands,
-                                 SmallVectorImpl<int64_t> &countPerRegion);
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def ParallelOp : SCF_Op<"parallel",
     [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,

mlir/lib/Dialect/SCF/SCF.cpp

if (auto condAttr = operands.front().dyn_cast_or_null<IntegerAttr>()) {
     countPerRegion.assign(1, cond ? 1 : 0);
     countPerRegion.push_back(cond ? 0 : 1);
   } else {
     // Non-constant condition: unknown invocations for both successors.
     countPerRegion.assign(2, kUnknownNumRegionInvocations);
   }
 }
 
+/// Each region of the op is invoked at most once. If the condition is constant,
+/// then the number of invocations can be narrowed to either 0 or 1.
+void IfOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands,
+    SmallVectorImpl<InvocationBounds> &invocationBounds) {
+  if (auto condAttr = operands.front().dyn_cast_or_null<IntegerAttr>()) {
+    // If the condition is true, `then` is executed once and `else` zero times,
+    // and vice-versa.
+    bool cond = condAttr.getValue().isOneValue();
+    invocationBounds.assign(1, {0, cond ? 1 : 0});
+    invocationBounds.emplace_back(0, cond ? 0 : 1);
+  } else {
+    // Non-constant condition: unknown invocations for both successors.
+    invocationBounds.assign(2, {0, 1});
+  }
+}
+
 namespace {
 // Pattern to remove unused IfOp results.
 struct RemoveUnusedResults : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   void transferBody(Block *source, Block *dest, ArrayRef<OpResult> usedResults,
                     PatternRewriter &rewriter) const {
     // Move all operations to the destination block.

mlir/test/Transforms/control-flow-sink.mlir

-// RUN: mlir-opt -control-flow-sink %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -control-flow-sink %s | FileCheck %s
 
 // CHECK-LABEL: @test_simple_sink
 // CHECK: %0 = arith.subi %arg2, %arg1 : i32
 // CHECK: %1 = test.region_if %arg0: i1 -> i32 then {
 // CHECK:   %2 = arith.subi %arg1, %arg2 : i32
 // CHECK:   test.region_if_yield %2 : i32
 // CHECK: } else {
 // CHECK:   %2 = arith.addi %arg1, %arg1 : i32
%result = test.region_if %cond : i1 -> i32 then {
     test.region_if_yield %v : i32
   } join {
     %v = arith.addi %3, %1 : i32
     test.region_if_yield %v : i32
   }
   return %result : i32
 }
 
+// -----
+
 // CHECK-LABEL: @test_region_sink
 // CHECK: %0 = test.region_if %arg0: i1 -> i32 then {
 // CHECK:   %1 = test.region_if %arg0: i1 -> i32 then {
 // CHECK:     test.region_if_yield %arg1 : i32
 // CHECK:   } else {
 // CHECK:     %2 = arith.subi %arg1, %arg2 : i32
 // CHECK:     test.region_if_yield %2 : i32
 // CHECK:   } join {
func @test_region_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
   } else {
     test.region_if_yield %arg0 : i32
   } join {
     test.region_if_yield %arg1 : i32
   }
   return %result1 : i32
 }
 
+// -----
+
 // CHECK-LABEL: @test_subgraph_sink
 // CHECK: %0 = test.region_if %arg0: i1 -> i32 then {
 // CHECK:   %1 = arith.subi %arg1, %arg2 : i32
 // CHECK:   %2 = arith.addi %arg1, %arg2 : i32
 // CHECK:   %3 = arith.subi %arg2, %arg1 : i32
 // CHECK:   %4 = arith.muli %3, %3 : i32
 // CHECK:   %5 = arith.muli %2, %1 : i32
 // CHECK:   %6 = arith.addi %5, %4 : i32
func @test_subgraph_sink(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
   } else {
     test.region_if_yield %arg0 : i32
   } join {
     test.region_if_yield %arg1 : i32
   }
   return %result : i32
 }
 
+// -----
+
 // CHECK-LABEL: @test_multiblock_region_sink
 // CHECK: %0 = arith.addi %arg1, %arg2 : i32
 // CHECK: %1 = "test.any_cond"() ( {
 // CHECK:   %3 = arith.addi %0, %arg2 : i32
 // CHECK:   %4 = arith.addi %3, %arg1 : i32
 // CHECK:   br ^bb1(%4 : i32)
 // CHECK: ^bb1(%5: i32):  // pred: ^bb0
 // CHECK:   %6 = arith.addi %5, %4 : i32
^bb0:
     br ^bb1(%2 : i32)
   ^bb1(%arg2 : i32):
     %4 = arith.addi %arg2, %2 : i32
     "test.yield"(%4) : (i32) -> ()
   }) : () -> i32
   %result = arith.addi %0, %3 : i32
   return %result : i32
 }
+
+// -----
+
+// CHECK-LABEL: @test_scf_if_sink
+// CHECK: %0 = scf.if %arg0 -> (i32) {
+// CHECK:   %1 = arith.addi %arg1, %arg1 : i32
+// CHECK:   scf.yield %1 : i32
+// CHECK: } else {
+// CHECK:   %1 = arith.muli %arg1, %arg1 : i32
+// CHECK:   scf.yield %1 : i32
+// CHECK: }
+// CHECK: return %0 : i32
+func @test_scf_if_sink(%cond : i1, %arg0 : i32) -> i32 {
+  %0 = arith.addi %arg0, %arg0 : i32
+  %1 = arith.muli %arg0, %arg0 : i32
+  %result = scf.if %cond -> i32 {
+    scf.yield %0 : i32
+  } else {
+    scf.yield %1 : i32
+  }
+  return %result : i32
+}
+
+// -----
+
+func private @consume(i32) -> ()
+
+// CHECK-LABEL: @test_scf_if_then_only_sink
+// CHECK: scf.if %arg0 {
+// CHECK:   %0 = arith.addi %arg1, %arg1 : i32
+// CHECK:   call @consume(%0) : (i32) -> ()
+// CHECK: }
+// CHECK: return
+func @test_scf_if_then_only_sink(%cond : i1, %arg0 : i32) {
+  %0 = arith.addi %arg0, %arg0 : i32
+  scf.if %cond {
+    call @consume(%0) : (i32) -> ()
+    scf.yield
+  }
+  return
+}