Auto generate StableHLO transform patterns #422
Conversation
|
just to make sure the names are what are expected, can we add them to Enzyme-JAX/src/enzyme_ad/jax/primitives.py Line 127 in 26356cc |
|
|
||
| // Naming convention: `<RootOp><ChildOp>Const`: | ||
|
|
||
| // Pattern: (x / cst_1) * cst_2 -> x * (cst_2 / cst_1) | ||
| def MulDivConst : BinOpLiftConstantComputationFlip<StableHLO_MulOp, StableHLO_DivOp>; | ||
| def MulDivConst : BinOpLiftConstantComputationFlip<"mul_div_const", StableHLO_MulOp, StableHLO_DivOp>; |
There was a problem hiding this comment.
How useful is it having a ton of individual patterns as transform ops?
I'm thinking a few possible outcomes here -- we offer individual patterns and grouped pattern ops for transform (i.e. {mul_div_const, mul_mul_const,...} and binop_const) -- or maybe this is something we resolve in the python layer with something like aliases?
I guess the real question is -- is there any issue with exploding the number of transform extension ops we're providing here?
I'd like to expand this to all StableHLO optim patterns so we can think about using some of the upstream ones (and upstreaming new ones) without impacting usability:
https://github.com/openxla/stablehlo/blob/main/stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td
There was a problem hiding this comment.
edit: fixed the link
So part of the original motivation for having the names was debugging both performance and issues.
For example, one can excise specific patterns to find bugs like #423
For performance it gives a search space since some patterns aren't always beneficial in practice (though from the hlo side it may seem strictly simplifying).
I don't think this means we need to have all things as individual patterns, but we can group things where convenient.
There was a problem hiding this comment.
The alternative would be to have one stablehlo_patterns op that is further parameterized by a list of patterns to include/exclude. But that means having (yet another) registration mechanism for that list whereas now we can use the op registration mechanism. Having logical groupings to reduce clutter makes sense to me as a short-term trade-off. Longer-term, we need to find a good way to manage this with e.g., adding ops on the fly or somehow referring to pattern name/location without creating a whole new op every time.
There was a problem hiding this comment.
Sounds good, agree. I think a few groupings are probably best in the short term, might remove the individual granularity and just use the StablehloPatternGroup from this PR.
I like the file locality of this approach in the meantime to keep things trivial for users contributing patterns / also lets you refer to identifiers in TD files for slightly-better-than-stringly-typed safety. Also will make it really easy to give existing StableHLO patterns bindings in transform dialect
The core of this PR is to make it trivial for patterns to have corresponding Transform dialect bindings.
Generates patterns and a populate method for each pattern, as well as a transform dialect op which wraps each pattern. Additionally there's some pattern-group ops as well.