Prepare move of CUDA compilation functionality

Currently PTX to SASS compilation functions are exported by the target `stream_executor/gpu:asm_compiler`, but implemented in `stream_executor/cuda:cuda_asm_compiler`.

I'm trying to clean this up by moving everything into `:cuda_asm_compiler`. Since we can't do atomic cross cutting changes between Jaxlib and XLA, I'm only preparing `:cuda_asm_compiler` as an alias for Jaxlib to depend on.

In a subsequent change I will move all the code over.

PiperOrigin-RevId: 622885343

third_party/xla/xla/stream_executor/cuda/BUILD

@@ -560,6 +560,22 @@ xla_cc_test(
 
 cuda_only_cc_library(
     name = "cuda_asm_compiler",
+    hdrs = ["cuda_asm_compiler.h"],
+    visibility = internal_visibility([
+        "//third_party/py/jax:__subpackages__",
+        "//tensorflow/compiler/mlir/tools/kernel_gen:__subpackages__",
+        "//xla/service/gpu:__subpackages__",
+        "//xla/stream_executor:__subpackages__",
+        "//tensorflow/core/kernels:__subpackages__",
+    ]),
+    deps = [
+        ":cuda_asm_compiler_legacy",
+        "//xla/stream_executor/gpu:asm_compiler",
+    ],
+)
+
+cuda_only_cc_library(
+    name = "cuda_asm_compiler_legacy",
     srcs = ["cuda_asm_compiler.cc"],
     copts = tf_additional_gpu_compilation_copts(),
     deps = [

third_party/xla/xla/stream_executor/cuda/cuda_asm_compiler.h

@@ -0,0 +1,22 @@
+/* Copyright 2020 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_CUDA_CUDA_ASM_COMPILER_H_
+#define XLA_STREAM_EXECUTOR_CUDA_CUDA_ASM_COMPILER_H_
+
+// TODO(hebecker): Move CUDA-specific functions from asm_compiler.h into here
+#include "xla/stream_executor/gpu/asm_compiler.h"  // IWYU pragma: export
+
+#endif  // XLA_STREAM_EXECUTOR_CUDA_CUDA_ASM_COMPILER_H_

third_party/xla/xla/stream_executor/gpu/BUILD

@@ -495,7 +495,7 @@ gpu_only_cc_library(
         "@local_tsl//tsl/platform:statusor",
         "@local_tsl//tsl/platform:subprocess",
     ] + if_cuda_is_configured([
-        "//xla/stream_executor/cuda:cuda_asm_compiler",
+        "//xla/stream_executor/cuda:cuda_asm_compiler_legacy",
         "//xla/stream_executor/cuda:cuda_driver",
         "//xla/stream_executor/cuda:ptxas_wrapper",
         "//xla/stream_executor/cuda:nvlink_wrapper",