-
Notifications
You must be signed in to change notification settings - Fork 451
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
[XLA:CPU][oneDNN] Alias result to addend when feasible #17637
base: main
Are you sure you want to change the base?
Conversation
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.
Thank you for the PR and sorry for the delay!
xla/service/cpu/onednn_matmul.cc
Outdated
case OneDnnFusionConfig::SUM: | ||
post_ops.append_sum(); | ||
break; |
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.
CreateMatMulPrimDesc
has been refactored to onednn_util.cc
in #18527. Please rebase this PR?
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.
Done
xla/service/cpu/onednn_config.proto
Outdated
@@ -50,6 +50,7 @@ message OneDnnFusionConfig { | |||
ELU = 8; | |||
RELU6 = 9; | |||
SIGMOID = 10; | |||
SUM = 11; |
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.
Could you please add a comment explaining that this is in-place accumulation? The name alone isn't enough to disambiguate it from BINARY_ADD
.
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.
Done
if (kind == OneDnnFusionConfig::SUM) { | ||
custom_call->set_output_to_operand_aliasing({{{}, {addend_idx, {}}}}); | ||
} |
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.
Could you please add a comment in the code explaining why we need to set the aliasing?
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.
Done
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.
Changes in this file seem self-contained. Can we refactor it to another PR?
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.
Thanks @penpornk for the review! I have added a commit to address your comments and created separate PRs for isolated changes. |
This PR adds a new oneDNN post-op type that will perform in-place addition whenever there is no broadcast overhead. It also includes matmul and convolution tests to verify the functionality.