-
Notifications
You must be signed in to change notification settings - Fork 344
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] Enable BMM+Mul+Add for bf16 #11241
[XLA:CPU] Enable BMM+Mul+Add for bf16 #11241
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 I'm sorry for the delay! I have a minor comment.
if (!IsBatchDot(*dot_instr)) { | ||
if (!IsRowMajor(lhs_shape) || !IsRowMajor(rhs_shape) || | ||
!IsRowMajor(output_shape)) { | ||
return false; | ||
} | ||
} |
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.
Nit: Let's avoid nested if
s. (Please apply code formatter.)
if (!IsBatchDot(*dot_instr)) { | |
if (!IsRowMajor(lhs_shape) || !IsRowMajor(rhs_shape) || | |
!IsRowMajor(output_shape)) { | |
return false; | |
} | |
} | |
if (!IsBatchDot(*dot_instr) && | |
(!IsRowMajor(lhs_shape) || !IsRowMajor(rhs_shape) || | |
!IsRowMajor(output_shape))) { | |
return false; | |
} |
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.
@penpornk I made the requested change. Thanks!
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 again!
Imported from GitHub PR openxla/xla#11241 This PR enable BatchMatmul + Mul + Add fusion for BF16 and also fixes a bug for the same. Copybara import of the project: -- abdbf9b89925c8553296122c846327b6dffa86ce by Kanvi Khanna <kanvi.khanna@intel.com>: Enable BMM+Mul+Add for bf16 -- 6674ac76b868f0399731d6185287eec76244df3a by Kanvi Khanna <kanvi.khanna@intel.com>: fix test -- 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 by Kanvi Khanna <kanvi.khanna@intel.com>: address review comment, fix test, format Merging this change closes #11241 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11241 from Intel-tensorflow:kanvi/bmm-mul-add_bf16 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 PiperOrigin-RevId: 631717784
Imported from GitHub PR openxla/xla#11241 This PR enable BatchMatmul + Mul + Add fusion for BF16 and also fixes a bug for the same. Copybara import of the project: -- abdbf9b89925c8553296122c846327b6dffa86ce by Kanvi Khanna <kanvi.khanna@intel.com>: Enable BMM+Mul+Add for bf16 -- 6674ac76b868f0399731d6185287eec76244df3a by Kanvi Khanna <kanvi.khanna@intel.com>: fix test -- 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 by Kanvi Khanna <kanvi.khanna@intel.com>: address review comment, fix test, format Merging this change closes #11241 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11241 from Intel-tensorflow:kanvi/bmm-mul-add_bf16 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 PiperOrigin-RevId: 631717784
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11241 from Intel-tensorflow:kanvi/bmm-mul-add_bf16 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 PiperOrigin-RevId: 631306324
Updates LLVM usage to match [aefad851672e](llvm/llvm-project@aefad851672e) FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11241 from Intel-tensorflow:kanvi/bmm-mul-add_bf16 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 PiperOrigin-RevId: 631745757
Imported from GitHub PR openxla/xla#11241 This PR enable BatchMatmul + Mul + Add fusion for BF16 and also fixes a bug for the same. Copybara import of the project: -- abdbf9b89925c8553296122c846327b6dffa86ce by Kanvi Khanna <kanvi.khanna@intel.com>: Enable BMM+Mul+Add for bf16 -- 6674ac76b868f0399731d6185287eec76244df3a by Kanvi Khanna <kanvi.khanna@intel.com>: fix test -- 8fb80082e8c545e007ab7a2d363eb7b7c251fd07 by Kanvi Khanna <kanvi.khanna@intel.com>: address review comment, fix test, format Merging this change closes #11241 PiperOrigin-RevId: 631745536
This PR enable BatchMatmul + Mul + Add fusion for BF16 and also fixes a bug for the same.