Skip to content

[opt] Support lifting operations through OneHotSelect#4394

Draft
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_930781718
Draft

[opt] Support lifting operations through OneHotSelect#4394
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_930781718

Conversation

@copybara-service

Copy link
Copy Markdown

[opt] Support lifting operations through OneHotSelect

This change introduces safety checks for lifting operations through OneHotSelect based on whether the operation preserves zero and distributes over bitwise OR. This allows more aggressive optimization of OneHotSelect nodes while maintaining correctness.

If the OneHotSelect's selector might be zero, lifting the operation is only safe if the operation preserves zero; i.e., replacing input X by 0 makes the output 0.

If the OneHotSelect's selector might have more than one bit set, lifting the operation is only safe if the operation distributes over bitwise OR; i.e., f(replace input X with A | B) produces the same result as replacing the result with f(replace input X with A) | f(replace input X with B).

This change introduces safety checks for lifting operations through OneHotSelect based on whether the operation preserves zero and distributes over bitwise OR. This allows more aggressive optimization of OneHotSelect nodes while maintaining correctness.

If the OneHotSelect's selector might be zero, lifting the operation is only safe if the operation preserves zero; i.e., replacing input X by 0 makes the output 0.

If the OneHotSelect's selector might have more than one bit set, lifting the operation is only safe if the operation distributes over bitwise OR; i.e., f(replace input X with A | B) produces the same result as replacing the result with f(replace input X with A) | f(replace input X with B).

PiperOrigin-RevId: 930781718
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.

1 participant