[mlir] Add integer range inference analysis

[krzysz00]

Summary

This commit defines a dataflow analysis for integer ranges, which
uses a newly-added InferIntRangeInterface to compute the lower and
upper bounds on the results of an operation from the bounds on the
arguments. The range inference is a flow-insensitive dataflow analysis
that can be used to simplify code, such as by statically identifying
bounds checks that cannot fail in order to eliminate them.

The InferIntRangeInterface has one method, inferResultRanges(), which
takes a vector of inferred ranges for each argument to an op
implementing the interface and a callback allowing the implementation
to define the ranges for each result. These ranges are stored as
ConstantIntRanges, which hold the lower and upper bounds for a
value. Bounds are tracked separately for the signed and unsigned
interpretations of a value, which ensures that the impact of
arithmetic overflows is correctly tracked during the analysis.

The commit also adds a -test-int-range-inference pass to test the
analysis until it is integrated into SCCP or otherwise exposed.

Finally, this commit fixes some bugs relating to the handling of
region iteration arguments and terminators in the data flow analysis
framework.

Depends on D124020

Depends on D124021

[rriddle]

mlir/lib/Dialect/Arithmetic/Analysis/IntRangeAnalysis.cpp

47–48

59–60

Can you add more newlines into these functions to better separate the functionality?

94–95

119–125

Can you split the functor into a separate variable? That would likely format this better.

120–124

?

138–141

Same comments here as above.

[krzysz00]

@rriddle To clarify, what types of tests do you think would be appropriate? Unit tests of the lattice implementation?

Secondly, what would the more generic implementation look like? That is, where's the missing generality?

Since it might help, I'll restate the context that this current set of patches is a general solution to.

If we take

func @example() -> index {
    // Is gpu.block_id, but I use our wrapper for clarity and because, unlike with gpu.block_id, we can know this is <= int32_max
    %0 = miopen.block_id
    %1 = affine.apply affine_map<(d0) -> (d0 floordiv 9 + ((d0 mod 9) floordiv 3) * 3 + (d0 mod 3) * 9)>(%0)
    return %1 : index
}

and run it through -lower-affine -canonicalize, we get

func @example() -> index {
    %c9 = arith.constant 9 : index
    %c0 = arith.constant 0 : index
    %c3 = arith.constant 3 : index
    %c-1 = arith.constant -1 : index
    %0 = miopen.block_id
    %1 = arith.cmpi slt, %0, %c0 : index
    %2 = arith.subi %c-1, %0 : index
    %3 = arith.select %1, %2, %0 : index
    %4 = arith.divsi %3, %c9 : index
    %5 = arith.subi %c-1, %4 : index
    %6 = arith.select %1, %5, %4 : index
    %7 = arith.remsi %0, %c9 : index
    %8 = arith.cmpi slt, %7, %c0 : index
    %9 = arith.addi %7, %c9 : index
    %10 = arith.select %8, %9, %7 : index
    %11 = arith.cmpi slt, %10, %c0 : index
    %12 = arith.subi %c-1, %10 : index                                                      %13 = arith.select %11, %12, %10 : index
    %14 = arith.divsi %13, %c3 : index
    %15 = arith.subi %c-1, %14 : index
    %16 = arith.select %11, %15, %14 : index
    %17 = arith.muli %16, %c3 : index
    %18 = arith.addi %6, %17 : index
    %19 = arith.remsi %0, %c3 : index
    %20 = arith.cmpi slt, %19, %c0 : index
    %21 = arith.addi %19, %c3 : index
    %22 = arith.select %20, %21, %19 : index
    %23 = arith.muli %22, %c9 : index
    %24 = arith.addi %18, %23 : index
    return %24 : index
}

Note how this code has a bunch of conditional logic for handling the case of the arguments to mod or division being negative when we can infer, from the fact that the map is being run on a non-negative value, that those conditions are statically false. (Note that it won't always make sense to drop those checks - we might have put a d0 - 1 somewhere in there to handle padding, for example)

Ideally (and -fold-inferred-constants -canonicalize -unsigned-when-equivalent will get us there) we want this simplified to

func @example() -> index {
    %c3 = arith.constant 3 : index
    %c9 = arith.constant 9 : index
    %0 = miopen.block_id
    %1 = arith.divui %0, %c9 : index
    %2 = arith.remui %0, %c9 : index
    %3 = arith.divui %2, %c3 : index
    %4 = arith.muli %3, %c3 : index
    %5 = arith.addi %0, %4 : index
    %6 = arith.remsi %0, %c3 : index
    %7 = arith.muli %6, %c9 : index
    %8 = arith.addi %5, %7 : index
    return %8 : index
}

This would make the IR we send to LLVM cleaner and enable certain optimizations (for example, when we switched some hardcoded remsi ..., %c64 to remui ..., %c64 in our generated code, LLVM could swap to bitwise arithmetic instead of actually emitting remainder instructions).

mlir/include/mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h

11

There's nothing stopping this from being somewhere more general, but I think that, because the passes that use all this are in arith, as are most of the implementors, having it in arith makes sense from a namespacing angle.

25

This is in a header because it'll be used by multiple passes - see https://reviews.llvm.org/D124024

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

30

I initially had that simplified model but found it hard to work with, especially since, unlike in LLVM, we don't have things like nuw/nsw.

Do you think it'd make sense for us to use this ConstantRange? I'll note its semantics are unclear, and I can't quite work out how they stuck both signed and unsigned ranges in there.

I could flip these IntegerAttrs to Optional<APInt>s - that's probably a better type.

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.td

45

What kind of callback? And I agree that this is missing a way to handle region arguments that aren't things like iter_args(), hence the changes to the data flow framework so I could add visitNonControlFlowArguments() in the analysis pass.

Do you see cases where we'd want to do interesting things with regions that aren't LoopLikeInterfaces or if/br-like?

mlir/include/mlir/Dialect/Arithmetic/Transforms/Passes.h

30

Should that patch be also disjoint from https://reviews.llvm.org/D124022 ?

[krzysz00]

Just wanted to follow up on my questions about next steps

[Mogball]

mlir/include/mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h

11

I don't have a strong opinion on the namespace other than "it should eventually not be in arith" but I don't think the pass and analysis should be located under mlir/Dialect/Arithmetic, since that's tied too closely to the dialect

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

76–83

These functions are more of an implementation detail of integer range analysis which I don't think should even be exposed publicly.

mlir/lib/Dialect/Arithmetic/Transforms/FoldInferredConstants.cpp

41

There's nothing about this pass that is specific to the arithmetic dialect.

[krzysz00]

mlir/include/mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h

11

On the other hand, ConstantOp is under Arithmetic ...

If y'all're happy with this all going into mlir/Analysis and mlir/Transforms and the main MLIR namespace, that's fine by me.

I'm just coming from the perspective where the passes and analyses are mostly about arith, and so mlir/Dialect/Arithmetic/ is a logical place to put them

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

76–83

Well, because of how data flow analysis work, these functions have to be methods on the struct.

And because there will be multiple passes that use this analysis, the struct in question has to be declared in a header.

And I'm not aware of how to put a private header together - I'd bo open to IntRangeAttrs being hidden from the public API, but I've got no idea how to do that.

[Mogball]

mlir/include/mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h

11

Yeah that's fine with me. SCCP is under mlir/Transforms as well

[krzysz00]

@rriddle Looking to follow up for more clarity about

• What sort of increased generality you're wanting to see in this patch
• What a test for the lattice functionality should look like (unit test? adding some sort of test pass?)
• What the callback-based version of InferIntRangeInterface you alluded to might look like

[krzysz00]

Move from integer attrs to Optional<APInt>

[Mogball]

So to keep the ball rolling on this one, I will pretend to be river and answer your questions :P

1. Just move things out of the arithmetic dialect and probably drop the arith namespace (mlir/Transforms). The analysis itself could go inside the source file like SCCPAnalysis.
2. Add 1 or 2 ops to the test dialect that, when combined with control-flow from scf or cf, cover all the "interesting" cases, i.e. integers ranges being overdefined, defined (and a range), and defined (collapsed to a constant).
3. The callback would be something like

void getIntegerRanges(ArrayRef<IntRangeAttrs> operandRanges, function_ref<void(Value, IntRangeAttrs)> setValueRange);

Which an implementing op could call on its results and any region arguments. This way you can implement visitNonControlFlowArguments using it too.

[rriddle]

Basically +1 to everything Jeff said. (It also seems like a lot of the comments are still unresolved?)

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

30

How does LLVM implement utilize ConstantRange then? Ideally I think we should try to avoid having two parallel ranges, given that now we have to implement things operating on both.

mlir/lib/Dialect/Arithmetic/Analysis/IntRangeAnalysis.cpp

30

Unresolved? Prefer using namespace blah; over namespace blah { in cpp files.

59–60

Unresolved?

94–95

Unresolved?

mlir/lib/Dialect/Arithmetic/IR/InferIntRangeInterface.cpp

16

Same comments about namespaces here.

mlir/lib/Dialect/Arithmetic/Transforms/FoldInferredConstants.cpp

25

Same comments here about static functions, and namespaces.

[krzysz00]

Move things around to address some review comments

[krzysz00]

Ok, so, to update the actions I've taken

1. I moved InferIntRangeInterface into Interfaces/
2. I moved IntRangeAnalysis into Analysis/. Having re-read y'all's comments on the topic, the claim I'm seeing is that, instead of directly exposing the fact that the integer range analysis is a dataflow analysis, there should be a wrapper around it that provides a more limited API. That makes sense, and I'll get on that next.
3. I left -fold-inferred-constants under arith because
   • It specifically creates arith.constant ops
   • It's not clear how it could be generalized past that
   • It looks like we're not supposed to make MLIRTransforms depend on any dialect in particular

If y'all have notes on how to make fold-inferred-constants dialect-independent, I'm all ears.

I haven't broken -fold-inferred-constants into its own commit yet, but I can get on that.

I'm also fixing to add some unit tests for IntRangeAttrs later.

[krzysz00]

Expose fewer details of IntRangeAnalysis

[krzysz00]

Fix building

[Mogball]

SCCP will defer to Dialect::materializeConstant. You could do the same with your pass.

[Mogball]

Mostly LGTM, just missing some tests

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

18 ↗

(On Diff #428189)

Put the includes together

30 ↗

(On Diff #428189)

These should be placed at the bottom of the class

33 ↗

(On Diff #428189)

Brief comments on the default constructor and other constructors?

mlir/lib/Analysis/IntRangeAnalysis.cpp

24 ↗

(On Diff #428189)

Typo

25 ↗

(On Diff #428189)

40 ↗

(On Diff #428189)

Operand lattices should never be uninitialized when an operation is visited.

106 ↗

(On Diff #428189)

ForwardDataFlowAnalysis will assert if any result lattices are uninitialized, so this check is not necessary

117 ↗

(On Diff #428189)

Please expand the name to result

125 ↗

(On Diff #428189)

Please expand mb

169 ↗

(On Diff #428189)

Operand lattices should never be uninitialized when an operation is visited.

172 ↗

(On Diff #428189)

191 ↗

(On Diff #428189)

This line is a no-op because markPessimisticFixpoint will overwrite the optimistic value

219 ↗

(On Diff #428189)

If you need to "undo" the correction, why not take it out of the helper function and perform it here (with an inverted condition)

mlir/lib/Dialect/Arithmetic/Transforms/FoldInferredConstants.cpp

30

I would call the dialect constant materializer, either the dialect of the defining op or the parent op if it's a block argument, then the pass would be totally dialect agnostic. You could also use a folder instance to unique constants as you go

[rriddle]

mlir/include/mlir/Analysis/IntRangeAnalysis.h

67 ↗

(On Diff #428189)

Can you use a std::unique_ptr here (i.e. use the pImpl technique) and move IntRangeAnalysisImpl to the cpp file? We shouldn't be exposing that at all within the header file.

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

18 ↗

(On Diff #428189)

+1. Also system includes should come last.

[krzysz00]

Address comments on code; tests will be added tomorrow-ish

[krzysz00]

Thanks for going over all this closely, y'all!

[Mogball]

mlir/lib/Analysis/IntRangeAnalysis.cpp

236 ↗

(On Diff #428506)

Add braces because if the comments

mlir/lib/Transforms/FoldInferredConstants.cpp

18 ↗

(On Diff #428506)

[Mogball]

MLIR is not entirely consistent with the namespace end comment but the preferred way is as such

mlir/include/mlir/Analysis/IntRangeAnalysis.h

25 ↗

(On Diff #428506)

42 ↗

(On Diff #428506)

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

100 ↗

(On Diff #428506)

mlir/lib/Analysis/IntRangeAnalysis.cpp

57–59 ↗

(On Diff #428506)

mlir/lib/Transforms/FoldInferredConstants.cpp

100 ↗

(On Diff #428506)

[krzysz00]

Add tests, fix typos

[Mogball]

LGTM!

Could you add one test op that sets the ranges of its block arguments and test that it works correctly?

mlir/lib/Analysis/IntRangeAnalysis.cpp

24–25 ↗

(On Diff #428801)

This class can just go in the anonymous namespace

105 ↗

(On Diff #428801)

The operand lattices should never be uninitialized.

148 ↗

(On Diff #428801)

braces

169 ↗

(On Diff #428801)

braces

222 ↗

(On Diff #428801)

braces

234 ↗

(On Diff #428801)

if any of the branches have braces, all of them should

mlir/test/lib/Dialect/Test/TestDialect.cpp

1411 ↗

(On Diff #428801)

mlir/test/lib/Dialect/Test/TestOps.td

2883–2886 ↗

(On Diff #428801)

2907–2911 ↗

(On Diff #428801)

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

1–2 ↗

(On Diff #428801)

what happened here lol

[rriddle]

Please wait for my approval before landing. I'd like to go over this more.

[rriddle]

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

30

Unresolved?

[krzysz00]

Fix braces, other such - add unit test for branches.

[krzysz00]

( @rriddle To bring the signed/unsigned/ConstantRange discussion out of inline comments)

I've got a few reasons I'd say it's reasonable to track the signed and unsigned interpretations of values separately, in contrast with LLVM's approach.

1. The LLVM approach appears to rely on having a bunch of flow-sensitive reasoning that generates things like nsw in order to make their analysis more precise. We don't have, and probably don't want to recreate, all that infrastructure, especially with how extensible MLIR is. So, we need another avenue to get analysis precision.
2. Having signed/unsigned bounds like this makes it easier to reason about what the range inference implementations are doing, especially when dealing with mixtures of signed and unsigned ops (like divui or remsi). Trying to precisely get bounds out of mixed code like that is tricky.
3. In most cases, we don't have double computation - the three ops that really have to do duplicate reasoning about the signed and unsigned result are addition, subtraction, and multiplication. Just about everything else can call fromSigned or fromUnsigned after computing one set of bounds to copy the bounds over if that's a valid thing to do.
4. Initial attempts at a one-range system were a confusing mess and I'm not looking to fight through making sure that's implemented correctly.

mlir/lib/Analysis/IntRangeAnalysis.cpp

24–25 ↗

(On Diff #428801)

Doesn't it need to have a defined namespace so it can be referenced in a header?

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

1–2 ↗

(On Diff #428801)

clang-format

[rriddle]

In D124023#3509442, @krzysz00 wrote:

( @rriddle To bring the signed/unsigned/ConstantRange discussion out of inline comments)

I've got a few reasons I'd say it's reasonable to track the signed and unsigned interpretations of values separately, in contrast with LLVM's approach.

1. The LLVM approach appears to rely on having a bunch of flow-sensitive reasoning that generates things like nsw in order to make their analysis more precise. We don't have, and probably don't want to recreate, all that infrastructure, especially with how extensible MLIR is. So, we need another avenue to get analysis precision.
2. Having signed/unsigned bounds like this makes it easier to reason about what the range inference implementations are doing, especially when dealing with mixtures of signed and unsigned ops (like divui or remsi). Trying to precisely get bounds out of mixed code like that is tricky.
3. In most cases, we don't have double computation - the three ops that really have to do duplicate reasoning about the signed and unsigned result are addition, subtraction, and multiplication. Just about everything else can call fromSigned or fromUnsigned after computing one set of bounds to copy the bounds over if that's a valid thing to do.
4. Initial attempts at a one-range system were a confusing mess and I'm not looking to fight through making sure that's implemented correctly.

I think this is underestimating some of the effects resulting from having both signed and unsigned:

• It seems like in the implementation of the interfaces for the arithmetic operations you are ignoring known information

E.g. if a signed operation has inputs that have known unsigned ranges, why shouldn't we try to preserve and utilize that information when possible?

• The duplication is not just about adding the interface to an op, but also when using the analysis

Bouncing off of the previous point, if a user doesn't want to ignore known information they will inevitably have to check both anyways.

• You aren't going to be able to get away from nsw and related flags

The arithmetic operations don't have those at this point, but that doesn't mean they won't grow them moving forward (the current modeling is not all encompassing). We also will want this analysis to support cases like the LLVM dialect, which do encode that information already.

I have some serious concerns about using multiple ranges, and I'd like to fully understand why a single range ended up being "a confusing mess". Especially given that LLVM has utilized it successfully, and hasn't revolted/reverted into a multi-range system.

[krzysz00]

It looks like, at the very least, I misread the comments on ConstantRange. Note that, for example, computeConstantRange at https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Analysis/ValueTracking.h#L548 has a forSigned parameter, and that the scalar evolution code has an enum requiring you to specify if you're computing the range for signed or unsigned values and corresponding utility functions getUnsignedRange() and getSignedRange().

This means my point 4, about things being a huge mess, doesn't apply to LLVM - what I initially tried to do is to use the same pair of bounds for both signed and unsigned values without any context about what I wanted, which I don't think is something that can be done precisely ... which is why LLVM doesn't do it either.

And to your points about ignoring information - that really shouldn't happen. The fromSigned and fromUnsigned constructors (which is what will be used when a particular operation has a signedness in its interpretation) will copy the bounds to the other range if possible. For example, if I say fromSigned(0, 15), I'll get the ranges unsigned = [0, 15], signed = [0, 15], but if I say fromSigned(-1, 1), I'll get unsigned = [-inf, +inf], signed = [-1, 1], because -1 <= x <= 1 doesn't correspond to a contiguous set of bit patterns from the unsigned perspective.

So, if you have a signed operation, it will use information about unsigned ranges implicitly - that's where fromUnsigned comes in. (In fact, if the unsigned view of that input has an upper bound <= int_max(type), the signed operation will know that argument is non-negative).

And so, if we used one range, we'd need to do what LLVM does and thread whether we want to know about each Value from a signed or unsigned point of view through the entire analysis. It seems more reasonable to me to just keep two ranges instead.

[krzysz00]

@rriddle I hope the above comment makes some sense

[rriddle]

mlir/include/mlir/Analysis/IntRangeAnalysis.h

30–31 ↗

(On Diff #428996)

33–35 ↗

(On Diff #428996)

Analyses are expected to be constructed from the operation, please change this to a constructor.

36 ↗

(On Diff #428996)

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

22 ↗

(On Diff #428996)

This doesn't look necessary.

28 ↗

(On Diff #428996)

The name of this is a bit weird, given that there aren't any attribute. Can we use ConstantIntRange instead?

32 ↗

(On Diff #428996)

Please add newlines between methods, applies throughout the commit.

39 ↗

(On Diff #428996)

Please drop trivial braces.

45 ↗

(On Diff #428996)

Where are these tuples being built? Can we drop this constructor?

71 ↗

(On Diff #428996)

Can we just pass in the values directly? Why go through tuple?

Same for the other method.

98 ↗

(On Diff #428996)

Why do we need to use optional here? Do we need the default constructor? If this is just for the analysis, we could wrap this class with an internal one that isn't outwardly exposed.

101 ↗

(On Diff #428996)

mlir/lib/Analysis/IntRangeAnalysis.cpp

53 ↗

(On Diff #428996)

Can you complete this comment?

mlir/test/lib/Dialect/Test/TestOps.td

2889–2891 ↗

(On Diff #428996)

2912–2914 ↗

(On Diff #428996)

2926–2928 ↗

(On Diff #428996)

[krzysz00]

Thanks for the comments on the range representation and catching the pointless complexity introduced by Optional<>

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

45 ↗

(On Diff #428996)

These tuples arise from methods in the implementation logic for arith that return various [min, max] ranges as a tuple. They're useful ergonomic constructors, IMO.

71 ↗

(On Diff #428996)

Utility methods that use tuples to return multiple values.

98 ↗

(On Diff #428996)

Having thought about it, we don't _need_ optional - ConstantRange down in LLVM doesn't need it. The case of unbounded ranges can be handled by [0, uint_max(width)] / [int_min(width), int_max(width)], respectively, and we could use width 0 APInts to handle things that aren't integers so we have a meaningful null value.

That would all, now that I think about it, simplify things substantially in many parts of the implementation at the cost of a few minor ergonomic annoyances. I'll make it happen.

[krzysz00]

Make the ints in ranges not optional, since that's not needed.

[krzysz00]

(submitting comments)

[krzysz00]

@rriddle Do you have any other comments?

[rriddle]

Nice, this is getting really close. Great work.

mlir/include/mlir/Analysis/IntRangeAnalysis.h

1 ↗

(On Diff #430446)

This looks a bit off.

18–20 ↗

(On Diff #430446)

Are all of these necessary?

31–35 ↗

(On Diff #430446)

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

21 ↗

(On Diff #430446)

attributes? This comment looks out of date.

51 ↗

(On Diff #430446)

61 ↗

(On Diff #430446)

63–65 ↗

(On Diff #430446)

Is this only for the pessimistic value state stuff? Could we just have a wrapper value for the analysis, and keep ConstantIntRanges clean from this? notAnInteger feels quite weird as an interface for a ConstantIntRange.

77–78 ↗

(On Diff #430446)

With above, can we please drop the tuple API until we need it?

93–95 ↗

(On Diff #430446)

Related to the comment above, if we used a wrapper class in the analysis we wouldn't need this to be public exposed.

107 ↗

(On Diff #430446)

nit: Can you hoist this near the constructors?

116 ↗

(On Diff #430446)

Can you document this?

45 ↗

(On Diff #428996)

Can we remove these until we need them then, I'm still not convinced yet of the value. The value may be more clear when there are uses, in which case they can be readded.

mlir/include/mlir/Transforms/Passes.h

54 ↗

(On Diff #430446)

mlir/lib/Transforms/FoldInferredConstants.cpp

1–2 ↗

(On Diff #430446)

This is a bit off.

88–97 ↗

(On Diff #430446)

SCCP in LLVM represents a range within the lattice, is that something we could do here (i.e. merge the necessary things into SCCP)? It would remove the need to have a duplicate pass with similar goals to SCCP.

[krzysz00]

Address review comments, move join/pessimistic stuff into wrapper class

[krzysz00]

Thanks, and I'm glad to hear this is all moving in a good direction

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

63–65 ↗

(On Diff #430446)

This becomes important for arith.constant and any other op we can think of whose return type might or might not be an integer. I think it'll improve code clarity if we keep this around as a helper method.

mlir/lib/Transforms/FoldInferredConstants.cpp

88–97 ↗

(On Diff #430446)

I agree that we might be able to merge this and SCCP, but I'd rather do that in a future commit once we've roped in the SCCP folks on here to make sure we've handled their use cases (ex - do they handle vector constants? We don't)

[krzysz00]

(and I did move the tuple methods over to the commit that defines InferIntRangeInterface for arithmetic)

[rriddle]

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

63–65 ↗

(On Diff #430446)

If the constant isn't an integer, can we just bail out? It seems like something to be handled outside of this class. For the interface I assume we would either:

• Not invoke the callback function
• Change the callback function to take an Optional<> and invoke it with None

What happens in the situations where an operation has one result that can have ranges inferred and one that doesn't? I assume we just wouldn't provide information for the result we can't infer, or am I wrong here? The case where constant doesn't return an int should fall into that same case (this isn't an int, so I'm not giving you anything).

mlir/lib/Transforms/FoldInferredConstants.cpp

88–97 ↗

(On Diff #430446)

Can we move this to being a test pass then? The premise of this pass is identical to SCCP, and I'd rather work to merge the necessary functionality into there than to rely on having a separate pass.

[krzysz00]

Integrate fold-inferred-constants into SCCP

[krzysz00]

Fix bug in dataflow analysis uncovered while testing arithmetic range implementations.

[krzysz00]

As to SCCP, I've gone and moved the invocation of the integer range analysis into SCCP. I couldn't merge it with the existing SCCP analysis due to their differences in how they detect loop-variant variables (the range-less SCCP was more conservative about it), but I did at least put both analyses in one pass.

I hope this change satisfies the objections you had.

(I also ran into a bug in how data flow analysis was handling terminators for stuff like scf.condition - this has now been fixed here.)

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

63–65 ↗

(On Diff #430446)

I've fixed the analysis so that failure to call the callback gives you [min(type), max(type)] and therefore removed notAnInteger(). Good call on how to clean all this up!

[krzysz00]

I will admit I'm not entirely happy with the SCCP integration - if you see any ways I could tighten it up, please do let me know.

[rriddle]

In D124023#3551660, @krzysz00 wrote:

I will admit I'm not entirely happy with the SCCP integration - if you see any ways I could tighten it up, please do let me know.

Ah sorry, didn't mean to imply you needed to do that in this patch. Can you split the SCCP out into a separate patch? Using a test pass for now to test the analysis? I'd like to get the base stuff landed, and I don't think that needs to be anchored on figuring out SCCP evolution.

[krzysz00]

I do want this functionality downstream for non-test reasons, though.

If it makes more sense, I can make -fold-inferred-constants a test pass in
the MLIR codebase and then temporarily pull it out of test/ in our fork
when I'm backporting this patch series. Would that work?

Or, if you don't think the SCCP evolution will be too contentious, I can
just split that into it's own patch and add it to the series - another few
days to make sure we get this right aren't the end of the world.

[rriddle]

In D124023#3553211, @krzysz00 wrote:

I do want this functionality downstream for non-test reasons, though.

If it makes more sense, I can make -fold-inferred-constants a test pass in
the MLIR codebase and then temporarily pull it out of test/ in our fork
when I'm backporting this patch series. Would that work?

Or, if you don't think the SCCP evolution will be too contentious, I can
just split that into it's own patch and add it to the series - another few
days to make sure we get this right aren't the end of the world.

Moving to test upstream, and then having a non-test pass downstream for now
would likely be the easiest. Given that you'll still be able to benefit from the work
you've done regardless of what we do with the SCCP integration.

[krzysz00]

Move inference-based folding to -test-int-range-inference

[krzysz00]

Rebase to more recent main, clean up remaining SCCP change

[rriddle]

Thanks! I really appreciate the patience during the review. This is really great work.

mlir/test/lib/Transforms/TestIntRangeInference.cpp

8 ↗

(On Diff #433828)

nit: Please drop usernames from TODO messages.

[krzysz00]

Adress final nit

River, Jeff, thank you so much for your thorough feedback and review!
This code wouldn't be as good as it is without y'all.

[mehdi_amini]

I revert, this broke the bot: https://lab.llvm.org/buildbot/#/builders/61/builds/27314

(you should have got an email)

[krzysz00]

I did get an email but the email only mentioned the warning, not the error.

I worked out the problem and put together a fix that I'll send in tomorrow.

[Mogball]

mlir/lib/Analysis/DataFlowAnalysis.cpp

579 ↗

(On Diff #433861)

https://github.com/llvm/llvm-project/issues/55873 was introduced by this code block here. I'm investigating it at the moment

mlir/include/mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h

+//===- IntRangeAnalysis.h - Infer Ranges Interfaces --*- 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 declares the dataflow analysis class for integer range inference
+// so that it can be used in transformations over the `arith` dialect such as
+// branch elimination or signed->unsigned rewriting
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITHMETIC_ANALYSIS_INTRANGEANALYSIS_H
+#define MLIR_DIALECT_ARITHMETIC_ANALYSIS_INTRANGEANALYSIS_H
+
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+
+namespace mlir {
+namespace arith {
+struct IntRangeAnalysis : public ForwardDataFlowAnalysis<IntRangeAttrs> {
+  using ForwardDataFlowAnalysis<IntRangeAttrs>::ForwardDataFlowAnalysis;
+
+  /// Define bounds on the results or block arguments of the operation
+  /// based on the bounds on the arguments given in `operands`
+  ChangeResult
+  visitOperation(Operation *op,
+                 ArrayRef<LatticeElement<IntRangeAttrs> *> operands) final;
+
+  /// Skip regions of branch ops when we can statically infer constant
+  /// values for operands to the branch op and said op tells us it's safe to do
+  /// so.
+  LogicalResult
+  getSuccessorsForOperands(BranchOpInterface branch,
+                           ArrayRef<LatticeElement<IntRangeAttrs> *> operands,
+                           SmallVectorImpl<Block *> &successors) final;
+
+  /// Skip regions of branch or loop ops when we can statically infer constant
+  /// values for operands to the branch op and said op tells us it's safe to do
+  /// so.
+  void
+  getSuccessorsForOperands(RegionBranchOpInterface branch,
+                           Optional<unsigned> sourceIndex,
+                           ArrayRef<LatticeElement<IntRangeAttrs> *> operands,
+                           SmallVectorImpl<RegionSuccessor> &successors) final;
+
+  /// Infer bounds on loop bounds
+  ChangeResult visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &region,
+      ArrayRef<LatticeElement<IntRangeAttrs> *> operands) final;
+};
+} // namespace arith
+} // namespace mlir
+
+#endif

mlir/include/mlir/Dialect/Arithmetic/IR/Arithmetic.h

 
 //===----------------------------------------------------------------------===//
 // Arithmetic Dialect Enum Attributes
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/ArithmeticOpsEnums.h.inc"
 
 //===----------------------------------------------------------------------===//
+// Arithmetic Dialect Interfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h"
+
+//===----------------------------------------------------------------------===//
 // Arithmetic Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Arithmetic/IR/ArithmeticOps.h.inc"
 
 namespace mlir {
 namespace arith {

mlir/include/mlir/Dialect/Arithmetic/IR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS ArithmeticOps.td)
 mlir_tablegen(ArithmeticOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(ArithmeticOpsEnums.cpp.inc -gen-enum-defs)
 add_mlir_dialect(ArithmeticOps arith)
 add_mlir_doc(ArithmeticOps ArithmeticOps Dialects/ -gen-dialect-doc)
+add_mlir_interface(InferIntRangeInterface)

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h

+//===- InferIntRangeInterface.h - Infer Ranges Interfaces --*- 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 definitions of the integer range inference interface
+// defined in `InferIntRange.td`
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE_H
+#define MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+#include <tuple>
+
+#include "llvm/ADT/APInt.h"
+
+namespace mlir {
+class Value;
+namespace arith {
+/// A collection of integer attributes representing minimum and maximum bounds
+/// on a scalar integer value. These bounds are inclusive on both ends, so
+/// bounds of [4, 5] mean 4 <= x <= 5. Separate bounds are tracked for
+/// the unsigned and signed interpretations of values.
+class IntRangeAttrs {
+  Optional<APInt> uminVal, umaxVal, sminVal, smaxVal;
+
+public:
+  IntRangeAttrs() : uminVal(), umaxVal(), sminVal(), smaxVal() {}
+  IntRangeAttrs(const Optional<APInt> &umin, const Optional<APInt> &umax,
+                const Optional<APInt> &smin, const Optional<APInt> &smax)
+      : uminVal(umin), umaxVal(umax), sminVal(smin), smaxVal(smax) {
+    // Unsigned integers that have upper bounds but no lower bound have a lower
+    // bound of 0.
+    if (umax && !umin) {
+      uminVal = APInt::getZero(umax->getBitWidth());
+    }
+  }
+
+  /// The minimum value of an integer when it is interpreted as unsigned.
+  const Optional<APInt> &umin() const;
+  /// The maximum value of an integer when it is interpreted as unsigned.
+  const Optional<APInt> &umax() const;
+  /// The minimum value of an integer when it is interpreted as signed;
+  const Optional<APInt> &smin() const;
+  /// The maximum value of an integer when it is interpreted as signed.
+  const Optional<APInt> &smax() const;
+
+  IntRangeAttrs(const std::tuple<Optional<APInt>, Optional<APInt>> &urange,
+                const std::tuple<Optional<APInt>, Optional<APInt>> &srange)
+      : IntRangeAttrs(std::get<0>(urange), std::get<1>(urange),
+                      std::get<0>(srange), std::get<1>(srange)) {}
+
+  /// Create an `IntRangeAttrs` where `min` is both the signed and unsigned
+  /// minimum and `max` is both the signed and unsigned maximum.
+  static IntRangeAttrs range(const Optional<APInt> &min,
+                             const Optional<APInt> &max);
+
+  /// Create an `IntRangeAttrs` with the signed minimum and maximum equal
+  /// to `smin` and `smax`, where the unsigned bounds are constructed from the
+  /// signed ones if the signed bounds correspond to a contiguous range of
+  /// values within [0, int_max] or are left unset otherwise.
+  static IntRangeAttrs fromSigned(const Optional<APInt> &smin,
+                                  const Optional<APInt> &smax);
+  static IntRangeAttrs
+  fromSigned(const std::tuple<Optional<APInt>, Optional<APInt>> &srange);
+
+  /// Create an `IntRangeAttrs` with the unsigned minimum and maximum equal
+  /// to `umin` and `umax` and the signed part equal to `umin` and `umax`
+  /// unless the sign bit changes between the minimum and maximum.
+  static IntRangeAttrs fromUnsigned(const Optional<APInt> &umin,
+                                    const Optional<APInt> &umax);
+  static IntRangeAttrs
+  fromUnsigned(const std::tuple<Optional<APInt>, Optional<APInt>> &urange);
+
+  /// Returns the union (computed separately for signed and unsigned bounds)
+  /// of `a` and `b`.
+  static IntRangeAttrs join(const IntRangeAttrs &a, const IntRangeAttrs &b);
+  static IntRangeAttrs getPessimisticValueState(MLIRContext *context);
+  static IntRangeAttrs getPessimisticValueState(Value v);
+
+  bool operator==(const IntRangeAttrs &other) const;
+
+  /// If either the signed or unsigned interpretations of the range
+  /// indicate that the value it bounds is a constant, return that constant
+  /// value.
+  Optional<APInt> getConstantValue() const;
+
+  friend raw_ostream &operator<<(raw_ostream &os, const IntRangeAttrs &range);
+};
+} // namespace arith
+} // namespace mlir
+
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h.inc"
+
+#endif // MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE_H

mlir/include/mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.td

+//===- InferIntRangeInterface.td - Integer Range Inference --*- 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
+//
+//===-----------------------------------------------------===//
+//
+// Defines the interface for range analysis on scalar integers
+//
+//===-----------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE
+#define MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def InferIntRangeInterface : OpInterface<"InferIntRangeInterface"> {
+  let description = [{
+    Allows operations to participate in range analysis for scalar integer values by
+    providing a methods that allows them to specify lower and upper bounds on their
+    result(s) given lower and upper bounds on their input(s) if known.
+  }];
+  let cppNamespace = "::mlir::arith";
+
+  let methods = [
+    InterfaceMethod<[{
+      Infer the bounds on the results of this op given the bounds on its arguments.
+      For each result value, the function should insert an `IntRangeAttrs`
+      representing the result's minimum and maximum value into `resultRanges`.
+      If a result does not have integer type, or if bounds can not be inferred,
+      the pair corresponding to it should be (nullptr, nullptr).
+
+      The returned attributes, if present, should have the same type as the
+      corresponding result.
+
+      This function will only be called when at least one result of the op is a
+      scalar integer value.
+
+      `argRanges` contains one `IntRangeAttrs` for each argument to the op in ODS
+       order. Non-integer arguments have `(nullptr, nullptr)` in their `argRanges`
+       element.
+    }],
+    "void", "inferResultRanges", (ins
+      "::llvm::ArrayRef<::mlir::arith::IntRangeAttrs>":$argRanges,
+      "::llvm::SmallVectorImpl<::mlir::arith::IntRangeAttrs> &":$resultRanges)
+  >];
+}
+#endif // MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE

mlir/include/mlir/Dialect/Arithmetic/Transforms/Passes.h

 std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
 
 /// Add patterns to expand Arithmetic ops for LLVM lowering.
 void populateArithmeticExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Create a pass to legalize Arithmetic ops for LLVM lowering.
 std::unique_ptr<Pass> createArithmeticExpandOpsPass();
 
+/// Create a pass to constant fold based on the results of range inferrence.
+std::unique_ptr<Pass> createArithmeticFoldInferredConstantsPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h.inc"
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Arithmetic/Transforms/Passes.td

def ArithmeticBufferize : Pass<"arith-bufferize", "ModuleOp"> {
   ];
 }
 
 def ArithmeticExpandOps : Pass<"arith-expand"> {
   let summary = "Legalize Arithmetic ops to be convertible to LLVM.";
   let constructor = "mlir::arith::createArithmeticExpandOpsPass()";
 }
 
+def ArithmeticFoldInferredConstants : Pass<"arith-fold-inferred-constants"> {
+  let summary = "Constant fold based on the results of integer range inference";
+  let constructor = "mlir::arith::createArithmeticFoldInferredConstantsPass()";
+}
+
 #endif // MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES

mlir/lib/Dialect/Arithmetic/Analysis/CMakeLists.txt

+add_mlir_dialect_library(MLIRArithmeticAnalysis
+  IntRangeAnalysis.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic
+
+  DEPENDS
+  mlir-headers
+
+  LINK_LIBS PUBLIC
+  MLIRArithmetic
+  MLIRAnalysis
+  MLIRControlFlowInterfaces
+  MLIRLoopLikeInterface
+  )

mlir/lib/Dialect/Arithmetic/Analysis/IntRangeAnalysis.cpp

+//===- IntRangeAnalysis.h - Infer Ranges Interfaces --*- 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 defines the dataflow analysis class for integer range inference
+// which is used in transformations over the `arith` dialect such as
+// branch elimination or signed->unsigned rewriting
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h"
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "int-range-analysis"
+
+namespace mlir {
+namespace arith {
+/// Given the ressults of getConstant{Lower,Upper}Bound()
+/// or getConstantStep() on a loopLikeInterface return the lower/upper bound for
+/// that result if possible.
+static Optional<APInt> getLoopBoundFromFold(Optional<OpFoldResult> loopBound,
+                                            IntRangeAnalysis &analysis,
+                                            bool getUpper) {
+  if (!loopBound.hasValue())
+    return {};
+  Optional<APInt> result;
+  if (loopBound->is<Attribute>()) {
+    if (auto bound =
+            loopBound->get<Attribute>().dyn_cast_or_null<IntegerAttr>())
+      result = bound.getValue();
+  } else if (loopBound->is<Value>()) {
+    LatticeElement<IntRangeAttrs> *lattice =
+        analysis.lookupLatticeElement(loopBound->get<Value>());
+    if (lattice != nullptr && !lattice->isUninitialized())
+      result =
+          getUpper ? lattice->getValue().smax() : lattice->getValue().smin();
+  }
+  // Loop bounds don't include the upper index, but integer range bounds do
+  if (result && getUpper) {
+    // Note: loops that don't execute (ex. %i = 0 to 0) will create bad bounds
+    // with this method, but they don't execute so it doesn't matter
+    result = *result - 1;
+  }
+  return result;
+}
+
+ChangeResult IntRangeAnalysis::visitOperation(
+    Operation *op, ArrayRef<LatticeElement<IntRangeAttrs> *> operands) {
+  ChangeResult ret = ChangeResult::NoChange;
+  // Ignore non-integer outputs - return early if the op has no scalar
+  // integer results
+  bool hasIntegerResult = false;
+  bool hasYieldedResult = false;
+  for (Value v : op->getResults()) {
+    if (v.getType().isIntOrIndex())
+      hasIntegerResult = true;
+    else
+      ret |= markAllPessimisticFixpoint(v);
+    for (Operation *user : v.getUsers())
+      hasYieldedResult |= user->hasTrait<OpTrait::IsTerminator>();
+  }
+  if (!hasIntegerResult)
+    return ret;
+  if (auto inferrable = dyn_cast<InferIntRangeInterface>(op)) {
+    LLVM_DEBUG(llvm::dbgs() << "Inferring ranges for ");
+    LLVM_DEBUG(inferrable->print(llvm::dbgs()));
+    LLVM_DEBUG(llvm::dbgs() << "\n");
+    SmallVector<IntRangeAttrs> argRanges(
+        llvm::map_range(operands, [](LatticeElement<IntRangeAttrs> *val) {
+          return val->getValue();
+        }));
+    SmallVector<IntRangeAttrs> resultRanges;
+    resultRanges.reserve(op->getNumResults());
+    inferrable.inferResultRanges(argRanges, resultRanges);
+    assert(resultRanges.size() == op->getNumResults() &&
+           "Range inference should provide one value per result");
+    for (auto pair : llvm::zip(op->getResults(), resultRanges)) {
+      LLVM_DEBUG(llvm::dbgs() << "Result range " << std::get<1>(pair) << "\n");
+      LatticeElement<IntRangeAttrs> &lattice =
+          getLatticeElement(std::get<0>(pair));
+      Optional<IntRangeAttrs> oldRange;
+      if (!lattice.isUninitialized())
+        oldRange = lattice.getValue();
+      ret |= lattice.join(std::get<1>(pair));
+      // Catch loop results with loop variant bounds and conservatively make
+      // them [-inf, inf] so we don't circle around infinitely often (because
+      // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+      // and often can't.
+      if (hasYieldedResult && oldRange.hasValue() &&
+          !(lattice.getValue() == *oldRange)) {
+        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+        ret |= lattice.join(IntRangeAttrs());
+        ret |= lattice.markPessimisticFixpoint();
+      }
+    }
+  } else if (op->getNumRegions() == 0) {
+    // No regions + no result inference method -> unbounded results (ex. memory
+    // ops)
+    ret |= markAllPessimisticFixpoint(op->getResults());
+  }
+  return ret;
+}
+
+LogicalResult IntRangeAnalysis::getSuccessorsForOperands(
+    BranchOpInterface branch,
+    ArrayRef<LatticeElement<IntRangeAttrs> *> operands,
+    SmallVectorImpl<Block *> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> mbConstValue =
+        enumPair.value()->getValue().getConstantValue();
+
+    if (mbConstValue)
+      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
+                              *mbConstValue);
+    return {};
+  };
+  SmallVector<Attribute> inferredConsts(
+      llvm::map_range(llvm::enumerate(operands), toConstantAttr));
+  if (Block *singleSucc = branch.getSuccessorForOperands(inferredConsts)) {
+    successors.push_back(singleSucc);
+    return success();
+  }
+  return failure();
+}
+
+void IntRangeAnalysis::getSuccessorsForOperands(
+    RegionBranchOpInterface branch, Optional<unsigned> sourceIndex,
+    ArrayRef<LatticeElement<IntRangeAttrs> *> operands,
+    SmallVectorImpl<RegionSuccessor> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> mbConstValue =
+        enumPair.value()->getValue().getConstantValue();
+
+    if (mbConstValue)
+      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
+                              *mbConstValue);
+    return {};
+  };
+  SmallVector<Attribute> inferredConsts(
+      llvm::map_range(llvm::enumerate(operands), toConstantAttr));
+  branch.getSuccessorRegions(sourceIndex, inferredConsts, successors);
+}
+
+ChangeResult IntRangeAnalysis::visitNonControlFlowArguments(
+    Operation *op, const RegionSuccessor &region,
+    ArrayRef<LatticeElement<IntRangeAttrs> *> operands) {
+  // Infer bounds for loop arguments that have static bounds
+  if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
+    Optional<Value> iv = loop.getSingleInductionVar();
+    if (!iv.hasValue())
+      return ForwardDataFlowAnalysis<
+          IntRangeAttrs>::visitNonControlFlowArguments(op, region, operands);
+    Optional<OpFoldResult> lowerBound = loop.getSingleLowerBound();
+    Optional<OpFoldResult> upperBound = loop.getSingleUpperBound();
+    Optional<OpFoldResult> step = loop.getSingleStep();
+    Optional<APInt> min =
+        getLoopBoundFromFold(lowerBound, *this, /*getUpper=*/false);
+    Optional<APInt> max =
+        getLoopBoundFromFold(upperBound, *this, /*getUpper=*/true);
+    Optional<APInt> stepVal =
+        getLoopBoundFromFold(step, *this, /*getUpper=*/false);
+    if (stepVal && stepVal.getValue().isNegative()) {
+      std::swap(min, max);
+      // Undo the upper bound correction we needed in the positive case.
+      min = *min + 1;
+    }
+
+    LatticeElement<IntRangeAttrs> &ivEntry = getLatticeElement(*iv);
+    return ivEntry.join(IntRangeAttrs::fromSigned(min, max));
+  }
+  return ForwardDataFlowAnalysis<IntRangeAttrs>::visitNonControlFlowArguments(
+      op, region, operands);
+}
+} // namespace arith
+} // namespace mlir

mlir/lib/Dialect/Arithmetic/CMakeLists.txt

+add_subdirectory(Analysis)
 add_subdirectory(IR)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

mlir/lib/Dialect/Arithmetic/IR/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS ArithmeticCanonicalization.td)
 mlir_tablegen(ArithmeticCanonicalization.inc -gen-rewriters)
 add_public_tablegen_target(MLIRArithmeticCanonicalizationIncGen)
 
 add_mlir_dialect_library(MLIRArithmetic
   ArithmeticOps.cpp
   ArithmeticDialect.cpp
+  InferIntRangeInterface.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic
 
   DEPENDS
   MLIRArithmeticOpsIncGen
+  MLIRInferIntRangeInterfaceIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   )

mlir/lib/Dialect/Arithmetic/IR/InferIntRangeInterface.cpp

+//===- InferIntRangeInterface.cpp -  Integer range inference interface ---===//
+//
+// 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/Arithmetic/IR/InferIntRangeInterface.h"
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.cpp.inc"
+#include "mlir/IR/BuiltinAttributes.h"
+
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+namespace arith {
+
+const Optional<APInt> &IntRangeAttrs::umin() const { return uminVal; }
+const Optional<APInt> &IntRangeAttrs::umax() const { return umaxVal; }
+const Optional<APInt> &IntRangeAttrs::smin() const { return sminVal; }
+const Optional<APInt> &IntRangeAttrs::smax() const { return smaxVal; }
+
+IntRangeAttrs IntRangeAttrs::range(const Optional<APInt> &min,
+                                   const Optional<APInt> &max) {
+  return {min, max, min, max};
+}
+
+IntRangeAttrs IntRangeAttrs::fromSigned(const Optional<APInt> &smin,
+                                        const Optional<APInt> &smax) {
+  Optional<APInt> umin, umax;
+  if (smin && smax) {
+    if (smin->isNonNegative() == smax->isNonNegative()) {
+      umin = smin->ult(*smax) ? smin : smax;
+      umax = smin->ugt(*smax) ? smin : smax;
+    }
+  } else if (smin) {
+    if (smin.getValue().isNonNegative()) {
+      umin = smin;
+      umax = APInt::getSignedMaxValue(smin->getBitWidth());
+    }
+  } else if (smax) {
+    if (smax.getValue().isNegative()) {
+      umin = smax;
+      umax = APInt::getSignedMinValue(smax->getBitWidth());
+    }
+  }
+  return {umin, umax, smin, smax};
+}
+
+IntRangeAttrs IntRangeAttrs::fromSigned(
+    const std::tuple<Optional<APInt>, Optional<APInt>> &srange) {
+  return IntRangeAttrs::fromSigned(std::get<0>(srange), std::get<1>(srange));
+}
+
+IntRangeAttrs IntRangeAttrs::fromUnsigned(const Optional<APInt> &umin,
+                                          const Optional<APInt> &umax) {
+  Optional<APInt> smin, smax;
+  // Note, unlike in fromSigned, we have no case for umin <= x < inf,
+  // since the upper bound is pontentially a negative number.
+  if (umin && umax) {
+    if (umin->isNonNegative() == umax->isNonNegative()) {
+      smin = umin->slt(*umax) ? umin : umax;
+      smax = umin->sgt(*umax) ? umin : umax;
+    }
+  } else if (umax) {
+    if (umax->isNonNegative()) {
+      smax = umax;
+      smin = APInt::getZero(umax->getBitWidth());
+    }
+  }
+  return {umin, umax, smin, smax};
+}
+
+IntRangeAttrs IntRangeAttrs::fromUnsigned(
+    const std::tuple<Optional<APInt>, Optional<APInt>> &urange) {
+  return IntRangeAttrs::fromUnsigned(std::get<0>(urange), std::get<1>(urange));
+}
+
+IntRangeAttrs IntRangeAttrs::join(const IntRangeAttrs &a,
+                                  const IntRangeAttrs &b) {
+  Optional<APInt> umin, umax, smin, smax;
+  if (a.umin() && b.umin())
+    umin = a.umin()->ult(*b.umin()) ? a.umin() : b.umin();
+  if (a.umax() && b.umax())
+    umax = a.umax()->ugt(*b.umax()) ? a.umax() : b.umax();
+  if (a.smin() && b.smin())
+    smin = a.smin()->slt(*b.smin()) ? a.smin() : b.smin();
+  if (a.smax() && b.smax())
+    smax = a.smax()->sgt(*b.smax()) ? a.smax() : b.smax();
+
+  return {umin, umax, smin, smax};
+}
+
+IntRangeAttrs IntRangeAttrs::getPessimisticValueState(MLIRContext *context) {
+  return {};
+}
+
+IntRangeAttrs IntRangeAttrs::getPessimisticValueState(Value v) {
+  return getPessimisticValueState(v.getContext());
+}
+
+Optional<APInt> IntRangeAttrs::getConstantValue() const {
+  if (umin() && umax() && umin() == umax())
+    return umin();
+  if (smin() && smax() && smin() == smax())
+    return smin();
+  return None;
+}
+
+bool IntRangeAttrs::operator==(const IntRangeAttrs &other) const {
+  return umin() == other.umin() && umax() == other.umax() &&
+         smin() == other.smin() && smax() == other.smax();
+}
+
+raw_ostream &operator<<(raw_ostream &os, const IntRangeAttrs &range) {
+  auto printValIfPresent = [&os](StringRef sign, const Optional<APInt> &val) {
+    if (val)
+      os << val.getValue();
+    else
+      os << sign << "inf";
+  };
+  os << "unsigned : [";
+  printValIfPresent("-", range.umin());
+  os << ", ";
+  printValIfPresent("+", range.umax());
+  os << "] signed : [";
+  printValIfPresent("-", range.smin());
+  os << ", ";
+  printValIfPresent("+", range.smax());
+  return os << "]";
+}
+} // namespace arith
+} // namespace mlir

mlir/lib/Dialect/Arithmetic/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArithmeticTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ExpandOps.cpp
+  FoldInferredConstants.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic/Transforms
 
   DEPENDS
   MLIRArithmeticTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
+  MLIRArithmeticAnalysis
   MLIRBufferization
   MLIRBufferizationTransforms
   MLIRIR
   MLIRMemRef
   MLIRPass
   MLIRTransforms
   )

mlir/lib/Dialect/Arithmetic/Transforms/FoldInferredConstants.cpp

+//===- FoldInferredConstants.cpp - Pass to materialize constants that can be
+// inferred by range 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Dialect/Arithmetic/Analysis/IntRangeAnalysis.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arithmetic/IR/InferIntRangeInterface.h"
+#include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Transforms/FoldUtils.h"
+
+using namespace mlir;
+using namespace mlir::arith;
+
+namespace {
+/// Patterend after mlir/lib/Transforms/SCCP.cpp
+LogicalResult replaceWithConstant(IntRangeAnalysis &analysis, OpBuilder &b,
+                                  Value value) {
+  LatticeElement<IntRangeAttrs> *mbInferredRange =
+      analysis.lookupLatticeElement(value);
+  if (!mbInferredRange)
+    return failure();
+  const IntRangeAttrs &inferredRange = mbInferredRange->getValue();
+  Optional<APInt> mbConstValue = inferredRange.getConstantValue();
+  if (!mbConstValue.hasValue())
+    return failure();
+  Value constant = b.createOrFold<ConstantOp>(
+      value.getLoc(), IntegerAttr::get(value.getType(), *mbConstValue));
+  value.replaceAllUsesWith(constant);
+  return success();
+}
+
+void rewrite(IntRangeAnalysis &analysis, MLIRContext *context,
+             MutableArrayRef<Region> initialRegions) {
+  SmallVector<Block *> worklist;
+  auto addToWorklist = [&](MutableArrayRef<Region> regions) {
+    for (Region &region : regions)
+      for (Block &block : llvm::reverse(region))
+        worklist.push_back(&block);
+  };
+
+  OpBuilder builder(context);
+
+  addToWorklist(initialRegions);
+  while (!worklist.empty()) {
+    Block *block = worklist.pop_back_val();
+
+    for (Operation &op : llvm::make_early_inc_range(*block)) {
+      builder.setInsertionPoint(&op);
+
+      // Replace any result with constants.
+      bool replacedAll = op.getNumResults() != 0;
+      for (Value res : op.getResults())
+        replacedAll &= succeeded(replaceWithConstant(analysis, builder, res));
+
+      // If all of the results of the operation were replaced, try to erase
+      // the operation completely.
+      if (replacedAll && wouldOpBeTriviallyDead(&op)) {
+        assert(op.use_empty() && "expected all uses to be replaced");
+        op.erase();
+        continue;
+      }
+
+      // Add any the regions of this operation to the worklist.
+      addToWorklist(op.getRegions());
+    }
+
+    // Replace any block arguments with constants.
+    builder.setInsertionPointToStart(block);
+    for (BlockArgument arg : block->getArguments())
+      (void)replaceWithConstant(analysis, builder, arg);
+  }
+}
+
+struct ArithmeticFoldInferredConstantsPass
+    : public ArithmeticFoldInferredConstantsBase<
+          ArithmeticFoldInferredConstantsPass> {
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+
+    IntRangeAnalysis analysis(ctx);
+    analysis.run(op);
+    rewrite(analysis, ctx, op->getRegions());
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::arith::createArithmeticFoldInferredConstantsPass() {
+  return std::make_unique<ArithmeticFoldInferredConstantsPass>();
+}