[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 ↗

(On Diff #392110)

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

[Mogball]

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

140 ↗

(On Diff #392110)

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

[rriddle]

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

138–140 ↗

(On Diff #392110)

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

5

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 ↗

(On Diff #392110)

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

1196

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

1196

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

1194

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

1194

Ah, right. I forgot about that

[Mogball]

Add back the constant cond

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

let extraClassDeclaration = [{
     /// value for `index`.
     OperandRange getSuccessorEntryOperands(unsigned index);
   }];
 
   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:
 

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

let methods = [
     InterfaceMethod<[{
         Populates `invocationBounds` with the minimum and maximum number of
         times this operation will invoke the attached regions (assuming the
         regions yield normally, i.e. do not abort or invoke an infinite loop).
         The minimum number of invocations is at least 0. If the maximum number
         of invocations cannot be statically determined, then it will not have a
         value (i.e., it is set to `llvm::None`).
 
-        `operands` is a set of optional attributes that either correspond to a
-        constant values for each operand of this operation, or null if that
+        `operands` is a set of optional attributes that either correspond to
+        constant values for each operand of this operation or null if that
         operand is not a constant.
+
+        This method may be called speculatively on operations where the provided
+        operands are not necessarily the same as the operation's current
+        operands. This may occur in analyses that wish to determine "what would
+        be the region invocations if these were the operands?"
       }],
       "void", "getRegionInvocationBounds",
       (ins "::mlir::ArrayRef<::mlir::Attribute>":$operands,
            "::llvm::SmallVectorImpl<::mlir::InvocationBounds> &"
              :$invocationBounds), [{}],
        [{ invocationBounds.append($_op->getNumRegions(),
                                   ::mlir::InvocationBounds::getUnknown()); }]
     >,

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult IfOp::fold(ArrayRef<Attribute> operands,
   // callbacks See: https://llvm.org/doxygen/ilist_8h_source.html#l00224
   getThenRegion().getBlocks().splice(getThenRegion().getBlocks().begin(),
                                      getElseRegion().getBlocks());
   getElseRegion().getBlocks().splice(getElseRegion().getBlocks().begin(),
                                      getThenRegion().getBlocks(), thenBlock);
   return success();
 }
 
+void IfOp::getRegionInvocationBounds(
+    ArrayRef<Attribute> operands,
+    SmallVectorImpl<InvocationBounds> &invocationBounds) {
+  if (auto cond = operands[0].dyn_cast_or_null<BoolAttr>()) {
+    // If the condition is known, then one region is known to be executed once
+    // and the other zero times.
+    invocationBounds.emplace_back(0, cond.getValue() ? 1 : 0);
+    invocationBounds.emplace_back(0, cond.getValue() ? 0 : 1);
+  } else {
+    // Non-constant condition. Each region may be executed 0 or 1 times.
+    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/Dialect/SCF/control-flow-sink.mlir

+// RUN: mlir-opt -split-input-file -control-flow-sink %s | FileCheck %s
+
+// CHECK-LABEL: @test_scf_if_sink
+// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
+// CHECK: %[[V0:.*]] = scf.if %[[ARG0]]
+// CHECK:   %[[V1:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
+// CHECK:   scf.yield %[[V1]]
+// CHECK: else
+// CHECK:   %[[V1:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
+// CHECK:   scf.yield %[[V1]]
+// CHECK: return %[[V0]]
+func @test_scf_if_sink(%arg0: i1, %arg1: i32) -> i32 {
+  %0 = arith.addi %arg1, %arg1 : i32
+  %1 = arith.muli %arg1, %arg1 : i32
+  %result = scf.if %arg0 -> 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-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
+// CHECK: scf.if %[[ARG0]]
+// CHECK:   %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
+// CHECK:   call @consume(%[[V0]])
+func @test_scf_if_then_only_sink(%arg0: i1, %arg1: i32) {
+  %0 = arith.addi %arg1, %arg1 : i32
+  scf.if %arg0 {
+    call @consume(%0) : (i32) -> ()
+    scf.yield
+  }
+  return
+}
+
+// -----
+
+func private @consume(i32) -> ()
+
+// CHECK-LABEL: @test_scf_if_double_sink
+// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
+// CHECK: scf.if %[[ARG0]]
+// CHECK:   scf.if %[[ARG0]]
+// CHECK:     %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
+// CHECK:     call @consume(%[[V0]])
+func @test_scf_if_double_sink(%arg0: i1, %arg1: i32) {
+  %0 = arith.addi %arg1, %arg1 : i32
+  scf.if %arg0 {
+    scf.if %arg0 {
+      call @consume(%0) : (i32) -> ()
+      scf.yield
+    }
+  }
+  return
+}