[NVIDIA XLA:GPU] Fused MHA Support in XLA GPU: Compiler lowerings and runtime
Hi @AyanmoI Can you please resolve conflicts? Thank you! Hi @AyanmoI Can you please resolve conflicts? Thank you! @cheshire Can u take a look. The comments are addressed.  @AyanmoI Could you mark the previous comments as "resolved" explicitly if they were? Some of the previous issues still seem to apply? DotDimensionMerger broke all tests. Thanks to swift action on our parts we have been able to resolve this quickly. Thanks to [Ilia Sergachev](mailto:sergachev@google.com) for moving the pass (https://github.com/openxla/xla/commit/0b3cc823a5dda412c324552e440e779212fb081b). The FMHA pipeline was moved to run before the Triton gemm and cublas gemm rewriter. All changes work successfully now. @cheshire I have addressed all pending comments and the CI is clean. Looks like this can be merged now? @cheshire the failures seem unrelated. @ddunl Is it possible to present the failures in a more UI-friendly way? It's very confusing to figure out what test is actually failing from "Py+CPP Test Suite - Ubuntu CPU, Python". I believe it is, that's how it is for XLA afaik. tagging @jakeharmon8 @MichaelHudgins  There are some linting errors, looking as to why they are not exposed (seems like opening PRs on XLA repo is much easier):

```
(other warnings about 80 cols)
...
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:487](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=487)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:881](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=881)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:884](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=884)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:910](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=910)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:913](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=913)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1182](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1182)
Lines should be <= 80 columns wide  [whitespace/line_length] [2]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1185](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1185)
Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1559](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1559)
header guard has wrong style, please use: TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSED_MHA_THUNK_H_  [build/header_guard] [5]
	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h:16](/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h?l=16)
 line should be "#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSED_MHA_THUNK_H_"  [build/header_guard] [5]
	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h:79](http://google3/third_party/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h?l=79)
Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
	//depot/[google3/third_party/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.cc:576](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.cc?l=576)
 header guard has wrong style, please use: TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_FUSED_MHA_RUNNER_H_  [build/header_guard] [5]
	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:16](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=16)
 line should be "#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_FUSED_MHA_RUNNER_H_"  [build/header_guard] [5]
	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:260](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=260)
Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:260](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=260)
Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.cc:98](/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.cc?l=98)
``` @AyanmoI it looks like the internal errors from the linter are due to long names (_e.g._, "GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_F16();") that are nearly 80 characters. There should be two options to fix this: 1) renaming with shorter names; or 2) adding "// NOLINT" to the end of each line reported above. I think the subsequent tests were blocked by this, so we can force those to run once the changes are made. > @AyanmoI it looks like the internal errors from the linter are due to long names (_e.g._, "GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_F16();") that are nearly 80 characters. There should be two options to fix this: 1) renaming with shorter names; or 2) adding "// NOLINT" to the end of each line reported above. I think the subsequent tests were blocked by this, so we can force those to run once the changes are made.

@gcforster Can you give me the complete log of all these func names/lines? I can add "// NOLINT" to those lines. I think it would make sense to have the long names for better description. @cheshire @gcforster I have eye-balled and added some `// NOLINT` to lines. Hope this works. Also, I had added these tests right from the beginning. Curious why these failures weren't showing up earlier? Sounds good. The errors I see are in file gpu_fused_mha_test.cc: (It looks like a couple may be missing based on the latest commits still)
Line 410: GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_F16
Line 413: GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_BF16
Line 487: GetModuleFMHABMM_BMM2_non_contracting_dim_stride_not_1_HloString_F16
Line 881: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_HloString_F16_smaller
Line 884: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_HloString_BF16_smaller
Line 910: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_arg_reversal_HloString_F16
Line 913: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_arg_reversal_HloString_BF16
Line 1182: GetModuleFMHABMM1_Scale_Mask_Softmax_BMM2_arg_reversal_HloString_F16
Line 1185: GetModuleFMHABMM1_Scale_Mask_Softmax_BMM2_arg_reversal_HloString_BF16
Line 1559: Needs new line at end of file > that are nearly 80 characters.

Ok. I was counting the no. of cols to <= 80. I'll add NOLINT to the end of all the lines having these function names. > Sounds good. The errors I see are in file gpu_fused_mha_test.cc: (It looks like a couple may be missing based on the latest commits still) Line 410: GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_F16 Line 413: GetModuleFMHABMM_BMM_arg_layout_manipulation_arg_reversal_HloString_BF16 Line 487: GetModuleFMHABMM_BMM2_non_contracting_dim_stride_not_1_HloString_F16 Line 881: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_HloString_F16_smaller Line 884: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_HloString_BF16_smaller Line 910: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_arg_reversal_HloString_F16 Line 913: GetModuleFMHABMM1_Scale_Bias_Mask_Softmax_BMM2_arg_reversal_HloString_BF16 Line 1182: GetModuleFMHABMM1_Scale_Mask_Softmax_BMM2_arg_reversal_HloString_F16 Line 1185: GetModuleFMHABMM1_Scale_Mask_Softmax_BMM2_arg_reversal_HloString_BF16 Line 1559: Needs new line at end of file

Done. @cheshire @gcforster Looks like these tests just timeout:
```
Failed targets (4)
[//tensorflow/examples/custom_ops_doc/multiplex_1:multiplex_1_test](https://source.cloud.google.com/results/invocations/efa7af75-200d-46f3-9fe4-5ac98ba01877/targets/%2F%2Ftensorflow%2Fexamples%2Fcustom_ops_doc%2Fmultiplex_1:multiplex_1_test)
Timed out
9:14:49 AM
03:46
[//tensorflow/python/data/experimental/kernel_tests:group_by_reducer_test](https://source.cloud.google.com/results/invocations/efa7af75-200d-46f3-9fe4-5ac98ba01877/targets/%2F%2Ftensorflow%2Fpython%2Fdata%2Fexperimental%2Fkernel_tests:group_by_reducer_test)
Timed out
9:14:49 AM
03:47
[//tensorflow/python/data/kernel_tests:interleave_test](https://source.cloud.google.com/results/invocations/efa7af75-200d-46f3-9fe4-5ac98ba01877/targets/%2F%2Ftensorflow%2Fpython%2Fdata%2Fkernel_tests:interleave_test)
Timed out
9:14:49 AM
04:29
[//tensorflow/python/framework:ops_test](https://source.cloud.google.com/results/invocations/efa7af75-200d-46f3-9fe4-5ac98ba01877/targets/%2F%2Ftensorflow%2Fpython%2Fframework:ops_test)
Timed out
9:14:49 AM
03:43
``` > There are some linting errors, looking as to why they are not exposed (seems like opening PRs on XLA repo is much easier):
> 
> ```
> (other warnings about 80 cols)
> ...
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:487](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=487)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:881](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=881)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:884](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=884)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:910](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=910)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:913](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=913)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1182](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1182)
> Lines should be <= 80 columns wide  [whitespace/line_length] [2]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1185](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1185)
> Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
> 	///tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:1559](/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=1559)
> header guard has wrong style, please use: TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSED_MHA_THUNK_H_  [build/header_guard] [5]
> 	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h:16](/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h?l=16)
>  line should be "#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSED_MHA_THUNK_H_"  [build/header_guard] [5]
> 	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h:79](http://google3/third_party/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.h?l=79)
> Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
> 	//depot/[google3/third_party/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.cc:576](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.cc?l=576)
>  header guard has wrong style, please use: TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_FUSED_MHA_RUNNER_H_  [build/header_guard] [5]
> 	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:16](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=16)
>  line should be "#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_FUSED_MHA_RUNNER_H_"  [build/header_guard] [5]
> 	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:260](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=260)
> Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
> 	///tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h:260](/tensorflow/compiler/xla/service/gpu/gpu_fused_mha_runner.h?l=260)
> Could not find a newline character at the end of the file.  [whitespace/ending_newline] [5]
> 	///tensorflow/compiler/xla/service/gpu/fused_mha_thunk.cc:98](/tensorflow/compiler/xla/service/gpu/fused_mha_thunk.cc?l=98)
> ```

@cheshire I see your point about moving this to OpenXLA since we are blocked unnecessarily by random TF test failures. I have moved TF stuff to OpenXLA internally and it needs some effort to do this manually. Maybe we can forced merge this since we know the failures are unrelated. I am going to open the next PR in OpenXLA. Otherwise, if you have a script to do the transfer automatically, that'll be really useful. I'd keep moving this work to OpenXLA manually as a last resort. The tests are currently pending. I'll keep you posted. @AyanmoI There are some build errors. 

```
[third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc:23](//depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc?l=23&ws=tap-presubmit-server/55736082&snapshot=2):10: error: module //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test does not depend on a module exporting 'third_party/tensorflow/compiler/xla/status_macros.h'
see [http://go/cpp-features#layering_check](https://www.google.com/url?q=http://go/cpp-features%23layering_check&sa=D); to fix run:
	build_cleaner //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test
#include "third_party/tensorflow/compiler/xla/status_macros.h"
         ^
[third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc:25](//depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc?l=25&ws=tap-presubmit-server/55736082&snapshot=2):10: error: module //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test does not depend on a module exporting 'third_party/tensorflow/compiler/xla/util.h'
see [http://go/cpp-features#layering_check](https://www.google.com/url?q=http://go/cpp-features%23layering_check&sa=D); to fix run:
	build_cleaner //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test
#include "third_party/tensorflow/compiler/xla/util.h"
         ^
[third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc:26](//depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc?l=26&ws=tap-presubmit-server/55736082&snapshot=2):10: error: module //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test does not directly depend on a module exporting 'third_party/tensorflow/compiler/xla/xla_data.proto.h', which is part of indirectly-used module blaze-out/k8-cuda11-opt/bin/third_party/tensorflow/compiler/xla/xla_data.proto.h
see [http://go/cpp-features#layering_check](https://www.google.com/url?q=http://go/cpp-features%23layering_check&sa=D); to fix run:
	build_cleaner //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test
#include "third_party/tensorflow/compiler/xla/xla_data.proto.h"
         ^
[third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc:71](//depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc?l=71&ws=tap-presubmit-server/55736082&snapshot=2):42: error: unused variable 'cc' [-Werror,-Wunused-variable]
  stream_executor::CudaComputeCapability cc = GetCudaComputeCapability();
                                         ^
[third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc:241](//depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter_test.cc?l=241&ws=tap-presubmit-server/55736082&snapshot=2):42: error: unused variable 'cc' [-Werror,-Wunused-variable]
  stream_executor::CudaComputeCapability cc = GetCudaComputeCapability();
                                         ^
5 errors generated.
```

https://github.com/tensorflow/tensorflow/pull/60233/files#diff-10e7b23ad8d7400d5a23bc499fa11b06bbac4271975b009930a68a3c124fa1f7R2880 needs deps added for includes, and there is an unused variable in the error message pasted above. @AyanmoI There is a merge conflict. Can you run 'git pull --rebase' and update the PR? @AyanmoI There is still a build failure. xla_cc_test "cudnn_fused_conv_rewriter_test" is missing //third_party/tensorflow/compiler/xla:status_macros and //third_party/tensorflow/compiler/xla:util from deps. @gcforster Still seems to have some unrelated test failures. I don't see any more build issues. Can you take a look? @AyanmoI I see a build failure still. The only message is about cc_library "cudnn_fused_mha_rewriter" (in the BUILD file) needing "//third_party/tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather" under deps. This is due to the include "third_party/tensorflow/compiler/xla/stream_executor/dnn.h" > @AyanmoI I see a build failure still. The only message is about cc_library "cudnn_fused_mha_rewriter" (in the BUILD file) needing "//third_party/tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather" under deps. This is due to the include "third_party/tensorflow/compiler/xla/stream_executor/dnn.h"

Let me see if I can remove the dependency of dnn.h. Otherwise, I'll add to the BUILD file. @gcforster What do u see now? @AyanmoI Build errors:

[third_party/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc:34](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/service/gpu/tests/gpu_fused_mha_test.cc?l=34&ws=tap-presubmit-server/56911180&snapshot=2):10: error: module //third_party/tensorflow/compiler/xla/service/gpu/tests:gpu_fused_mha_test_gpu does not depend on a module exporting 'third_party/tensorflow/compiler/xla/service/gpu/cublas_cudnn.h'

and

In file included from [third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.cc:16](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.cc?l=16&ws=tap-presubmit-server/56911180&snapshot=2):
[./third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.h:21](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/service/gpu/cudnn_fused_mha_rewriter.h?l=21&ws=tap-presubmit-server/56911180&snapshot=2):10: error: module //third_party/tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter does not depend on a module exporting 'third_party/tensorflow/compiler/xla/stream_executor/dnn.h'

Changes needed for these errors:
third_party/tensorflow/compiler/xla/service/gpu/BUILD
cc_library "cudnn_fused_mha_rewriter" deps
(remove) "//third_party/tensorflow/compiler/xla/stream_executor:dnn_proto_cc"
(add) "//third_party/tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather"

third_party/tensorflow/compiler/xla/service/gpu/tests/BUILD
xla_test "gpu_fused_mha_test" deps
(remove) "//third_party/absl/memory"
(add) "//third_party/tensorflow/compiler/xla/service/gpu:cublas_cudnn" One thing that might be worth trying - you can add `--features=layering_check` to your bazel command. I'm not sure that this will work how we want in TF (there's lots of BUILD files that disable layering_check) but layering_check is what's producing the build errors internally. (XLA external CI is updated to use layering_check as well as of today) > 

@gcforster I can see `stream_executor_install_hdrs_gather` defined anywhere. I get this error locally when I include `"//tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather",` in the deps for `cudnn_fused_mha_rewriter`
```
ERROR: /workspace/tensorflow/tensorflow/compiler/xla/service/gpu/BUILD:2866:11: in deps attribute of cc_library rule //tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter: _transitive_hdrs rule '//tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather' is misplaced here (expected genrule, cc_library, cc_inc_library, cc_embed_data, go_library, objc_library, cc_import, cc_proto_library, gentpl, gentplvars, genantlr, sh_library, cc_binary or cc_test) and '//tensorflow/compiler/xla/stream_executor:stream_executor_install_hdrs_gather' does not have mandatory providers: 'CcInfo'
ERROR: /workspace/tensorflow/tensorflow/compiler/xla/service/gpu/BUILD:2866:11: Analysis of target '//tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter' failed
ERROR: Analysis of target '//tensorflow/compiler/xla/service/gpu:cudnn_fused_mha_rewriter_test' failed; build aborted:
INFO: Elapsed time: 5.771s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (177 packages loaded, 16067 targets configured)
ERROR: Couldn't start the build. Unable to run tests
``` > One thing that might be worth trying - you can add `--features=layering_check` to your bazel command. I'm not sure that this will work how we want in TF (there's lots of BUILD files that disable layering_check) but layering_check is what's producing the build errors internally. (XLA external CI is updated to use layering_check as well as of today)

@ddunl `--features=layering_check` doesn't seem to reveal these build issues locally.
cuDNN graph API will support the runtime generation of a fused kernel for various flavors of Fused Dot Attention modules. Currently XLA fuser does not have the capability to perform such fusions and integration of cuDNN graph API support makes it a great opportunity to optimize the graph further and enhance performance.

This PR adds the plumbing for the following patterns for inference use case only:
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 - BMM2
BMM1 - Softmax - Dropout - BMM2
BMM1 - Scale - Bias - Softmax - BMM2
BMM1 - Scale - Bias - Softmax - Dropout - BMM2

<img width="386" alt="image" src="https://github.com/tensorflow/tensorflow/assets/42984676/c575bba3-532f-483d-a9f4-9cf427ce0fa5">

The stream executor changes have been merged via https://github.com/tensorflow/tensorflow/pull/60145
