[XLA:CPU][oneDNN] : Add OneDNN softmax rewriter
Hi @sachinmuradi Can you please rebase your branch and resolve conflicts? Thank you! This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. Hi @sachinmuradi Can you please rebase your branch and resolve conflicts? Thank you! This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. Hi @sachinmuradi Can you please rebase your branch and resolve conflicts? Thank you! Hi @sachinmuradi Any update on this PR? Please. Thank you! Closing this PR since its replacement PR https://github.com/openxla/xla/pull/7540 has been merged. Thank you again for the changes!
This PR rewrites the below mentioned HLO pattern, into softmax pattern and executes it through custom call to oneDNN softmax library function.
This pattern is mainly seen in flax/jax [softmax ](https://github.com/google/jax/blob/a259df0d76e50e0e54fa5a69a7c6b78975cde10a/jax/_src/random.py#L1025)operation, with Float32 datatype. The rewrite uses 'divide' operation as root node to match the pattern and is currently called before any layout assignment / float normalization passes to avoid more convert ops or layout normalization ops during pattern matching.

**HLO Pattern :**
![image](https://github.com/tensorflow/tensorflow/assets/43043975/7875cf46-49c1-45d6-9055-8fa32ec0cae7)
