[mlir][scf] Add getNumRegionInvocations to IfOp

[Mogball]

Summary

Implements the RegionBranchOpInterface method getNumRegionInvocations to scf::IfOp so that, when the condition is constant, the number of region executions can be analyzed by NumberOfExecutions.

[rriddle]

Can you add some justification to the patch description?

[Mogball]

Fixing the build and updating description

[jpienaar]

I don't this interfaces well, but this makes sense to me

mlir/unittests/Dialect/SCF/SCFOps.cpp

30

R"mlir( ? (Given we could then, as mentioned in chat, look at injecting some syntax highlighting)

51

EXPECT_EQ ? (Expect is preferred except where it isn't safe to continue if assertion failed)

[Mogball]

review comments

[ftynse]

Shouldn't we just expect that ifs with known conditions are canonicalized away? I.e., is this ever necessary?

[rriddle]

In D115087#3194277, @ftynse wrote:

Shouldn't we just expect that ifs with known conditions are canonicalized away? I.e., is this ever necessary?

It depends on how the interface gets used. If used in a speculative fashion, the provided operands are not necessarily the same as the current operands of the if. Some algorithm/analysis could ask "what would be the region invocations if these were the operands". Not entirely sure how much that applies in this situation, but it does for e.g. the branch interfaces (e.g. during constant propagation).

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

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

void ForOp::getNumRegionInvocations(ArrayRef<Attribute> operands,
   countPerRegion.resize(1);
 
   auto lb = operands[0].dyn_cast_or_null<IntegerAttr>();
   auto ub = operands[1].dyn_cast_or_null<IntegerAttr>();
   auto step = operands[2].dyn_cast_or_null<IntegerAttr>();
 
   // Loop bounds are not known statically.
   if (!lb || !ub || !step || step.getValue().getSExtValue() == 0) {
-    countPerRegion[0] = -1;
+    countPerRegion[0] = kUnknownNumRegionInvocations;
     return;
   }
 
   countPerRegion[0] =
       ceilDiv(ub.getValue().getSExtValue() - lb.getValue().getSExtValue(),
               step.getValue().getSExtValue());
 }
 
if (elseRegion)
       regions.push_back(RegionSuccessor(elseRegion));
     return;
   }
 
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
 }
 
+/// If the condition is a constant, returns 1 for the executed block and 0 for
+/// the other. Otherwise, returns `kUnknownNumRegionInvocations` for both
+/// successors.
+void IfOp::getNumRegionInvocations(ArrayRef<Attribute> operands,
+                                   SmallVectorImpl<int64_t> &countPerRegion) {
+  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();
+    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);
+  }
+}
+
 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/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(Quant)
+add_subdirectory(SCF)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Utils)

mlir/unittests/Dialect/SCF/CMakeLists.txt

+add_mlir_unittest(MLIRSCFTests
+  SCFOps.cpp
+  )
+target_link_libraries(MLIRSCFTests
+  PRIVATE
+  MLIRIR
+  MLIRParser
+  MLIRSCF
+  MLIRStandard
+  )

mlir/unittests/Dialect/SCF/SCFOps.cpp

+//===- SCFOps.cpp - SCF Op Unit Tests -------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Parser.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+class SCFOpsTest : public testing::Test {
+public:
+  SCFOpsTest() {
+    context.getOrLoadDialect<scf::SCFDialect>();
+    context.getOrLoadDialect<StandardOpsDialect>();
+  }
+
+protected:
+  MLIRContext context;
+};
+
+TEST_F(SCFOpsTest, IfOpNumRegionInvocations) {
+  const char *const code = R"mlir(
+func @test(%cond : i1) -> () {
+  scf.if %cond {
+    scf.yield
+  } else {
+    scf.yield
+  }
+  return
+}
+)mlir";
+  Builder builder(&context);
+
+  auto module = parseSourceString(code, &context);
+  ASSERT_TRUE(module);
+  scf::IfOp op;
+  module->walk([&](scf::IfOp ifOp) { op = ifOp; });
+  ASSERT_TRUE(op);
+
+  SmallVector<int64_t> countPerRegion;
+  op.getNumRegionInvocations({Attribute()}, countPerRegion);
+  EXPECT_EQ(countPerRegion.size(), 2u);
+  EXPECT_EQ(countPerRegion[0], kUnknownNumRegionInvocations);
+  EXPECT_EQ(countPerRegion[1], kUnknownNumRegionInvocations);
+
+  countPerRegion.clear();
+  op.getNumRegionInvocations(
+      {builder.getIntegerAttr(builder.getI1Type(), true)}, countPerRegion);
+  EXPECT_EQ(countPerRegion.size(), 2u);
+  EXPECT_EQ(countPerRegion[0], 1);
+  EXPECT_EQ(countPerRegion[1], 0);
+
+  countPerRegion.clear();
+  op.getNumRegionInvocations(
+      {builder.getIntegerAttr(builder.getI1Type(), false)}, countPerRegion);
+  EXPECT_EQ(countPerRegion.size(), 2u);
+  EXPECT_EQ(countPerRegion[0], 0);
+  EXPECT_EQ(countPerRegion[1], 1);
+}
+} // end anonymous namespace