Fused MHA Support in XLA GPU: Stream Executor changes
@cheshire There were some issues rebasing the other [PR](https://github.com/tensorflow/tensorflow/pull/60002) branch. So I created a this PR with the same changes  and closed the other PR. Please approve this PR. Sorry about the confusion. I am resolving the build issue in the tests. This is due to introduction of new cudnn frontend. I think @akuegel is also fixing this internally. > I think @akuegel is also fixing this internally.

I just checked, and there are differences in the fixes. Let's import again and see whether additional fixes will be needed, and if not, we can take it as is. This has been merged. I had to also apply some fixes suggested by ClangTidy as it was also blocking the submit. Unfortunately this change needs to be rolled back, it seems it breaks JAX build under CUDA 11.4 and CuDNN 8.2 > Unfortunately this change needs to be rolled back, it seems it breaks JAX build under CUDA 11.4 and CuDNN 8.2

@akuegel what are the issues? We can work on fixing them and re-merge. Can we re-open this PR to work on fixing the issues? I don't see a way to re-open this.  It fails to build under CUDA 11.4 /CuDNN 8.2:

```
external/xla/xla/stream_executor/cuda/cuda_dnn.cc: At global scope:
external/xla/xla/stream_executor/cuda/cuda_dnn.cc:3500:5: error: 'cudnnBackendTensorReordering_t' has not been declared
 3500 |     cudnnBackendTensorReordering_t cudnn_tensor_order_type =
      |     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
external/xla/xla/stream_executor/cuda/cuda_dnn.cc:3501:9: error: 'CUDNN_TENSOR_REORDERING_NONE' was not declared in this scope; did you mean 'cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_NONE'?
 3501 |         CUDNN_TENSOR_REORDERING_NONE,
      |         ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
      |         cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_NONE
In file included from bazel-out/k8-opt/bin/external/cudnn_frontend_archive/_virtual_includes/cudnn_frontend/third_party/cudnn_frontend/include/cudnn_frontend_ConvDesc.h:35,
                 from bazel-out/k8-opt/bin/external/cudnn_frontend_archive/_virtual_includes/cudnn_frontend/third_party/cudnn_frontend/include/cudnn_frontend.h:100,
                 from external/xla/xla/stream_executor/cuda/cuda_dnn.cc:56:
bazel-out/k8-opt/bin/external/cudnn_frontend_archive/_virtual_includes/cudnn_frontend/third_party/cudnn_frontend/include/cudnn_frontend_utils.h:376:5: note: 'cudnn_frontend::cudnnBackendTensorReordering_t::CUDNN_TENSOR_REORDERING_NONE' declared here
  376 |     CUDNN_TENSOR_REORDERING_NONE,
      |     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~

etc.
```

 @akuegel After reopening the PR, why does it now have only one file? Was this partially merged? It looks like only the last commit got reverted. I don't think the issue (of build failing with CUDA 11.4 and CuDNN 8.2) will have resolved with reverting the last commit only. > 

Actually all of it was reverted:

https://github.com/tensorflow/tensorflow/commit/7a4d2e8b1a503a45ff86d40b0387fc832302379f

Why it doesn't show properly merged in this PR, I don't know :-(
Maybe it is easier to create a new PR? > > 
> 
> Actually all of it was reverted:
> 
> [7a4d2e8](https://github.com/tensorflow/tensorflow/commit/7a4d2e8b1a503a45ff86d40b0387fc832302379f)
> 
> Why it doesn't show properly merged in this PR, I don't know :-( Maybe it is easier to create a new PR?

Yeah..I guess. Let me do that. @akuegel @cheshire opened a new PR https://github.com/tensorflow/tensorflow/pull/60139
cuDNN currently supports the following patterns for Multi-headed attention:

BMM1 - BMM2
BMM1 - Scale - Bias - Mask - Softmax - BMM2
BMM1 - Scale - Bias - Mask - Softmax - Dropout - BMM2
BMM1 - Scale - Mask - Softmax - BMM2
BMM1 - Scale - Mask - Softmax - Dropout - BMM2
BMM1 - Softmax - Dropout - BMM2

This PR adds support for the stream executor for these patterns.