Skip to content
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

Merged
merged 2 commits into from
Jan 27, 2025

Conversation

mmanzoorTT
Copy link
Contributor

@mmanzoorTT mmanzoorTT commented Jan 9, 2025

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.

@mmanzoorTT mmanzoorTT force-pushed the mmanzoor/reduce-and-op branch 3 times, most recently from ca387ed to 5b5c37e Compare January 9, 2025 21:44
@mmanzoorTT mmanzoorTT changed the title Reduce and implementation Add support for stablehlo.reduce op for logical and Jan 9, 2025
@mmanzoorTT mmanzoorTT changed the title Add support for stablehlo.reduce op for logical and Add support for stablehlo.reduce op for logical and operator Jan 9, 2025
@mmanzoorTT mmanzoorTT force-pushed the mmanzoor/reduce-and-op branch from 5b5c37e to e800cae Compare January 9, 2025 21:46
@mmanzoorTT mmanzoorTT requested review from azecevicTT and AleksKnezevic and removed request for azecevicTT January 9, 2025 21:46
@mmanzoorTT mmanzoorTT marked this pull request as ready for review January 9, 2025 21:46
Copy link
Contributor

@azecevicTT azecevicTT left a 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.

@mmanzoorTT mmanzoorTT force-pushed the mmanzoor/reduce-and-op branch 6 times, most recently from 6b2428d to 4a3083f Compare January 24, 2025 17:37
@mmanzoorTT
Copy link
Contributor Author

mmanzoorTT commented Jan 24, 2025

@azecevicTT Thanks a lot for the feedback. I have updated my decomposition to use ttnn.prod (as suggested). ttnn.min can be problematic in case of negative numbers. This change does not rely of any value/type assumption. The only assumption used here is that zero is considered as false and non-zero numbers are considered as true. I have added an example and improve the description for TTIR op to avoid any ambiguity.

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.

@mmanzoorTT mmanzoorTT force-pushed the mmanzoor/reduce-and-op branch from 4a3083f to 608fcaa Compare January 24, 2025 17:58
Copy link
Contributor

@azecevicTT azecevicTT left a 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!

test/ttmlir/Silicon/StableHLO/reduction/reduce_and_op.mlir Outdated Show resolved Hide resolved
test/ttmlir/Decomposition/TTIR/reduction/reduce_and.mlir Outdated Show resolved Hide resolved
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.
@mmanzoorTT mmanzoorTT force-pushed the mmanzoor/reduce-and-op branch from 608fcaa to cc18937 Compare January 27, 2025 15:12
@mmanzoorTT mmanzoorTT merged commit 2a944ff into main Jan 27, 2025
23 checks passed
@mmanzoorTT mmanzoorTT deleted the mmanzoor/reduce-and-op branch January 27, 2025 20:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants