[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

46–47 ↗

(On Diff #425014)

58–59 ↗

(On Diff #425014)

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

93–94 ↗

(On Diff #425014)

118–124 ↗

(On Diff #425014)

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

119–123 ↗

(On Diff #425014)

?

137–140 ↗

(On Diff #425014)

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

10 ↗

(On Diff #425014)

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.

24 ↗

(On Diff #425014)

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

29 ↗

(On Diff #425014)

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

44 ↗

(On Diff #424573)

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 ↗

(On Diff #425014)

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

10 ↗

(On Diff #425014)

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

75–82 ↗

(On Diff #424573)

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

40 ↗

(On Diff #425014)

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

[krzysz00]

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

10 ↗

(On Diff #425014)

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

75–82 ↗

(On Diff #424573)

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

10 ↗

(On Diff #425014)

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

29 ↗

(On Diff #425014)

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

29 ↗

(On Diff #425014)

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

58–59 ↗

(On Diff #425014)

Unresolved?

93–94 ↗

(On Diff #425014)

Unresolved?

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

16 ↗

(On Diff #426780)

Same comments about namespaces here.

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

25 ↗

(On Diff #426780)

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

19

Put the includes together

31

These should be placed at the bottom of the class

34

Brief comments on the default constructor and other constructors?

mlir/lib/Analysis/IntRangeAnalysis.cpp

25

Typo

26

41

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

107

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

118

Please expand the name to result

126

Please expand mb

170

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

173

192

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

220

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

29 ↗

(On Diff #428189)

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

68

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

19

+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

237

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

26

43

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

101

mlir/lib/Analysis/IntRangeAnalysis.cpp

58–60

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

25–26

This class can just go in the anonymous namespace

106

The operand lattices should never be uninitialized.

149

braces

170

braces

223

braces

235

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

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

1407

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

2895–2898

2919–2923

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

2–3

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

29 ↗

(On Diff #425014)

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

25–26

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

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

2–3

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

31–32

34–36

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

37

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

23

This doesn't look necessary.

29

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

33

Please add newlines between methods, applies throughout the commit.

40

Please drop trivial braces.

46

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

72

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

Same for the other method.

99

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.

102

mlir/lib/Analysis/IntRangeAnalysis.cpp

54

Can you complete this comment?

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

2901–2903

2924–2926

2938–2940

[krzysz00]

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

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

46

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.

72

Utility methods that use tuples to return multiple values.

99

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

2

This looks a bit off.

19–21

Are all of these necessary?

32–36

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

22

attributes? This comment looks out of date.

46

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.

52

62

64–66

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.

78–79

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

94–96

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

108

nit: Can you hoist this near the constructors?

117

Can you document this?

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

64–66

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

64–66

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

64–66

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

9

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

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/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_ANALYSIS_INTRANGEANALYSIS_H
+#define MLIR_ANALYSIS_INTRANGEANALYSIS_H
+
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+
+namespace mlir {
+namespace detail {
+class IntRangeAnalysisImpl;
+} // end namespace detail
+
+class IntRangeAnalysis {
+public:
+  /// Analyze all operations rooted under (but not including)
+  /// `topLevelOperation`.
+  IntRangeAnalysis(Operation *topLevelOperation);
+  IntRangeAnalysis(IntRangeAnalysis &&other);
+  ~IntRangeAnalysis();
+
+  /// Get inferred range for value `v` if one exists.
+  Optional<ConstantIntRanges> getResult(Value v);
+
+private:
+  std::unique_ptr<detail::IntRangeAnalysisImpl> impl;
+};
+} // end namespace mlir
+
+#endif

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(CastInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
+add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(SideEffectInterfaces)
 add_mlir_interface(TilingInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
 set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)

mlir/include/mlir/Interfaces/InferIntRangeInterface.h

+//===- InferIntRangeInterface.h - Integer Range Inference --*- 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_INTERFACES_INFERINTRANGEINTERFACE_H
+#define MLIR_INTERFACES_INFERINTRANGEINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+/// A set of arbitrary-precision integers representing bounds on a given 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 in order to enable more
+/// precice inference of the interplay between operations with signed and
+/// unsigned semantics.
+class ConstantIntRanges {
+public:
+  /// Bound umin <= (unsigned)x <= umax and smin <= signed(x) <= smax.
+  /// Non-integer values should be bounded by APInts of bitwidth 0.
+  ConstantIntRanges(const APInt &umin, const APInt &umax, const APInt &smin,
+                    const APInt &smax)
+      : uminVal(umin), umaxVal(umax), sminVal(smin), smaxVal(smax) {
+    assert(uminVal.getBitWidth() == umaxVal.getBitWidth() &&
+           umaxVal.getBitWidth() == sminVal.getBitWidth() &&
+           sminVal.getBitWidth() == smaxVal.getBitWidth() &&
+           "All bounds in the ranges must have the same bitwidth");
+  }
+
+  bool operator==(const ConstantIntRanges &other) const;
+
+  /// The minimum value of an integer when it is interpreted as unsigned.
+  const APInt &umin() const;
+
+  /// The maximum value of an integer when it is interpreted as unsigned.
+  const APInt &umax() const;
+
+  /// The minimum value of an integer when it is interpreted as signed.
+  const APInt &smin() const;
+
+  /// The maximum value of an integer when it is interpreted as signed.
+  const APInt &smax() const;
+
+  /// Return the bitwidth that should be used for integer ranges describing
+  /// `type`. For concrete integer types, this is their bitwidth, for `index`,
+  /// this is the internal storage bitwidth of `index` attributes, and for
+  /// non-integer types this is 0.
+  static unsigned getStorageBitwidth(Type type);
+
+  /// Create an `IntRangeAttrs` where `min` is both the signed and unsigned
+  /// minimum and `max` is both the signed and unsigned maximum.
+  static ConstantIntRanges range(const APInt &min, const 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 they correspond to a contigious range of bit patterns when
+  /// viewed as unsigned values and are left at [0, int_max()] otherwise.
+  static ConstantIntRanges fromSigned(const APInt &smin, const APInt &smax);
+
+  /// 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 ConstantIntRanges fromUnsigned(const APInt &umin, const APInt &umax);
+
+  /// Returns the union (computed separately for signed and unsigned bounds)
+  /// of `a` and `b`.
+  ConstantIntRanges rangeUnion(const ConstantIntRanges &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 ConstantIntRanges &range);
+
+private:
+  APInt uminVal, umaxVal, sminVal, smaxVal;
+};
+
+/// The type of the `setResultRanges` callback provided to ops implementing
+/// InferIntRangeInterface. It should be called once for each integer result
+/// value and be passed the ConstantIntRanges corresponding to that value.
+using SetIntRangeFn = function_ref<void(Value, const ConstantIntRanges &)>;
+} // end namespace mlir
+
+#include "mlir/Interfaces/InferIntRangeInterface.h.inc"
+
+#endif // MLIR_INTERFACES_INFERINTRANGEINTERFACE_H

mlir/include/mlir/Interfaces/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_INTERFACES_INFERINTRANGEINTERFACE
+#define MLIR_INTERFACES_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";
+
+  let methods = [
+    InterfaceMethod<[{
+      Infer the bounds on the results of this op given the bounds on its arguments.
+      For each result value or block argument (that isn't a branch argument,
+      since the dataflow analysis handles those case), the method should call
+      `setValueRange` with that `Value` as an argument. When `setValueRange`
+      is not called for some value, it will recieve a default value of the mimimum
+      and maximum values forits type (the unbounded range).
+
+      When called on an op that also implements the RegionBranchOpInterface
+      or BranchOpInterface, this method should not attempt to infer the values
+      of the branch results, as this will be handled by the analyses that use
+      this interface.
+
+      This function will only be called when at least one result of the op is a
+      scalar integer value or the op has a region.
+
+      `argRanges` contains one `IntRangeAttrs` for each argument to the op in ODS
+       order. Non-integer arguments will have the an unbounded range of width-0
+       APInts in their `argRanges` element.
+    }],
+    "void", "inferResultRanges", (ins
+      "::llvm::ArrayRef<::mlir::ConstantIntRanges>":$argRanges,
+      "::mlir::SetIntRangeFn":$setResultRanges)
+  >];
+}
+#endif // MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
+  IntRangeAnalysis.cpp
   Liveness.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
   )
 
 add_mlir_library(MLIRAnalysis
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
+  IntRangeAnalysis.cpp
   Liveness.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
+  MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Presburger)

mlir/lib/Analysis/DataFlowAnalysis.cpp

void ForwardDataFlowSolver::visitOperation(Operation *op) {
 
   // Process region holding operations.
   if (op->getNumRegions()) {
     // Check to see if we can reason about the internal control flow of this
     // region operation.
     if (auto branch = dyn_cast<RegionBranchOpInterface>(op))
       return visitRegionBranchOperation(branch, operandLattices);
 
-    // If we can't, conservatively mark all regions as executable.
-    // TODO: Let the `visitOperation` method decide how to propagate
-    // information to the block arguments.
-    for (Region &region : op->getRegions())
-      markEntryBlockExecutable(&region, /*markPessimisticFixpoint=*/true);
+    for (Region &region : op->getRegions()) {
+      analysis.visitNonControlFlowArguments(op, RegionSuccessor(&region),
+                                            operandLattices);
+      // `visitNonControlFlowArguments` is required to define all of the region
+      // argument lattices.
+      assert(llvm::none_of(
+                 region.getArguments(),
+                 [&](Value value) {
+                   return analysis.getLatticeElement(value).isUninitialized();
+                 }) &&
+             "expected `visitNonControlFlowArguments` to define all argument "
+             "lattices");
+      markEntryBlockExecutable(&region, /*markPessimisticFixpoint=*/false);
+    }
   }
 
   // If this op produces no results, it can't produce any constants.
   if (op->getNumResults() == 0)
     return;
 
   // If all of the results of this operation are already resolved, bail out
   // early.
if (op->getNumSuccessors() == 0) {
     if (isa<CallableOpInterface>(parentOp))
       return visitCallableTerminatorOperation(parentOp, op, operandLattices);
 
     // Otherwise, check to see if the parent tracks region control flow.
     auto regionInterface = dyn_cast<RegionBranchOpInterface>(parentOp);
     if (!regionInterface || !isBlockExecutable(parentOp->getBlock()))
       return;
 
+    // If the branch is a RegionBranchTerminatorOpInterface,
+    // construct the set of operand lattices as the set of non control-flow
+    // arguments of the parent and the values this op returns. This allows
+    // for the correct lattices to be passed to getSuccessorsForOperands()
+    // in cases such as scf.while.
+    ArrayRef<AbstractLatticeElement *> branchOpLattices = operandLattices;
+    SmallVector<AbstractLatticeElement *, 0> parentLattices;
+    if (auto regionTerminator =
+            dyn_cast<RegionBranchTerminatorOpInterface>(op)) {
+      parentLattices.reserve(regionInterface->getNumOperands());
+      for (Value parentOperand : regionInterface->getOperands()) {
+        AbstractLatticeElement *operandLattice =
+            analysis.lookupLatticeElement(parentOperand);
+        if (!operandLattice || operandLattice->isUninitialized())
+          return;
+        parentLattices.push_back(operandLattice);
+      }
+      unsigned regionNumber = parentRegion->getRegionNumber();
+      OperandRange iterArgs =
+          regionInterface.getSuccessorEntryOperands(regionNumber);
+      OperandRange terminatorArgs =
+          regionTerminator.getSuccessorOperands(regionNumber);
+      assert(iterArgs.size() == terminatorArgs.size() &&
+             "Number of iteration arguments for region should equal number of "
+             "those arguments defined by terminator");
+      if (!iterArgs.empty()) {
+        unsigned iterStart = iterArgs.getBeginOperandIndex();
+        unsigned terminatorStart = terminatorArgs.getBeginOperandIndex();
+        for (unsigned i = 0, e = iterArgs.size(); i < e; ++i)
+          parentLattices[iterStart + i] = operandLattices[terminatorStart + i];
+      }
+      branchOpLattices = parentLattices;
+    }
     // Query the set of successors of the current region using the current
     // optimistic lattice state.
     SmallVector<RegionSuccessor, 1> regionSuccessors;
     analysis.getSuccessorsForOperands(regionInterface,
                                       parentRegion->getRegionNumber(),
-                                      operandLattices, regionSuccessors);
+                                      branchOpLattices, regionSuccessors);
     if (regionSuccessors.empty())
       return;
 
     // Try to get "region-like" successor operands if possible in order to
     // propagate the operand states to the successors.
     if (isRegionReturnLike(op)) {
       auto getOperands = [&](Optional<unsigned> regionIndex) {
         // Determine the individual region  successor operands for the given
         // region index (if any).
         return *getRegionBranchSuccessorOperands(op, regionIndex);
       };
-      return visitRegionSuccessors(parentOp, regionSuccessors, operandLattices,
+      return visitRegionSuccessors(parentOp, regionSuccessors, branchOpLattices,
                                    getOperands);
     }
 
     // If this terminator is not "region-like", conservatively mark all of the
     // successor values as having reached the pessimistic fixpoint.
     for (auto &it : regionSuccessors) {
       // If the successor is a region, mark the entry block as executable so
       // that we visit operations defined within. If the successor is the

mlir/lib/Analysis/IntRangeAnalysis.cpp

+//===- IntRangeAnalysis.cpp - 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/Analysis/IntRangeAnalysis.h"
+#include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "int-range-analysis"
+
+using namespace mlir;
+
+namespace {
+/// A wrapper around ConstantIntRanges that provides the lattice functions
+/// expected by dataflow analysis.
+struct IntRangeLattice {
+  IntRangeLattice(const ConstantIntRanges &value) : value(value){};
+  IntRangeLattice(ConstantIntRanges &&value) : value(value){};
+
+  bool operator==(const IntRangeLattice &other) const {
+    return value == other.value;
+  }
+
+  /// wrapper around rangeUnion()
+  static IntRangeLattice join(const IntRangeLattice &a,
+                              const IntRangeLattice &b) {
+    return a.value.rangeUnion(b.value);
+  }
+
+  /// Creates a range with bitwidth 0 to represent that we don't know if the
+  /// value being marked overdefined is even an integer.
+  static IntRangeLattice getPessimisticValueState(MLIRContext *context) {
+    APInt noIntValue = APInt::getZeroWidth();
+    return ConstantIntRanges::range(noIntValue, noIntValue);
+  }
+
+  /// Create a maximal range ([0, uint_max(t)] / [int_min(t), int_max(t)])
+  /// range that is used to mark the value v as unable to be analyzed further,
+  /// where t is the type of v.
+  static IntRangeLattice getPessimisticValueState(Value v) {
+    unsigned int width = ConstantIntRanges::getStorageBitwidth(v.getType());
+    APInt umin = APInt::getMinValue(width);
+    APInt umax = APInt::getMaxValue(width);
+    APInt smin = width != 0 ? APInt::getSignedMinValue(width) : umin;
+    APInt smax = width != 0 ? APInt::getSignedMaxValue(width) : umax;
+    return ConstantIntRanges{umin, umax, smin, smax};
+  }
+
+  ConstantIntRanges value;
+};
+} // end anonymous namespace
+
+namespace mlir {
+namespace detail {
+class IntRangeAnalysisImpl : public ForwardDataFlowAnalysis<IntRangeLattice> {
+  using ForwardDataFlowAnalysis<IntRangeLattice>::ForwardDataFlowAnalysis;
+
+public:
+  /// 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<IntRangeLattice> *> 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<IntRangeLattice> *> 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<IntRangeLattice> *> operands,
+                           SmallVectorImpl<RegionSuccessor> &successors) final;
+
+  /// Call the InferIntRangeInterface implementation for region-using ops
+  /// that implement it, and infer the bounds of loop induction variables
+  /// for ops that implement LoopLikeOPInterface.
+  ChangeResult visitNonControlFlowArguments(
+      Operation *op, const RegionSuccessor &region,
+      ArrayRef<LatticeElement<IntRangeLattice> *> operands) final;
+};
+} // end namespace detail
+} // end namespace mlir
+
+/// Given the results of getConstant{Lower,Upper}Bound()
+/// or getConstantStep() on a LoopLikeInterface return the lower/upper bound for
+/// that result if possible.
+static APInt getLoopBoundFromFold(Optional<OpFoldResult> loopBound,
+                                  Type boundType,
+                                  detail::IntRangeAnalysisImpl &analysis,
+                                  bool getUpper) {
+  unsigned int width = ConstantIntRanges::getStorageBitwidth(boundType);
+  if (loopBound.hasValue()) {
+    if (loopBound->is<Attribute>()) {
+      if (auto bound =
+              loopBound->get<Attribute>().dyn_cast_or_null<IntegerAttr>())
+        return bound.getValue();
+    } else if (loopBound->is<Value>()) {
+      LatticeElement<IntRangeLattice> *lattice =
+          analysis.lookupLatticeElement(loopBound->get<Value>());
+      if (lattice != nullptr)
+        return getUpper ? lattice->getValue().value.smax()
+                        : lattice->getValue().value.smin();
+    }
+  }
+  return getUpper ? APInt::getSignedMaxValue(width)
+                  : APInt::getSignedMinValue(width);
+}
+
+ChangeResult detail::IntRangeAnalysisImpl::visitOperation(
+    Operation *op, ArrayRef<LatticeElement<IntRangeLattice> *> operands) {
+  ChangeResult result = ChangeResult::NoChange;
+  // Ignore non-integer outputs - return early if the op has no scalar
+  // integer results
+  bool hasIntegerResult = false;
+  for (Value v : op->getResults()) {
+    if (v.getType().isIntOrIndex())
+      hasIntegerResult = true;
+    else
+      result |= markAllPessimisticFixpoint(v);
+  }
+  if (!hasIntegerResult)
+    return result;
+
+  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<ConstantIntRanges> argRanges(
+        llvm::map_range(operands, [](LatticeElement<IntRangeLattice> *val) {
+          return val->getValue().value;
+        }));
+
+    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(v);
+      Optional<IntRangeLattice> oldRange;
+      if (!lattice.isUninitialized())
+        oldRange = lattice.getValue();
+      result |= lattice.join(IntRangeLattice(attrs));
+
+      // 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).
+      bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
+        return op->hasTrait<OpTrait::IsTerminator>();
+      });
+      if (isYieldedResult && oldRange.hasValue() &&
+          !(lattice.getValue() == *oldRange)) {
+        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+        result |= lattice.markPessimisticFixpoint();
+      }
+    };
+
+    inferrable.inferResultRanges(argRanges, joinCallback);
+    for (Value opResult : op->getResults()) {
+      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(opResult);
+      // setResultRange() not called, make pessimistic.
+      if (lattice.isUninitialized())
+        result |= lattice.markPessimisticFixpoint();
+    }
+  } else if (op->getNumRegions() == 0) {
+    // No regions + no result inference method -> unbounded results (ex. memory
+    // ops)
+    result |= markAllPessimisticFixpoint(op->getResults());
+  }
+  return result;
+}
+
+LogicalResult detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
+    BranchOpInterface branch,
+    ArrayRef<LatticeElement<IntRangeLattice> *> operands,
+    SmallVectorImpl<Block *> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> maybeConstValue =
+        enumPair.value()->getValue().value.getConstantValue();
+
+    if (maybeConstValue) {
+      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
+                              *maybeConstValue);
+    }
+    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 detail::IntRangeAnalysisImpl::getSuccessorsForOperands(
+    RegionBranchOpInterface branch, Optional<unsigned> sourceIndex,
+    ArrayRef<LatticeElement<IntRangeLattice> *> operands,
+    SmallVectorImpl<RegionSuccessor> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> maybeConstValue =
+        enumPair.value()->getValue().value.getConstantValue();
+
+    if (maybeConstValue) {
+      return IntegerAttr::get(branch->getOperand(enumPair.index()).getType(),
+                              *maybeConstValue);
+    }
+    return {};
+  };
+  SmallVector<Attribute> inferredConsts(
+      llvm::map_range(llvm::enumerate(operands), toConstantAttr));
+  branch.getSuccessorRegions(sourceIndex, inferredConsts, successors);
+}
+
+ChangeResult detail::IntRangeAnalysisImpl::visitNonControlFlowArguments(
+    Operation *op, const RegionSuccessor &region,
+    ArrayRef<LatticeElement<IntRangeLattice> *> operands) {
+  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<ConstantIntRanges> argRanges(
+        llvm::map_range(operands, [](LatticeElement<IntRangeLattice> *val) {
+          return val->getValue().value;
+        }));
+
+    ChangeResult result = ChangeResult::NoChange;
+    auto joinCallback = [&](Value v, const ConstantIntRanges &attrs) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(v);
+      Optional<IntRangeLattice> oldRange;
+      if (!lattice.isUninitialized())
+        oldRange = lattice.getValue();
+      result |= lattice.join(IntRangeLattice(attrs));
+
+      // 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).
+      bool isYieldedValue = llvm::any_of(v.getUsers(), [](Operation *op) {
+        return op->hasTrait<OpTrait::IsTerminator>();
+      });
+      if (isYieldedValue && oldRange.hasValue() &&
+          !(lattice.getValue() == *oldRange)) {
+        LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+        result |= lattice.markPessimisticFixpoint();
+      }
+    };
+
+    inferrable.inferResultRanges(argRanges, joinCallback);
+    for (Value regionArg : region.getSuccessor()->getArguments()) {
+      LatticeElement<IntRangeLattice> &lattice = getLatticeElement(regionArg);
+      // setResultRange() not called, make pessimistic.
+      if (lattice.isUninitialized())
+        result |= lattice.markPessimisticFixpoint();
+    }
+
+    return result;
+  }
+
+  // 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<
+          IntRangeLattice>::visitNonControlFlowArguments(op, region, operands);
+    }
+    Optional<OpFoldResult> lowerBound = loop.getSingleLowerBound();
+    Optional<OpFoldResult> upperBound = loop.getSingleUpperBound();
+    Optional<OpFoldResult> step = loop.getSingleStep();
+    APInt min = getLoopBoundFromFold(lowerBound, iv->getType(), *this,
+                                     /*getUpper=*/false);
+    APInt max = getLoopBoundFromFold(upperBound, iv->getType(), *this,
+                                     /*getUpper=*/true);
+    // Assume positivity for uniscoverable steps by way of getUpper = true.
+    APInt stepVal =
+        getLoopBoundFromFold(step, iv->getType(), *this, /*getUpper=*/true);
+
+    if (stepVal.isNegative()) {
+      std::swap(min, max);
+    } else {
+      // Correct the upper bound by subtracting 1 so that it becomes a <= bound,
+      // because loops do not generally include their upper bound.
+      max -= 1;
+    }
+
+    LatticeElement<IntRangeLattice> &ivEntry = getLatticeElement(*iv);
+    return ivEntry.join(ConstantIntRanges::fromSigned(min, max));
+  }
+  return ForwardDataFlowAnalysis<IntRangeLattice>::visitNonControlFlowArguments(
+      op, region, operands);
+}
+
+IntRangeAnalysis::IntRangeAnalysis(Operation *topLevelOperation) {
+  impl = std::make_unique<mlir::detail::IntRangeAnalysisImpl>(
+      topLevelOperation->getContext());
+  impl->run(topLevelOperation);
+}
+
+IntRangeAnalysis::~IntRangeAnalysis() = default;
+IntRangeAnalysis::IntRangeAnalysis(IntRangeAnalysis &&other) = default;
+
+Optional<ConstantIntRanges> IntRangeAnalysis::getResult(Value v) {
+  LatticeElement<IntRangeLattice> *result = impl->lookupLatticeElement(v);
+  if (result == nullptr || result->isUninitialized())
+    return llvm::None;
+  return result->getValue().value;
+}

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   CastInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DataLayoutInterfaces.cpp
   DerivedAttributeOpInterface.cpp
+  InferIntRangeInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   SideEffectInterfaces.cpp
   TilingInterface.cpp
   VectorInterfaces.cpp
   ViewLikeInterface.cpp
   )
 

 
 
 add_mlir_interface_library(CallInterfaces)
 add_mlir_interface_library(CastInterfaces)
 add_mlir_interface_library(ControlFlowInterfaces)
 add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DataLayoutInterfaces)
 add_mlir_interface_library(DerivedAttributeOpInterface)
+add_mlir_interface_library(InferIntRangeInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 add_mlir_interface_library(LoopLikeInterface)

mlir/lib/Interfaces/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/Interfaces/InferIntRangeInterface.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/InferIntRangeInterface.cpp.inc"
+
+using namespace mlir;
+
+bool ConstantIntRanges::operator==(const ConstantIntRanges &other) const {
+  return umin().getBitWidth() == other.umin().getBitWidth() &&
+         umin() == other.umin() && umax() == other.umax() &&
+         smin() == other.smin() && smax() == other.smax();
+}
+
+const APInt &ConstantIntRanges::umin() const { return uminVal; }
+
+const APInt &ConstantIntRanges::umax() const { return umaxVal; }
+
+const APInt &ConstantIntRanges::smin() const { return sminVal; }
+
+const APInt &ConstantIntRanges::smax() const { return smaxVal; }
+
+unsigned ConstantIntRanges::getStorageBitwidth(Type type) {
+  if (type.isIndex())
+    return IndexType::kInternalStorageBitWidth;
+  if (auto integerType = type.dyn_cast<IntegerType>())
+    return integerType.getWidth();
+  // Non-integer types have their bounds stored in width 0 `APInt`s.
+  return 0;
+}
+
+ConstantIntRanges ConstantIntRanges::range(const APInt &min, const APInt &max) {
+  return {min, max, min, max};
+}
+
+ConstantIntRanges ConstantIntRanges::fromSigned(const APInt &smin,
+                                                const APInt &smax) {
+  unsigned int width = smin.getBitWidth();
+  APInt umin, umax;
+  if (smin.isNonNegative() == smax.isNonNegative()) {
+    umin = smin.ult(smax) ? smin : smax;
+    umax = smin.ugt(smax) ? smin : smax;
+  } else {
+    umin = APInt::getMinValue(width);
+    umax = APInt::getMaxValue(width);
+  }
+  return {umin, umax, smin, smax};
+}
+
+ConstantIntRanges ConstantIntRanges::fromUnsigned(const APInt &umin,
+                                                  const APInt &umax) {
+  unsigned int width = umin.getBitWidth();
+  APInt smin, smax;
+  if (umin.isNonNegative() == umax.isNonNegative()) {
+    smin = umin.slt(umax) ? umin : umax;
+    smax = umin.sgt(umax) ? umin : umax;
+  } else {
+    smin = APInt::getSignedMinValue(width);
+    smax = APInt::getSignedMaxValue(width);
+  }
+  return {umin, umax, smin, smax};
+}
+
+ConstantIntRanges
+ConstantIntRanges::rangeUnion(const ConstantIntRanges &other) const {
+  // "Not an integer" poisons everything and also cannot be fed to comparison
+  // operators.
+  if (umin().getBitWidth() == 0)
+    return *this;
+  if (other.umin().getBitWidth() == 0)
+    return other;
+
+  const APInt &uminUnion = umin().ult(other.umin()) ? umin() : other.umin();
+  const APInt &umaxUnion = umax().ugt(other.umax()) ? umax() : other.umax();
+  const APInt &sminUnion = smin().slt(other.smin()) ? smin() : other.smin();
+  const APInt &smaxUnion = smax().sgt(other.smax()) ? smax() : other.smax();
+
+  return {uminUnion, umaxUnion, sminUnion, smaxUnion};
+}
+
+Optional<APInt> ConstantIntRanges::getConstantValue() const {
+  // Note: we need to exclude the trivially-equal width 0 values here.
+  if (umin() == umax() && umin().getBitWidth() != 0)
+    return umin();
+  if (smin() == smax() && smin().getBitWidth() != 0)
+    return smin();
+  return None;
+}
+
+raw_ostream &mlir::operator<<(raw_ostream &os, const ConstantIntRanges &range) {
+  return os << "unsigned : [" << range.umin() << ", " << range.umax()
+            << "] signed : [" << range.smin() << ", " << range.smax() << "]";
+}

mlir/test/Interfaces/InferIntRangeInterface/infer-int-range-test-ops.mlir

+// RUN: mlir-opt -test-int-range-inference %s | FileCheck %s
+
+// CHECK-LABEL: func @constant
+// CHECK: %[[cst:.*]] = "test.constant"() {value = 3 : index}
+// CHECK: return %[[cst]]
+func.func @constant() -> index {
+  %0 = test.with_bounds { umin = 3 : index, umax = 3 : index,
+                               smin = 3 : index, smax = 3 : index}
+  func.return %0 : index
+}
+
+// CHECK-LABEL: func @increment
+// CHECK: %[[cst:.*]] = "test.constant"() {value = 4 : index}
+// CHECK: return %[[cst]]
+func.func @increment() -> index {
+  %0 = test.with_bounds { umin = 3 : index, umax = 3 : index, smin = 0 : index, smax = 0x7fffffffffffffff : index }
+  %1 = test.increment %0
+  func.return %1 : index
+}
+
+// CHECK-LABEL: func @maybe_increment
+// CHECK: test.reflect_bounds {smax = 4 : index, smin = 3 : index, umax = 4 : index, umin = 3 : index}
+func.func @maybe_increment(%arg0 : i1) -> index {
+  %0 = test.with_bounds { umin = 3 : index, umax = 3 : index,
+                               smin = 3 : index, smax = 3 : index}
+  %1 = scf.if %arg0 -> index {
+    scf.yield %0 : index
+  } else {
+    %2 = test.increment %0
+    scf.yield %2 : index
+  }
+  %3 = test.reflect_bounds %1
+  func.return %3 : index
+}
+
+// CHECK-LABEL: func @maybe_increment_br
+// CHECK: test.reflect_bounds {smax = 4 : index, smin = 3 : index, umax = 4 : index, umin = 3 : index}
+func.func @maybe_increment_br(%arg0 : i1) -> index {
+  %0 = test.with_bounds { umin = 3 : index, umax = 3 : index,
+                               smin = 3 : index, smax = 3 : index}
+  cf.cond_br %arg0, ^bb0, ^bb1
+^bb0:
+    %1 = test.increment %0
+    cf.br ^bb2(%1 : index)
+^bb1:
+    cf.br ^bb2(%0 : index)
+^bb2(%2 : index):
+  %3 = test.reflect_bounds %2
+  func.return %3 : index
+}
+
+// CHECK-LABEL: func @for_bounds
+// CHECK: test.reflect_bounds {smax = 1 : index, smin = 0 : index, umax = 1 : index, umin = 0 : index}
+func.func @for_bounds() -> index {
+  %c0 = test.with_bounds { umin = 0 : index, umax = 0 : index,
+                                smin = 0 : index, smax = 0 : index}
+  %c1 = test.with_bounds { umin = 1 : index, umax = 1 : index,
+                                smin = 1 : index, smax = 1 : index}
+  %c2 = test.with_bounds { umin = 2 : index, umax = 2 : index,
+                                smin = 2 : index, smax = 2 : index}
+
+  %0 = scf.for %arg0 = %c0 to %c2 step %c1 iter_args(%arg2 = %c0) -> index {
+    scf.yield %arg0 : index
+  }
+  %1 = test.reflect_bounds %0
+  func.return %1 : index
+}
+
+// CHECK-LABEL: func @no_analysis_of_loop_variants
+// CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -1 : index, umin = 0 : index}
+func.func @no_analysis_of_loop_variants() -> index {
+  %c0 = test.with_bounds { umin = 0 : index, umax = 0 : index,
+                                smin = 0 : index, smax = 0 : index}
+  %c1 = test.with_bounds { umin = 1 : index, umax = 1 : index,
+                                smin = 1 : index, smax = 1 : index}
+  %c2 = test.with_bounds { umin = 2 : index, umax = 2 : index,
+                                smin = 2 : index, smax = 2 : index}
+
+  %0 = scf.for %arg0 = %c0 to %c2 step %c1 iter_args(%arg2 = %c0) -> index {
+    %1 = test.increment %arg2
+    scf.yield %1 : index
+  }
+  %2 = test.reflect_bounds %0
+  func.return %2 : index
+}
+
+// CHECK-LABEL: func @region_args
+// CHECK: test.reflect_bounds {smax = 4 : index, smin = 3 : index, umax = 4 : index, umin = 3 : index}
+func.func @region_args() {
+  test.with_bounds_region { umin = 3 : index, umax = 4 : index,
+                            smin = 3 : index, smax = 4 : index } %arg0 {
+    %0 = test.reflect_bounds %arg0
+  }
+  func.return
+}
+
+// CHECK-LABEL: func @func_args_unbound
+// CHECK: test.reflect_bounds {smax = 9223372036854775807 : index, smin = -9223372036854775808 : index, umax = -1 : index, umin = 0 : index}
+func.func @func_args_unbound(%arg0 : index) -> index {
+  %0 = test.reflect_bounds %arg0
+  func.return %0 : index
+}

mlir/test/lib/Dialect/Test/CMakeLists.txt

add_mlir_library(MLIRTestDialect
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRDerivedAttributeOpInterface
   MLIRDialect
   MLIRDLTI
   MLIRFunc
   MLIRFuncTransforms
   MLIRIR
+  MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRLinalg
   MLIRLinalgTransforms
   MLIRLLVMIR
   MLIRPass
   MLIRReduce
   MLIRTensor
   MLIRTransformUtils
   MLIRTransforms
 )

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

 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 class DLTIDialect;
 class RewritePatternSet;

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

 //===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//
 //
 // 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 "TestDialect.h"
 #include "TestAttributes.h"
 #include "TestInterfaces.h"
 #include "TestTypes.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
 
 using namespace mlir;
for (Operation &op : block)
       if (failed(mlir::verify(&op)))
         return emitOpError("nested op hasn't been verified");
 
   emitRemark("success run of region verifier");
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// Test InferIntRangeInterface
+//===----------------------------------------------------------------------===//
+
+void TestWithBoundsOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                         SetIntRangeFn setResultRanges) {
+  setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});
+}
+
+ParseResult TestWithBoundsRegionOp::parse(OpAsmParser &parser,
+                                          OperationState &result) {
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  // Parse the input argument
+  OpAsmParser::Argument argInfo;
+  argInfo.type = parser.getBuilder().getIndexType();
+  parser.parseArgument(argInfo);
+
+  // Parse the body region, and reuse the operand info as the argument info.
+  Region *body = result.addRegion();
+  return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/false);
+}
+
+void TestWithBoundsRegionOp::print(OpAsmPrinter &p) {
+  p.printOptionalAttrDict((*this)->getAttrs());
+  p << ' ';
+  p.printRegionArgument(getRegion().getArgument(0), /*argAttrs=*/{},
+                        /*omitType=*/true);
+  p << ' ';
+  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
+}
+
+void TestWithBoundsRegionOp::inferResultRanges(
+    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
+  Value arg = getRegion().getArgument(0);
+  setResultRanges(arg, {getUmin(), getUmax(), getSmin(), getSmax()});
+}
+
+void TestIncrementOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
+                                        SetIntRangeFn setResultRanges) {
+  const ConstantIntRanges &range = argRanges[0];
+  APInt one(range.umin().getBitWidth(), 1);
+  setResultRanges(getResult(),
+                  {range.umin().uadd_sat(one), range.umax().uadd_sat(one),
+                   range.smin().sadd_sat(one), range.smax().sadd_sat(one)});
+}
+
+void TestReflectBoundsOp::inferResultRanges(
+    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
+  const ConstantIntRanges &range = argRanges[0];
+  MLIRContext *ctx = getContext();
+  Builder b(ctx);
+  setUminAttr(b.getIndexAttr(range.umin().getZExtValue()));
+  setUmaxAttr(b.getIndexAttr(range.umax().getZExtValue()));
+  setSminAttr(b.getIndexAttr(range.smin().getSExtValue()));
+  setSmaxAttr(b.getIndexAttr(range.smax().getSExtValue()));
+  setResultRanges(getResult(), range);
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

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

 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/PatternBase.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 
 // Include the attribute definitions.
 include "TestAttrDefs.td"
 // Include the type definitions.

 }
 
 
 // This is used to test encoding of a string attribute into an SSA name of a
 // pretty printed value name.
 def CustomResultsNameOp
  : TEST_Op<"custom_result_name",
            [DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
-  let arguments = (ins 
+  let arguments = (ins
     Variadic<AnyInteger>:$optional,
     StrArrayAttr:$names
   );
   let results = (outs Variadic<AnyInteger>:$r);
 }
 
 // This is used to test the OpAsmOpInterface::getDefaultDialect() feature:
 // operations nested in a region under this op will drop the "test." dialect
let assemblyFormat = [{
     $args $body attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     mlir::Region &getLoopBody() { return getBody(); }
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test InferIntRangeInterface
+//===----------------------------------------------------------------------===//
+def TestWithBoundsOp : TEST_Op<"with_bounds",
+                          [DeclareOpInterfaceMethods<InferIntRangeInterface>,
+                           NoSideEffect]> {
+  let arguments = (ins IndexAttr:$umin,
+                       IndexAttr:$umax,
+                       IndexAttr:$smin,
+                       IndexAttr:$smax);
+  let results = (outs Index:$fakeVal);
+
+  let assemblyFormat = "attr-dict";
+}
+
+def TestWithBoundsRegionOp : TEST_Op<"with_bounds_region",
+                          [DeclareOpInterfaceMethods<InferIntRangeInterface>,
+                           SingleBlock, NoTerminator]> {
+  let arguments = (ins IndexAttr:$umin,
+                       IndexAttr:$umax,
+                       IndexAttr:$smin,
+                       IndexAttr:$smax);
+  // The region has one argument of index type
+  let regions = (region SizedRegion<1>:$region);
+  let hasCustomAssemblyFormat = 1;
+}
+
+def TestIncrementOp : TEST_Op<"increment",
+                         [DeclareOpInterfaceMethods<InferIntRangeInterface>,
+                         NoSideEffect]> {
+  let arguments = (ins Index:$value);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = "attr-dict $value";
+}
+
+def TestReflectBoundsOp : TEST_Op<"reflect_bounds",
+                         [DeclareOpInterfaceMethods<InferIntRangeInterface>]> {
+  let arguments = (ins Index:$value,
+                       OptionalAttr<IndexAttr>:$umin,
+                       OptionalAttr<IndexAttr>:$umax,
+                       OptionalAttr<IndexAttr>:$smin,
+                       OptionalAttr<IndexAttr>:$smax);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = "attr-dict $value";
+}
 #endif // TEST_OPS

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestControlFlowSink.cpp
   TestInlining.cpp
+  TestIntRangeInference.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
+  MLIRAnalysis
+  MLIRInferIntRangeInterface
   MLIRTestDialect
   MLIRTransforms
   )
 
 target_include_directories(MLIRTestTransforms
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test
   )

mlir/test/lib/Transforms/TestIntRangeInference.cpp

+//===- TestIntRangeInference.cpp - Create consts from range inference ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// TODO: This pass is needed to test integer range inference until that
+// functionality has been integrated into SCCP.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/IntRangeAnalysis.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "mlir/Support/TypeID.h"
+#include "mlir/Transforms/FoldUtils.h"
+
+using namespace mlir;
+
+/// Patterned after SCCP
+static LogicalResult replaceWithConstant(IntRangeAnalysis &analysis,
+                                         OpBuilder &b, OperationFolder &folder,
+                                         Value value) {
+  Optional<ConstantIntRanges> maybeInferredRange = analysis.getResult(value);
+  if (!maybeInferredRange)
+    return failure();
+  const ConstantIntRanges &inferredRange = maybeInferredRange.getValue();
+  Optional<APInt> maybeConstValue = inferredRange.getConstantValue();
+  if (!maybeConstValue.hasValue())
+    return failure();
+
+  Operation *maybeDefiningOp = value.getDefiningOp();
+  Dialect *valueDialect =
+      maybeDefiningOp ? maybeDefiningOp->getDialect()
+                      : value.getParentRegion()->getParentOp()->getDialect();
+  Attribute constAttr = b.getIntegerAttr(value.getType(), *maybeConstValue);
+  Value constant = folder.getOrCreateConstant(b, valueDialect, constAttr,
+                                              value.getType(), value.getLoc());
+  if (!constant)
+    return failure();
+
+  value.replaceAllUsesWith(constant);
+  return success();
+}
+
+static 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);
+  OperationFolder folder(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, folder, 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, folder, arg);
+  }
+}
+
+namespace {
+struct TestIntRangeInference
+    : PassWrapper<TestIntRangeInference, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestIntRangeInference)
+
+  StringRef getArgument() const final { return "test-int-range-inference"; }
+  StringRef getDescription() const final {
+    return "Test integer range inference analysis";
+  }
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    IntRangeAnalysis analysis(op);
+    rewrite(analysis, op->getContext(), op->getRegions());
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestIntRangeInference() {
+  PassRegistration<TestIntRangeInference>();
+}
+} // end namespace test
+} // end namespace mlir

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

 void registerTestDataLayoutQuery();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDiagnosticsPass();
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
 void registerTestExpandMathPass();
 void registerTestComposeSubView();
 void registerTestMultiBuffering();
+void registerTestIntRangeInference();
 void registerTestIRVisitorsPass();
 void registerTestGenericIRVisitorsPass();
 void registerTestGenericIRVisitorsInterruptPass();
 void registerTestInterfaces();
 void registerTestLinalgCodegenStrategy();
 void registerTestLinalgElementwiseFusion();
 void registerTestLinalgFusionTransforms();
 void registerTestLinalgTensorFusionTransforms();
#endif
   mlir::test::registerTestDecomposeCallGraphTypes();
   mlir::test::registerTestDataFlowPass();
   mlir::test::registerTestDataLayoutQuery();
   mlir::test::registerTestDominancePass();
   mlir::test::registerTestDynamicPipelinePass();
   mlir::test::registerTestExpandMathPass();
   mlir::test::registerTestComposeSubView();
   mlir::test::registerTestMultiBuffering();
+  mlir::test::registerTestIntRangeInference();
   mlir::test::registerTestIRVisitorsPass();
   mlir::test::registerTestGenericIRVisitorsPass();
   mlir::test::registerTestInterfaces();
   mlir::test::registerTestLinalgCodegenStrategy();
   mlir::test::registerTestLinalgElementwiseFusion();
   mlir::test::registerTestLinalgFusionTransforms();
   mlir::test::registerTestLinalgTensorFusionTransforms();
   mlir::test::registerTestLinalgTiledLoopFusionTransforms();

mlir/unittests/Interfaces/CMakeLists.txt

 add_mlir_unittest(MLIRInterfacesTests
   ControlFlowInterfacesTest.cpp
   DataLayoutInterfacesTest.cpp
+  InferIntRangeInterfaceTest.cpp
   InferTypeOpInterfaceTest.cpp
 )
 
 target_link_libraries(MLIRInterfacesTests
   PRIVATE
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRDLTI
   MLIRFunc
+  MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRParser
 )

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

+//===- InferIntRangeInterfaceTest.cpp - Unit Tests for InferIntRange... --===//
+//
+// 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/Interfaces/InferIntRangeInterface.h"
+#include "llvm/ADT/APInt.h"
+#include <limits>
+
+#include <gtest/gtest.h>
+
+using namespace mlir;
+
+TEST(IntRangeAttrs, BasicConstructors) {
+  APInt zero = APInt::getZero(64);
+  APInt two(64, 2);
+  APInt three(64, 3);
+  ConstantIntRanges boundedAbove(zero, two, zero, three);
+  EXPECT_EQ(boundedAbove.umin(), zero);
+  EXPECT_EQ(boundedAbove.umax(), two);
+  EXPECT_EQ(boundedAbove.smin(), zero);
+  EXPECT_EQ(boundedAbove.smax(), three);
+}
+
+TEST(IntRangeAttrs, FromUnsigned) {
+  APInt zero = APInt::getZero(64);
+  APInt maxInt = APInt::getSignedMaxValue(64);
+  APInt minInt = APInt::getSignedMinValue(64);
+  APInt minIntPlusOne = minInt + 1;
+
+  ConstantIntRanges canPortToSigned =
+      ConstantIntRanges::fromUnsigned(zero, maxInt);
+  EXPECT_EQ(canPortToSigned.smin(), zero);
+  EXPECT_EQ(canPortToSigned.smax(), maxInt);
+
+  ConstantIntRanges cantPortToSigned =
+      ConstantIntRanges::fromUnsigned(zero, minInt);
+  EXPECT_EQ(cantPortToSigned.smin(), minInt);
+  EXPECT_EQ(cantPortToSigned.smax(), maxInt);
+
+  ConstantIntRanges signedNegative =
+      ConstantIntRanges::fromUnsigned(minInt, minIntPlusOne);
+  EXPECT_EQ(signedNegative.smin(), minInt);
+  EXPECT_EQ(signedNegative.smax(), minIntPlusOne);
+}
+
+TEST(IntRangeAttrs, FromSigned) {
+  APInt zero = APInt::getZero(64);
+  APInt one = zero + 1;
+  APInt negOne = zero - 1;
+  APInt intMax = APInt::getSignedMaxValue(64);
+  APInt intMin = APInt::getSignedMinValue(64);
+  APInt uintMax = APInt::getMaxValue(64);
+
+  ConstantIntRanges noUnsignedBound =
+      ConstantIntRanges::fromSigned(negOne, one);
+  EXPECT_EQ(noUnsignedBound.umin(), zero);
+  EXPECT_EQ(noUnsignedBound.umax(), uintMax);
+
+  ConstantIntRanges positive = ConstantIntRanges::fromSigned(one, intMax);
+  EXPECT_EQ(positive.umin(), one);
+  EXPECT_EQ(positive.umax(), intMax);
+
+  ConstantIntRanges negative = ConstantIntRanges::fromSigned(intMin, negOne);
+  EXPECT_EQ(negative.umin(), intMin);
+  EXPECT_EQ(negative.umax(), negOne);
+
+  ConstantIntRanges preserved = ConstantIntRanges::fromSigned(zero, one);
+  EXPECT_EQ(preserved.umin(), zero);
+  EXPECT_EQ(preserved.umax(), one);
+}
+
+TEST(IntRangeAttrs, Join) {
+  APInt zero = APInt::getZero(64);
+  APInt one = zero + 1;
+  APInt two = zero + 2;
+  APInt intMin = APInt::getSignedMinValue(64);
+  APInt intMax = APInt::getSignedMaxValue(64);
+  APInt uintMax = APInt::getMaxValue(64);
+
+  ConstantIntRanges maximal(zero, uintMax, intMin, intMax);
+  ConstantIntRanges zeroOne(zero, one, zero, one);
+
+  EXPECT_EQ(zeroOne.rangeUnion(maximal), maximal);
+  EXPECT_EQ(maximal.rangeUnion(zeroOne), maximal);
+
+  EXPECT_EQ(zeroOne.rangeUnion(zeroOne), zeroOne);
+
+  ConstantIntRanges oneTwo(one, two, one, two);
+  ConstantIntRanges zeroTwo(zero, two, zero, two);
+  EXPECT_EQ(zeroOne.rangeUnion(oneTwo), zeroTwo);
+
+  ConstantIntRanges zeroOneUnsignedOnly(zero, one, intMin, intMax);
+  ConstantIntRanges zeroOneSignedOnly(zero, uintMax, zero, one);
+  EXPECT_EQ(zeroOneUnsignedOnly.rangeUnion(zeroOneSignedOnly), maximal);
+}