load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm") load("@xla//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm") load( "@xla//xla/tsl:tsl.bzl", "check_deps", "internal_visibility", ) load("//tensorflow:tensorflow.bzl", "if_libtpu", "if_with_tpu_support", "tf_cc_test", "tf_copts", "tf_cuda_cc_test", "tf_cuda_only_cc_test") load("//tensorflow:tensorflow.default.bzl", "cc_header_only_library", "filegroup", "tf_custom_op_py_strict_library") load("//tensorflow/compiler/jit:package_groups.bzl", "legacy_jit_users_package_group") load("//tensorflow/core/platform:build_config.bzl", "tf_additional_all_protos", "tf_proto_library") load( "//tensorflow/core/platform:build_config_root.bzl", "tf_cuda_tests_tags", ) load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], default_visibility = [ ":internal", "//third_party/cloud_tpu/inference_converter:__pkg__", "//waymo/ml/deploy/system/aot_compilation:__pkg__", ], licenses = ["notice"], ) legacy_jit_users_package_group(name = "legacy_jit_users") package_group( name = "internal", includes = [ "//tensorflow/compiler/tf2xla:internal", ], packages = [ "//tensorflow/c/...", "//tensorflow/compiler/tests/...", "//tensorflow/core/tfrt/saved_model/...", "//tensorflow/python/...", ], ) package_group( name = "friends", includes = [ "//tensorflow/compiler/tf2xla:friends", ], ) # defs.cc/h only contains string constants, and can be included in mobile # builds. filegroup( name = "mobile_srcs_no_runtime", srcs = [ "defs.cc", ], visibility = [":friends"], ) filegroup( name = "mobile_hdrs_no_runtime", srcs = [ "defs.h", ], visibility = [":friends"], ) # Target that bundles up the XLA CPU and GPU JIT devices. # This target is deprecated as it pulls in all the XLA devices and bloats the dependencies. # Please use the individual targets in the deps list as needed. See b/336889334. cc_library( name = "jit", visibility = internal_visibility([ ":legacy_jit_users", "//tensorflow/python/profiler:__pkg__", "//tensorflow/python/kernel_tests/nn_ops:__pkg__", ]), deps = [ ":xla_cpu_device", ":xla_cpu_jit", "//tensorflow/compiler/plugin", ] + if_cuda_or_rocm([ ":xla_gpu_device", ":xla_gpu_jit", ]) + if_with_tpu_support([ ":xla_tpu_device", ":xla_tpu_jit", ]), alwayslink = 1, ) cc_library( name = "xla_cpu_jit", visibility = ["//visibility:public"], deps = [ ":jit_compilation_passes", ":xla_kernel_creator", # buildcleaner: keep "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core/tfrt/common:pjrt_cpu_client_registration", ] + if_with_tpu_support( if_false = ["@xla//xla/service:cpu_plugin"], if_true = [], ), alwayslink = 1, ) cc_library( name = "xla_gpu_jit", visibility = ["//visibility:public"], deps = if_cuda_or_rocm([ ":jit_compilation_passes", ":xla_kernel_creator", # buildcleaner: keep "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "@xla//xla/pjrt/plugin/xla_gpu:gpu_static_registration", "@xla//xla/service:gpu_plugin", "//tensorflow/core/tfrt/common:pjrt_gpu_client_registration", ]) + if_cuda([ "@xla//xla/stream_executor/cuda:all_runtime", # buildcleaner: keep ]) + if_rocm([ "@xla//xla/stream_executor/rocm:all_runtime", # buildcleaner: keep ]), alwayslink = 1, ) cc_library( name = "xla_tpu_jit", visibility = ["//visibility:public"], deps = if_libtpu([ "//tensorflow/core/tpu/graph_rewrite:tpu_rewrite_pass_registration", "//tensorflow/core/tpu/graph_rewrite:configure_tpu_embedding_rewrite_registration", ]), alwayslink = 1, ) cc_library( name = "xla_cpu_device", srcs = ["xla_cpu_device.cc"], visibility = [":friends"], deps = [ ":common", ":flags", ":jit_compilation_passes", ":xla_device", ":xla_kernel_creator", # buildcleaner: keep "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "@com_google_absl//absl/memory", "@xla//xla/stream_executor:platform_manager", ] + if_libtpu( if_false = [ "@xla//xla/service:cpu_plugin", # buildcleaner: keep ], if_true = [], ), alwayslink = 1, ) cc_library( name = "xla_gpu_device", srcs = ["xla_gpu_device.cc"], visibility = [":friends"], deps = [ ":common", ":flags", ":jit_compilation_passes", ":xla_device", ":xla_device_no_jit_rewrite_registration", ":xla_kernel_creator", # buildcleaner: keep "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor/gpu:gpu_init", ] + if_libtpu( if_false = [ "@xla//xla/service:gpu_plugin", # buildcleaner: keep ], if_true = [], ), alwayslink = 1, ) cc_library( name = "xla_tpu_device", srcs = ["xla_tpu_device.cc"], visibility = [":friends"], deps = [ ":xla_device", "//tensorflow/compiler/jit/kernels:xla_ops_no_jit_rewrite_registration", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework_internal", "//tensorflow/core/framework:kernel_def_proto_cc", "//tensorflow/core/tpu:tpu_defs", ], alwayslink = 1, ) cc_library( name = "xla_tensor", srcs = ["xla_tensor.cc"], hdrs = ["xla_tensor.h"], visibility = [":friends"], deps = [ "//tensorflow/compiler/tf2xla:common", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/memory", "@xla//xla:shape_util", "@xla//xla/client:local_client", "@xla//xla/service:shaped_buffer", ], ) cc_library( name = "xla_device_context", srcs = ["xla_device_context.cc"], hdrs = ["xla_device_context.h"], visibility = ["//visibility:public"], deps = [ ":xla_tensor", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_proto_parsing", "//tensorflow/core/common_runtime:dma_helper", "//tensorflow/core/framework:allocator", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "@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", "@xla//xla:literal", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla/client:local_client", "@xla//xla/service:stream_pool", "@xla//xla/stream_executor:allocator_stats", "@xla//xla/stream_executor:event", "@xla//xla/stream_executor:stream", "@xla//xla/stream_executor:stream_executor_h", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], ) cc_library( name = "xla_device_no_jit_rewrite_registration", srcs = [ "xla_compile_on_demand_op.cc", "xla_compiler_options_util.cc", "xla_device.cc", "xla_device_ops.cc", "xla_ops_on_regular_devices.cc", "xla_platform_info.cc", ], hdrs = [ "xla_compile_on_demand_op.h", "xla_compiler_options_util.h", "xla_device.h", "xla_device_ops.h", "xla_platform_info.h", ], # Public visibility is needed for external TF/XLA backends. visibility = ["//visibility:public"], deps = [ ":common", ":device_compilation_cluster_signature", ":device_compilation_profiler", ":device_compiler", ":device_executable_persistor", ":flags_headers", ":pjrt_base_device", ":pjrt_device_compiler_client", ":pjrt_device_context", ":variable_info", ":variable_info_util", ":xla_compile_util", ":xla_device_compiler_client", ":xla_device_context", ":xla_launch_util", ":xla_tensor", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:array_ops_op_lib", "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:dataset_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:math_ops_op_lib", "//tensorflow/core:nn_ops_op_lib", "//tensorflow/core:no_op_op_lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:sendrecv_ops_op_lib", "//tensorflow/core:state_ops_op_lib", "//tensorflow/core/kernels:constant_op", "//tensorflow/core/kernels:fifo_queue", "//tensorflow/core/kernels:function_ops", "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:resource_variable_ops", "//tensorflow/core/kernels:shape_ops", "//tensorflow/core/kernels:variable_ops", "//tensorflow/core/kernels/data:finalize_dataset_op", "//tensorflow/core/kernels/data:generator_dataset_op", "//tensorflow/core/kernels/data:iterator_ops", "//tensorflow/core/kernels/data:optional_ops", "//tensorflow/core/kernels/data:options_dataset_op", "//tensorflow/core/kernels/data:prefetch_dataset_op", "//tensorflow/core/platform:stream_executor_no_cuda", "//tensorflow/core/profiler/lib:traceme", "//tensorflow/core/tfrt/common:async_value_tensor", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:global_state", "//tensorflow/core/tfrt/common:pjrt_util", "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@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/types:optional", "@com_google_absl//absl/types:span", "@xla//xla:executable_run_options", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt:tf_pjrt_client", "@xla//xla/service:compiler", "@xla//xla/service:executable", "@xla//xla/service:stream_pool", "@xla//xla/service/gpu:gpu_executable_run_options", "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor:stream", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/stream_executor/integrations:tf_allocator_adapter", "@xla//xla/tsl/framework:device_id_utils", ], alwayslink = 1, ) check_deps( name = "xla_device_check_deps", disallowed_deps = [ "@xla//xla/service/gpu:gpu_compiler", "@xla//xla/service/cpu:cpu_compiler", ], deps = [":xla_device"], ) cc_library( name = "xla_device", hdrs = [ "xla_compile_on_demand_op.h", "xla_device.h", "xla_device_ops.h", ], # Public visibility is needed for external TF/XLA backends. visibility = ["//visibility:public"], deps = [ ":device_compilation_cluster_signature", ":device_compilation_profiler", ":jit_compilation_passes", ":variable_info", ":variable_info_util", ":xla_device_no_jit_rewrite_registration", ":xla_launch_util", ":xla_tensor", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:array_ops_op_lib", "//tensorflow/core:control_flow_ops_op_lib", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:dataset_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:functional_ops_op_lib", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:math_ops_op_lib", "//tensorflow/core:nn_ops_op_lib", "//tensorflow/core:no_op_op_lib", "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:sendrecv_ops_op_lib", "//tensorflow/core:state_ops_op_lib", "//tensorflow/core/kernels:constant_op", "//tensorflow/core/kernels:fifo_queue", "//tensorflow/core/kernels:function_ops", "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:resource_variable_ops", "//tensorflow/core/kernels:shape_ops", "//tensorflow/core/kernels:variable_ops", "//tensorflow/core/kernels/data:finalize_dataset_op", "//tensorflow/core/kernels/data:generator_dataset_op", "//tensorflow/core/kernels/data:iterator_ops", "//tensorflow/core/kernels/data:optional_ops", "//tensorflow/core/kernels/data:options_dataset_op", "//tensorflow/core/kernels/data:prefetch_dataset_op", "//tensorflow/core/platform:stream_executor_no_cuda", "//tensorflow/core/tfrt/common:async_value_tensor", "@com_google_absl//absl/types:optional", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", ], ) cc_library( name = "shape_inference_helpers", srcs = ["shape_inference_helpers.cc"], hdrs = ["shape_inference_helpers.h"], visibility = [":friends"], deps = select({ "//tensorflow:android": [ "//tensorflow/core:portable_tensorflow_lib", ], "//conditions:default": [ "//tensorflow/core:graph", ], }), ) cc_library( name = "flags", srcs = ["flags.cc"], hdrs = ["flags.h"], visibility = [":friends"], deps = [ "//tensorflow/compiler/mlir/tensorflow:dump_graph", "//tensorflow/core:framework_internal", "//tensorflow/core:framework_types_hdr", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/tpu/kernels:sparse_core_xla_flags_defaults", "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@xla//xla:parse_flags_from_env", ], ) # Header-only version of "flags" library, for linking from the shared object # without ODR violations. cc_library( name = "flags_headers", hdrs = ["flags.h"], visibility = [":friends"], deps = [ "//tensorflow/core:framework_internal", "//tensorflow/core:framework_types_hdr", "//tensorflow/core:lib", "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/types:optional", ], ) cc_header_only_library( name = "flags_headers_only", features = [ "-parse_headers", # buildifier: disable=no-parse-headers ], deps = [":flags_headers"], ) cc_library( name = "common", srcs = [ "defs.cc", ], hdrs = [ "defs.h", ], visibility = [":friends"], ) # Internal targets below this point. cc_library( name = "variable_info", srcs = ["variable_info.cc"], hdrs = ["variable_info.h"], visibility = [ ":internal", # We reuse VariableInfo in TFRT's implementation of TpuExecuteOp. "//learning/brain/tfrt/tf_tpu:__pkg__", "//learning/brain/tfrt/tpu_common:__pkg__", "//tensorflow/core/common_runtime/next_pluggable_device:__pkg__", ], deps = [ "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", ], ) cc_library( name = "variable_info_util", srcs = ["variable_info_util.cc"], hdrs = ["variable_info_util.h"], visibility = [ ":internal", # We reuse VariableInfo in TFRT's implementation of TpuExecuteOp. "//learning/brain/tfrt/tf_tpu:__pkg__", ], deps = [ ":variable_info", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/algorithm:container", ], ) cc_library( name = "pjrt_compile_util", srcs = ["pjrt_compile_util.cc"], hdrs = ["pjrt_compile_util.h"], visibility = [ ":internal", "//tensorflow/core/tfrt/gpu/kernel:__pkg__", ], deps = [ ":device_compilation_profiler", ":device_compiler", ":xla_compile_util", ":xla_device_no_jit_rewrite_registration", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:lib", "@tsl//tsl/platform:statusor", "@xla//xla/pjrt:pjrt_client", ], ) tf_cuda_cc_test( name = "pjrt_compile_util_test", srcs = ["pjrt_compile_util_test.cc"], tags = [ "config-cuda-only", "no_oss", # This test only runs with GPU. "requires-gpu-nvidia", "xla", ], deps = [ ":pjrt_compile_util", ":test_util", ":xla_gpu_jit", "//tensorflow/cc:function_ops", "//tensorflow/cc:math_ops", "//tensorflow/cc:scope", "//tensorflow/core:framework", "//tensorflow/core:test", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "//tensorflow/core/platform:status_matchers", "//tensorflow/core/platform:statusor", "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_googletest//:gtest_main", ], ) cc_library( name = "xla_launch_util", srcs = ["xla_launch_util.cc"], hdrs = ["xla_launch_util.h"], visibility = [ ":internal", # We reuse VariableInfo in TFRT's implementation of TpuExecuteOp. "//learning/brain/tfrt/tf_tpu:__pkg__", "//learning/brain/tfrt/tpu_plugin:__pkg__", "//tensorflow/core/tfrt/gpu/kernel:__pkg__", ], deps = [ ":pjrt_tensor_buffer", ":pjrt_tensor_buffer_util", ":variable_info", ":variable_info_util", ":xla_tensor", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:dma_helper", "//tensorflow/core/common_runtime:gpu_device_context", "//tensorflow/core/common_runtime/gpu:gpu_serving_device_selector", "//tensorflow/core/tfrt/common:async_value_tensor", "@com_google_absl//absl/algorithm:container", "@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/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", "@xla//xla:future", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:util", "@xla//xla/client:local_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt:pjrt_common", "@xla//xla/pjrt:pjrt_executable", "@xla//xla/service:executable", "@xla//xla/service:maybe_owning_device_memory", "@xla//xla/service:shaped_buffer", "@xla//xla/service:transfer_manager", "@xla//xla/stream_executor:device_memory", "@xla//xla/stream_executor:device_memory_allocator", "@xla//xla/stream_executor:event", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/framework:device_id_utils", "@xla//xla/tsl/framework:serving_device_selector_policies", ], ) tf_cc_test( name = "xla_launch_util_test", srcs = ["xla_launch_util_test.cc"], deps = [ ":device_compilation_profiler", ":device_compiler", ":device_executable_persistor", ":flags_headers", ":pjrt_device_compiler_client", ":variable_info", ":variable_info_util", ":xla_cpu_device", ":xla_cpu_jit", ":xla_device_no_jit_rewrite_registration", ":xla_launch_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:portable_gif_internal", "//tensorflow/core:session_options", "//tensorflow/core/framework:fake_input", "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:refcount", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt:pjrt_common", "@xla//xla/pjrt:pjrt_executable", "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/tests:literal_test_util", "@xla//xla/tsl/framework:device_id_utils", "@xla//xla/tsl/lib/core:status_test_util", "@xla//xla/tsl/platform:statusor", ], ) tf_cuda_only_cc_test( name = "xla_launch_util_gpu_test", srcs = ["xla_launch_util_gpu_test.cc"], deps = [ ":device_compiler", ":flags_headers", ":pjrt_device_compiler_client", ":pjrt_device_context", ":variable_info", ":variable_info_util", ":xla_device_no_jit_rewrite_registration", ":xla_gpu_jit", ":xla_launch_util", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "//tensorflow/core/framework:fake_input", "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:refcount", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_absl//absl/container:flat_hash_set", "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:status", "@tsl//tsl/platform:statusor", "@xla//xla/pjrt:pjrt_client", "@xla//xla/tests:literal_test_util", "@xla//xla/tsl/lib/core:status_test_util", ], ) cc_library( name = "xla_compile_util", srcs = ["xla_compile_util.cc"], hdrs = ["xla_compile_util.h"], visibility = [ ":internal", "//tensorflow/core/tfrt/gpu/kernel:__pkg__", ], deps = [ ":flags_headers", "//tensorflow/compiler/tf2xla:xla_argument", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core/tfrt/common:global_state", "//tensorflow/core/util:determinism", "@com_google_absl//absl/status", ], ) tf_cc_test( name = "xla_compile_util_test", srcs = [ "xla_compile_util_test.cc", ], deps = [ ":flags_headers", ":xla_compile_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:test_main", "//tensorflow/core/framework:fake_input", "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/tpu:tpu_defs", "@com_google_googletest//:gtest", ], ) tf_proto_library( name = "xla_compilation_cache_proto", srcs = ["xla_compilation_cache.proto"], protodeps = tf_additional_all_protos() + ["@xla//xla/service:hlo_proto"], visibility = ["//visibility:public"], ) tf_proto_library( name = "tf_graph_to_hlo_compiler_proto", srcs = ["tf_graph_to_hlo_compiler.proto"], protodeps = tf_additional_all_protos(), visibility = ["//visibility:public"], deps = ["//tensorflow/compiler/tf2xla:xla_argument_proto"], ) cc_library( name = "device_compiler", hdrs = ["device_compiler.h"], copts = tf_copts(), visibility = [":internal"], deps = [ ":device_compilation_cache", ":device_compilation_cluster_signature", ":device_compilation_profiler", ":device_compiler_client", ":device_executable_persistor", ":flags_headers", ":tf_graph_to_hlo_compiler", ":xla_compile_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:framework_lite", "//tensorflow/core:lib", "//tensorflow/core/framework:resource_base", "//tensorflow/core/platform:thread_annotations", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", "@com_google_absl//absl/types:span", ], ) tf_cc_test( name = "device_compiler_disable_test", srcs = [ "device_compiler_disable_test.cc", ], tags = ["no_cuda_on_cpu_tap"], deps = [ ":device_compilation_profiler", ":device_compiler", ":flags", ":xla_cpu_jit", ":xla_device_compiler_client", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:test", "//tensorflow/core:test_main", "@xla//xla/client:client_library", ], ) cc_library( name = "jit_compilation_passes", srcs = ["jit_compilation_pass_registration.cc"], deps = [ ":compilation_passes", ":xla_activity_logging_listener", "//tensorflow/compiler/tf2xla:functionalize_control_flow_pass_registration", "//tensorflow/compiler/tf2xla:mlir_bridge_pass_registration", "//tensorflow/core:core_cpu_internal", ], alwayslink = 1, ) cc_library( name = "get_compiler_ir", srcs = ["get_compiler_ir.cc"], hdrs = ["get_compiler_ir.h"], visibility = [":internal"], deps = [ ":compilability_check_util", ":device_compiler", ":variable_info", ":variable_info_util", ":xla_device_no_jit_rewrite_registration", ":xla_launch_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:core_cpu_internal", "//tensorflow/core/common_runtime/eager:tensor_handle", "//tensorflow/core/platform:refcount", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@xla//xla:status_macros", "@xla//xla/client:executable_build_options", "@xla//xla/client:local_client", "@xla//xla/hlo/translate:portable_api", "@xla//xla/service:hlo_graph_dumper", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor/host:host_platform_id", ], alwayslink = 1, ) # Header-only version of "flags" library, for linking from the shared object # without ODR violations. cc_library( name = "get_compiler_ir_hdrs", textual_hdrs = ["get_compiler_ir.h"], visibility = [":internal"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:statusor", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], ) cc_header_only_library( name = "get_compiler_ir_hdrs_only", features = [ "-parse_headers", # buildifier: disable=no-parse-headers ], deps = [":get_compiler_ir_hdrs"], ) # This target can be used by XLA device plugins to prevent circular dependencies, and provides access to all of the required headers for building a device library. cc_header_only_library( name = "xla_jit_headers_lib", visibility = ["//visibility:public"], deps = [ ":xla_cpu_device", ":xla_cpu_jit", ":xla_gpu_device", ":xla_gpu_jit", ], ) cc_library( name = "xla_kernel_creator", srcs = [ "xla_kernel_creator.cc", ], hdrs = [ "xla_kernel_creator.h", ], visibility = [ ":internal", "//tensorflow/core/common_runtime/eager:__pkg__", ], deps = [ ":common", ":compilability_check_util", ":compilation_passes", ":flags", ":jit_compilation_passes", "//tensorflow/compiler/jit/kernels:xla_ops_no_jit_rewrite_registration", "//tensorflow/compiler/tf2xla:mlir_bridge_pass", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", ], alwayslink = 1, ) tf_cc_test( name = "xla_kernel_creator_test", srcs = [ "xla_kernel_creator.h", "xla_kernel_creator_test.cc", ], deps = [ ":xla_kernel_creator", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:session_options", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", ], ) cc_library( name = "resource_operation_safety_analysis", srcs = ["resource_operation_safety_analysis.cc"], hdrs = ["resource_operation_safety_analysis.h"], deps = [ ":xla_cluster_util", "//tensorflow/compiler/tf2xla:resource_operation_table", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@xla//xla/service/graphcycles", ], ) tf_cc_test( name = "resource_operation_safety_analysis_test", srcs = ["resource_operation_safety_analysis_test.cc"], deps = [ ":common", ":resource_operation_safety_analysis", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:sendrecv_ops", "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_absl//absl/strings", ], ) cc_library( name = "shape_inference", srcs = ["shape_inference.cc"], hdrs = ["shape_inference.h"], visibility = [":friends"], deps = [ ":shape_inference_helpers", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], ) cc_library( name = "test_util", testonly = 1, srcs = ["test_util.cc"], hdrs = ["test_util.h"], deps = [ ":shape_inference", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@xla//xla:status_macros", ], ) tf_cc_test( name = "shape_inference_test", srcs = ["shape_inference_test.cc"], deps = [ ":shape_inference", ":test_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:scope", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:ops", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/kernels:constant_op", "@com_google_absl//absl/log:check", ], ) cc_library( name = "encapsulate_util", srcs = ["encapsulate_util.cc"], hdrs = ["encapsulate_util.h"], deps = [ ":shape_inference", "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", ], ) tf_cc_test( name = "encapsulate_util_test", srcs = ["encapsulate_util_test.cc"], deps = [ ":encapsulate_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:ops", "//tensorflow/cc:scope", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log:check", ], ) cc_library( name = "compilation_passes", srcs = [ "build_xla_ops_pass.cc", "clone_constants_for_better_clustering.cc", "cluster_scoping_pass.cc", "deadness_analysis.cc", "deadness_analysis_internal.h", "encapsulate_subgraphs_pass.cc", "encapsulate_xla_computations_pass.cc", "extract_outside_compilation_pass.cc", "force_xla_constants_on_host_pass.cc", "increase_dynamism_for_auto_jit_pass.cc", "mark_for_compilation_pass.cc", "mark_for_compilation_pass_test_helper.cc", "partially_decluster_pass.cc", "report_clustering_info_pass.cc", ], hdrs = [ "build_xla_ops_pass.h", "clone_constants_for_better_clustering.h", "cluster_scoping_pass.h", "deadness_analysis.h", "encapsulate_subgraphs_pass.h", "encapsulate_xla_computations_pass.h", "extract_outside_compilation_pass.h", "force_xla_constants_on_host_pass.h", "increase_dynamism_for_auto_jit_pass.h", "mark_for_compilation_pass.h", "mark_for_compilation_pass_test_helper.h", "partially_decluster_pass.h", "report_clustering_info_pass.h", ], visibility = [ ":internal", "//third_party/cloud_tpu/inference_converter:__pkg__", "//waymo/onboard/ml/chauffeur_net:__pkg__", ], deps = [ "compilability_check_util", ":common", ":device_util", ":encapsulate_util", ":flags", ":resource_operation_safety_analysis", ":shape_inference_helpers", ":xla_activity_listener", ":xla_cluster_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:scope", "//tensorflow/cc:scope_internal", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:resource_operation_table", "//tensorflow/compiler/tf2xla:side_effect_util", "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/cc:xla_jit_ops", "//tensorflow/compiler/tf2xla/cc:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core/framework:bounds_check", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@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/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/types:optional", "@xla//xla:status_macros", "@xla//xla:union_find", "@xla//xla:util", "@xla//xla:xla_data_proto_cc", "@xla//xla/service/graphcycles", ], ) cc_library( name = "xla_cluster_util", srcs = ["xla_cluster_util.cc"], hdrs = ["xla_cluster_util.h"], deps = [ ":flags", ":xla_activity_proto_cc", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:function_body", "//tensorflow/core/framework:bounds_check", "//tensorflow/core/platform:hash", "@com_google_absl//absl/algorithm:container", "@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/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:optional", "@xla//xla:status_macros", "@xla//xla/service/graphcycles", ], ) cc_library( name = "device_util", srcs = ["device_util.cc"], hdrs = ["device_util.h"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/numeric:bits", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@xla//xla:status_macros", ], ) tf_cc_test( name = "device_util_test", srcs = ["device_util_test.cc"], deps = [ ":device_util", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", ], ) tf_cc_test( name = "deadness_analysis_test", size = "small", srcs = [ "deadness_analysis_internal.h", "deadness_analysis_test.cc", ], deps = [ ":common", ":compilation_passes", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", "//tensorflow/cc:sendrecv_ops", "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_absl//absl/container:flat_hash_map", ], ) cc_library( name = "compilation_passes_test_main", testonly = True, srcs = ["compilation_passes_test_main.cc"], visibility = ["//visibility:public"], deps = [ ":flags", "//tensorflow/core:lib", "//tensorflow/core:test", "@com_google_absl//absl/strings", ], ) tf_cc_test( name = "compilation_passes_test", size = "small", srcs = [ "build_xla_ops_pass_test.cc", "clone_constants_for_better_clustering_test.cc", "cluster_scoping_pass_test.cc", "encapsulate_subgraphs_pass_test.cc", "encapsulate_xla_computations_pass_test.cc", "extract_outside_compilation_pass_test.cc", "force_xla_constants_on_host_pass_test.cc", "increase_dynamism_for_auto_jit_pass_test.cc", "mark_for_compilation_pass_test.cc", "partially_decluster_pass_test.cc", "rearrange_function_argument_pass_test.cc", ], tags = [ # TODO(b/141643254) Re-enable msan after fixing # use-of-uninitialized-value error. "nomsan", ] + tf_cuda_tests_tags(), deps = [ ":common", ":compilability_check_util", ":compilation_passes", ":compilation_passes_test_main", ":encapsulate_util", ":flags", ":node_matchers", ":test_util", ":xla_cluster_util", ":xla_cpu_device", ":xla_gpu_device", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:scope", "//tensorflow/cc:sendrecv_ops", "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:rearrange_function_argument", "//tensorflow/compiler/tf2xla:side_effect_util", "//tensorflow/compiler/tf2xla:test_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/cc:xla_jit_ops", "//tensorflow/compiler/tf2xla/cc:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:all_kernels", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:session_options", "//tensorflow/core:test", "//tensorflow/core:testlib", "//tensorflow/core/common_runtime:device_set", "@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/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@xla//xla/hlo/testlib:test", ], ) tf_cc_test( name = "xla_cluster_util_test", size = "small", srcs = [ "xla_cluster_util_test.cc", ], deps = [ ":common", ":xla_cluster_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_lib", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/strings", "@xla//xla:status_macros", ], ) cc_library( name = "node_matchers", testonly = True, srcs = ["node_matchers.cc"], hdrs = ["node_matchers.h"], deps = [ "//tensorflow/cc:ops", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@xla//xla/hlo/testlib:test", ], ) tf_cc_test( name = "node_matchers_test", srcs = ["node_matchers_test.cc"], deps = [ ":node_matchers", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:ops", "//tensorflow/core:ops", "//tensorflow/core:test_main", ], ) cc_library( name = "compilability_check_util", srcs = ["compilability_check_util.cc"], hdrs = ["compilability_check_util.h"], visibility = [ ":friends", ], deps = [ ":common", ":device_util", ":flags", ":resource_operation_safety_analysis", ":xla_activity_listener", ":xla_activity_proto_cc", ":xla_cluster_util", "//tensorflow/compiler/tf2xla:resource_operation_table", "//tensorflow/compiler/tf2xla:tf2xla_defs", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:function_body", "//tensorflow/core/common_runtime:function_utils", "@com_google_absl//absl/algorithm:container", "@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:optional", "@com_google_absl//absl/types:span", "@xla//xla:union_find", "@xla//xla:util", "@xla//xla/service/graphcycles", ], ) tf_cc_test( name = "compilability_check_util_test", srcs = ["compilability_check_util_test.cc"], deps = [ ":common", ":compilability_check_util", ":xla_cpu_device", ":xla_cpu_jit", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:scope", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], ) tf_cc_test( name = "xla_activity_listener_test", srcs = ["xla_activity_listener_test.cc"], deps = [ ":flags", ":xla_activity_listener", ":xla_cpu_device", ":xla_cpu_jit", "//tensorflow/cc:cc_ops", "//tensorflow/cc:ops", "//tensorflow/core:all_kernels", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:ops", "//tensorflow/core:test", "//tensorflow/core/common_runtime:direct_session_internal", "//tensorflow/core/kernels:cwise_op", "//tensorflow/core/kernels:matmul_op", "//tensorflow/core/kernels:partitioned_function_ops", ], ) tf_custom_op_py_strict_library( name = "xla_ops_py", kernels = ["//tensorflow/compiler/jit/ops:xla_ops"], visibility = [ ":friends", ], deps = [ "//tensorflow/compiler/jit/ops:xla_ops_grad", "//tensorflow/compiler/jit/ops:xla_ops_wrapper_py", ], ) cc_library( name = "xla_activity_listener", srcs = ["xla_activity_listener.cc"], hdrs = ["xla_activity_listener.h"], visibility = ["//visibility:public"], deps = [ ":xla_activity_proto_cc", "//tensorflow/core:lib", "@com_google_absl//absl/synchronization", ], ) tf_proto_library( name = "xla_activity_proto", srcs = ["xla_activity.proto"], protodeps = tf_additional_all_protos(), ) cc_library( name = "xla_activity_logging_listener", srcs = ["xla_activity_logging_listener.cc"], deps = [ ":xla_activity_listener", ":xla_activity_proto_cc", "//tensorflow/core:lib", "@com_google_absl//absl/memory", ], alwayslink = 1, ) cc_library( name = "tf_to_hlo_compiler", hdrs = ["tf_to_hlo_compiler.h"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", ], ) cc_library( name = "tf_graph_to_hlo_compiler", srcs = ["tf_graph_to_hlo_compiler.cc"], hdrs = ["tf_graph_to_hlo_compiler.h"], deps = [ ":tf_graph_to_hlo_compiler_proto_cc", ":tf_to_hlo_compiler", "//tensorflow/compiler/tf2xla:xla_argument", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:lib", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:path", "@xla//xla/tsl/lib/strings:proto_serialization", ], ) cc_library( name = "device_compilation_profiler", srcs = ["device_compilation_profiler.cc"], hdrs = ["device_compilation_profiler.h"], deps = [ ":xla_activity_listener", ":xla_activity_proto_cc", ":xla_compile_util", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "@com_google_absl//absl/strings", "@tsl//tsl/platform:mutex", ], ) cc_library( name = "device_compiler_client", srcs = ["device_compiler_client.cc"], hdrs = ["device_compiler_client.h"], visibility = [":internal"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core/util:determinism", "@xla//xla/client:executable_build_options", ], ) tf_cc_test( name = "device_compiler_client_test", srcs = ["device_compiler_client_test.cc"], deps = [ ":device_compiler_client", "@com_google_googletest//:gtest_main", ], ) cc_library( name = "xla_device_compiler_client", srcs = ["xla_device_compiler_client.cc"], hdrs = ["xla_device_compiler_client.h"], deps = [ ":device_compiler_client", "@xla//xla/client:local_client", ], ) cc_library( name = "device_executable_persistor", srcs = ["device_executable_persistor.cc"], hdrs = ["device_executable_persistor.h"], deps = [ ":device_compiler_client", ":xla_compilation_cache_proto_cc", ":xla_device_compiler_client", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib_headers_for_pybind", "//tensorflow/core:portable_gif_internal", "//tensorflow/core/platform:path", "//tensorflow/core/platform:status", "//tensorflow/core/platform:statusor", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@tsl//tsl/platform:errors", "@tsl//tsl/platform:statusor", "@xla//xla:util", "@xla//xla/pjrt:pjrt_client", "@xla//xla/service:hlo_proto_cc", ], ) cc_library( name = "device_compilation_cache", hdrs = ["device_compilation_cache.h"], deps = [ ":device_compilation_cluster_signature", ":xla_compile_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework_lite", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", ], ) cc_library( name = "device_compilation_cluster_signature", srcs = ["device_compilation_cluster_signature.cc"], hdrs = ["device_compilation_cluster_signature.h"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core/platform:types", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", ], ) cc_library( name = "pjrt_device_compiler_client", srcs = ["pjrt_device_compiler_client.cc"], hdrs = ["pjrt_device_compiler_client.h"], deps = [ ":device_compiler_client", "@xla//xla/pjrt:pjrt_client", ], ) cc_library( name = "pjrt_base_device", srcs = ["pjrt_base_device.cc"], hdrs = ["pjrt_base_device.h"], # Public visibility is needed for external TF/XLA backends. visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/core:framework", "//tensorflow/core/common_runtime:local_device", ], ) cc_library( name = "pjrt_tensor_buffer", hdrs = ["pjrt_tensor_buffer.h"], visibility = ["//visibility:public"], deps = [ "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "@xla//xla/pjrt:pjrt_client", ], ) cc_library( name = "pjrt_tensor_buffer_util", srcs = ["pjrt_tensor_buffer_util.cc"], hdrs = ["pjrt_tensor_buffer_util.h"], visibility = ["//visibility:public"], deps = [ ":pjrt_tensor_buffer", "//tensorflow/core:framework", "//tensorflow/core/common_runtime:dma_helper", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@tsl//tsl/platform:statusor", "@xla//xla/pjrt:pjrt_client", ], ) tf_cuda_only_cc_test( name = "pjrt_tensor_buffer_util_test", srcs = ["pjrt_tensor_buffer_util_test.cc"], tags = tf_cuda_tests_tags() + [ "config-cuda-only", "no_oss", ], deps = [ ":pjrt_tensor_buffer_util", ":test_util", ":xla_gpu_jit", "//tensorflow/core:framework", "//tensorflow/core:gpu_runtime", # Needed to register DEVICE_GPU "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:statusor", "@xla//xla:shape_util", "@xla//xla/pjrt:pjrt_client", ], ) cc_library( name = "pjrt_device_context", srcs = [ "pjrt_device_context.cc", ], hdrs = [ "pjrt_device_context.h", ], visibility = ["//visibility:public"], deps = [ ":pjrt_tensor_buffer", ":pjrt_tensor_buffer_util", "//tensorflow/c/experimental/next_pluggable_device:tensor_pjrt_buffer_util", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/core:framework", "//tensorflow/core/common_runtime:dma_helper", "//tensorflow/core/common_runtime/next_pluggable_device:next_pluggable_device_api", "//tensorflow/core/platform:status", "//tensorflow/core/profiler/lib:traceme", "//tensorflow/core/tfrt/common:async_value_tensor", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "@com_google_absl//absl/status", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt:pjrt_common", "@xla//xla/tsl/c:tsl_status_internal", "@xla//xla/tsl/framework:device_id_utils", ], ) cc_library( name = "xla_host_recv_device_context", srcs = [ "xla_host_recv_device_context.cc", ], hdrs = [ "xla_host_recv_device_context.h", ], visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:common", "//tensorflow/core:framework", "@tf_runtime//:async_value", "@xla//xla:shape_util", "@xla//xla/stream_executor:device_memory", "@xla//xla/stream_executor:stream", ], ) cc_library( name = "xla_host_send_device_context", srcs = [ "xla_host_send_device_context.cc", ], hdrs = [ "xla_host_send_device_context.h", ], visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:common", "//tensorflow/core:framework", "@tf_runtime//:async_value", "@xla//xla:shape_util", "@xla//xla/stream_executor:device_memory", "@xla//xla/stream_executor:stream", ], ) tf_cuda_only_cc_test( name = "xla_host_send_recv_device_context_test", srcs = ["xla_host_send_recv_device_context_test.cc"], tags = tf_cuda_tests_tags() + [ "config-cuda-only", "no_oss", # Temporarily disable OSS. ], deps = [ ":flags", ":xla_device", ":xla_gpu_device", ":xla_host_recv_device_context", ":xla_host_send_device_context", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework_internal", "//tensorflow/core:gpu_runtime", # Needed to register DEVICE_GPU "//tensorflow/core:test", "//tensorflow/core/framework:tensor_testutil", "@com_google_googletest//:gtest_main", "@xla//xla:shape_util", "@xla//xla/stream_executor:device_memory", "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor:stream", "@xla//xla/stream_executor:stream_executor_h", ], ) tf_cc_test( name = "device_compilation_cluster_signature_test", srcs = [ "device_compilation_cluster_signature_test.cc", ], deps = [ ":device_compilation_cluster_signature", ":flags", "//tensorflow/compiler/tf2xla:common", "//tensorflow/core:test", "//tensorflow/core:test_main", "@xla//xla/client:client_library", ], ) tf_cc_test( name = "device_compilation_profiler_test", srcs = ["device_compilation_profiler_test.cc"], tags = [ "nomsan", # TODO(b/284492454) ], deps = [ ":device_compilation_profiler", ":xla_activity_proto_cc", "//tensorflow/compiler/jit/tests:device_compiler_test_helper", "//tensorflow/core:protos_all_cc", "@com_google_googletest//:gtest_main", ], ) tf_cc_test( name = "device_executable_persistor_test", srcs = ["device_executable_persistor_test.cc"], tags = ["no_cuda_on_cpu_tap"], deps = [ ":device_compiler_client", ":device_executable_persistor", ":pjrt_device_compiler_client", ":xla_compilation_cache_proto_cc", ":xla_cpu_device", ":xla_cpu_jit", ":xla_device_compiler_client", "//tensorflow/cc:function_ops", "//tensorflow/cc:math_ops", "//tensorflow/cc:scope", "//tensorflow/core:test", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status_matchers", "//tensorflow/core/platform:statusor", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_googletest//:gtest_main", "@xla//xla/client:client_library", "@xla//xla/client:executable_build_options", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", ], ) tf_cc_test( name = "device_compilation_cache_test", srcs = ["device_compilation_cache_test.cc"], deps = [ ":device_compilation_cache", ":xla_compile_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:test", "//tensorflow/core/platform:errors", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest_main", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc", ], ) # TODO(b/261212343): Support running this test on CPU and OSS. tf_cuda_cc_test( name = "device_compiler_test", srcs = ["device_compiler_test.cc"], tags = [ "config-cuda-only", "no_oss", # This test only runs with GPU. "requires-gpu-nvidia", "xla", ], deps = [ ":device_compilation_cluster_signature", ":device_compiler", ":device_compiler_client", ":tf_graph_to_hlo_compiler_proto_cc", ":xla_compile_util", ":xla_device_compiler_client", ":xla_gpu_device", ":xla_gpu_jit", "//tensorflow/cc:function_ops", "//tensorflow/cc:math_ops", "//tensorflow/cc:scope", "//tensorflow/compiler/jit/tests:device_compiler_test_helper", "//tensorflow/compiler/tf2xla:xla_argument", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:test", "//tensorflow/core/framework:fake_input", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:env", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status", "//tensorflow/core/platform:status_matchers", "//tensorflow/core/platform:statusor", "//tensorflow/core/platform:test", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:path", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/lib/strings:proto_serialization", "@xla//xla/tsl/platform:statusor", ], ) tf_cuda_cc_test( name = "device_context_test", srcs = ["device_context_test.cc"], tags = tf_cuda_tests_tags() + [ "config-cuda-only", "no_oss", # Temporarily disable OSS. ], deps = [ ":flags", ":xla_device", ":xla_gpu_device", ":xla_gpu_jit", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework_internal", "//tensorflow/core:test", "//tensorflow/core/framework:tensor_testutil", "@com_google_googletest//:gtest_main", ], ) tf_cuda_cc_test( name = "xla_compiler_options_util_test", srcs = ["xla_compiler_options_util_test.cc"], tags = tf_cuda_tests_tags(), deps = [ ":device_compiler", ":device_executable_persistor", ":flags", ":pjrt_base_device", ":pjrt_device_compiler_client", ":test_util", ":xla_device_compiler_client", ":xla_device_no_jit_rewrite_registration", ":xla_gpu_device", ":xla_gpu_jit", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_argument", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "//tensorflow/core:framework_types_hdr", "//tensorflow/core/platform:refcount", "//tensorflow/core/tpu:tpu_defs", "@com_google_googletest//:gtest_main", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", "@xla//xla/stream_executor:device_memory_allocator", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/platform:statusor", ], ) tf_cuda_cc_test( name = "xla_platform_info_test", srcs = ["xla_platform_info_test.cc"], tags = tf_cuda_tests_tags() + ["config-cuda-only"], deps = [ ":device_compilation_profiler", ":device_compiler", ":flags_headers", ":test_util", ":xla_device_no_jit_rewrite_registration", ":xla_gpu_device", ":xla_gpu_jit", "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "//tensorflow/core:framework_types_hdr", "//tensorflow/core:lib_proto_parsing", "//tensorflow/core:test", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status_matchers", "//tensorflow/core/platform:statusor", "//tensorflow/core/protobuf:error_codes_proto_impl_cc", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:pjrt_util", "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/log:check", "@com_google_googletest//:gtest_main", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/tsl/platform:statusor", ], )