[Util] Implement InferIntDivisibilityOpInterface for affine ops #22723
Conversation
Max191
requested review from
MaheshRavishankar,
benvanik,
hanhanW,
jerryyin,
krzysz00 and
nirvedhmeshram
compiler/src/iree/compiler/Dialect/Util/Transforms/test/test_integer_divisibility_analysis.mlir
Show resolved
Hide resolved
nirvedhmeshram
left a comment
nirvedhmeshram
left a comment
There was a problem hiding this comment.
LGTM, I do have a non-blocking but related question
compiler/src/iree/compiler/Dialect/Util/Transforms/test/test_integer_divisibility_analysis.mlir
Outdated
Show resolved
Hide resolved
|
So, For affine ops we have traditionally just converted them to arith ops and done that. Not sure how tricky it is to add support for all affine expressions |
compiler/src/iree/compiler/Dialect/Util/Transforms/TestIntegerDivisibilityAnalysis.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Util/Transforms/test/test_integer_divisibility_analysis.mlir
Show resolved
Hide resolved
This PR should support arbitrary affine expressions. It implements the basic binary operators, and then the full divisibility is just a composition of them. I'm not sure if lowering to arith ops is a great solution, since this analysis will be used early on in compilation (e.g., sometimes before all tiling and distribution is done), and we do a lot of affine apply/min/max composition during these compilation phases. The affine divisibility implementation turned out to not be too bad, since I mainly just had to implement the core operators. |
Max191
force-pushed
the
affine-divisibility-analysis
branch
from
b51e734 to
463e4ef
Compare
MaheshRavishankar
left a comment
MaheshRavishankar
left a comment
There was a problem hiding this comment.
Nice! Looks good to me.
Maybe @bjacob has more nuanced feedback.
compiler/src/iree/compiler/Dialect/Util/Transforms/test/test_integer_divisibility_analysis.mlir
Show resolved
Hide resolved
| IREE::Util::ConstantIntDivisibility | ||
| visitConstantExpr(AffineConstantExpr expr) { | ||
| // Constant expressions are trivial, since they are always static. | ||
| uint64_t constValue = std::abs(expr.getValue()); |
There was a problem hiding this comment.
So this is going to make sdiv positive. I dont fully understand how this was set up, but just wanted to see if you had a rationale for this, or (like me) this is not fully reasoned about.
There was a problem hiding this comment.
Browsing the code (which is new to me), the ConstantIntDivisibility constructor at the next line surprisingly takes (uint64_t udiv, uint64_t sdiv):
I guess that this motivates uint64_t constValue above, as it is going to be casted to uint64_t anyway at the next line.
I wonder if that constructor is a typo or intentional? Noticing that it is at least consistent with this method a few lines down:
To me, not knowing anything about the specifics here, this looks suspisciously close to a violation of "don't use unsigned to denote non-negative values" rule,
https://google.github.io/styleguide/cppguide.html#Integer_Types
In particular, do not use unsigned types to say a number will never be negative.
Happy to dive into this (I still don't know anything about this stuff) BUT this looks like in order to be useful, we need to dive into the first principles here -- what are the exact semantics that we want here, why are unsigned types involved, should they be involved etc.
There was a problem hiding this comment.
I am not entirely sure. Maybe @stellaraccident has more context on those, but I would propose first landing this and then iterating on the "signed" vs "unsigned" parts here. +1 on nailing this down a bit more though. This seems like a good thing to address soon before it gets out of hand.
There was a problem hiding this comment.
Dropping a summary of what Benoit and I discussed here so it is not lost: https://discord.com/channels/689900678990135345/1442894583624106036/1442903483014578317
Full thread starts here: https://discord.com/channels/689900678990135345/1442894583624106036
| } | ||
| int64_t constValue = std::abs(constRhs.getValue()); | ||
| IREE::Util::ConstantIntDivisibility lhsDiv = visit(expr.getLHS()); | ||
| uint64_t divUDiv = lhsDiv.udiv() % static_cast<uint64_t>(constValue) == 0 |
There was a problem hiding this comment.
Nit: do cast once for readability?
Max191
force-pushed
the
affine-divisibility-analysis
branch
from
463e4ef to
d1f4fc3
Compare
Implements divisibility inference for affine.apply, affine.min, and affine.max operations through AffineExprDivisibilityFinder visitor. The visitor recursively computes divisibility for affine expressions: - Multiplication: product of operand divisibilities - Addition/Subtraction: GCD (union) of operand divisibilities - Division (floor/ceil): quotient when evenly divisible, else 1 - Modulo: falls back to minimum divisibility (1,1) - Min/Max: GCD of all branch divisibilities Adds TestIntegerDivisibilityAnalysis pass to more directly test divisibility analysis without relying on IR optimizations. The pass probes values with iree_unregistered.test_int_divisibility ops and annotates them with computed divisibility attributes. Also updates arith::DivUIOp divisibility calculation to fallback to minimum divisibility when there is a remainder division, because we can't infer the divisibility when there is a remainder. Signed-off-by: Max Dawkins <[email protected]>
Signed-off-by: Max Dawkins <[email protected]>
Signed-off-by: Max Dawkins <[email protected]>
Signed-off-by: Max Dawkins <[email protected]>
Max191
force-pushed
the
affine-divisibility-analysis
branch
from
097a923 to
a013ceb
Compare
…ps (#22723)" This reverts commit fa24a3a. Signed-off-by: Bangtian Liu <[email protected]>
#22860) …ps (#22723)" This PR breaks compilation of torch_models for llama 8b fp8 with data tiling. This is post submit CI run [before](https://github.com/iree-org/iree/actions/runs/20032789150/job/57449215649) and [after](https://github.com/iree-org/iree/actions/runs/20037515165/job/57463661192) this PR was merged. This reverts commit fa24a3a. ci-extra: test_torch
…ops (#22860)" (#23137) Reapply #22723 now that the torch model failures were fixed by #23118. ci-extra: test_torch Signed-off-by: Max Dawkins <[email protected]>
This PR implements the InferIntDivisibilityOpInterface for affine.apply, affine.min, and affine.max operations. Affine apply gets the divisibility of its result expression, and affine.min/max gets the GCD of all result expression divisibilities. The implementation supports the following divisibilities and any compositions of them: - Multiplication: product of operand divisibilities - Addition/Subtraction: GCD (union) of operand divisibilities - Division (floor/ceil): quotient when evenly divisible, else 1 - Modulo: falls back to minimum divisibility (1,1) This PR also adds the TestIntegerDivisibilityAnalysis pass to more directly test divisibility analysis without relying on IR optimizations. The pass probes values consumed by `"iree_unregistered.test_int_divisibility"` ops and annotates them with computed divisibility attributes. There is a small change to the arith.divui divisibility implementation to fallback to minimum divisibility when there is a remainder division, because we can't infer the divisibility when there is a remainder. --------- Signed-off-by: Max Dawkins <[email protected]> Signed-off-by: Keshav Vinayak Jha <[email protected]>
#22860) …ps (#22723)" This PR breaks compilation of torch_models for llama 8b fp8 with data tiling. This is post submit CI run [before](https://github.com/iree-org/iree/actions/runs/20032789150/job/57449215649) and [after](https://github.com/iree-org/iree/actions/runs/20037515165/job/57463661192) this PR was merged. This reverts commit fa24a3a. ci-extra: test_torch Signed-off-by: Keshav Vinayak Jha <[email protected]>
…ops (#22860)" (#23137) Reapply #22723 now that the torch model failures were fixed by #23118. ci-extra: test_torch Signed-off-by: Max Dawkins <[email protected]> Signed-off-by: Keshav Vinayak Jha <[email protected]>
7 participants
This PR implements the InferIntDivisibilityOpInterface for affine.apply, affine.min, and affine.max operations. Affine apply gets the divisibility of its result expression, and affine.min/max gets the GCD of all result expression divisibilities. The implementation supports the following divisibilities and any compositions of them:
This PR also adds the TestIntegerDivisibilityAnalysis pass to more directly test divisibility analysis without relying on IR optimizations. The pass probes values consumed by
"iree_unregistered.test_int_divisibility"ops and annotates them with computed divisibility attributes.There is a small change to the arith.divui divisibility implementation to fallback to minimum divisibility when there is a remainder division, because we can't infer the divisibility when there is a remainder.