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

extend --convert-elementwise-to-affine to support ops with scalar operands #769

Conversation

AlexanderViand-Intel
Copy link
Collaborator

@AlexanderViand-Intel AlexanderViand-Intel commented Jul 5, 2024

As originally pointed out in #143 and recently mentioned in #763, the --convert-elementwise-to-affine pass did not handle ops with "mixed" tensor/scalar operands, such as polynomial.mul_scalar.

It turns out this was a much easier fix than anticipated, at least for mul_scalar as ElementwiseMappable does not (as I wrongly assumed) imply SameOperandsAndResultType. In fact, there was already a ToDo (#534) that suggested the three-character change necessary to make this work.
However, this still does not work for polynomial.ntt/intt, as ElementwiseMappable does require all tensor operands to agree, and that wouldn't hold for a "tensorized" ntt/intt (cf. #143 (comment))

With this change, any ElementwiseMappable operation with at least one tensor operand is converted into an affine.for, and tensor operands are replaced with tensor.extracts, while scalar operands are left unmodified. This result is probably what one would expect (despite the semantics for this not being formally noted anywhere upstream, afaik), effectively "broadcasting" the scalar to the tensor operand(s).

PS: I noticed the tests for this pass also had some mostly meaningless "expected output" "tests" in it from development so I used the opportunity to remove those.


Resolves #534
(Mostly, see #763 (comment)) Resolves #143

@j2kun
Copy link
Collaborator

j2kun commented Jul 6, 2024

LGTM! How convenient that you did it right the first time :)

@j2kun j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Jul 6, 2024
@copybara-service copybara-service bot merged commit 633609a into google:main Jul 8, 2024
10 checks passed
@AlexanderViand-Intel AlexanderViand-Intel deleted the partial-elementwise-to-affine branch July 9, 2024 06:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing
Projects
None yet
3 participants