load("//xla:xla.default.bzl", "xla_cc_test")
load("//xla/backends/cpu/nanort:package_groups.bzl", "xla_cpu_nanort_packages")
load("//xla/tsl:tsl.bzl", "internal_visibility")
load("//xla/tsl/platform/default:rules_cc.bzl", "cc_library")

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

# Required to load package group definitions.
xla_cpu_nanort_packages()

cc_library(
    name = "nanort_client",
    srcs = ["nanort_client.cc"],
    hdrs = ["nanort_client.h"],
    visibility = internal_visibility([
        "//xla/backends/cpu/nanort:nanort_users",
    ]),
    deps = [
        ":nanort_executable",
        "//xla:debug_options_flags",
        "//xla:shape_util",
        "//xla:util",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/pjrt:utils",
        "//xla/service:compiler",
        "//xla/service:dump",
        "//xla/service:executable",
        "//xla/service:hlo_module_config",
        "//xla/service/cpu:cpu_compiler_pure",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/status:statusor",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_tsl//tsl/profiler/lib:traceme_encode",
    ],
)

xla_cc_test(
    name = "nanort_client_test",
    srcs = ["nanort_client_test.cc"],
    deps = [
        ":nanort_client",
        ":nanort_executable",
        "//xla:array2d",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu:alignment",
        "//xla/ffi",
        "//xla/ffi:execution_context",
        "//xla/ffi:ffi_api",
        "//xla/hlo/builder:xla_builder",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/parser:hlo_parser",
        "//xla/pjrt:pjrt_client",
        "//xla/pjrt:pjrt_executable",
        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
        "//xla/runtime:device_id",
        "//xla/service:computation_placer_hdr",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "nanort_executable",
    srcs = ["nanort_executable.cc"],
    hdrs = ["nanort_executable.h"],
    visibility = internal_visibility([
        "//xla/backends/cpu/nanort:nanort_users",
    ]),
    deps = [
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/cpu:alignment",
        "//xla/backends/cpu/runtime:buffer_allocations",
        "//xla/backends/cpu/runtime:function_library",
        "//xla/backends/cpu/runtime:thread_pool_task_runner",
        "//xla/backends/cpu/runtime:thunk",
        "//xla/ffi:execution_context",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:device_id",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_layout",
        "//xla/service:computation_placer_hdr",
        "//xla/service:executable",
        "//xla/service:hlo_value",
        "//xla/service/cpu:cpu_executable",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:fixed_array",
        "@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/strings:str_format",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
        "@local_tsl//tsl/platform:casts",
        "@local_tsl//tsl/platform:platform_port",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_tsl//tsl/profiler/lib:traceme_encode",
    ],
)

cc_library(
    name = "ifrt_client",
    srcs = ["ifrt_client.cc"],
    hdrs = ["ifrt_client.h"],
    copts = [
        "-Wno-unused-variable",
    ],
    visibility = internal_visibility([
        "//xla/backends/cpu/nanort:nanort_users",
    ]),
    deps = [
        ":nanort_client",
        ":nanort_executable",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu:alignment",
        "//xla/hlo/builder:xla_computation",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/ir:hlo_sharding",
        "//xla/pjrt:mlir_to_hlo",
        "//xla/pjrt:pjrt_compiler",
        "//xla/pjrt:pjrt_executable",
        "//xla/pjrt:pjrt_layout",
        "//xla/pjrt:utils",
        "//xla/python/ifrt",
        "//xla/python/ifrt:attribute_map",
        "//xla/python/ifrt:basic_device_list",
        "//xla/python/ifrt:client_impl_util",
        "//xla/python/ifrt:user_context",
        "//xla/python/ifrt/hlo:hlo_program",
        "//xla/python/pjrt_ifrt:pjrt_dtype",
        "//xla/python/pjrt_ifrt:xla_ifrt",
        "//xla/service:hlo_module_config",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/base:dynamic_annotations",
        "@com_google_absl//absl/base:nullability",
        "@com_google_absl//absl/container:btree",
        "@com_google_absl//absl/container:inlined_vector",
        "@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",
        "@llvm-project//llvm:Support",
        "@local_tsl//tsl/platform:fingerprint",
        "@local_tsl//tsl/platform:platform_port",
    ],
)

cc_library(
    name = "register_nanort_for_ifrt_tests",
    testonly = True,
    srcs = ["register_nanort_for_ifrt_tests.cc"],
    visibility = internal_visibility([
        "//xla/backends/cpu/nanort:nanort_users",
    ]),
    deps = [
        ":ifrt_client",
        "//xla/python/ifrt:test_util",
    ],
    alwayslink = True,
)

xla_cc_test(
    name = "ifrt_client_test",
    srcs = ["ifrt_client_test.cc"],
    deps = [
        ":ifrt_client",
        ":register_nanort_for_ifrt_tests",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:xla_data_proto_cc",
        "//xla/pjrt:mlir_to_hlo",
        "//xla/pjrt:pjrt_executable",
        "//xla/python/ifrt",
        "//xla/python/ifrt:array_impl_test_lib",
        "//xla/python/ifrt:client_impl_test_lib",
        "//xla/python/ifrt:test_util",
        "//xla/python/ifrt:tuple_impl_test_lib",
        "//xla/python/ifrt:user_context",
        "//xla/python/ifrt/hlo:hlo_program",
        "//xla/python/pjrt_ifrt:xla_executable_impl_test_lib",
        "//xla/python/pjrt_ifrt:xla_ifrt",
        "//xla/tsl/concurrency:ref_count",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
    ],
)
