How to enable oneDNN optimization？

[ZhongYFeng]

❓ Questions and Help

I am doing training and inference on XLA_CPU, but I find that the training speed is particularly slow. Compared with pytorch, the training speed is about 10 times slower.
According to the log, I found that mklcnn acceleration is enabled by default during pytorch training, but when I use xla training, mklcnn is not enabled.

2021-11-12 16:12:44.683410: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.

How can I enable mkldnn on xla to get training acceleration?

[ZhongYFeng]

When I was compiling xla, I added the --config=mkl option to the bazel build of the build_torch_xla_libs.sh file, but there was still the same warning, indicating that mkl did not take effect

[miaoshasha]

Thanks Zhong! We'll be looking at it and get back to you, thanks for bring the question up!

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.