load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library")
load(
    "@local_config_rocm//rocm:build_defs.bzl",
    "if_rocm_is_configured",
)
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal", "xla_py_proto_library")
load("//xla/tests:build_defs.bzl", "xla_test")
load("//xla/tsl:tsl.bzl", "internal_visibility", "nvtx_headers")
load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
load("//xla/tsl/platform:build_config.bzl", "tf_proto_library")
load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = internal_visibility([":friends"]),
    licenses = ["notice"],
)

package_group(
    name = "friends",
    includes = ["//xla:friends"],
)

#===-------------------------------------------------------------------------------------------===//
# Runtime tracing libraries
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "thunk_runtime_dependencies",
    # Register GPU collectives plugin.
    deps = ["//xla/backends/gpu/collectives:gpu_collectives_plugin"],
)

cc_library(
    name = "annotation",
    srcs = ["annotation.cc"],
    hdrs = ["annotation.h"],
    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
    deps = [
        "//xla:printer",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/utils:hlo_longest_prefix",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@local_tsl//tsl/platform:errors",
        "@local_tsl//tsl/profiler/lib:nvtx_utils",
        "@local_tsl//tsl/profiler/lib:scoped_annotation",
    ] + if_cuda_is_configured(nvtx_headers()),
)

#===-------------------------------------------------------------------------------------------===//
# Command Buffer Integration
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "command_state",
    srcs = ["command_state.cc"],
    hdrs = ["command_state.h"],
    deps = [
        "//xla/tsl/lib/gtl:int_type",
        "@com_google_absl//absl/base:no_destructor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
    ],
)

xla_cc_test(
    name = "command_state_test",
    srcs = ["command_state_test.cc"],
    deps = [
        ":command_state",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "command",
    srcs = ["command.cc"],
    hdrs = ["command.h"],
    deps = [
        ":command_state",
        ":thunk",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/stream_executor:command_buffer",
    ],
)

cc_library(
    name = "command_buffer_cmd",
    srcs = ["command_buffer_cmd.cc"],
    hdrs = ["command_buffer_cmd.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":annotation",
        ":collective_broadcast_thunk",
        ":collective_execution",
        ":collective_permute_thunk",
        ":collective_thunk",
        ":command",
        ":command_state",
        ":copy_thunk",
        ":custom_call_thunk",
        ":dynamic_slice_thunk",
        ":gpublas_lt_matmul_thunk",
        ":p2p_thunk_common",
        ":recv_thunk",
        ":send_thunk",
        ":shaped_slice",
        ":thunk",
        ":while_thunk",
        "//xla:debug_options_flags",
        "//xla:executable_run_options",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:reduction_kind",
        "//xla/ffi:attribute_map",
        "//xla/ffi:call_frame",
        "//xla/ffi:execution_state",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/hlo/evaluator:hlo_evaluator",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/runtime:device_id",
        "//xla/runtime:execution_graph",
        "//xla/runtime:object_pool",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer",
        "//xla/service:custom_call_status_internal",
        "//xla/service:custom_call_status_public_headers",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu:stream_executor_util",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:tensor_map",
        "//xla/stream_executor:trace_command_buffer_factory",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/profiler/lib:scoped_annotation",
    ],
)

xla_test(
    name = "command_buffer_cmd_test",
    srcs = ["command_buffer_cmd_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command",
        ":command_buffer_cmd",
        ":command_state",
        ":copy_thunk",
        ":shaped_slice",
        ":thunk",
        "//xla:shape_util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor:stream_executor_address_allocator",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:gpu_test_kernels_fatbin",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "//xla/tsl/util:safe_reinterpret_cast",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "command_buffer_cmd_emitter",
    srcs = ["command_buffer_cmd_emitter.cc"],
    hdrs = ["command_buffer_cmd_emitter.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":collective_broadcast_thunk",
        ":collective_permute_thunk",
        ":collective_thunk",
        ":command",
        ":command_buffer_cmd",
        ":command_state",
        ":conditional_thunk",
        ":copy_thunk",
        ":cudnn_thunk",
        ":custom_call_thunk",
        ":custom_kernel_thunk",
        ":dynamic_slice_thunk",
        ":gemm_thunk",
        ":gpublas_lt_matmul_thunk",
        ":kernel_thunk",
        ":memset_thunk",
        ":recv_thunk",
        ":replica_id_thunk",
        ":send_thunk",
        ":sequential_thunk",
        ":thunk",
        ":while_thunk",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

#===-------------------------------------------------------------------------------------------===//
# XLA Thunks Runtime
#===-------------------------------------------------------------------------------------------===//

cc_library(
    name = "dynamic_slice_thunk",
    srcs = ["dynamic_slice_thunk.cc"],
    hdrs = ["dynamic_slice_thunk.h"],
    deps = [
        ":dynamic_slice_thunk_proto_cc",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/evaluator:hlo_evaluator",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
    ],
)

xla_test(
    name = "dynamic_slice_thunk_test",
    srcs = ["dynamic_slice_thunk_test.cc"],
    backends = [
        "gpu",
    ],
    deps = [
        ":collective_multimem_registry",
        ":custom_call_thunk",
        ":dynamic_slice_thunk",
        ":dynamic_slice_thunk_proto_cc",
        ":gemm_thunk",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_deserialization",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu:ffi",
        "//xla/ffi",
        "//xla/ffi:attribute_map",
        "//xla/ffi:ffi_api",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/parser:hlo_parser",
        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "command_buffer_thunk",
    srcs = ["command_buffer_thunk.cc"],
    hdrs = ["command_buffer_thunk.h"],
    deps = [
        ":command",
        ":command_buffer_cmd",
        ":command_state",
        ":sequential_thunk",  # build_cleaner: keep
        ":thunk",
        "//xla/service:buffer_assignment",  # build_cleaner: keep
        "//xla/service/gpu:buffer_allocations",  # build_cleaner: keep
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/profiler/lib:profiler_lock",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_tsl//tsl/profiler/lib:traceme_encode",
    ] + xla_internal(["service:execution_graph_visualizer_google"]),
)

xla_test(
    name = "command_buffer_thunk_test",
    srcs = ["command_buffer_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command",
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":dynamic_slice_thunk",
        ":gpublas_lt_matmul_thunk",
        ":memset_thunk",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:gpu_test_kernels",
        "//xla/stream_executor/gpu:gpu_test_kernels_fatbin",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/profiler/lib:profiler_lock",
    ],
)

xla_test(
    name = "cuda_command_buffer_thunk_test",
    srcs = ["cuda_command_buffer_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["cuda-only"],
    deps = [
        ":command",
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":thunk",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:command_buffer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:engine_options",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/cuda:cudnn_plugin",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@cudnn_frontend_archive//:cudnn_frontend",
        "@local_config_cuda//cuda:cuda_headers",
    ],
)

cc_library(
    name = "conditional_thunk",
    srcs = ["conditional_thunk.cc"],
    hdrs = ["conditional_thunk.h"],
    deps = [
        ":host_memory_pool",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/functional:overload",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "conditional_thunk_test",
    srcs = ["conditional_thunk_test.cc"],
    deps = [
        ":conditional_thunk",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla/service:buffer_assignment",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

cc_library(
    name = "convolution_thunk",
    srcs = ["convolution_thunk.cc"],
    hdrs = ["convolution_thunk.h"],
    deps = [
        ":shaped_slice",
        ":thunk",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:gpu_conv_runner",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:lazy_op_runner",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/protobuf:dnn_proto_cc",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "convolution_thunk_test",
    srcs = ["convolution_thunk_test.cc"],
    deps = [
        ":convolution_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "convolution_reorder_thunk",
    srcs = ["convolution_reorder_thunk.cc"],
    hdrs = ["convolution_reorder_thunk.h"],
    deps = [
        ":convolution_filter_thunk_proto_cc",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:buffer_assignment_proto_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:dnn",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "convolution_reorder_thunk_test",
    srcs = ["convolution_reorder_thunk_test.cc"],
    deps = [
        ":convolution_reorder_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

tf_proto_library(
    name = "convolution_filter_thunk_proto",
    srcs = ["convolution_filter_thunk.proto"],
    deps = [
        "//xla/service:shaped_slice_proto",
    ],
)

tf_proto_library(
    name = "copy_thunk_proto",
    srcs = ["copy_thunk.proto"],
    protodeps = [
        # keep sorted
        "//xla/service:buffer_assignment_proto",
    ],
)

cc_library(
    name = "copy_thunk",
    srcs = ["copy_thunk.cc"],
    hdrs = ["copy_thunk.h"],
    deps = [
        ":copy_thunk_proto_cc",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla:shape_util",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "copy_thunk_test",
    srcs = ["copy_thunk_test.cc"],
    deps = [
        ":copy_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "cub_sort_thunk",
    srcs = ["cub_sort_thunk.cc"],
    hdrs = ["cub_sort_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/ffi:call_frame",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "cub_sort_thunk_test",
    srcs = ["cub_sort_thunk_test.cc"],
    tags = ["gpu"],
    deps = [
        ":cub_sort_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/service:buffer_assignment",
        "//xla/service:platform_util",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ] + if_cuda_is_configured([
        "//xla/stream_executor/cuda:all_runtime",
    ]) + if_rocm_is_configured([
        "//xla/stream_executor/rocm:all_runtime",
    ]),
)

cc_library(
    name = "custom_call_thunk",
    srcs = ["custom_call_thunk.cc"],
    hdrs = ["custom_call_thunk.h"],
    deps = [
        ":collective_clique_requests",
        ":collective_cliques",
        ":collective_params",
        ":custom_call_target",
        ":shaped_slice",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:util",
        "//xla/ffi",
        "//xla/ffi:attribute_map",
        "//xla/ffi:call_frame",
        "//xla/ffi:execution_context",
        "//xla/ffi:execution_state",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/runtime:object_pool",
        "//xla/service:buffer_assignment",
        "//xla/service:custom_call_status",
        "//xla/service:custom_call_status_internal",
        "//xla/service:custom_call_target_registry",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform",
    ],
)

xla_test(
    name = "custom_call_thunk_test",
    srcs = ["custom_call_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":collective_multimem_registry",
        ":custom_call_thunk",
        ":shaped_slice",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla/ffi",
        "//xla/ffi:attribute_map",
        "//xla/ffi:execution_state",
        "//xla/ffi:ffi_api",
        "//xla/ffi:type_registry",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:custom_call_status_public_headers",
        "//xla/service:custom_call_target_registry",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_address_allocator",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "fft_thunk",
    srcs = ["fft_thunk.cc"],
    hdrs = ["fft_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:buffer_assignment_proto_cc",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:fft",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "fft_thunk_test",
    srcs = ["fft_thunk_test.cc"],
    deps = [
        ":fft_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "gemm_thunk",
    srcs = ["gemm_thunk.cc"],
    hdrs = ["gemm_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/profiler/lib:nvtx_utils",
    ],
)

xla_cc_test(
    name = "gemm_thunk_test",
    srcs = ["gemm_thunk_test.cc"],
    deps = [
        ":gemm_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "gpublas_lt_matmul_thunk",
    srcs = ["gpublas_lt_matmul_thunk.cc"],
    hdrs = ["gpublas_lt_matmul_thunk.h"],
    deps = [
        ":thunk",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "gpublas_lt_matmul_thunk_test",
    srcs = ["gpublas_lt_matmul_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["pjrt_migration_candidate"],
    use_legacy_runtime = True,
    deps = [
        ":gpublas_lt_matmul_thunk",
        ":thunk",
        ":thunk_id",
        "//xla:error_spec",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:cublas_cudnn",
        "//xla/service/gpu:matmul_utils",
        "//xla/service/gpu/transforms:gemm_rewriter",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:semantic_version",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:gpu_blas_lt",
        "//xla/tests:hlo_test_base",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/time",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

cc_library(
    name = "infeed_thunk",
    srcs = ["infeed_thunk.cc"],
    hdrs = ["infeed_thunk.h"],
    deps = [
        ":shaped_slice",
        ":thunk",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_transfer_manager",
        "//xla/service/gpu:io_feed_manager",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "infeed_thunk_test",
    srcs = ["infeed_thunk_test.cc"],
    deps = [
        ":infeed_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

cc_library(
    name = "kernel_thunk",
    srcs = ["kernel_thunk.cc"],
    hdrs = ["kernel_thunk.h"],
    deps = [
        ":print_buffer_contents",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_tsl//tsl/profiler/lib:traceme_encode",
    ],
)

xla_test(
    name = "kernel_thunk_test",
    srcs = ["kernel_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":kernel_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_id",
        ":thunk_proto_cc",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_test_kernels",
        "//xla/stream_executor/gpu:tma_metadata",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cuda_library(
    name = "select_k_exec_raft",
    srcs = ["select_k_exec_raft.cc"],
    hdrs = ["select_k_exec.h"],
    copts = [
        "-fexceptions",
        "-DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE",
    ],
    tags = ["cuda-only"],
    deps = [
        "//xla:status_macros",
        "//xla:types",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:scratch_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@local_config_cuda//cuda:cuda_headers",
        "@raft//:raft_matrix",
    ],
)

cc_library(
    name = "select_k_exec_stub",
    srcs = ["select_k_exec_stub.cc"],
    hdrs = ["select_k_exec.h"],
    deps = [
        "//xla:types",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/status",
    ],
)

xla_test(
    name = "select_k_exec_raft_test",
    srcs = ["select_k_exec_raft_test.cc"],
    backends = [
        "gpu",
    ],
    tags = [
        "cuda-only",
    ],
    deps = [
        ":select_k_exec_raft",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "select_k_thunk",
    srcs = ["select_k_thunk.cc"],
    hdrs = ["select_k_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_id",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:types",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_allocator",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:statusor",
    ] + if_cuda_is_configured(
        [":select_k_exec_raft"],
        no_cuda = [":select_k_exec_stub"],
    ),
)

xla_cc_test(
    name = "select_k_thunk_test",
    srcs = ["select_k_thunk_test.cc"],
    deps = [
        ":select_k_thunk",
        ":thunk",
        ":thunk_id",
        ":thunk_proto_cc",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "memset_thunk",
    srcs = ["memset_thunk.cc"],
    hdrs = ["memset_thunk.h"],
    deps = [
        ":shaped_slice",
        ":thunk",
        "//xla:shape_util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "memset_thunk_test",
    srcs = ["memset_thunk_test.cc"],
    deps = [
        ":memset_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_gather_thunk",
    srcs = ["all_gather_thunk.cc"],
    hdrs = ["all_gather_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:future",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

xla_cc_test(
    name = "all_gather_thunk_test",
    srcs = ["all_gather_thunk_test.cc"],
    deps = [
        ":all_gather_thunk",
        ":collective_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "collective_kernel_thunk",
    srcs = ["collective_kernel_thunk.cc"],
    hdrs = ["collective_kernel_thunk.h"],
    deps = [
        ":all_reduce",
        ":collective_cliques",
        ":collective_metadata_thunk",
        ":collective_multimem",
        ":collective_params",
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:rank_id",
        "//xla/core/collectives:reduction_kind",
        "//xla/runtime:device_id",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/stream_executor/gpu:collective_kernel_metadata",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "collective_kernel_thunk_test",
    srcs = ["collective_kernel_thunk_test.cc"],
    backend_tags = {
        "gpu": [
            "multi_gpu",
        ],
    },
    backends = ["h100"],
    deps = [
        ":collective_clique_requests",
        ":collective_kernel_thunk",
        ":collective_multimem_registry",
        ":collective_params",
        ":collective_thunk",
        ":thunk",
        "//xla:array",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/core/collectives:reduction_kind",
        "//xla/pjrt:worker_thread",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer_hdr",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/service/gpu:ptx_compile_options_from_debug_options",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/cuda:assemble_compilation_provider",
        "//xla/stream_executor/cuda:compilation_options",
        "//xla/stream_executor/cuda:compilation_provider",
        "//xla/stream_executor/cuda:compilation_provider_options",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/cuda:cuda_platform",
        "//xla/stream_executor/gpu:gpu_init",
        "//xla/tsl/concurrency:future",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_reduce_thunk",
    srcs = ["all_reduce_thunk.cc"],
    hdrs = ["all_reduce_thunk.h"],
    deps = [
        ":collective_kernel_thunk",
        ":collective_thunk",
        ":thunk",
        "//xla:future",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:reduction_kind",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

xla_cc_test(
    name = "all_reduce_thunk_test",
    srcs = ["all_reduce_thunk_test.cc"],
    deps = [
        ":all_reduce_thunk",
        ":collective_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_to_all_thunk",
    srcs = ["all_to_all_thunk.cc"],
    hdrs = ["all_to_all_thunk.h"],
    deps = [
        ":collective_execution",
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:rendezvous",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:event",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "all_to_all_thunk_test",
    srcs = ["all_to_all_thunk_test.cc"],
    deps = [
        ":all_to_all_thunk",
        ":collective_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "ragged_all_to_all_thunk",
    srcs = ["ragged_all_to_all_thunk.cc"],
    hdrs = ["ragged_all_to_all_thunk.h"],
    deps = [
        ":collective_thunk",
        ":ragged_all_to_all",
        ":thunk",
        "//xla:future",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:rendezvous",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:event",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

xla_cc_test(
    name = "ragged_all_to_all_thunk_test",
    srcs = ["ragged_all_to_all_thunk_test.cc"],
    deps = [
        ":collective_thunk",
        ":ragged_all_to_all_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "collective_broadcast_thunk",
    srcs = ["collective_broadcast_thunk.cc"],
    hdrs = ["collective_broadcast_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:future",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:status_macros",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

xla_test(
    name = "collective_broadcast_thunk_test",
    srcs = ["collective_broadcast_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["pjrt_migration_candidate"],
    use_legacy_runtime = True,
    deps = [
        ":collective_broadcast_thunk",
        ":collective_thunk",
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:verified_hlo_module",
        "//xla/service:backend",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:gpu_executable",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tests:hlo_test_base",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_permute_thunk",
    srcs = ["collective_permute_thunk.cc"],
    hdrs = ["collective_permute_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer",
        "//xla/service:rendezvous",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:node_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

xla_test(
    name = "collective_permute_thunk_test",
    srcs = ["collective_permute_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["pjrt_migration_candidate"],
    use_legacy_runtime = True,
    deps = [
        ":collective_permute_thunk",
        ":collective_thunk",
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:verified_hlo_module",
        "//xla/service:backend",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:gpu_executable",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tests:hlo_test_base",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_clique_requests",
    srcs = ["collective_clique_requests.cc"],
    hdrs = ["collective_clique_requests.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
    ],
)

xla_cc_test(
    name = "collective_clique_requests_test",
    srcs = ["collective_clique_requests_test.cc"],
    deps = [
        ":collective_clique_requests",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/runtime:device_id",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
    ],
)

cc_library(
    name = "collective_multimem",
    srcs = ["collective_multimem.cc"],
    hdrs = ["collective_multimem.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:rank_id",
        "//xla/runtime:device_id",
        "//xla/service:rendezvous",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_executor_header",
        "//xla/stream_executor/gpu:multicast_memory",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "collective_multimem_registry",
    srcs = ["collective_multimem_registry.cc"],
    hdrs = ["collective_multimem_registry.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        ":collective_multimem",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/runtime:device_id",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "collective_params",
    srcs = ["collective_params.cc"],
    hdrs = ["collective_params.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//xla:executable_run_options",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/runtime:device_id",
        "//xla/service:computation_placer",
        "//xla/service:executable",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "collective_cliques",
    srcs = ["collective_cliques.cc"],
    hdrs = ["collective_cliques.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        ":collective_clique_requests",
        ":collective_params",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_cliques",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/runtime:device_id",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/base:no_destructor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "collective_execution",
    srcs = ["collective_execution.cc"],
    hdrs = ["collective_execution.h"],
    deps = [
        ":collective_cliques",
        ":collective_params",
        "//xla:debug_options_flags",
        "//xla:status_macros",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:communicator",
        "//xla/runtime:device_id",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "collective_thunk",
    srcs = ["collective_thunk.cc"],
    hdrs = ["collective_thunk.h"],
    deps = [
        ":collective_execution",
        ":collective_params",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:debug_options_flags",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer",
        "//xla/service:rendezvous",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/llvm_ir:llvm_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "collective_thunk_test",
    srcs = ["collective_thunk_test.cc"],
    deps = [
        ":collective_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "p2p_thunk_common",
    srcs = ["p2p_thunk_common.cc"],
    hdrs = ["p2p_thunk_common.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/parser:hlo_parser",
        "//xla/runtime:device_id",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer_hdr",
        "//xla/service:source_target_pairs",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@llvm-project//mlir:IR",
    ],
)

cc_library(
    name = "recv_thunk",
    srcs = ["recv_thunk.cc"],
    hdrs = ["recv_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "recv_thunk_test",
    srcs = ["recv_thunk_test.cc"],
    deps = [
        ":collective_thunk",
        ":recv_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "send_thunk",
    srcs = ["send_thunk.cc"],
    hdrs = ["send_thunk.h"],
    deps = [
        ":collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "send_thunk_test",
    srcs = ["send_thunk_test.cc"],
    deps = [
        ":collective_thunk",
        ":send_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

xla_test(
    name = "send_recv_thunk_test",
    srcs = ["send_recv_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["pjrt_migration_candidate"],
    use_legacy_runtime = True,
    deps = [
        ":collective_thunk",
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":recv_thunk",
        ":send_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:verified_hlo_module",
        "//xla/service:backend",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:gpu_executable",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tests:hlo_test_base",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "collective_group_thunk",
    srcs = ["collective_group_thunk.cc"],
    hdrs = ["collective_group_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:future",
        "//xla/backends/gpu/collectives:gpu_communicator",
        "//xla/core/collectives:communicator",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:status_macros",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "norm_thunk",
    srcs = ["norm_thunk.cc"],
    hdrs = ["norm_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:gpu_norm_runner",
        "//xla/service/gpu:gpu_norm_runner_proto_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:lazy_op_runner",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "norm_thunk_test",
    srcs = ["norm_thunk_test.cc"],
    deps = [
        ":norm_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "outfeed_thunk",
    srcs = ["outfeed_thunk.cc"],
    hdrs = ["outfeed_thunk.h"],
    deps = [
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:shaped_slice_proto_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_transfer_manager",
        "//xla/service/gpu:io_feed_manager",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "outfeed_thunk_test",
    srcs = ["outfeed_thunk_test.cc"],
    deps = [
        ":outfeed_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "replica_id_thunk",
    srcs = ["replica_id_thunk.cc"],
    hdrs = ["replica_id_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer_hdr",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "collective_metadata_thunk",
    srcs = ["collective_metadata_thunk.cc"],
    hdrs = ["collective_metadata_thunk.h"],
    deps = [
        ":collective_multimem",
        ":collective_multimem_registry",
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_clique_rendezvous",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:collective_kernel_metadata",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_protobuf//:protobuf_lite",
    ],
)

cc_library(
    name = "sequential_thunk",
    srcs = ["sequential_thunk.cc"],
    hdrs = ["sequential_thunk.h"],
    deps = [
        ":annotation",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@local_tsl//tsl/profiler/lib:scoped_annotation",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

xla_cc_test(
    name = "sequential_thunk_test",
    srcs = ["sequential_thunk_test.cc"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        ":thunk_id",
        ":thunk_proto_cc",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "host_send_recv_thunk",
    srcs = ["host_send_recv_thunk.cc"],
    hdrs = ["host_send_recv_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "thunk_id",
    hdrs = ["thunk_id.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        "//xla/tsl/lib/gtl:int_type",
    ],
)

xla_cc_test(
    name = "thunk_id_test",
    srcs = ["thunk_id_test.cc"],
    deps = [
        ":thunk_id",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "thunk",
    srcs = ["thunk.cc"],
    hdrs = ["thunk.h"],
    deps = [
        ":collective_clique_requests",
        ":collective_cliques",
        ":collective_multimem_registry",
        ":collective_params",
        ":thunk_id",
        ":thunk_proto_cc",
        "//xla:executable_run_options",
        "//xla:status_macros",
        "//xla:util",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/ffi:execution_context",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/service/gpu:gpu_executable_run_options",
        "//xla/service/gpu:ir_emission_utils",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/lib/gtl:int_type",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "thunk_test",
    srcs = ["thunk_test.cc"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "shaped_slice",
    hdrs = ["shaped_slice.h"],
    deps = [
        "//xla/service:shaped_slice",
    ],
)

xla_cc_test(
    name = "for_all_thunks_test",
    srcs = ["for_all_thunks_test.cc"],
    tags = ["gpu"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_thunk",
        ":conditional_thunk",
        ":dynamic_slice_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla:shape_util",
        "//xla/service:buffer_assignment",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "triangular_solve_thunk",
    srcs = ["triangular_solve_thunk.cc"],
    hdrs = ["triangular_solve_thunk.h"],
    deps = [
        ":make_batch_pointers",
        ":shaped_slice",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "triangular_solve_thunk_test",
    srcs = ["triangular_solve_thunk_test.cc"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        ":triangular_solve_thunk",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

xla_cc_test(
    name = "replica_id_thunk_test",
    srcs = ["replica_id_thunk_test.cc"],
    deps = [
        ":replica_id_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

cc_library(
    name = "while_thunk",
    srcs = ["while_thunk.cc"],
    hdrs = ["while_thunk.h"],
    deps = [
        ":host_memory_pool",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "while_thunk_test",
    srcs = ["while_thunk_test.cc"],
    backends = ["gpu"],
    tags = ["test_migrated_to_hlo_runner_pjrt"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":while_thunk",
        "//xla:executable_run_options",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:verified_hlo_module",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tests:hlo_pjrt_test_base",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "wait_for_streams_thunk",
    srcs = ["wait_for_streams_thunk.cc"],
    hdrs = ["wait_for_streams_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
    ],
)

xla_cc_test(
    name = "wait_for_streams_thunk_test",
    srcs = ["wait_for_streams_thunk_test.cc"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        ":wait_for_streams_thunk",
        "//xla/tsl/platform:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "cudnn_thunk",
    srcs = ["cudnn_thunk.cc"],
    hdrs = ["cudnn_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:buffer_assignment_proto_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:dnn",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/profiler/lib:nvtx_utils",
    ],
)

xla_cc_test(
    name = "cudnn_thunk_test",
    srcs = ["cudnn_thunk_test.cc"],
    deps = [
        ":cudnn_thunk",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
        "@com_google_protobuf//:protobuf",
    ],
)

cc_library(
    name = "custom_call_target",
    hdrs = ["custom_call_target.h"],
    deps = [
        "//xla/service:custom_call_status",
    ],
)

cc_library(
    name = "host_memory_pool",
    srcs = ["host_memory_pool.cc"],
    hdrs = ["host_memory_pool.h"],
    deps = [
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "buffer_comparator",
    srcs = ["buffer_comparator.cc"],
    hdrs = ["buffer_comparator.h"],
    deps = [
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:buffer_comparator_kernel",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@eigen_archive//:eigen3",
    ],
)

xla_test(
    name = "buffer_comparator_test",
    srcs = ["buffer_comparator_test.cc"],
    backends = ["gpu"],
    disabled_backends = [],
    deps = [
        ":buffer_comparator",
        "//xla:shape_util",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service:platform_util",
        "//xla/service/gpu:stream_executor_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/cleanup",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/platform:ml_dtypes",
    ],
)

cc_library(
    name = "make_batch_pointers",
    srcs = ["make_batch_pointers.cc"],
    hdrs = ["make_batch_pointers.h"],
    deps = [
        "//xla:types",
        "//xla:util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:make_batch_pointers_kernel",
        "//xla/stream_executor/rocm:rocm_platform_id",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
    ],
)

xla_test(
    name = "make_batch_pointers_test",
    srcs = ["make_batch_pointers_test.cc"],
    backends = ["gpu"],
    use_legacy_runtime = True,
    deps = [
        ":make_batch_pointers",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_reduce",
    srcs = ["all_reduce.cc"],
    hdrs = ["all_reduce.h"],
    deps = [
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/core/collectives:rank_id",
        "//xla/service:collective_ops_utils",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/stream_executor/gpu:collective_kernel_metadata",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util:safe_reinterpret_cast",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "all_reduce_test",
    srcs = ["all_reduce_test.cc"],
    backend_tags = {
        "gpu": [
            "multi_gpu",
            "xla_amdgpu_any",
            "no_oss",
        ],
    },
    backends = ["gpu"],
    disabled_backends = [],
    use_legacy_runtime = True,
    deps = [
        ":all_reduce",
        "//xla:array",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/core/collectives:rank_id",
        "//xla/core/collectives:reduction_kind",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
        "//xla/service:collective_ops_utils",
        "//xla/service:hlo_runner",
        "//xla/service:platform_util",
        "//xla/service/gpu:gpu_constants",
        "//xla/service/gpu:launch_dimensions",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:all_reduce_kernel",
        "//xla/stream_executor/gpu:collective_kernel_metadata",
        "//xla/stream_executor/gpu:gpu_executor_header",
        "//xla/stream_executor/gpu:gpu_init",
        "//xla/stream_executor/gpu:multicast_memory",
        "//xla/stream_executor/host:host_platform",
        "//xla/tests:literal_test_util",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest",
    ],
)

cc_library(
    name = "ragged_all_to_all",
    srcs = ["ragged_all_to_all.cc"],
    hdrs = ["ragged_all_to_all.h"],
    deps = [
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:ragged_all_to_all_kernel",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "ragged_all_to_all_test",
    srcs = ["ragged_all_to_all_test.cc"],
    backends = ["gpu"],
    disabled_backends = [],
    deps = [
        ":ragged_all_to_all",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_address_handle",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:gpu_init",
        "//xla/stream_executor/host:host_platform",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest",
    ],
)

cc_library(
    name = "topk",
    srcs = ["topk.cc"],
    hdrs = ["topk.h"],
    compatible_with = [],
    tags = ["gpu"],
    deps = [
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:kernel_args_packing_spec",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/stream_executor/gpu:topk_kernel",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/numeric:bits",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
    ],
)

xla_test(
    name = "topk_test",
    srcs = ["topk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":topk",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/gpu:kernel_serialization_check",
        "//xla/stream_executor/host:host_platform",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest",
    ],
)

tf_proto_library(
    name = "thunk_proto",
    srcs = [
        "thunk.proto",
    ],
    deps = [
        ":convolution_filter_thunk_proto",
        ":copy_thunk_proto",
        ":dynamic_slice_thunk_proto",
        "//xla:xla_data_proto",
        "//xla/core/host_offloading:host_offloading_executable_proto",
        "//xla/ffi:attribute_map_proto",
        "//xla/ffi:execution_state_proto",
        "//xla/service:buffer_assignment_proto",
        "//xla/service:hlo_proto",
        "//xla/service:shaped_slice_proto",
        "//xla/service/gpu:gpu_conv_runner_proto",
        "//xla/service/gpu:gpu_norm_runner_proto",
        "//xla/service/gpu:launch_dimensions_proto",
        "//xla/service/gpu/kernels:custom_kernel_proto",
        "//xla/stream_executor:launch_dim_proto",
        "//xla/stream_executor/gpu:gpu_blas_lt_proto",
        "//xla/stream_executor/gpu:tma_metadata_proto",
    ],
)

xla_py_proto_library(
    name = "thunk_proto_py",
    deps = [":thunk_proto"],
)

tf_proto_library(
    name = "dynamic_slice_thunk_proto",
    srcs = [
        "dynamic_slice_thunk.proto",
    ],
    deps = [
        "//xla:xla_data_proto",
        "//xla:xla_proto",
        "//xla/service:buffer_assignment_proto",
    ],
)

cc_library(
    name = "thunk_proto_deserialization",
    srcs = ["thunk_proto_deserialization.cc"],
    hdrs = ["thunk_proto_deserialization.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":collective_broadcast_thunk",
        ":collective_group_thunk",
        ":collective_permute_thunk",
        ":collective_thunk",
        ":conditional_thunk",
        ":convolution_reorder_thunk",
        ":convolution_thunk",
        ":copy_thunk",
        ":cub_sort_thunk",
        ":cudnn_thunk",
        ":custom_call_thunk",
        ":custom_kernel_thunk",
        ":dynamic_slice_thunk",
        ":fft_thunk",
        ":gemm_thunk",
        ":gpublas_lt_matmul_thunk",
        ":host_execute_thunk",
        ":host_send_recv_thunk",
        ":infeed_thunk",
        ":kernel_thunk",
        ":memset_thunk",
        ":norm_thunk",
        ":outfeed_thunk",
        ":ragged_all_to_all_thunk",
        ":recv_thunk",
        ":replica_id_thunk",
        ":send_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":triangular_solve_thunk",
        ":wait_for_streams_thunk",
        ":while_thunk",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:kernel_spec",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@com_google_protobuf//:protobuf",
    ],
)

xla_cc_test(
    name = "thunk_proto_deserialization_test",
    srcs = ["thunk_proto_deserialization_test.cc"],
    deps = [
        ":conditional_thunk",
        ":copy_thunk",
        ":custom_kernel_thunk",
        ":host_execute_thunk",
        ":host_send_recv_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_proto_cc",
        ":thunk_proto_deserialization",
        ":while_thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:hlo_module_config",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util:safe_reinterpret_cast",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "nvshmem_collective_thunk",
    srcs = ["nvshmem_collective_thunk.cc"],
    hdrs = ["nvshmem_collective_thunk.h"],
    deps = [
        ":collective_execution",
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_clique_key",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives",
        "//xla/core/collectives:collectives_registry",
        "//xla/service:computation_placer",
        "//xla/stream_executor:event",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "nvshmem_collective_permute_thunk",
    srcs = ["nvshmem_collective_permute_thunk.cc"],
    hdrs = ["nvshmem_collective_permute_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_permute_thunk",
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "nvshmem_send_thunk",
    srcs = ["nvshmem_send_thunk.cc"],
    hdrs = ["nvshmem_send_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "nvshmem_recv_thunk",
    srcs = ["nvshmem_recv_thunk.cc"],
    hdrs = ["nvshmem_recv_thunk.h"],
    tags = ["gpu"],
    deps = [
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":p2p_thunk_common",
        ":thunk",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/hlo/ir:collective_op_group_mode",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "nvshmem_all_reduce_thunk",
    srcs = ["nvshmem_all_reduce_thunk.cc"],
    hdrs = ["nvshmem_all_reduce_thunk.h"],
    deps = [
        ":all_reduce_thunk",
        ":collective_kernel_thunk",
        ":collective_thunk",
        ":nvshmem_collective_thunk",
        ":thunk",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu/collectives:gpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/hlo/ir:hlo",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service/gpu:backend_configs_cc",
        "//xla/service/gpu/transforms/collectives:collective_ops_utils",
        "//xla/stream_executor:stream",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "host_execute_thunk",
    srcs = ["host_execute_thunk.cc"],
    hdrs = ["host_execute_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_proto_cc",
        "//xla:executable_run_options",
        "//xla:shape_tree",
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla/backends/gpu/host_offloading:gpu_host_offloading_allocator",
        "//xla/core/host_offloading:host_offloading_allocator",
        "//xla/core/host_offloading:host_offloading_buffer",
        "//xla/core/host_offloading:host_offloading_executable",
        "//xla/core/host_offloading:host_offloading_executable_proto_cc",
        "//xla/core/host_offloading:host_offloading_nanort_executable",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:shaped_slice_proto_cc",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform:platform_port",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "host_execute_thunk_test",
    srcs = ["host_execute_thunk_test.cc"],
    backends = ["gpu"],
    deps = [
        ":host_execute_thunk",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/cpu:alignment",
        "//xla/backends/cpu/nanort:nanort_client",
        "//xla/backends/cpu/nanort:nanort_executable",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/parser:hlo_parser",
        "//xla/service:buffer_assignment",
        "//xla/service:compiler",
        "//xla/service:executable",
        "//xla/service:platform_util",
        "//xla/service/cpu:cpu_aot_compilation_result",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/tests:literal_test_util",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
        "@local_tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "thunk_pass_pipeline",
    srcs = ["thunk_pass_pipeline.cc"],
    hdrs = ["thunk_pass_pipeline.h"],
    deps = [
        ":sequential_thunk",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_description",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
    ],
)

xla_cc_test(
    name = "thunk_pass_pipeline_test",
    srcs = ["thunk_pass_pipeline_test.cc"],
    deps = [
        ":sequential_thunk",
        ":thunk",
        ":thunk_pass_pipeline",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_description",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "command_buffer_conversion_pass",
    srcs = ["command_buffer_conversion_pass.cc"],
    hdrs = ["command_buffer_conversion_pass.h"],
    deps = [
        ":command_buffer_cmd",
        ":command_buffer_cmd_emitter",
        ":command_buffer_thunk",
        ":conditional_thunk",
        ":copy_thunk",
        ":custom_call_thunk",
        ":sequential_thunk",
        ":thunk",
        ":thunk_pass_pipeline",
        ":while_thunk",
        "//xla:util",
        "//xla/ffi:ffi_api",
        "//xla/hlo/ir:hlo",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:semantic_version",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/platform",
        "@local_tsl//tsl/profiler/lib:profiler_lock",
        "@local_tsl//tsl/profiler/lib:traceme",
    ],
)

xla_test(
    name = "command_buffer_conversion_pass_test",
    srcs = ["command_buffer_conversion_pass_test.cc"],
    backends = ["gpu"],
    deps = [
        ":all_gather_thunk",
        ":collective_thunk",
        ":command_buffer_conversion_pass",
        ":command_buffer_thunk",
        ":conditional_thunk",
        ":copy_thunk",
        ":cudnn_thunk",
        ":custom_call_thunk",
        ":gemm_thunk",
        ":replica_id_thunk",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_pass_pipeline",
        ":while_thunk",
        "//xla:shape_util",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
        "//xla/service:hlo_module_config",
        "//xla/service:platform_util",
        "//xla/service/gpu:gpu_device_info_for_tests",
        "//xla/service/gpu:matmul_utils",
        "//xla/stream_executor:blas",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream_executor_h",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "thunk_buffer_debug_pass",
    srcs = [
        "thunk_buffer_debug_checksum.cc",
        "thunk_buffer_debug_filter.cc",
        "thunk_buffer_debug_float_check.cc",
        "thunk_buffer_debug_pass.cc",
        "thunk_buffer_debug_saver_inserter.cc",
    ],
    hdrs = [
        "thunk_buffer_debug_checksum.h",
        "thunk_buffer_debug_filter.h",
        "thunk_buffer_debug_float_check.h",
        "thunk_buffer_debug_pass.h",
        "thunk_buffer_debug_saver_inserter.h",
    ],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_proto_cc",
        ":buffer_debug_log_structs",
        ":buffers_checksum_thunk",
        ":buffers_float_check_thunk",
        ":custom_call_thunk",
        ":runtime_intrinsics",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_id",
        ":thunk_pass_pipeline",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/backends/gpu:ffi",
        "//xla/ffi",
        "//xla/ffi:attribute_map",
        "//xla/ffi/api:c_api",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:dump",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:stream",
        "//xla/stream_executor/gpu:buffer_debug_log",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/functional:bind_front",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_googlesource_code_re2//:re2",
    ],
)

xla_cc_test(
    name = "thunk_buffer_debug_pass_test",
    srcs = ["thunk_buffer_debug_pass_test.cc"],
    deps = [
        ":buffers_checksum_thunk",
        ":buffers_float_check_thunk",
        ":conditional_thunk",
        ":custom_call_thunk",
        ":runtime_intrinsics",
        ":sequential_thunk",
        ":shaped_slice",
        ":thunk",
        ":thunk_buffer_debug_pass",
        ":thunk_id",
        ":thunk_pass_pipeline",
        ":while_thunk",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:hlo_module_config",
        "//xla/stream_executor:device_description",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "runtime_intrinsics",
    srcs = ["runtime_intrinsics.cc"],
    hdrs = ["runtime_intrinsics.h"],
    deps = [
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/gpu:ffi",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/service:collective_ops_utils",
        "//xla/service:custom_call_status",
        "//xla/service:custom_call_target_registry",
        "//xla/service:platform_util",
        "//xla/stream_executor:memory_allocation",
        "//xla/stream_executor:stream",
        "//xla/tsl/lib/io:record_writer",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:platform_port",
    ],
    alwayslink = 1,
)

xla_test(
    name = "runtime_intrinsics_test",
    srcs = ["runtime_intrinsics_test.cc"],
    backends = ["gpu"],
    tags = ["test_migrated_to_hlo_runner_pjrt"],
    deps = [
        ":runtime_intrinsics",
        "//xla:literal",
        "//xla:literal_util",
        "//xla/hlo/ir:hlo",
        "//xla/tests:hlo_pjrt_test_base",
        "//xla/tests:xla_internal_test_main",
        "//xla/tsl/lib/io:record_reader",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/testing:temporary_directory",
        "@com_google_absl//absl/base:log_severity",
        "@com_google_absl//absl/log:scoped_mock_log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_googletest//:gtest",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:tstring",
    ],
)

cc_library(
    name = "buffers_checksum_thunk",
    srcs = ["buffers_checksum_thunk.cc"],
    hdrs = ["buffers_checksum_thunk.h"],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_structs",
        ":thunk",
        ":thunk_id",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/cuda:cuda_platform_id",
        "//xla/stream_executor/gpu:buffer_debug_log",
        "//xla/stream_executor/gpu:buffer_debug_xor_checksum_kernel",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
    ],
)

xla_test(
    name = "buffers_checksum_thunk_test",
    srcs = ["buffers_checksum_thunk_test.cc"],
    backend_tags = {
        "gpu": [
            "multi_gpu",
        ],
    },
    backends = ["gpu"],
    tags = [
        "cuda-only",
        "gpu",
    ],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_structs",
        ":buffers_checksum_thunk",
        ":collective_clique_requests",
        ":collective_multimem_registry",
        ":collective_params",
        ":thunk",
        ":thunk_id",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:mock_stream",
        "//xla/stream_executor:mock_stream_executor",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:buffer_debug_log",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "buffers_float_check_thunk",
    srcs = ["buffers_float_check_thunk.cc"],
    hdrs = ["buffers_float_check_thunk.h"],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_structs",
        ":thunk",
        "//xla:types",
        "//xla:util",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/stream_executor/cuda:cuda_platform_id",
        "//xla/stream_executor/gpu:buffer_debug_float_check_kernel",
        "//xla/stream_executor/gpu:buffer_debug_log",
        "//xla/stream_executor/gpu:gpu_kernel_registry",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
    ],
)

xla_test(
    name = "buffers_float_check_thunk_test",
    srcs = ["buffers_float_check_thunk_test.cc"],
    backend_tags = {
        "gpu": [
            "multi_gpu",
        ],
    },
    backends = ["gpu"],
    tags = [
        "cuda-only",
        "gpu",
    ],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_structs",
        ":buffers_float_check_thunk",
        ":collective_clique_requests",
        ":collective_multimem_registry",
        ":collective_params",
        ":thunk",
        ":thunk_id",
        "//xla:types",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:executable",
        "//xla/service/gpu:buffer_allocations",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_memory_allocator",
        "//xla/stream_executor/gpu:buffer_debug_log",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

tf_proto_library(
    name = "buffer_debug_log_proto",
    srcs = ["buffer_debug_log.proto"],
)

xla_py_proto_library(
    name = "buffer_debug_log_proto_py",
    deps = [":buffer_debug_log_proto"],
)

cc_library(
    name = "buffer_debug_log_structs",
    hdrs = ["buffer_debug_log_structs.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        ":buffer_debug_log_proto_cc",
        "//xla/tsl/lib/gtl:int_type",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "buffer_debug_log_entry_metadata_store",
    srcs = ["buffer_debug_log_entry_metadata_store.cc"],
    hdrs = ["buffer_debug_log_entry_metadata_store.h"],
    compatible_with = get_compatible_with_portable(),
    deps = [
        ":buffer_debug_log_proto_cc",
        ":buffer_debug_log_structs",
        ":thunk_id",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "print_buffer_contents",
    srcs = ["print_buffer_contents.cc"],
    hdrs = ["print_buffer_contents.h"],
    deps = [
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:stream",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_test(
    name = "print_buffer_contents_test",
    srcs = ["print_buffer_contents_test.cc"],
    backends = ["gpu"],
    deps = [
        ":print_buffer_contents",
        "//xla/service:platform_util",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:platform",
        "//xla/stream_executor:platform_manager",
        "//xla/stream_executor:stream",
        "//xla/tsl/lib/core:status_test_util",
        "@com_google_absl//absl/log:scoped_mock_log",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "custom_kernel_thunk",
    srcs = ["custom_kernel_thunk.cc"],
    hdrs = ["custom_kernel_thunk.h"],
    deps = [
        ":print_buffer_contents",
        ":thunk",
        ":thunk_proto_cc",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu:launch_dimensions",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:device_address",
        "//xla/stream_executor:kernel",
        "//xla/stream_executor:kernel_args",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "custom_kernel_thunk_test",
    srcs = ["custom_kernel_thunk_test.cc"],
    deps = [
        ":custom_kernel_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla/codegen/emitters:kernel_arguments",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/gpu/kernels:custom_kernel",
        "//xla/stream_executor:launch_dim",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/util/proto:parse_text_proto",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

xla_cc_test(
    name = "buffer_debug_log_entry_metadata_store_test",
    srcs = ["buffer_debug_log_entry_metadata_store_test.cc"],
    deps = [
        ":buffer_debug_log_entry_metadata_store",
        ":buffer_debug_log_structs",
        ":thunk_id",
        "//xla/tsl/util/proto:proto_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)
