-
Notifications
You must be signed in to change notification settings - Fork 14
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add support for stablehlo.reduce op for logical and operator #1741
Conversation
ca387ed
to
5b5c37e
Compare
5b5c37e
to
e800cae
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You should write a much more detailed description about this decomposition, ideally as a comment before class/struct of that decomposition. If I understand this correctly assumption is that all of the tensor values are either exactly 0.0 or 1.0. Those are very strong assumptions that should be well documented at least. I know we are currently lowering only stablehlo.reduce
to ttir.reduce_and
, but bundled with type conversion involved (i1 -> bfloat16) we are making a strong TTIR -> StableHLO dependency, and reduce_and
should exist independently of StableHLO. Nothing stops us from calling ttir.reduce_and
with float32 with arbitrary values, there is no way to check this during compilation (values part) and worst of all this will just silently fail as sum
and equal
on arbitrary values are perfectly fine.
Even if we ignore all of this and rely on developers discipline and knowledge that only exists in the heads of a few people instead of a compiler, bfloat16
is notorious for its small precision. It has only 7 precision bits, so even some fairly small numbers that are perfectly fine as dimension size in practice can't be represented exactly with bfloat16
, so this approach will have a lot of false negatives.
With assumption that all of the elements have values of 0.0 or 1.0 better approach would be reduction with min
or prod
. Both of these exist in TTNN.
But the question about types and value assumptions remains open. If we want TTIR to be a public interface that any frontend can lower to, we should think about the way how we can represent those assumptions in code, especially as I said these things will fail just silently. If I saw an op that accepts bfloat16
tensor that's called reduce_and
and I gave it [2, 1] and it returns 0 I would be very surprised for a good reason.
6b2428d
to
4a3083f
Compare
@azecevicTT Thanks a lot for the feedback. I have updated my decomposition to use TTNN does not support boolean data type, so stablehlo->ttir conversion uses type conversion of (i1 -> bf16). Other frontends/dialects may use different approach. The decomposed op in updated PR has same input/output type as of the original op removing any dependency on stablehlo. |
4a3083f
to
608fcaa
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks much better!
TTNN does not support reduction for logical and operator. So stablehlo.reduce for stablehlo.and operator is decomposed into reduction min op along given dimension and then comparing with zero.
608fcaa
to
cc18937
Compare
TTNN does not support reduction for logical and operator. So stablehlo.reduce
for stablehlo.and operator is decomposed into reduction sum op along give
dimension and then compared with the size of given dimension.