[XLA:CPU] OneDNN matmul library call for XLA HLO Dot operation
@d0k Could you please review this? We would like to get some feedback and suggestions. @d0k Thanks for the comments. I have addressed those. Please check. @d0k Can we move FloatNormalization passes for CPU (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L655-L669) right before  `CpuInstructionFusion`, for example, before https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L832 > @d0k Can we move FloatNormalization passes for CPU (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L655-L669) right before `CpuInstructionFusion`, for example, before https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L832

I'm not strictly opposed to that, but if your problem can be solved by configuring FloatNormalization (i.e. overriding FloatSupport) I'd prefer that solution. > > @d0k Can we move FloatNormalization passes for CPU (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L655-L669) right before `CpuInstructionFusion`, for example, before https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L832
> 
> I'm not strictly opposed to that, but if your problem can be solved by configuring FloatNormalization (i.e. overriding FloatSupport) I'd prefer that solution.

@d0k Thanks! I think we can solve by configuring FloatNormalization. I have one more question. Can we move `TreeReductionRewriter`, (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L740C22-L740C43),
right before `CpuInstructionFusion`? > @d0k Thanks! I think we can solve by configuring FloatNormalization. I have one more question. Can we move `TreeReductionRewriter`, (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc#L740C22-L740C43), right before `CpuInstructionFusion`?

I think right before layout assignment is the latest you can run it, feel free to move it there. @penpornk Now that all dependencies are merged in master branch, this PR is ready for review. Thanks! > There are 8 [Py+CPP Ubuntu CPU](https://source.cloud.google.com/results/invocations/50666fb9-baa8-4277-adf2-457e5bc701cd/log) test failures. Could you please help take a look? + Since you need to make changes anyway, could you please address my previous comments (that I said I would do from internally) as well? Thank you very much!
> 
> ```
> //tensorflow/compiler/tests:matrix_inverse_op_test_cpu                   FAILED in 3 out of 3 in 9.5s
> //tensorflow/compiler/tests:matrix_inverse_op_test_cpu_mlir_bridge_test  FAILED in 3 out of 3 in 8.5s
> //tensorflow/compiler/tests:matrix_solve_op_test_cpu                     FAILED in 3 out of 3 in 5.1s
> //tensorflow/compiler/xla/client/lib:qr_test_cpu                         FAILED in 3 out of 3 in 8.3s
> //tensorflow/compiler/xla/service/cpu/tests:cpu_eigen_dot_operation_test FAILED in 3 out of 3 in 7.1s
> //tensorflow/python/eager/polymorphic_function:polymorphic_function_xla_jit_test_cpu FAILED in 3 out of 3 in 13.9s
> //tensorflow/python/eager/polymorphic_function:polymorphic_function_xla_jit_test_cpu_mlir_bridge_test FAILED in 3 out of 3 in 13.7s
> //tensorflow/compiler/xla/tests:triangular_solve_test_cpu                FAILED in 9 out of 9 in 28.0s
> ```

Thanks @penpornk for the review. I will look into those failures and address your comments. @penpornk I think I have addressed your comments and fixed the UT failures. Please check. Thanks! @penpornk Now that 2.14 branch-cut is done, let us know if there is anything else needed for PR?  @penpornk We noticed that this PR got reverted. Could you please tell us more and a reproducer for jaxlib crash. @penpornk  We noticed that this PR got reverted. Could you please tell us more and a reproducer for jaxlib crash. @mdfaijul Sorry for the late reply! The error message reported by the JAX team was 
```
[ RUN      ] LaxBackedNumpyTests.testCorrCoef0 (shape=(5,), dtype=<class 'ml_dtypes.bfloat16'>, rowvar=True)
could not create a primitive descriptor for a matmul primitive
Fatal Python error: Illegal instruction
```
I'm thinking this could be because oneDNN tried to create a bf16 matmul primitive on a machine that doesn't have AVX512_BF16. (I just realized that jaxlib builds with `--config=mkl_open_source_only`.) I'm trying to build jaxlib reproduce the issue. Will get back to you once I have the commands. @mdfaijul Here is the repro instruction. Please run this on an AVX2 machine. If you don't have one, constraining `ONEDNN_MAX_CPU_ISA=AVX2` might do the trick too.

```bash
# Use TF docker image to save time. JAX currently requires Python 3.9+ and NumPy 1.22+
sudo docker run --name jax-dbg -w /root -it -d \  
  -v "${HOME}:/mnt" \
  tensorflow/build:latest-python3.9 \
  bash
sudo docker attach jax-dbg  # Goes into the docker image
pip install numpy wheel build  # Prereqs for building JAX
git clone https://github.com/google/jax.git
cd jax

# This JAX commit uses the XLA commit that still has the changes from your PR.
git checkout a1d8787e15914a24dc3e2f8cc46ce5542e602570  

# The run failed test. Please run on an AVX2 machine.
bazel test //tests:lax_numpy_test_cpu --test_filter=LaxBackedNumpyTests.testCorrCoef0
``` @penpornk Thanks for the reproducer. I have a fix and opened a new PR https://github.com/tensorflow/tensorflow/pull/61686 @mdfaijul Thank you very much! How much was the PR changed? If it's not too much, maybe it's better that I reverted my rollback and add those addition fixes internally (I made some internal changes while merging your PR). If it's a big changes I can reapply my changes onto your new PR instead. @penpornk It's only small change. You can see the in this commit https://github.com/tensorflow/tensorflow/pull/61686/commits/2abff97921a566d143be67451aad85e735e9f1ab
This PR enables OneDNN library call for XLA HLO dot operation through custom_call instruction. In particular,

- We insert an HLO pass for pattern matching. This pass is after DotDecomposer in the pass pipeline for CPU. The pass rewrites dot operation to custom_call operation.
- IR emitter generates LLVM IR for the custom_call for the given target function
- The symbol for the target function is registered with CustomCallTargetRegistry
- OneDNN matmul primitive is used for the computation