[mlir] Add NumberOfExecutions analysis + update RegionBranchOpInterface interface to query number of region invocations

[ezhulenev]

Summary

Implements RFC discussed in: https://llvm.discourse.group/t/rfc-operationinstancesinterface-or-any-better-name/2158/10

[ezhulenev]

Follow block parents to compute number of block executions + compute number of executions for ForOp with constant bounds

[ftynse]

Very nice!

I have a bunch of nits and several generalization suggestions.

mlir/include/mlir/Analysis/NumberOfExecutions.h

38

Nit: explicit?

69

There doesn't seem to be an accessor for this field, so IMO it can be dropped (unless there is an anticipated need for it). At which point, we can also consider using BlockNumberOfExecutionsInfo = Optional<int64_t>, which can be nicely constructed from an int64_t or from llvm::None.

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

86–88

It should be possible to do DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getNumRegionInvocations"]> in the def and avoid this explicit declaration.

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

135

Nit: I'd be explicit and say something like "Populates countPerRegion with the number..."

137

Nit: could we have a named constant instead of -1?

149

Nit: countPerRegion.resize(numRegions, -1) avoids the loop

mlir/lib/Analysis/NumberOfExecutions.cpp

45–46

I wonder if this can be generalized to work on all constant-like operations. Currently, we only have a ConstantLike trait, so we cannot get the constant value as attribute without knowing the concrete op type. Maybe we can turn ConstantLike into an interface with value being an interface method. Out of scope for this commit probably.

53

Could we rather query the op? Hardcoding FuncOp makes this analysis non-applicable to custom function types (we have some, e.g., in LLVM dialect and in Flang).

58

Let's also be future-proof and guard against graph regions (https://mlir.llvm.org/docs/LangRef/#graph-regions). They don't have blocks (yet), but if they do, the logic of following block successors below will not necessarily apply. An assertion could do.

66

Nit: we usually don't declare function stack variables as const in MLIR

71

Nit: I would have just used auto and kept the original lambda type here, it's not like we are passing it around.

mlir/lib/Dialect/SCF/SCF.cpp

303–306

Could we rather extend the one we already have in MathExtras.h to support negative RHS?

[ezhulenev]

Address PR comments

mlir/include/mlir/Analysis/NumberOfExecutions.h

69

I added a custom non-default constructible type to catch accidental default constructions of None instead of a failure. And maybe more advanced analysis can add some other fields in the future.

Added this to class documentation for the next reader and removed unused Block*.

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

86–88

Cool, learning something new every day :)

mlir/lib/Analysis/NumberOfExecutions.cpp

45–46

Yeah, that would be useful also in the context of TF graphs with tf.Const ops, but I did not find any relevant examples and chose the easiest path.

53

It guess in can just implement RegionBranchOpInterface, I decided to special case it because it's "built in" and not defined in any of the dialect .td files. I expect all other FuncOpLike things to implement RegionBranchOpInterface.

71

Unfortunately lambdas do not support recursion without ugly tricks.

[ftynse]

LGTM. Please wait for one more reviewer to double-check the logic according to the forum discussion.

mlir/include/mlir/Analysis/NumberOfExecutions.h

69

Ack. (Phabricator doesn't let anybody other than the author to mark comments as done).

mlir/lib/Analysis/NumberOfExecutions.cpp

53

ODS is not mandatory to implement an interface, ops defined in C++ can derive Traits<InterfaceName::Trait> and get the same functionality.

I'm fine with your current implementation given the explanation provided. We should look into having FuncOp implement the interface separately and consult with @rriddle about future plans regarding it being built-in.

[ezhulenev]

mlir/lib/Analysis/NumberOfExecutions.cpp

53

Yeah, I know, I just don't know is it "ok" to add it to built in type :) Assuming FuncOp is as fundamental as Region I think it's ok to have a special case, but I'd like to hear opinion from someone with a better judgement.

[mehdi_amini]

mlir/lib/Analysis/NumberOfExecutions.cpp

53

FuncOp isn't as fundamental as Region, I rather treat FuncOp as any other op.

For example LLVM, SPIRV, GPU dialects aren't using FuncOp.

[ezhulenev]

mlir/lib/Analysis/NumberOfExecutions.cpp

53

I've looked into adding RegionBranchOpInterface to FuncOp`, but it feels weird because the FuncOp does not execute the region, only the CallOp does.

For example: getSuccessorEntryOperands: "Returns the operands of this operation used as the entry arguments", but the FuncOp is OpTrait::ZeroOperands.

Maybe add a special trait, something like OpTrait::InvokesRegionsOnce, an additional piece of information just like OpTrait::OneRegion.

Or just rely on OpTrait::FunctionLike here, assuming that all functions execute attached region once.

[ezhulenev]

Use FunctionLike trait to to get number of region invocations

[rriddle]

mlir/lib/Analysis/NumberOfExecutions.cpp

14

llvm includes go after MLIR includes.

19

Is the dependency on standard ops necessary?

23

Please trim some of these includes, several are unnecessary.

36

Missing static here?

66

This is not correct and seems like it would result in incorrect usages of the analysis. If you want to estimate function executions, you'll need to use the callgraph.

75

Drop trivial braces.

122

Nit: Flip the condition of this if, and remove the trivial braces.

if (!...)
  blockInfo.insert

mlir/lib/Dialect/Async/IR/Async.cpp

125

nit: Seems like push_back(1) is better than an explicit resize and assignment.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

78

Prefer explicit qualification, i.e. mlir::kUnknownNumRegionInvocations

[rriddle]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

Why should functions be default to executed once? That doesn't seem like a conservatively safe default.

[silvas]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

I agree, there seems to be a code smell here.

One solution is for queries to be of the form Optional<int64_t> getNumberOfExecutions(Block *block, Region *assumingThisRegionIsEnteredOnce).

So the analysis could hold a DenseMap<Block *, NumberOfExecutionsAssumingDirectParentRegionEnteredOnce>. Then on a query, you would walk up parents, multiplying NumberOfExecutionsAssumingDirectParentRegionEnteredOnce at each step until we reach assumingThisRegionIsEnteredOnce (and until we implement a call graph analysis, we will return None if assumingThisRegionIsEnteredOnce is not an ancestor of block)

And the body of a module should probably be an illegal Region to pass as assumingThisRegionIsEnteredOnce, because it is a "declarative region" and never "executes" (we might need a new IR property to represent this?)

test-print-number-of-executions would have some hardcoded policy to print values assuming assumingThisRegionIsEnteredOnce == enclosingFunc.getBody()

[ezhulenev]

Address PR comments

mlir/lib/Analysis/NumberOfExecutions.cpp

19

It's needed for dyn_cast<ConstantOp>

66

This is roughly ~= "when a function called it executes its entry block once". And this seems like a reasonable assumption. This is not about how many call operations are in the graph.

[silvas]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

what about a function like

def f(i):
  print("here") # How many executions will this be assigned?
  if i > 0:
    f(i - 1)

I think we need to be more precise about what we mean by "execution" here.

Also, does this pass correctly handle a case like?

def g(i):
  if i == 0:
    abort()
  print("here") # How many executions?

[rriddle]

At the moment it is somewhat unclear at what scope this analysis is intended to be constructed and used. It seems like this currently doesn't handle things related to the callgraph, doesn't handle control flow well, and assumes that it only operates locally within a single function? I'm slightly concerned that we are building an analysis that is very tied with the constraints on async tokens as determined by D90716, without much reusability outside of those constraints. Along with Sean's point, I don't think the concept of "Execution" is very well scoped here and leads to some weird code smells with functions/modules/etc.

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

201

Replace -1 here with kUnknownNumRegionInvocations.

mlir/lib/Analysis/NumberOfExecutions.cpp

54

Use m_Constant here instead, we should avoid hardcoding operation types whenever possible.

[ezhulenev]

Use matchPattern to match constants

[ezhulenev]

In D90922#2379863, @rriddle wrote:

At the moment it is somewhat unclear at what scope this analysis is intended to be constructed and used. It seems like this currently doesn't handle things related to the callgraph ...

I'd say that call graph is irrelevant for this analysis. Call operations "consumes" SSA values and then it is the responsibility of the callee to do the right thing. But I'm speaking from the point of view of async-ref-counting, I'd love to hear what @ftynse can tell about his need for register assignment.

Maybe there is a better terminology then "number of executions" that is less confusing.

mlir/lib/Analysis/NumberOfExecutions.cpp

66

#1. 1 execution of the basic block that contains print operations. Function call is irrelevant here.

#2. Unknown, because it involves dynamic control flow

CFG will look something like this:

^bb0
  %cond = cmp i == 0
  cond_br %cond ^abort, ^print
^abort
  br ^print
^print

def g(i):
  abort()
  print("here") # is a bit more tricky, but let's assume this is an 
                # estimate given that no operations can terminate a program

I think we need to be more precise about what we mean by "execution" here.

"How many times this operation will be executed in this region assuming region invoked N times".

Call operation is about completely different region (different instance of a region).

[ezhulenev]

FWIW Original async ref counting PR uses "number of operation instances in the region"

[silvas]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

I see that the current code actually does walk parents in the way I described. Sorry.

I thought somewhat about the definition of "execution" that we mean here. It's really "for *each* time that Region is entered, how many times will given operation be executed". That makes the definition work for my case f above, since it yields getNumberOfExecutions(thePrintOp, f.getBody()) == 1.

Also, because of the possibility of aborting or calling a function that cannot be proven to terminate in the middle of a block, I think it is best for getNumberOfExecutions to take an Operation * instead of a Block, so that cases like this behave as intended, and users don't have to do their own within-a-block check:

def h():
  print("here1") # Number of executions is 1 per entry to h.getBody()
  call @weProvedThatItAbortsOrInfiniteLoops()
  print("here") # Number of executions is 0 per entry to h.getBody()

As such, could we do something like an interface to this analysis that looks like:

Optional<IntegerRange> getNumberOfExecutions(Operation *op, Region *perEntryOfThisRegion)

IntegerRange is just a pair of integers that represent a lower and upper bound on the number of executions, which composes nicely to handle cases where we can't prove that the statement executes at all (like those tricky abort/infinite-loop cases) by having a lower bound of 0 (we could do something more precise than an interval, such as an arbitrary set of integers, but that might get unwieldy).

To handle this correctly, getNumRegionInvocations will need to be wordsmithed to something like "number of times the body will be invoked if the body yields normally (i.e. doesn't abort or invoke an infinite loop)". Then this analysis can do a bottom up traversal of the region tree, doing the op-level analysis within each block looking for possibly aborting/infiniteLooping operations and correctly handling any ops that follow ops that might abort/infiniteLoop, and transitively propagating that up the region tree.

We can make this more precise with interprocedural analysis that does the equivalent of inferring attributes like LLVM's norecurse/willreturn/mustprogress (see https://llvm.org/docs/LangRef.html#function-attributes)

Under this definition, we have:

def f(i):
  print("here") # Number of executions is 1 per entry to f.getBody()
  if i > 0:
    f(i - 1)

def g(i):
  if i == 0:
    abort()
  print("here") # Number of executions is 0 or 1 per entry to g.getBody()

def h():
  print("here1") # Number of executions is 1 per entry to h.getBody()
  call @weProvedThatItAbortsOrInfiniteLoops()
  print("here") # Number of executions is 0 per entry to h.getBody()

def i():
  print("here1") # Number of executions is 1 per entry to i.getBody()
  call @cannotProveItDoesntAbortOrInfiniteLoop()
  print("here") # Number of executions is 0 or 1 per entry to i.getBody()

[ezhulenev]

Reworderd NumberOfExecutions documentation

[ezhulenev]

Update getNumRegionInvocations documentation to mention "yields normally" property

[ezhulenev]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

I updated documentation for the interface and NumberOfExecutions to mention "yields normally" and "each time" properties.

I like this API ...

Optional<IntegerRange> getNumberOfExecutions(Operation *op, Region *perEntryOfThisRegion)

... but for now computing the IntegerRange is not really needed and it adds a lot of complexity.

I plan to add API with Optional<int64_t> to represent "known fixed number" vs "everything else"

Optional<int64_t> getNumberOfExecutions(Operation *op, Region *perEntryOfThisRegion);

and assume that all operations are "normal": do not abort, or go into infinite loop == "completed in finite amount of time".

[ezhulenev]

Implment Optional<int64_t> getNumberOfExecutions(Operation *op, Region *perEntryOfThisRegion) API

[ezhulenev]

Check that block lies inside the perEntryOfThisRegion

[silvas]

mlir/lib/Analysis/NumberOfExecutions.cpp

66

Thanks.

I think with the new API, we don't need these special cases. FunctionLike already didn't make sense because an op is not "executed" by the module region (we should probably return "unknown" (or an error) for getNumberOfExecutions(someOpInAFunction, module.getBody()). OpTrait::OneRegion doesn't guarantee anything about the number of invocations. E.g. scf::ForOp implements OpTrait::OneRegion.

136

The note that all operates are assumed to not abort/infloop should be in the documentation comment for users of the API; it is not an implementation detail because it makes the analysis incorrect for some potential users.

[ezhulenev]

Do not special case functions for number of executions computation

mlir/lib/Analysis/NumberOfExecutions.cpp

66

Done. It required few more changes, but the special case for functions is gone.

136

Added as a top-level documentation because it is also relevant to blocks.

[silvas]

Thanks Eugene. All I have left are nits. This LGTM.

We probably want to wait on River to give final sign-off. In particular surrounding this analysis not being conservative w.r.t. potentially aborting/nonterminating computations.

mlir/include/mlir/Support/MathExtras.h

32

nit: do we have a unittest you could extend? (or add? :) )

[ezhulenev]

Add a unit test for Support/MathExtras

[rriddle]

The new wording makes more sense to me and clears the confusion I had previously, thanks! LGTM after the remaining comments.

mlir/include/mlir/Analysis/NumberOfExecutions.h

38

nit: in a finite amount

98

nit: Use /// for comments here.

mlir/lib/Analysis/NumberOfExecutions.cpp

16

Can this function include be removed?

39

This seems like it would break when a graph region is nested within a CFG region. Can we avoid the direct assert here and just use kUnknownNumRegionInvocations instead? I could see an assert on the top level operation, but is there a reason to assert on nested operations?

51–60

129

nit: in a finite amount of time

209–222

[ezhulenev]

Address PR comments

mlir/lib/Analysis/NumberOfExecutions.cpp

209–222

I had to walk block and then operations to get sensible output for filecheck, otherwise it prints inner operation before its parent.

[ezhulenev]

Rebased

mlir/include/mlir/Analysis/NumberOfExecutions.h

+//===- NumberOfExecutions.h - Number of executions analysis -----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains an analysis for computing how many times a block within a
+// region is executed. The analysis iterates over all associated regions that
+// are attached to the given top-level operation.
+//
+// It is possible to query number of executions information on block level.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_NUMBER_OF_EXECUTIONS_H
+#define MLIR_ANALYSIS_NUMBER_OF_EXECUTIONS_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
+
+namespace mlir {
+
+class Block;
+class BlockNumberOfExecutionsInfo;
+class Operation;
+
+/// Represents an analysis for computing how many times a block within a region
+/// is executed. The analysis iterates over all associated regions that are
+/// attached to the given top-level operation.
+class NumberOfExecutions {
+public:
+  /// Creates a new NumberOfExecutions analysis that computes how many times a
+  /// block within a region is executed for all associated regions.
+  explicit NumberOfExecutions(Operation *op);
+
+  /// Returns the number of times block `block` inside an `operation` will be
+  /// executed. Returns empty optional if this is not known statically.
+  Optional<int64_t> getNumberOfExecutions(Block *block) const;
+
+  /// Dumps the number of executions information to the given stream.
+  void print(raw_ostream &os) const;
+
+private:
+  /// The operation this analysis was constructed from.
+  Operation *operation;
+
+  /// A mapping from blocks to number of executions information.
+  DenseMap<Block *, BlockNumberOfExecutionsInfo> blockNumbersOfExecution;
+};
+
+/// Represents number of block executions information. This class is not default
+/// constructible to catch accidental constuction working in operations with
+/// `blockNumbersOfExecution` map defined in NumberOfExecutions analysis.
+class BlockNumberOfExecutionsInfo {
+public:
+  /// Returns the number of times this block will be executed within a parent
+  /// operation. This does not take into account how many times the parent
+  /// operation itself will be executed.
+  Optional<int64_t> getNumberOfExecutions() const;
+
+  static BlockNumberOfExecutionsInfo known(int64_t value);
+  static BlockNumberOfExecutionsInfo unknown();
+
+private:
+  BlockNumberOfExecutionsInfo(Optional<int64_t> numberOfExecutions);
+
+  Optional<int64_t> numberOfExecutions;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_ANALYSIS_NUMBER_OF_EXECUTIONS_H

mlir/include/mlir/Dialect/Async/IR/Async.h

 #ifndef MLIR_DIALECT_ASYNC_IR_ASYNC_H
 #define MLIR_DIALECT_ASYNC_IR_ASYNC_H
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace async {
 
 namespace detail {
 struct ValueTypeStorage;
 } // namespace detail

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

 //===- AsyncOps.td - Async operations definition -----------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the operation definition file for Async dialect operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ASYNC_OPS
 #define ASYNC_OPS
 
 include "mlir/Dialect/Async/IR/AsyncBase.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // Async op definitions
 //===----------------------------------------------------------------------===//
 
 // Base class for the operation in this dialect
 class Async_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<AsyncDialect, mnemonic, traits>;
 
 def Async_ExecuteOp :
   Async_Op<"execute", [SingleBlockImplicitTerminator<"YieldOp">,
+                       DeclareOpInterfaceMethods<RegionBranchOpInterface,
+                                                 ["getNumRegionInvocations"]>,
                        AttrSizedOperandSegments]> {
   let summary = "Asynchronous execute operation";
   let description = [{
     The `body` region attached to the `async.execute` operation semantically
     can be executed concurrently with the successor operation. In the followup
     example "compute0" can be executed concurrently with "compute1".
 
     The actual concurrency semantics depends on the dialect lowering to the
let results = (outs Async_TokenType:$token,
                       Variadic<Async_AnyValueType>:$results);
   let regions = (region SizedRegion<1>:$body);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let verifier = [{ return ::verify(*this); }];
 }
 
 def Async_YieldOp :
     Async_Op<"yield", [HasParent<"ExecuteOp">, NoSideEffect, Terminator]> {
   let summary = "terminator for Async execute operation";
   let description = [{
     The `async.yield` is a special terminator operation for the block inside
     `async.execute` operation.
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let assemblyFormat = "attr-dict ($operands^ `:` type($operands))?";

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

unsigned getNumIterOperands() {
       return getOperation()->getNumOperands() - getNumControlOperands();
     }
 
     /// Return operands used when entering the region at 'index'. These operands
     /// correspond to the loop iterator operands, i.e., those exclusing the
     /// induction variable. LoopOp only has one region, so 0 is the only valid
     /// value for `index`.
     OperandRange getSuccessorEntryOperands(unsigned index);
+
+    /// Returns the number of invocations of the body block if the loop bounds
+    /// are constants. Returns `-1` (unknown) otherwise.
+    void getNumRegionInvocations(ArrayRef<Attribute> operands,
+                                 SmallVectorImpl<int64_t> &countPerRegion);
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def IfOp : SCF_Op<"if",
       [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveSideEffects,

mlir/include/mlir/Interfaces/ControlFlowInterfaces.h

 LogicalResult verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
                                             Optional<OperandRange> operands);
 } // namespace detail
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
+// A constant value to represent unknown number of region invocations.
+extern const int64_t kUnknownNumRegionInvocations;
+
 namespace detail {
 /// Verify that types match along control flow edges described the given op.
 LogicalResult verifyTypesAlongControlFlowEdges(Operation *op);
 } //  namespace detail
 
 /// This class represents a successor of a region. A region successor can either
 /// be another region, or the parent operation. If the successor is a region,
 /// this class represents the destination region, as well as a set of arguments

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

InterfaceMethod<[{
         operation as a successor. This method allows for describing which
         regions may be executed when entering an operation, and which regions
         are executed after having executed another region of the parent op. The
         successor region must be non-empty.
       }],
       "void", "getSuccessorRegions",
       (ins "Optional<unsigned>":$index, "ArrayRef<Attribute>":$operands,
            "SmallVectorImpl<RegionSuccessor> &":$regions)
+    >,
+    InterfaceMethod<[{
+        Populates countPerRegion with the number of times this operation will
+        invoke the attached regions. If the number of region invocations is not
+        known statically it will set the number of invocations to
+        `kUnknownNumRegionInvocations`.
+
+        `operands` is a set of optional attributes that either correspond to a
+        constant values for each operand of this operation, or null if that
+        operand is not a constant.
+      }],
+      "void", "getNumRegionInvocations",
+      (ins "ArrayRef<Attribute>":$operands,
+           "SmallVectorImpl<int64_t> &":$countPerRegion), [{}],
+      /*defaultImplementation=*/[{
+        unsigned numRegions = this->getOperation()->getNumRegions();
+        assert(countPerRegion.empty());
+        countPerRegion.resize(numRegions, kUnknownNumRegionInvocations);
+      }]
     >
   ];
 
   let verify = [{
     static_assert(!ConcreteOpType::template hasTrait<OpTrait::ZeroRegion>(),
                   "expected operation to have non-zero regions");
     return success();
   }];

mlir/include/mlir/Support/MathExtras.h

 //===- MathExtras.h - Math functions relevant to MLIR -----------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains math functions relevant to MLIR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_MATHEXTRAS_H_
 #define MLIR_SUPPORT_MATHEXTRAS_H_
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/APInt.h"
 
 namespace mlir {
 
 /// Returns the result of MLIR's ceildiv operation on constants. The RHS is
-/// expected to be positive.
+/// expected to be non-zero.
 inline int64_t ceilDiv(int64_t lhs, int64_t rhs) {
-  assert(rhs >= 1);
+  assert(rhs != 0);
   // C/C++'s integer division rounds towards 0.
-  return lhs % rhs > 0 ? lhs / rhs + 1 : lhs / rhs;
+  int64_t x = (rhs > 0) ? -1 : 1;
+  return (lhs * rhs > 0) ? ((lhs + x) / rhs) + 1 : -(-lhs / rhs);
 }
 
 /// Returns the result of MLIR's floordiv operation on constants. The RHS is
-/// expected to be positive.
+/// expected to be non-zero.
 inline int64_t floorDiv(int64_t lhs, int64_t rhs) {
-  assert(rhs >= 1);
+  assert(rhs != 0);
   // C/C++'s integer division rounds towards 0.
-  return lhs % rhs < 0 ? lhs / rhs - 1 : lhs / rhs;
+  int64_t x = (rhs < 0) ? 1 : -1;
+  return (lhs * rhs < 0) ? -((-lhs + x) / rhs) - 1 : lhs / rhs;
 }
 
 /// Returns MLIR's mod operation on constants. MLIR's mod operation yields the
 /// remainder of the Euclidean division of 'lhs' by 'rhs', and is therefore not
 /// C's % operator.  The RHS is always expected to be positive, and the result
 /// is always non-negative.
 inline int64_t mod(int64_t lhs, int64_t rhs) {
   assert(rhs >= 1);

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AffineAnalysis.cpp
   AffineStructures.cpp
   BufferAliasAnalysis.cpp
   CallGraph.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
+  NumberOfExecutions.cpp
   PresburgerSet.cpp
   SliceAnalysis.cpp
   Utils.cpp
   )
 
 add_mlir_library(MLIRAnalysis
   BufferAliasAnalysis.cpp
   CallGraph.cpp
   Liveness.cpp
+  NumberOfExecutions.cpp
   SliceAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers
 
add_mlir_library(MLIRLoopAnalysis
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
   MLIRPresburger
   MLIRSCF
   )
-  
 add_subdirectory(Presburger)

mlir/lib/Analysis/NumberOfExecutions.cpp

+//===- NumberOfExecutions.cpp - Number of executions analysis -------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of the number of executions analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "mlir/Analysis/NumberOfExecutions.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Region.h"
+#include "mlir/IR/RegionKindInterface.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+
+#define DEBUG_TYPE "number-of-executions-analysis"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// NumberOfExecutions
+//===----------------------------------------------------------------------===//
+
+/// Computes blocks number of executions information for the given region.
+void computeRegionBlockNumberOfExecutions(
+    Region &region, DenseMap<Block *, BlockNumberOfExecutionsInfo> &blockInfo) {
+  Operation *parentOp = region.getParentOp();
+  int regionId = region.getRegionNumber();
+
+  // CFG analysis does not make sense for Graph regions.
+  if (auto kindInterface = dyn_cast<RegionKindInterface>(parentOp)) {
+    auto kind = kindInterface.getRegionKind(regionId);
+    assert(kind == RegionKind::SSACFG && "only SSA CFG regions are supported");
+    (void)kind;
+  }
+
+  SmallVector<int64_t, 4> numRegionsInvocations;
+  if (auto regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp)) {
+    SmallVector<Attribute, 2> operands;
+    operands.reserve(parentOp->getNumOperands());
+
+    for (Value operand : parentOp->getOperands()) {
+      ConstantOp cst = operand.getDefiningOp<ConstantOp>();
+      operands.push_back(cst ? cst.value() : Attribute());
+    }
+
+    regionInterface.getNumRegionInvocations(operands, numRegionsInvocations);
+  }
+
+  // Functions will always execute attached region once.
+  if (auto func = dyn_cast<FuncOp>(parentOp))
+    numRegionsInvocations.push_back(1);
+
+  // If we do not have region invocation information then we can't compute
+  // execution number for the blocks inside the region.
+  if (numRegionsInvocations.empty() ||
+      numRegionsInvocations[regionId] == kUnknownNumRegionInvocations) {
+    for (Block &block : region) {
+      blockInfo.insert({&block, BlockNumberOfExecutionsInfo::unknown()});
+    }
+    return;
+  }
+
+  int64_t numRegionInvocations = numRegionsInvocations[regionId];
+
+  // DFS traversal looking for loops in the CFG.
+  llvm::SmallSet<Block *, 4> loopStart;
+
+  llvm::unique_function<void(Block *, llvm::SmallSet<Block *, 4> &)> dfs =
+      [&](Block *block, llvm::SmallSet<Block *, 4> &visited) {
+        // Found a loop in the CFG that starts at the `block`.
+        if (visited.contains(block)) {
+          loopStart.insert(block);
+          return;
+        }
+
+        // Continue DFS traversal.
+        visited.insert(block);
+        for (Block *successor : block->getSuccessors())
+          dfs(successor, visited);
+        visited.erase(block);
+      };
+
+  llvm::SmallSet<Block *, 4> visited;
+  dfs(&region.front(), visited);
+
+  // Start from the entry block and follow only blocks with single succesor.
+  Block *block = &region.front();
+  while (block && !loopStart.contains(block)) {
+    // Block will be executed exactly once.
+    blockInfo.insert(
+        {block, BlockNumberOfExecutionsInfo::known(numRegionInvocations)});
+
+    // We reached the exit block or block with multiple successors.
+    if (block->getNumSuccessors() != 1)
+      break;
+
+    // Continue traversal.
+    block = block->getSuccessor(0);
+  }
+
+  // For all blocks that we did not visit set the executions number to unknown.
+  for (Block &block : region) {
+    if (blockInfo.count(&block))
+      continue;
+    blockInfo.insert({&block, BlockNumberOfExecutionsInfo::unknown()});
+  }
+}
+
+/// Creates a new NumberOfExecutions analysis that computes how many times a
+/// block within a region is executed for all associated regions.
+NumberOfExecutions::NumberOfExecutions(Operation *op) : operation(op) {
+  operation->walk([&](Region *region) {
+    computeRegionBlockNumberOfExecutions(*region, blockNumbersOfExecution);
+  });
+}
+
+Optional<int64_t>
+NumberOfExecutions::getNumberOfExecutions(Block *block) const {
+  auto blockIt = blockNumbersOfExecution.find(block);
+  const auto &blockInfo = blockIt->getSecond();
+
+  // Immediately return None if we do not know block number of executions
+  // within its parent operation.
+  if (!blockInfo.getNumberOfExecutions().hasValue())
+    return None;
+
+  // Follow parent operations until we reach the `operation`.
+  int64_t numberOfExecutions = *blockInfo.getNumberOfExecutions();
+  Operation *parent = block->getParentOp();
+
+  while (parent != operation) {
+    Block *parentBlock = parent->getBlock();
+    auto parentIt = blockNumbersOfExecution.find(parentBlock);
+    const auto &parentBlockInfo = parentIt->getSecond();
+
+    // We stumbled upon an operation with unknown number of executions.
+    if (!parentBlockInfo.getNumberOfExecutions().hasValue())
+      return None;
+
+    // Number of block executions is a product of all parent blocks executions.
+    numberOfExecutions *= *parentBlockInfo.getNumberOfExecutions();
+    parent = parent->getParentOp();
+  }
+
+  return numberOfExecutions;
+}
+
+void NumberOfExecutions::print(raw_ostream &os) const {
+  unsigned blockId = 0;
+
+  operation->walk([&](Block *block) {
+    llvm::errs() << "Block: " << blockId++ << "\n";
+    llvm::errs() << "Number of executions: ";
+    if (auto n = getNumberOfExecutions(block))
+      llvm::errs() << *n << "\n";
+    else
+      llvm::errs() << "<unknown>\n";
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// BlockNumberOfExecutionsInfo
+//===----------------------------------------------------------------------===//
+
+BlockNumberOfExecutionsInfo::BlockNumberOfExecutionsInfo(
+    Optional<int64_t> numberOfExecutions)
+    : numberOfExecutions(numberOfExecutions) {}
+
+Optional<int64_t> BlockNumberOfExecutionsInfo::getNumberOfExecutions() const {
+  return numberOfExecutions;
+}
+
+BlockNumberOfExecutionsInfo BlockNumberOfExecutionsInfo::known(int64_t value) {
+  return {value};
+}
+BlockNumberOfExecutionsInfo BlockNumberOfExecutionsInfo::unknown() {
+  return {None};
+}

mlir/lib/Dialect/Async/IR/Async.cpp

 }
 
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
 
 constexpr char kOperandSegmentSizesAttr[] = "operand_segment_sizes";
 
+void ExecuteOp::getNumRegionInvocations(
+    ArrayRef<Attribute> operands, SmallVectorImpl<int64_t> &countPerRegion) {
+  (void)operands;
+  assert(countPerRegion.empty());
+  countPerRegion.resize(1);
+  countPerRegion[0] = 1;
+}
+
+void ExecuteOp::getSuccessorRegions(Optional<unsigned> index,
+                                    ArrayRef<Attribute> operands,
+                                    SmallVectorImpl<RegionSuccessor> &regions) {
+  // The `body` region branch back to the parent operation.
+  if (index.hasValue()) {
+    assert(*index == 0);
+    regions.push_back(RegionSuccessor(getResults()));
+    return;
+  }
+
+  // Otherwise the successor is the body region.
+  regions.push_back(RegionSuccessor(&body()));
+}
+
 static void print(OpAsmPrinter &p, ExecuteOp op) {
   p << op.getOperationName();
 
   // [%tokens,...]
   if (!op.dependencies().empty())
     p << " [" << op.dependencies() << "]";
 
   // (%value as %unwrapped: !async.value<!arg.type>, ...)

mlir/lib/Dialect/SCF/SCF.cpp

 //===- SCF.cpp - Structured Control Flow Operations -----------------------===//
 //
 // 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/IR/BlockAndValueMapping.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace mlir::scf;
 
 //===----------------------------------------------------------------------===//
 // SCFDialect Dialect Interfaces
 //===----------------------------------------------------------------------===//
void ForOp::getSuccessorRegions(Optional<unsigned> index,
   }
 
   // Otherwise, the loop may branch back to itself or the parent operation.
   assert(index.getValue() == 0 && "expected loop region");
   regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
   regions.push_back(RegionSuccessor(getResults()));
 }
 
+void ForOp::getNumRegionInvocations(ArrayRef<Attribute> operands,
+                                    SmallVectorImpl<int64_t> &countPerRegion) {
+  assert(countPerRegion.empty());
+  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;
+    return;
+  }
+
+  countPerRegion[0] =
+      ceilDiv(ub.getValue().getSExtValue() - lb.getValue().getSExtValue(),
+              step.getValue().getSExtValue());
+}
+
 ValueVector mlir::scf::buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps, ValueRange iterArgs,
     function_ref<ValueVector(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuilder) {
   assert(lbs.size() == ubs.size() &&
          "expected the same number of lower and upper bounds");
   assert(lbs.size() == steps.size() &&

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionBranchOpInterface
 //===----------------------------------------------------------------------===//
 
+namespace mlir {
+// A constant value to represent unknown number of region invocations.
+const int64_t kUnknownNumRegionInvocations = -1;
+} // namespace mlir
+
 /// Verify that types match along all region control flow edges originating from
 /// `sourceNo` (region # if source is a region, llvm::None if source is parent
 /// op). `getInputsTypesForRegion` is a function that returns the types of the
 /// inputs that flow from `sourceIndex' to the given region, or llvm::None if
 /// the exact type match verification is not necessary (e.g., if the Op verifies
 /// the match itself).
 static LogicalResult
 verifyTypesAlongAllEdges(Operation *op, Optional<unsigned> sourceNo,

mlir/test/Analysis/test-number-of-executions.mlir

+// RUN: mlir-opt %s -test-print-number-of-executions -split-input-file 2>&1    \
+// RUN:   | FileCheck %s --dump-input=always
+
+// CHECK-LABEL: Number of executions: empty
+func @empty() {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: sequential
+func @sequential() {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  br ^bb1
+^bb1:
+  // CHECK: Block: 1
+  // CHECK: Number of executions: 1
+  br ^bb2
+^bb2:
+  // CHECK: Block: 2
+  // CHECK: Number of executions: 1
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: conditional
+func @conditional(%cond : i1) {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  br ^bb1
+^bb1:
+  // CHECK: Block: 1
+  // CHECK: Number of executions: 1
+  cond_br %cond, ^bb2, ^bb3
+^bb2:
+  // CHECK: Block: 2
+  // CHECK: Number of executions: <unknown>
+  br ^bb4
+^bb3:
+  // CHECK: Block: 3
+  // CHECK: Number of executions: <unknown>
+  br ^bb4
+^bb4:
+  // CHECK: Block: 4
+  // CHECK: Number of executions: <unknown>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: loop
+func @loop(%cond : i1) {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  br ^bb1
+^bb1:
+  // CHECK: Block: 1
+  // CHECK: Number of executions: <unknown>
+  br ^bb2
+^bb2:
+  // CHECK: Block: 2
+  // CHECK: Number of executions: <unknown>
+  br ^bb3
+^bb3:
+  // CHECK: Block: 3
+  // CHECK: Number of executions: <unknown>
+  cond_br %cond, ^bb1, ^bb4
+^bb4:
+  // CHECK: Block: 4
+  // CHECK: Number of executions: <unknown>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: scf_if_dynamic_branch
+func @scf_if_dynamic_branch(%cond : i1) {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  scf.if %cond {
+    // CHECK: Block: 1
+    // CHECK: Number of executions: <unknown>
+  } else {
+    // CHECK: Block: 2
+    // CHECK: Number of executions: <unknown>
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: async_execute
+func @async_execute() {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  async.execute {
+   // CHECK: Block: 1
+   // CHECK: Number of executions: 1
+    async.yield
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: async_execute_with_scf_if
+func @async_execute_with_scf_if(%cond : i1) {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  async.execute {
+    // CHECK: Block: 1
+    // CHECK: Number of executions: 1
+    scf.if %cond {
+    // CHECK: Block: 2
+    // CHECK: Number of executions: <unknown>
+    } else {
+    // CHECK: Block: 3
+    // CHECK: Number of executions: <unknown>
+    }
+    async.yield
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: scf_for_constant_bounds
+func @scf_for_constant_bounds() {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+
+  scf.for %i = %c0 to %c2 step %c1 {
+    // CHECK: Block: 1
+    // CHECK: Number of executions: 2
+  }
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: propagate_parent_num_executions
+func @propagate_parent_num_executions() {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+
+  scf.for %i = %c0 to %c2 step %c1 {
+    // CHECK: Block: 1
+    // CHECK: Number of executions: 2
+    async.execute {
+      // CHECK: Block: 2
+      // CHECK: Number of executions: 2
+      async.yield
+    }
+  }
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: Number of executions: clear_num_executions
+func @clear_num_executions(%step : index) {
+  // CHECK: Block: 0
+  // CHECK: Number of executions: 1
+  %c0 = constant 0 : index
+  %c2 = constant 2 : index
+
+  scf.for %i = %c0 to %c2 step %step {
+    // CHECK: Block: 1
+    // CHECK: Number of executions: <unknown>
+    async.execute {
+      // CHECK: Block: 2
+      // CHECK: Number of executions: <unknown>
+      async.yield
+    }
+  }
+
+  return
+}

mlir/test/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTestTransforms
   TestLinalgCodegenStrategy.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestNumberOfExecutions.cpp
   TestOpaqueLoc.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSCFUtils.cpp
   TestVectorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR

mlir/test/lib/Transforms/TestNumberOfExecutions.cpp

+//===- TestNumberOfExecutions.cpp - Test number of executions analysis ----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for constructing and resolving number of
+// executions information.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/NumberOfExecutions.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestNumberOfExecutionsPass
+    : public PassWrapper<TestNumberOfExecutionsPass, FunctionPass> {
+  void runOnFunction() override {
+    llvm::errs() << "Number of executions: " << getFunction().getName() << "\n";
+    getAnalysis<NumberOfExecutions>().print(llvm::errs());
+  }
+};
+
+} // end anonymous namespace
+
+namespace mlir {
+void registerTestNumberOfExecutionsPass() {
+  PassRegistration<TestNumberOfExecutionsPass>(
+      "test-print-number-of-executions",
+      "Print the contents of a constructed number of executions analysis.");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopPermutationPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchers();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
+void registerTestNumberOfExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestPrintDefUsePass();
 void registerTestPrintNestingPass();
 void registerTestRecursiveTypesPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSpirvModuleCombinerPass();
#endif
   registerTestLinalgTransforms();
   registerTestLivenessPass();
   registerTestLoopFusion();
   registerTestLoopMappingPass();
   registerTestLoopUnrollingPass();
   registerTestMatchers();
   registerTestMemRefDependenceCheck();
   registerTestMemRefStrideCalculation();
+  registerTestNumberOfExecutionsPass();
   registerTestOpaqueLoc();
   registerTestPreparationPassWithAllowedMemrefResults();
   registerTestPrintDefUsePass();
   registerTestPrintNestingPass();
   registerTestRecursiveTypesPass();
   registerTestReducer();
   registerTestGpuParallelLoopMappingPass();
   registerTestSpirvEntryPointABIPass();