load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")
load("//xla:xla.default.bzl", "xla_cc_test")
load(
    "//xla/stream_executor:build_defs.bzl",
    "if_cuda_or_rocm_is_configured",
)
load("//xla/tsl:tsl.bzl", "internal_visibility")
load("//xla/tsl/platform:rules_cc.bzl", "cc_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(["//xla:internal"]),
    licenses = ["notice"],
)

package_group(
    name = "friends",
    includes = [
        "//xla:friends",
        "//xla:internal",
    ],
    packages = [
        "//tensorflow/core/tfrt/ifrt/...",
        "//third_party/australis/...",
        "//third_party/gxlang/...",
        "//third_party/openxla_pjrt_plugin/...",
        "//third_party/py/jax/...",
        "//third_party/py/torch_tpu/...",
    ],
)

cc_library(
    name = "worker_thread",
    srcs = ["worker_thread.cc"],
    hdrs = ["worker_thread.h"],
    deps = [
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:env",
    ],
)

xla_cc_test(
    name = "worker_thread_test",
    srcs = ["worker_thread_test.cc"],
    deps = [
        ":worker_thread",
        "//xla/hlo/testlib:test",
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "pjrt_c_api_client",
    hdrs = ["pjrt_c_api_client.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/pjrt/c_api_client:pjrt_c_api_client",
    ],
)

cc_library(
    name = "semaphore",
    srcs = ["semaphore.cc"],
    hdrs = ["semaphore.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:types",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:logging",
    ],
)

xla_cc_test(
    name = "semaphore_test",
    srcs = ["semaphore_test.cc"],
    deps = [
        ":semaphore",
        "//xla/hlo/testlib:test",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:env",
    ],
)

cc_library(
    name = "dynamic_shapes",
    srcs = ["dynamic_shapes.cc"],
    hdrs = ["dynamic_shapes.h"],
    deps = [
        ":device_event",
        ":raw_buffer",
        "//xla:shape_util",
        "//xla/tsl/concurrency:async_value",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:platform_port",
    ],
)

cc_library(
    name = "common_pjrt_client",
    srcs = [
        "abstract_tracked_device_buffer.cc",
        "common_pjrt_client.cc",
        "host_to_device_transfer_manager.cc",
    ],
    hdrs = [
        "abstract_tracked_device_buffer.h",
        "common_pjrt_client.h",
        "host_to_device_transfer_manager.h",
    ],
    visibility = internal_visibility([
        ":friends",
        "//third_party/py/torch_tpu:__pkg__",
    ]),
    deps = [
        ":async_work_runner",
        ":device_event",
        ":dynamic_shapes",
        ":host_callback",
        ":pjrt_client",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":raw_buffer",
        ":staging_buffer",
        ":transpose",
        ":utils",
        "//xla:executable_run_options",
        "//xla:future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/error:error_codes",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/c:pjrt_c_api_device_event_hdrs",
        "//xla/runtime:device_id",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:future",
        "//xla/tsl/concurrency:ref_count",
        "//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/base:nullability",
        "@com_google_absl//absl/cleanup",
        "@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:any_invocable",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/log:vlog_is_on",
        "@com_google_absl//absl/random",
        "@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",
        "@llvm-project//llvm:Support",
        "@tsl//tsl/platform:casts",
        "@tsl//tsl/platform:context",
        "@tsl//tsl/profiler/lib:connected_traceme",
        "@tsl//tsl/profiler/lib:context_types_hdrs",
        "@tsl//tsl/profiler/lib:scoped_memory_debug_annotation",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

cc_library(
    name = "infer_dispatch_info",
    srcs = ["infer_dispatch_info.cc"],
    hdrs = ["infer_dispatch_info.h"],
    visibility = internal_visibility([
        ":friends",
    ]),
    deps = [
        ":common_pjrt_client",
        ":pjrt_client",
        ":pjrt_executable",
        ":pjrt_layout",
        ":utils",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/ir:hlo_sharding",
        "//xla/mlir/utils:type_util",
        "//xla/service/spmd/shardy/stablehlo_round_trip:export_shardings",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@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/types:span",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@shardy//shardy/dialect/sdy/ir:dialect",
    ],
)

cc_library(
    name = "pjrt_api",
    srcs = ["pjrt_api.cc"],
    hdrs = ["pjrt_api.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_status_utils",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "@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",
    ],
)

xla_cc_test(
    name = "pjrt_api_test",
    srcs = ["pjrt_api_test.cc"],
    deps = [
        ":pjrt_api",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_wrapper_impl",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/protobuf:error_codes_proto_impl_cc",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "maybe_owning_mlir_module",
    hdrs = ["maybe_owning_mlir_module.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@llvm-project//mlir:IR",
    ],
)

cc_library(
    name = "pjrt_client",
    srcs = ["pjrt_client.cc"],
    hdrs = ["pjrt_client.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":host_memory_allocator",
        ":maybe_owning_mlir_module",
        ":pjrt_abi_version",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        ":pjrt_layout",
        ":scoped_async_tracking_event",
        ":utils",
        "//xla:future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/distributed:key_value_store_interface",
        "//xla/pjrt/distributed/coordination:coordination_service_proto_cc",
        "//xla/service:computation_placer_hdr",
        "//xla/service:hlo_cost_analysis",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/lib/gtl:int_type",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//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/functional:any_invocable",
        "@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",
        "@tsl//tsl/platform:errors",
    ],
)

cc_library(
    name = "mock_pjrt_client",
    testonly = 1,
    hdrs = ["mock_pjrt_client.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_executable",
        "//xla/hlo/builder:xla_computation",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_for_library",
    ],
)

cc_library(
    name = "pjrt_client_utils",
    srcs = ["pjrt_client_utils.cc"],
    hdrs = ["pjrt_client_utils.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":pjrt_client",
        "//xla:shape_util",
        "//xla:util",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "pjrt_client_test_common",
    testonly = 1,
    srcs = ["pjrt_client_test.cc"],
    hdrs = ["pjrt_client_test.h"],
    deps = [
        ":pjrt_client",
        ":pjrt_compiler",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:types",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu:alignment",
        "//xla/hlo/builder:xla_builder",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/parser:hlo_parser",
        "//xla/hlo/testlib:test",
        "//xla/tests:literal_test_util",
        "//xla/tsl/framework:allocator",
        "//xla/tsl/lib/core:status_test_util",
        "@com_google_absl//absl/base:core_headers",
        "@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",
        "@tsl//tsl/platform:statusor",
    ],
    alwayslink = 1,
)

cc_library(
    name = "pjrt_device_dimensions",
    srcs = ["pjrt_device_dimensions.cc"],
    hdrs = ["pjrt_device_dimensions.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        "//xla/pjrt/proto:pjrt_device_dimensions_proto_cc",
        "@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/types:span",
    ],
)

xla_cc_test(
    name = "pjrt_device_dimensions_test",
    srcs = ["pjrt_device_dimensions_test.cc"],
    deps = [
        ":pjrt_device_dimensions",
        "//xla/pjrt/proto:pjrt_device_dimensions_proto_cc",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "pjrt_executable",
    srcs = ["pjrt_executable.cc"],
    hdrs = ["pjrt_executable.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":compiled_memory_stats",
        ":pjrt_abi_version",
        ":pjrt_common",
        ":pjrt_device_dimensions",
        ":pjrt_layout",
        ":utils",
        "//xla:debug_options_flags",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla:xla_proto_cc",
        "//xla/backends/gpu/target_config",
        "//xla/client:executable_build_options",
        "//xla/ffi:execution_context",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:executable_metadata_proto_cc",
        "//xla/pjrt/proto:execute_options_proto_cc",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:compiler",
        "//xla/service:computation_layout",
        "//xla/service:hlo_cost_analysis",
        "//xla/service:hlo_proto_cc",
        "//xla/service:hlo_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@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",
        "@com_google_absl//absl/types:span",
        "@com_google_protobuf//:protobuf",
    ],
)

xla_cc_test(
    name = "pjrt_executable_test",
    srcs = ["pjrt_executable_test.cc"],
    deps = [
        ":pjrt_executable",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:executable_metadata_proto_cc",
        "//xla/pjrt/proto:execute_options_proto_cc",
        "//xla/service:computation_placer_hdr",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "pjrt_device_description",
    hdrs = ["pjrt_device_description.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_common",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "pjrt_compiler",
    srcs = ["pjrt_compiler.cc"],
    hdrs = ["pjrt_compiler.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":maybe_owning_mlir_module",
        ":pjrt_abi_version",
        ":pjrt_common",
        ":pjrt_device_description",
        ":pjrt_device_dimensions",
        ":pjrt_executable",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/pjrt/proto:topology_description_proto_cc",
        "//xla/runtime:chip_id",
        "//xla/runtime:device_id",
        "//xla/runtime:process_id",
        "//xla/tsl/platform:status_macros",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/base:no_destructor",
        "@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/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:fingerprint",
    ],
)

xla_cc_test(
    name = "pjrt_compiler_test",
    srcs = ["pjrt_compiler_test.cc"],
    deps = [
        ":maybe_owning_mlir_module",
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_device_description",
        ":pjrt_executable",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "@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:status_matchers",
        "@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 = "pjrt_common",
    srcs = ["pjrt_common.cc"],
    hdrs = ["pjrt_common.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/pjrt/proto:pjrt_value_type_proto_cc",
        "//xla/runtime:chip_id",
        "//xla/runtime:device_id",
        "//xla/runtime:process_id",
        "@com_google_absl//absl/container:inlined_vector",
    ],
)

cc_library(
    name = "pjrt_abi_version",
    hdrs = ["pjrt_abi_version.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_common",
        "//xla/pjrt/proto:pjrt_abi_version_proto_cc",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
    ],
)

xla_cc_test(
    name = "pjrt_common_test",
    srcs = ["pjrt_common_test.cc"],
    deps = [
        ":pjrt_common",
        "//xla/pjrt/proto:pjrt_value_type_proto_cc",
        "@com_google_googletest//:gtest_main",
    ],
)

xla_cc_test(
    name = "utils_test",
    srcs = ["utils_test.cc"],
    deps = [
        ":layout_mode",
        ":utils",
        "//xla:shape_util",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "utils",
    srcs = ["utils.cc"],
    hdrs = ["utils.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        ":layout_mode",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/client:executable_build_options",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/service:computation_layout",
        "//xla/service:computation_placer",
        "//xla/service:hlo_proto_cc",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@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/strings",
        "@com_google_absl//absl/types:span",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@tsl//tsl/platform:logging",
        "@tsl//tsl/platform:platform_port",
    ],
)

cc_library(
    name = "pjrt_layout",
    hdrs = ["pjrt_layout.h"],
    visibility = internal_visibility(["//xla:friends"]),
    deps = [
        "//xla:shape_util",
        "//xla/hlo/parser:hlo_parser",
        "//xla/tsl/platform:status_macros",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@tsl//tsl/platform:statusor",
    ],
)

cc_library(
    name = "layout_mode",
    srcs = ["layout_mode.cc"],
    hdrs = ["layout_mode.h"],
    visibility = ["//xla:friends"],
    deps = [
        "//xla:shape_util",
        "//xla/hlo/parser:hlo_parser",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
    ],
)

cc_library(
    name = "metrics",
    srcs = ["metrics.cc"],
    hdrs = ["metrics.h"],
    deps = ["//xla/tsl/lib/monitoring:counter"],
)

cc_library(
    name = "mlir_to_hlo",
    srcs = ["mlir_to_hlo.cc"],
    hdrs = ["mlir_to_hlo.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:util",
        "//xla/client:executable_build_options",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/translate:stablehlo",
        "//xla/mlir/utils:error_util",
        "//xla/mlir_hlo:hlo_dialect_registration",
        "//xla/mlir_hlo:mhlo_passes",
        "//xla/mlir_hlo:stablehlo_extension_passes",
        "//xla/service/spmd/shardy:constants",
        "//xla/service/spmd/shardy:utils",
        "//xla/service/spmd/shardy/sdy_round_trip:pipelines",
        "//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_export",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:vlog_is_on",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:ArithDialect",
        "@llvm-project//mlir:BytecodeWriter",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:FuncExtensions",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:MLProgramDialect",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:ReconcileUnrealizedCasts",
        "@llvm-project//mlir:ShapeDialect",
        "@llvm-project//mlir:Support",
        "@llvm-project//mlir:Transforms",
        "@shardy//shardy/dialect/sdy/ir:dialect",
        "@shardy//shardy/dialect/sdy/ir:register",
        "@stablehlo//:chlo_ops",
        "@stablehlo//:register",
        "@stablehlo//:stablehlo_ops",
        "@stablehlo//:stablehlo_passes",
        "@stablehlo//:stablehlo_portable_api",
        "@stablehlo//:stablehlo_serialization",
        "@stablehlo//:version",
    ],
)

xla_cc_test(
    name = "mlir_to_hlo_test",
    srcs = ["mlir_to_hlo_test.cc"],
    deps = [
        ":mlir_to_hlo",
        "//xla/hlo/testlib:test",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@stablehlo//:stablehlo_portable_api",
    ],
)

cc_library(
    name = "pjrt_future",
    hdrs = ["pjrt_future.h"],
    visibility = internal_visibility([":friends"]),
    deps = ["//xla:future"],
)

cc_library(
    name = "scoped_async_tracking_event",
    hdrs = ["scoped_async_tracking_event.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:future",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
    ],
)

cc_library(
    name = "host_memory_spaces",
    srcs = ["host_memory_spaces.cc"],
    hdrs = ["host_memory_spaces.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_client",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:fingerprint",
    ],
)

cc_library(
    name = "tpu_constants",
    hdrs = ["tpu_constants.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@com_google_absl//absl/strings:string_view",
    ],
)

# Transitional forwarding target. Use pjrt/plugin/xla_cpu:xla_cpu_pjrt_client instead.
cc_library(
    name = "tfrt_cpu_pjrt_client",
    hdrs = ["tfrt_cpu_pjrt_client.h"],
    visibility = internal_visibility([
        "//xla:friends",
    ]),
    deps = [
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
    ],
)

cc_library(
    name = "lru_cache",
    hdrs = ["lru_cache.h"],
    visibility = internal_visibility([
        "//xla:friends",
    ]),
    deps = [
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:hash_container_defaults",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
    ],
)

xla_cc_test(
    name = "lru_cache_test",
    srcs = ["lru_cache_test.cc"],
    deps = [
        ":lru_cache",
        "//xla/hlo/testlib:test",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/random",
        "@com_google_googletest//:gtest_main",
    ],
)

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

cc_library(
    name = "transpose",
    srcs = [
        "transpose.cc",
    ],
    hdrs = [
        "transpose.h",
    ],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":lru_cache",
        ":transpose_kernels",
        "//xla:ef57",
        "//xla:permutation_util",
        "//xla:util",
        "//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:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/log:absl_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",
        "@tsl//tsl/profiler/lib:traceme",
    ],
)

xla_cc_test(
    name = "transpose_test",
    srcs = ["transpose_test.cc"],
    shard_count = 10,
    deps = [
        ":transpose",
        ":transpose_kernels",
        "//xla:array",
        "//xla:ef57",
        "//xla:permutation_util",
        "//xla:types",
        "//xla:util",
        "//xla/hlo/testlib:test",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "//xla/tsl/protobuf:error_codes_proto_impl_cc",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/flags:flag",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/log:log_streamer",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "tf_pjrt_client",
    srcs = ["tf_pjrt_client.cc"],
    hdrs = ["tf_pjrt_client.h"],
    visibility = ["//visibility:public"],
    deps = [
        ":maybe_owning_mlir_module",
        ":pjrt_client",
        ":pjrt_common",
        ":pjrt_compiler",
        ":pjrt_executable",
        "//xla:future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/service:computation_placer_hdr",
        "//xla/service:hlo_cost_analysis",
        "//xla/tsl/platform:status_macros",
        "@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/functional:any_invocable",
        "@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:string_view",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:casts",
    ],
)

xla_cc_test(
    name = "tf_pjrt_client_test",
    srcs = ["tf_pjrt_client_test.cc"],
    deps = [
        ":tf_pjrt_client",
        "//xla:literal_util",
        "//xla:xla_data_proto_cc",
        "//xla/hlo/parser:hlo_parser",
        "//xla/pjrt/plugin/xla_cpu:cpu_client_options",
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
        "//xla/service:hlo_proto_cc",
        "@com_google_absl//absl/base",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:env",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "host_callback",
    srcs = ["host_callback.cc"],
    hdrs = ["host_callback.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":pjrt_client",
        ":pjrt_executable",
        "//xla:future",
        "//xla:shape_util",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:ffi",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@tsl//tsl/platform:logging",
    ],
)

xla_cc_test(
    name = "host_callback_test",
    srcs = ["host_callback_test.cc"],
    deps = [
        ":host_callback",
        ":pjrt_client",
        "//xla:future",
        "//xla:xla_data_proto_cc",
        "//xla/tests:literal_test_util",
        "//xla/tsl/lib/core:status_test_util",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "exceptions",
    hdrs = ["exceptions.h"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:cord",
    ],
)

cc_library(
    name = "status_casters",
    hdrs = ["status_casters.h"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":exceptions",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@tsl//tsl/platform:macros",
    ],
)

cc_library(
    name = "triton_stub",
    srcs = [
        "triton.h",
        "triton_stub.cc",
    ],
    tags = ["manual"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@tsl//tsl/platform:path",
    ],
)

cc_library(
    name = "triton_cuda",
    srcs = [
        "triton.h",
        "triton_cuda.cc",
    ],
    compatible_with = [],
    tags = [
        "cuda-only",
        "gpu",
        "manual",
    ],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:xla_proto_cc",
        "//xla/backends/gpu/codegen/triton:compilation_pipeline",
        "//xla/service/gpu/llvm_gpu_backend",
        "//xla/service/gpu/llvm_gpu_backend:nvptx_backend",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor/cuda:cuda_compute_capability",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:status_macros",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@llvm-project//llvm:Core",
        "@llvm-project//llvm:IRReader",
        "@llvm-project//llvm:Linker",
        "@llvm-project//llvm:MC",
        "@llvm-project//llvm:Support",
        "@llvm-project//llvm:Target",
        "@llvm-project//llvm:TargetParser",
        "@llvm-project//mlir:ArithDialect",
        "@llvm-project//mlir:ArithToLLVM",
        "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
        "@llvm-project//mlir:ExecutionEngineUtils",
        "@llvm-project//mlir:FuncExtensions",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:IndexToLLVM",
        "@llvm-project//mlir:LLVMCommonConversion",
        "@llvm-project//mlir:LLVMDialect",
        "@llvm-project//mlir:LLVMIRTransforms",
        "@llvm-project//mlir:LLVMToLLVMIRTranslation",
        "@llvm-project//mlir:NVVMToLLVMIRTranslation",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:ToLLVMIRTranslation",
        "@triton//:TritonDialects",
        "@triton//:TritonTransforms",
        "@triton//third_party/nvidia:NVGPUToLLVM",
        "@triton//third_party/nvidia:TritonNVIDIAGPUToLLVM",
        "@tsl//tsl/platform:path",
    ],
)

cc_library(
    name = "triton_rocm",
    srcs = [
        "triton.h",
        "triton_rocm.cc",
    ],
    compatible_with = [],
    tags = [
        "gpu",
        "manual",
        "rocm-only",
    ],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:debug_options_flags",
        "//xla:xla_proto_cc",
        "//xla/backends/gpu/codegen/triton:compilation_pipeline",
        "//xla/service/gpu:target_constants",
        "//xla/service/gpu/llvm_gpu_backend:amdgpu_backend",
        "//xla/stream_executor:device_description",
        "//xla/stream_executor/rocm:rocm_compute_capability",
        "//xla/tsl/platform:env",
        "//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",
        "@llvm-project//llvm:Core",
        "@llvm-project//llvm:TargetParser",
        "@llvm-project//mlir:AffineDialect",
        "@llvm-project//mlir:AffineToStandard",
        "@llvm-project//mlir:ArithDialect",
        "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:FuncExtensions",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:LLVMDialect",
        "@llvm-project//mlir:LLVMIRTransforms",
        "@llvm-project//mlir:LLVMToLLVMIRTranslation",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:ROCDLToLLVMIRTranslation",
        "@llvm-project//mlir:TensorDialect",
        "@llvm-project//mlir:ToLLVMIRTranslation",
        "@llvm-project//mlir:Transforms",
        "@triton//:TritonDialects",
        "@tsl//tsl/platform:path",
        "@tsl//tsl/platform:random",
    ],
)

cc_library(
    name = "triton",
    srcs = [],
    hdrs = ["triton.h"],
    visibility = internal_visibility([":friends"]),
    deps = if_cuda_or_rocm_is_configured(
        [],
        [":triton_stub"],
    ) + if_rocm_is_configured(
        [":triton_rocm"],
    ) + if_cuda_is_configured(
        [":triton_cuda"],
    ) + [
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
    ],
)

cc_library(
    name = "device_event",
    srcs = [
        "device_event.cc",
        "device_event_utils.cc",
    ],
    hdrs = [
        "device_event.h",
        "device_event_utils.h",
    ],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:future",
        "//xla/pjrt/c:pjrt_c_api_device_event_hdrs",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_status_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:executor",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "device_event_utils_test",
    srcs = ["device_event_utils_test.cc"],
    deps = [
        ":device_event",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:executor",
        "//xla/tsl/concurrency:future",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/status",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "staging_buffer",
    srcs = ["staging_buffer.cc"],
    hdrs = ["staging_buffer.h"],
    visibility = internal_visibility([":friends"]),
    deps = ["@com_google_absl//absl/types:span"],
)

cc_library(
    name = "raw_buffer",
    srcs = ["raw_buffer.cc"],
    hdrs = ["raw_buffer.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        ":async_work_runner",
        ":device_event",
        ":pjrt_client",
        ":staging_buffer",
        "//xla:future",
        "//xla:literal",
        "//xla:shape_util",
        "//xla/pjrt/c:pjrt_c_api_device_event_hdrs",
        "//xla/pjrt/c:pjrt_c_api_raw_buffer_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_status_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:status_macros",
        "@com_google_absl//absl/functional:any_invocable",
        "@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/types:span",
    ],
)

xla_cc_test(
    name = "raw_buffer_test",
    srcs = ["raw_buffer_test.cc"],
    deps = [
        ":raw_buffer",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "async_work_runner",
    hdrs = ["async_work_runner.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:executor",
        "//xla/tsl/concurrency:ref_count",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "thread_pool_async_work_runner",
    srcs = ["thread_pool_async_work_runner.cc"],
    hdrs = ["thread_pool_async_work_runner.h"],
    deps = [
        ":async_work_runner",
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/strings:string_view",
        "@tsl//tsl/platform:unbounded_work_queue",
    ],
)

cc_library(
    name = "pjrt_phase_compile_sample_plugin",
    testonly = True,
    srcs = ["pjrt_phase_compile_sample_plugin.cc"],
    hdrs = ["pjrt_phase_compile_sample_plugin.h"],
    visibility = ["//visibility:private"],
    deps = [
        ":maybe_owning_mlir_module",
        ":pjrt_client",
        ":pjrt_compiler",
        ":pjrt_executable",
        "//xla/hlo/builder:xla_computation",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_helpers",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/c:pjrt_c_api_status_utils",
        "//xla/pjrt/c:pjrt_c_api_wrapper_impl",
        "//xla/pjrt/plugin/xla_cpu:cpu_client_options",
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Transforms",
        "@stablehlo//:stablehlo_passes_optimization",
        "@stablehlo//:stablehlo_portable_api",
        "@stablehlo//:stablehlo_serialization",
    ],
)

cc_library(
    name = "errors",
    srcs = ["errors.cc"],
    hdrs = ["errors.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:cord",
        "@com_google_absl//absl/strings:string_view",
    ],
)

xla_cc_test(
    name = "pjrt_phase_compiler_test",
    srcs = ["pjrt_phase_compiler_test.cc"],
    deps = [
        ":maybe_owning_mlir_module",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_phase_compile_sample_plugin",
        "//xla:debug_options_flags",
        "//xla/backends/cpu:target_machine_options",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/testlib:test",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/plugin/xla_cpu:cpu_topology",
        "//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/tsl/lib/core:status_test_util",
        "//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/strings:string_view",
        "@com_google_googletest//:gtest",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@stablehlo//:reference_api",
    ],
)

xla_cc_test(
    name = "pjrt_phase_compile_extension_test",
    srcs = ["pjrt_phase_compile_extension_test.cc"],
    deps = [
        ":pjrt_api",
        ":pjrt_compiler",
        ":pjrt_executable",
        ":pjrt_phase_compile_sample_plugin",
        ":string_utils",
        "//xla/pjrt/c:pjrt_c_api_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
        "//xla/pjrt/c:pjrt_c_api_phase_compile_internal",
        "//xla/pjrt/c_api_client:pjrt_c_api_client",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "//xla/tsl/lib/core:status_test_util",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@stablehlo//:reference_api",
    ],
)

cc_library(
    name = "string_utils",
    srcs = ["string_utils.cc"],
    hdrs = ["string_utils.h"],
    deps = [
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "partial_program_utils",
    srcs = ["partial_program_utils.cc"],
    hdrs = ["partial_program_utils.h"],
    deps = [
        "//xla/pjrt/proto:compile_options_proto_cc",
        "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
        "@com_google_absl//absl/cleanup",
        "@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 = "host_memory_allocator",
    srcs = ["host_memory_allocator.cc"],
    hdrs = ["host_memory_allocator.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla/runtime:device_id",
        "//xla/tsl/framework:allocator",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@tsl//tsl/platform:platform_port",
    ],
)

cc_library(
    name = "never_run_on_fiber",
    hdrs = ["never_run_on_fiber.h"],
    deps = [
        ":async_work_runner",
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/synchronization",
    ],
)

xla_cc_test(
    name = "errors_test",
    srcs = ["errors_test.cc"],
    deps = [
        ":errors",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:status_matchers",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "compiled_memory_stats",
    srcs = ["compiled_memory_stats.cc"],
    hdrs = ["compiled_memory_stats.h"],
    visibility = internal_visibility([":friends"]),
    deps = [
        "//xla:shape_util",
        "//xla/pjrt/proto:executable_metadata_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/service:hlo_value",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "exceptions_test",
    srcs = ["exceptions_test.cc"],
    deps = [
        ":exceptions",
        "//xla/tsl/platform:env",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:cord",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "local_device_state",
    hdrs = ["local_device_state.h"],
    visibility = internal_visibility([":friends"]),
    deps = ["//xla/pjrt/se:local_device_state"],
    alwayslink = 1,
)

cc_library(
    name = "pjrt_stream_executor_client",
    hdrs = [
        "pjrt_stream_executor_client.h",
    ],
    visibility = internal_visibility([":friends"]),
    deps = ["//xla/pjrt/se:pjrt_stream_executor_client"],
    alwayslink = 1,
)

cc_library(
    name = "stream_executor_executable",
    hdrs = ["stream_executor_executable.h"],
    visibility = internal_visibility([":friends"]),
    deps = ["//xla/pjrt/se:stream_executor_executable"],
    alwayslink = 1,
)

alias(
    name = "stream_executor_executable_proto_cc",
    actual = "//xla/pjrt/se:stream_executor_executable_proto_cc",
    visibility = ["//visibility:public"],
)
