Move CUDA specific functions from asm_compiler to cuda_asm_compiler target

This avoids:
- a forward declaration of `GpuContext`
- the `:asm_compiler_header` header only target

The moved code is unchanged - I just move it from one
file to another and fix up includes and dependencies.

Note that this is adding just another `#ifdef` to the redzone allocator code. I will clean this up in a subsequent change.

PiperOrigin-RevId: 623285804

Author: beckerhe

jaxlib/cuda/BUILD

@@ -399,7 +399,7 @@ cc_library(
         ":triton_utils",
         "//jaxlib/gpu:triton_cc_proto",
         "@xla//xla/service:custom_call_status",
-        "@xla//xla/stream_executor/gpu:asm_compiler",
+        "@xla//xla/stream_executor/cuda:cuda_asm_compiler",
         "@xla//xla/tsl/cuda:cudart",
         "@tsl//tsl/platform:env",
         "@com_google_absl//absl/base:core_headers",

jaxlib/gpu/triton_kernels.cc

@@ -30,9 +30,12 @@
 #include "jaxlib/gpu/triton_utils.h"
 #include "jaxlib/gpu/vendor.h"
 #include "xla/service/custom_call_status.h"
-#include "xla/stream_executor/gpu/asm_compiler.h"
 #include "tsl/platform/env.h"
 
+#ifdef JAX_GPU_CUDA
+#include "xla/stream_executor/cuda/cuda_asm_compiler.h"
+#endif
+
 #define GPU_RETURN_IF_ERROR(expr) JAX_RETURN_IF_ERROR(JAX_AS_STATUS(expr))
 
 namespace jax::JAX_GPU_NAMESPACE {