feat: propagate bounds in whileop #1453
Conversation
avik-pal
force-pushed
the
ap/while_noop
branch
from
68b2367 to
37e44ff
Compare
| auto step = info.getConstantStep().value(); | ||
|
|
||
| // currently only removes remainder of induction variable | ||
| whileBody.walk([&](stablehlo::RemOp remOp) -> WalkResult { |
There was a problem hiding this comment.
instead could we just look at all users of the induction var (that likely will be faster than a walk of the body)
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
EnzymeJAX Benchmarks
| Benchmark suite | Current: e358839 | Previous: 54edf38 | Ratio |
|---|---|---|---|
scatter_sum / JaX / cpu / Primal |
0.000004446296999958577 s |
0.000004485886999464128 s |
0.99 |
scatter_sum / JaXPipe / cpu / Primal |
0.000004340780000347877 s |
0.000004439690998697188 s |
0.98 |
scatter_sum / JaX / tpu / Primal |
0.0001532494204999 s |
0.0001577955517001 s |
0.97 |
scatter_sum / JaXPipe / tpu / Primal |
0.0001514654375001 s |
0.0001410659988992 s |
1.07 |
This comment was automatically generated by workflow using github-action-benchmark.
avik-pal
force-pushed
the
ap/while_noop
branch
from
efaec18 to
fc223de
Compare
| } | ||
|
|
||
| // Initialize bounds map with induction variable bounds | ||
| DenseMap<Value, Bounds> boundsMap; |
There was a problem hiding this comment.
min/max only applies if step >= 0, if negative this needs to be changed
There was a problem hiding this comment.
the induction var is initialized as
APInt minBound(bitWidth, std::min(start, limit), true);
APInt maxBound(bitWidth, std::max(start, limit), true);
boundsMap[inductionVar] = Bounds(minBound, maxBound);There was a problem hiding this comment.
hm, is min/max inclusive or exclusive then. If step > 0, range is [start, limit). If negative it's (limit, start]
There was a problem hiding this comment.
inclusive in the subsequent usage, needs to be fixed here
| return false; | ||
| }; | ||
|
|
||
| bool rewriteCompareOp(PatternRewriter &rewriter, |
There was a problem hiding this comment.
this seems somewhat (but not wholly) similar to:
Enzyme-JAX/src/enzyme_ad/jax/Passes/AffineCFG.cpp
Line 5941 in 9861ed8
avik-pal
changed the title
| DenseMap<Value, Bounds> boundsMap; | ||
| if (step > 0) { | ||
| APInt minBound(bitWidth, start, true); | ||
| APInt maxBound(bitWidth, limit - 1, true); |
There was a problem hiding this comment.
this isn't necessarily correct, since start == limit is a valid case [where the loop body isn't executed]
There was a problem hiding this comment.
though I suppose if the body is never executed, what is done in the after region doesn't matter anyways
There was a problem hiding this comment.
added a check for loop count >= 1
There was a problem hiding this comment.
I'm not seeing where since if step > limit - start, loop count is still zero [and step only checked for sign]
There was a problem hiding this comment.
A bit above
if (!info.isValid() || !info.isConstant() ||
info.getConstantNumIters() <= 0)
return failure();
600c987 to
e116b87
Compare
avik-pal
force-pushed
the
ap/while_noop
branch
from
f0beb10 to
e358839
Compare
2 participants
No description provided.