tensorflow / tensorflow

An Open Source Machine Learning Framework for Everyone
https://tensorflow.org
Apache License 2.0
182.91k stars 73.92k forks source link

PR #11241: [XLA:CPU] Enable BMM+Mul+Add for bf16 #67157

Closed copybara-service[bot] closed 1 week ago

copybara-service[bot] commented 1 week ago

PR #11241: [XLA:CPU] Enable BMM+Mul+Add for bf16

Imported from GitHub PR https://github.com/openxla/xla/pull/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=https://github.com/openxla/xla/pull/11241 from Intel-tensorflow:kanvi/bmm-mul-add_bf16 8fb80082e8c545e007ab7a2d363eb7b7c251fd07