[XLA:CPU] Layer Norm XLA HLO Pattern Matcher with oneDNN custom call rewrite
Hi @akhilgoe Can you please rebase your branch and resolve conflicts? Thank you! Hi @gbaned...the conflicts have been resolved. Thanks! Hi @akhilgoe Can you please check @d0k's comments and keep us posted ? Thank you! Hi @penpornk, could you help to review/merge this PR? Hi @akhilgoe Can you please rebase your branch and resolve conflicts? Thank you! Hi @akhilgoe Can you please rebase your branch and resolve conflicts? Thank you! I have resolved the merge conflicts. Hi @akhilgoe Can you please rebase your branch and resolve the conflicts? Thank you! Replaced by https://github.com/openxla/xla/pull/7485. please close this PR. Closing this PR as its replacement (https://github.com/openxla/xla/pull/7485) has been merged.
This PR enables OneDNN library call for the matched XLA HLO Layer Norm pattern through custom_call instruction. In particular, this PR adds:

1. oneDNN ops rewriter pass that will rewrite custom oneDNN ops like Layer Norm and Softmax. Currently, this pass is located before the BF16 all-reduce to F32 promotion.
2. Layer Norm pattern matcher as seen in some Flax models
3. oneDNN Layer Norm custom call rewrite
4. Variable number of arguments for oneDNN custom call targets (by @mdfaijul )
5. A test for XLA HLO Layer Norm pattern match and rewrite verification