[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

19

[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

101

[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

This class can just go in the anonymous namespace

105

The operand lattices should never be uninitialized.

148

braces

169

braces

222

braces

234

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

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

1411

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

2883–2886

2907–2911

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

1–2

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

24–25

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

mlir/unittests/Interfaces/InferIntRangeInterfaceTest.cpp

1–2

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

2889–2891

2912–2914

2926–2928

[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

mlir/lib/Transforms/FoldInferredConstants.cpp

2–3

This is a bit off.

89–98

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

89–98

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

89–98

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

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/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/Analysis/DataFlowAnalysis.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
+
+namespace mlir {
+namespace detail {
+class IntRangeAnalysisImpl;
+} // end namespace detail
+
+class IntRangeAnalysis {
+public:
+  IntRangeAnalysis(MLIRContext *ctx);
+  ~IntRangeAnalysis();
+  IntRangeAnalysis(IntRangeAnalysis &&other);
+
+  /// Analyze all operations rooted under (but not including)
+  /// `topLevelOperation`.
+  void run(Operation *topLevelOperation);
+  /// Get inferred range for value `v` if one exists.
+  Optional<IntRangeAttrs> 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 - 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_INTERFACES_INFERINTRANGEINTERFACE_H
+#define MLIR_INTERFACES_INFERINTRANGEINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/APInt.h"
+#include <tuple>
+
+namespace mlir {
+class Value;
+
+/// 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 {
+public:
+  /// Return a range with no bounds for the signed and unsigned interpretations.
+  IntRangeAttrs() : uminVal(), umaxVal(), sminVal(), smaxVal() {}
+  /// Bound umin <= (unsigned)x <= umax and smin <= signed(x) <= smax,
+  /// with llvm::None being considered +-infinity.
+  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());
+    }
+  }
+
+  /// Convenience wrapper that allows providing [min, max] ranges.
+  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)) {}
+
+  /// 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;
+
+  /// 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);
+  /// Helper methods for when IntRangeAttrs is used in a dataflow analysis.
+  static IntRangeAttrs getPessimisticValueState(MLIRContext *context);
+  static IntRangeAttrs getPessimisticValueState(Value v);
+
+  /// 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;
+
+  bool operator==(const IntRangeAttrs &other) const;
+
+  friend raw_ostream &operator<<(raw_ostream &os, const IntRangeAttrs &range);
+
+private:
+  Optional<APInt> uminVal, umaxVal, sminVal, smaxVal;
+};
+} // 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 [-inf, +inf]
+      (the unbounded range).
+
+      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 unbounded range in their
+       `argRanges` element.
+    }],
+    "void", "inferResultRanges", (ins
+      "::llvm::ArrayRef<::mlir::IntRangeAttrs>":$argRanges,
+      "::llvm::function_ref<void(::mlir::Value, const ::mlir::IntRangeAttrs &)>":$setResultRanges)
+  >];
+}
+#endif // MLIR_DIALECT_ARITHMETIC_IR_INFERINTRANGEINTERFACE

mlir/include/mlir/Transforms/Passes.h

createCanonicalizerPass(const GreedyRewriteConfig &config,
                         ArrayRef<std::string> enabledPatterns = llvm::None);
 
 /// Creates a pass to perform control-flow sinking.
 std::unique_ptr<Pass> createControlFlowSinkPass();
 
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
+/// Create a pass to constant fold based on the results of range inferrence.
+std::unique_ptr<Pass> createFoldInferredConstantsPass();
+
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
 std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
 
 /// Creates a pass to strip debug information from a function.
 std::unique_ptr<Pass> createStripDebugInfoPass();
 
 /// Creates a pass which prints the list of ops and the number of occurrences in

mlir/include/mlir/Transforms/Passes.td

def CSE : Pass<"cse"> {
   }];
   let constructor = "mlir::createCSEPass()";
   let statistics = [
     Statistic<"numCSE", "num-cse'd", "Number of operations CSE'd">,
     Statistic<"numDCE", "num-dce'd", "Number of operations DCE'd">
   ];
 }
 
+def FoldInferredConstants : Pass<"fold-inferred-constants"> {
+  let summary = "Constant fold based on the results of integer range inference";
+  let constructor = "mlir::createFoldInferredConstantsPass()";
+}
+
 def Inliner : Pass<"inline"> {
   let summary = "Inline function calls";
   let constructor = "mlir::createInlinerPass()";
   let options = [
     Option<"defaultPipelineStr", "default-pipeline", "std::string",
            /*default=*/"", "The default optimizer pipeline used for callables">,
     ListOption<"opPipelineList", "op-pipelines", "OpPassManager",
                "Callable operation specific optimizer pipelines (in the form "

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/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/Interfaces/InferIntRangeInterface.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "int-range-analysis"
+
+using namespace mlir;
+
+namespace mlir {
+namespace detail {
+class IntRangeAnalysisImpl : public ForwardDataFlowAnalysis<IntRangeAttrs> {
+  using ForwardDataFlowAnalysis<IntRangeAttrs>::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<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;
+};
+} // 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 Optional<APInt>
+getLoopBoundFromFold(Optional<OpFoldResult> loopBound,
+                     detail::IntRangeAnalysisImpl &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)
+      result =
+          getUpper ? lattice->getValue().smax() : lattice->getValue().smin();
+  }
+  return result;
+}
+
+ChangeResult detail::IntRangeAnalysisImpl::visitOperation(
+    Operation *op, ArrayRef<LatticeElement<IntRangeAttrs> *> 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<IntRangeAttrs> argRanges(
+        llvm::map_range(operands, [](LatticeElement<IntRangeAttrs> *val) {
+          return val->isUninitialized() ? IntRangeAttrs() : val->getValue();
+        }));
+
+    auto joinCallback = [&](Value v, const IntRangeAttrs &attrs) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+      LatticeElement<IntRangeAttrs> &lattice = getLatticeElement(v);
+      Optional<IntRangeAttrs> oldRange;
+      if (!lattice.isUninitialized())
+        oldRange = lattice.getValue();
+      result |= lattice.join(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);
+  } 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<IntRangeAttrs> *> operands,
+    SmallVectorImpl<Block *> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> maybeConstValue =
+        enumPair.value()->getValue().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<IntRangeAttrs> *> operands,
+    SmallVectorImpl<RegionSuccessor> &successors) {
+  auto toConstantAttr = [&branch](auto enumPair) -> Attribute {
+    Optional<APInt> maybeConstValue =
+        enumPair.value()->getValue().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<IntRangeAttrs> *> 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<IntRangeAttrs> argRanges(
+        llvm::map_range(operands, [](LatticeElement<IntRangeAttrs> *val) {
+          return val->getValue();
+        }));
+
+    ChangeResult result = ChangeResult::NoChange;
+    auto joinCallback = [&](Value v, const IntRangeAttrs &attrs) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferred range " << attrs << "\n");
+      LatticeElement<IntRangeAttrs> &lattice = getLatticeElement(v);
+      Optional<IntRangeAttrs> oldRange;
+      if (!lattice.isUninitialized())
+        oldRange = lattice.getValue();
+      result |= lattice.join(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);
+    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<
+          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);
+    else if (max) {
+      // Correct the upper bound by subtracting 1 so that it becomes a <= bound,
+      // because loops do not generally include their upper bound.
+      max = *max - 1;
+    }
+
+    LatticeElement<IntRangeAttrs> &ivEntry = getLatticeElement(*iv);
+    return ivEntry.join(IntRangeAttrs::fromSigned(min, max));
+  }
+  return ForwardDataFlowAnalysis<IntRangeAttrs>::visitNonControlFlowArguments(
+      op, region, operands);
+}
+
+IntRangeAnalysis::IntRangeAnalysis(MLIRContext *ctx) {
+  impl = std::make_unique<mlir::detail::IntRangeAnalysisImpl>(ctx);
+}
+
+IntRangeAnalysis::~IntRangeAnalysis() = default;
+IntRangeAnalysis::IntRangeAnalysis(IntRangeAnalysis &&other) = default;
+
+void IntRangeAnalysis::run(Operation *topLevelOperation) {
+  impl->run(topLevelOperation);
+}
+
+Optional<IntRangeAttrs> IntRangeAnalysis::getResult(Value v) {
+  LatticeElement<IntRangeAttrs> *result = impl->lookupLatticeElement(v);
+  if (result == nullptr || result->isUninitialized())
+    return llvm::None;
+  return result->getValue();
+}

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/Interfaces/InferIntRangeInterface.cpp.inc"
+
+using namespace mlir;
+
+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 = APInt::getSignedMinValue(smax->getBitWidth());
+      umax = smax;
+    }
+  }
+  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 &mlir::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 << "]";
+}

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
+  FoldInferredConstants.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   OpStats.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   SymbolPrivatize.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRCopyOpInterface
+  MLIRInferIntRangeInterface
   MLIRLoopLikeInterface
+  MLIRSideEffectInterfaces
   MLIRPass
   MLIRSupport
   MLIRTransformUtils
   )

mlir/lib/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/IntRangeAnalysis.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Transforms/FoldUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+/// Patterned after SCCP.cpp
+static LogicalResult replaceWithConstant(IntRangeAnalysis &analysis,
+                                         OpBuilder &b, OperationFolder &folder,
+                                         Value value) {
+  Optional<IntRangeAttrs> maybeInferredRange = analysis.getResult(value);
+  if (!maybeInferredRange)
+    return failure();
+  const IntRangeAttrs &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 FoldInferredConstantsPass
+    : public FoldInferredConstantsBase<FoldInferredConstantsPass> {
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+
+    IntRangeAnalysis analysis(ctx);
+    analysis.run(op);
+    rewrite(analysis, ctx, op->getRegions());
+  }
+};
+} // end namespace
+
+std::unique_ptr<Pass> mlir::createFoldInferredConstantsPass() {
+  return std::make_unique<FoldInferredConstantsPass>();
+}

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

+// RUN: mlir-opt -fold-inferred-constants %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 }
+  %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 %{{.*}}
+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 @func_args_unbound
+// CHECK: test.reflect_bounds %{{.*}}
+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/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.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<IntRangeAttrs> argRanges,
+    llvm::function_ref<void(Value, const IntRangeAttrs &)> setResultRanges) {
+  setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});
+}
+
+void TestIncrementOp::inferResultRanges(
+    ArrayRef<IntRangeAttrs> argRanges,
+    llvm::function_ref<void(Value, const IntRangeAttrs &)> setResultRanges) {
+  const IntRangeAttrs &range = argRanges[0];
+  auto inc = [](const APInt &value) -> APInt { return value + 1; };
+  setResultRanges(getResult(), {range.umin().map(inc), range.umax().map(inc),
+                                range.smin().map(inc), range.smax().map(inc)});
+}
+
+void TestReflectBoundsOp::inferResultRanges(
+    ArrayRef<IntRangeAttrs> argRanges,
+    llvm::function_ref<void(Value, const IntRangeAttrs &)> setResultRanges) {
+  const IntRangeAttrs &range = argRanges[0];
+  MLIRContext *ctx = getContext();
+  Builder b(ctx);
+  if (range.umin().hasValue())
+    setUminAttr(b.getIndexAttr(range.umin()->getZExtValue()));
+  else
+    removeUminAttr();
+
+  if (range.umax().hasValue())
+    setUmaxAttr(b.getIndexAttr(range.umax()->getZExtValue()));
+  else
+    removeUmaxAttr();
+
+  if (range.smin().hasValue())
+    setSminAttr(b.getIndexAttr(range.smin()->getSExtValue()));
+  else
+    removeSminAttr();
+
+  if (range.smax().hasValue())
+    setSmaxAttr(b.getIndexAttr(range.smax()->getSExtValue()));
+  else
+    removeSmaxAttr();
+
+  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 OptionalAttr<IndexAttr>:$umin,
+                          OptionalAttr<IndexAttr>:$umax,
+                          OptionalAttr<IndexAttr>:$smin,
+                          OptionalAttr<IndexAttr>:$smax);
+  let results = (outs Index:$fakeVal);
+
+  let assemblyFormat = [{
+    attr-dict
+  }];
+}
+
+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/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 InferIntRangeInterf.
+//===//
+//
+// 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) {
+  IntRangeAttrs blank;
+  EXPECT_FALSE(blank.umin().hasValue());
+  EXPECT_FALSE(blank.umax().hasValue());
+  EXPECT_FALSE(blank.smin().hasValue());
+  EXPECT_FALSE(blank.smax().hasValue());
+
+  llvm::APInt zero = llvm::APInt::getZero(64);
+  llvm::APInt two(64, 2);
+  llvm::APInt three(64, 3);
+  IntRangeAttrs boundedAbove({}, two, {}, three);
+  EXPECT_EQ(boundedAbove.umin(), zero);
+  EXPECT_EQ(boundedAbove.umax(), two);
+  EXPECT_FALSE(boundedAbove.smin().hasValue());
+  EXPECT_EQ(boundedAbove.smax(), three);
+}
+
+TEST(IntRangeAttrs, FromUnsigned) {
+  llvm::APInt zero = llvm::APInt::getZero(64);
+  llvm::APInt maxInt = llvm::APInt::getSignedMaxValue(64);
+  llvm::APInt minInt = llvm::APInt::getSignedMinValue(64);
+  llvm::APInt minIntPlusOne = minInt + 1;
+
+  IntRangeAttrs canPortToSigned = IntRangeAttrs::fromUnsigned(zero, maxInt);
+  EXPECT_EQ(canPortToSigned.smin(), zero);
+  EXPECT_EQ(canPortToSigned.smax(), maxInt);
+
+  IntRangeAttrs cantPortToSigned = IntRangeAttrs::fromUnsigned(zero, minInt);
+  EXPECT_FALSE(cantPortToSigned.smin().hasValue());
+  EXPECT_FALSE(cantPortToSigned.smax().hasValue());
+
+  IntRangeAttrs signedNegative =
+      IntRangeAttrs::fromUnsigned(minInt, minIntPlusOne);
+  EXPECT_EQ(signedNegative.smin(), minInt);
+  EXPECT_EQ(signedNegative.smax(), minIntPlusOne);
+}
+
+TEST(IntRangeAttrs, FromSigned) {
+  llvm::APInt zero = llvm::APInt::getZero(64);
+  llvm::APInt one = zero + 1;
+  llvm::APInt negOne = zero - 1;
+  llvm::APInt intMax = llvm::APInt::getSignedMaxValue(64);
+  llvm::APInt intMin = llvm::APInt::getSignedMinValue(64);
+
+  IntRangeAttrs noUnsignedBound = IntRangeAttrs::fromSigned(negOne, one);
+  EXPECT_FALSE(noUnsignedBound.umin().hasValue());
+  EXPECT_FALSE(noUnsignedBound.umax().hasValue());
+
+  IntRangeAttrs positive = IntRangeAttrs::fromSigned(one, {});
+  EXPECT_EQ(positive.umin(), one);
+  EXPECT_EQ(positive.umax(), intMax);
+
+  IntRangeAttrs negative = IntRangeAttrs::fromSigned({}, negOne);
+  EXPECT_EQ(negative.umin(), intMin);
+  EXPECT_EQ(negative.umax(), negOne);
+
+  IntRangeAttrs preserved = IntRangeAttrs::fromSigned(zero, one);
+  EXPECT_EQ(preserved.umin(), zero);
+  EXPECT_EQ(preserved.umax(), one);
+}
+
+TEST(IntRangeAttrs, Join) {
+  llvm::APInt zero = llvm::APInt::getZero(64);
+  llvm::APInt one = zero + 1;
+  llvm::APInt two = zero + 2;
+
+  IntRangeAttrs blank;
+  IntRangeAttrs zeroOne(zero, one, zero, one);
+
+  EXPECT_EQ(IntRangeAttrs::join(zeroOne, blank), blank);
+  EXPECT_EQ(IntRangeAttrs::join(blank, zeroOne), blank);
+
+  EXPECT_EQ(IntRangeAttrs::join(zeroOne, zeroOne), zeroOne);
+
+  IntRangeAttrs oneTwo(one, two, one, two);
+  IntRangeAttrs zeroTwo(zero, two, zero, two);
+  EXPECT_EQ(IntRangeAttrs::join(zeroOne, oneTwo), zeroTwo);
+
+  IntRangeAttrs zeroOneUnsigned(zero, one, {}, {});
+  IntRangeAttrs zeroOneSigned({}, {}, zero, one);
+  EXPECT_EQ(IntRangeAttrs::join(zeroOneUnsigned, zeroOneSigned), blank);
+}