load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@xla//third_party/rules_python/python:py_library.bzl", "py_library") # copybara:uncomment_begin(google-only) # load("//devtools/deps/check:deps_check.bzl", "check_dependencies") # copybara:uncomment_end load("@xla//xla:xla.default.bzl", "xla_py_proto_library") load("@xla//xla/service/cpu:build_defs.bzl", "runtime_copts") load("@xla//xla/tsl/mkl:build_defs.bzl", "mkl_deps") load("@xla//xla/tsl/platform:build_config_root.bzl", "if_static") load( "@xla//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", ) load("//tensorflow:tensorflow.bzl", "if_google", "if_libtpu", "tf_cc_binary", "tf_cc_test", "tf_copts", "tf_cuda_cc_test", "tf_gen_op_wrapper_py", "tf_openmp_copts") load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "filegroup", "get_compatible_with_portable") load( "//tensorflow/core/platform:build_config.bzl", "tf_additional_tensor_coding_deps", "tf_proto_library", ) load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], default_visibility = [":internal"], licenses = ["notice"], ) package_group( name = "internal", packages = [ "//tensorflow/compiler/aot/...", "//tensorflow/compiler/jit/...", "//tensorflow/compiler/mlir/...", "//tensorflow/compiler/tests/...", "//tensorflow/compiler/tf2xla/...", "//tensorflow/core/tpu/...", "//tensorflow/python/compiler/...", "//tensorflow/python/util/...", ], ) package_group( name = "friends", includes = [":internal"], packages = [ "//platforms/performance/automl/...", "//tensorflow/...", "//tensorflow_federated/cc/core/impl/executors/...", "//tensorflow_models/...", "//third_party/deepmind/deepmind_research/density_functional_approximation_dm21/...", "//third_party/mlir_edge/model_curriculum/iree/...", "//third_party/mlperf/submissions/training/v0_7/models/...", "//third_party/py/keras_cv/...", "//third_party/py/litert_torch/...", "//third_party/py/orbax/...", "//third_party/py/tf_keras/...", "//third_party/sparse_conv/ops/...", "//waymo/ml/deploy/benchmark/...", ], ) cc_library( name = "tf2xla_supported_ops_lib", srcs = ["tf2xla_supported_ops.cc"], hdrs = ["tf2xla_supported_ops.h"], visibility = ["//visibility:public"], deps = [ ":xla_compiler", ":xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/strings", ], ) tf_cc_binary( name = "tf2xla_supported_ops", srcs = ["tf2xla_supported_ops_main.cc"], visibility = ["//visibility:public"], deps = [":tf2xla_supported_ops_lib"], ) tf_proto_library( name = "tf2xla_proto", srcs = ["tf2xla.proto"], protodeps = [ "//tensorflow/core:protos_all", ], visibility = ["//visibility:public"], ) xla_py_proto_library( name = "tf2xla_py", has_services = False, visibility = ["//visibility:public"], deps = [":tf2xla_proto"], ) tf_proto_library( name = "host_compute_metadata_proto", srcs = ["host_compute_metadata.proto"], protodeps = [ "//tensorflow/core:protos_all", ], visibility = ["//visibility:public"], ) cc_library( name = "graph_compiler_util", srcs = ["graph_compiler_util.cc"], hdrs = ["graph_compiler_util.h"], visibility = [":friends"], deps = [ ":functionalize_control_flow", ":sharding_util", ":tf2xla_proto_cc", ":tf2xla_util", ":xla_compiler", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:graph", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", ], ) cc_library( name = "encoded_buffer_allocation_info", hdrs = ["encoded_buffer_allocation_info.h"], visibility = [":friends"], deps = [ "@xla//xla/backends/cpu:buffer_allocation_info", ], ) tf_cc_test( name = "encoded_buffer_allocation_info_test", srcs = ["encoded_buffer_allocation_info_test.cc"], deps = [ ":encoded_buffer_allocation_info", "@com_google_googletest//:gtest_main", "@xla//xla/backends/cpu:buffer_allocation_info", ], ) cc_library( name = "tf2xla", srcs = ["tf2xla.cc"], hdrs = ["tf2xla.h"], visibility = ["//visibility:public"], deps = [ ":common", ":graph_compiler_util", ":tf2xla_proto_cc", ":tf2xla_util", ":xla_compiler", ":xla_op_registry", "//tensorflow/compiler/aot:aot_only_var_handle_op", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:status", "@com_google_absl//absl/strings", "@xla//xla/client", "@xla//xla/hlo/builder:xla_computation", ], ) cc_library( name = "mlir_tf2xla", srcs = ["mlir_tf2xla.cc"], hdrs = ["tf2xla.h"], visibility = ["//visibility:public"], deps = [ ":tf2xla_proto_cc", ":tf2xla_util", ":xla_compiler", "//tensorflow/compiler/jit", "//tensorflow/compiler/mlir/tensorflow:device_util", "//tensorflow/compiler/mlir/tensorflow:import_utils", "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags", "//tensorflow/compiler/mlir/tensorflow/transforms:bridge", "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes", "//tensorflow/compiler/mlir/tensorflow/translate/tools:parsers", "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util", "//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:framework_types_hdr", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:device_set", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "//tensorflow/core/platform:types", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@llvm-project//mlir:IR", "@tsl//tsl/platform:errors", "@tsl//tsl/platform:statusor", "@xla//xla/client", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/tsl/platform:statusor", ], ) # The filegroups below are explicitly used by # tensorflow/tools/pip_package:build_pip_package to ensure we include the proper # sources for the XLA AOT CPU runtime; as these are necessary outside of bazel # when linking tfcompile objects using saved_model_cli (e.g. using the # tensorflow pip package). The associated .cc files are included in tensorflow # pip package's xla_aot_runtime_srcs/ subdirectory. All necessary headers are # also included in the pip package's include/tensorflow/ and include/external/ # subdirectories. Note however that sometimes additional object files may need # to be linked when linking aot xla objects, e.g. abseil libraries. See the deps # attribute of the "xla_compiled_cpu_runtime_standalone" target below for an # exhaustive list. filegroup( name = "xla_compiled_cpu_runtime_hdrs", srcs = [ "allocator.h", "encoded_buffer_allocation_info.h", "xla_compiled_cpu_function.h", "//tensorflow/core/kernels:xla_cpu_runtime_hdrs", "//tensorflow/core/platform:xla_cpu_runtime_srcs", "@tsl//tsl/platform:xla_cpu_runtime_srcs", "@xla//xla:cpu_runtime_hdrs", "@xla//xla/backends/cpu:xla_cpu_runtime_hdrs", "@xla//xla/service:custom_call_status_hdrs", "@xla//xla/service/cpu:runtime_hdrs", "@xla//xla/tsl/concurrency:xla_cpu_runtime_hdrs", "@xla//xla/tsl/framework:xla_cpu_runtime_hdrs", "@xla//xla/tsl/framework/fixedpoint:xla_cpu_runtime_hdrs", "@xla//xla/tsl/lib/math:xla_cpu_runtime_hdrs", "@xla//xla/tsl/util:xla_cpu_runtime_hdrs", ], visibility = [ "//tensorflow/tools/pip_package:__pkg__", "//tensorflow/tools/pip_package:__subpackages__", ], ) filegroup( name = "xla_compiled_cpu_runtime_srcs", srcs = [ "allocator.cc", "xla_compiled_cpu_function.cc", "//tensorflow/core/kernels:xla_cpu_runtime_srcs", "//tensorflow/core/platform:xla_cpu_runtime_srcs", "@tsl//tsl/platform:xla_cpu_runtime_srcs", "@xla//xla:cpu_runtime_srcs", "@xla//xla/service:custom_call_status_srcs", "@xla//xla/service/cpu:runtime_srcs", "@xla//xla/tsl/concurrency:xla_cpu_runtime_srcs", ], visibility = [ "//tensorflow/tools/pip_package:__pkg__", "//tensorflow/tools/pip_package:__subpackages__", ], ) # This stand-alone target is used to ensure that we can build tf_library type # targets against the subset of sources declared in # xla_compiled_cpu_runtime_{srcs,hdrs}. # # The macros in tensorflow/python/tools/tools.bzl produce AOT compiled binaries # that rely on this target, as do unit tests in tensorflow/python/tools. # # See above for the significance of the source filegroups. cc_library( name = "xla_compiled_cpu_runtime_standalone", srcs = [ ":xla_compiled_cpu_runtime_srcs", ], hdrs = [ ":xla_compiled_cpu_runtime_hdrs", ], copts = runtime_copts() + tf_openmp_copts(), defines = [ "EIGEN_NEON_GEBP_NR=4", ], features = [ "fully_static_link", "-parse_headers", ], linkstatic = 1, visibility = [":friends"], # Note, we specifically removed MKL and multithreaded dependencies so the # standalone does not require the MKL binary blob or threading libraries. # # TODO(ebrevdo): Remove tf_additoinal_tensor_coding_deps in favor of # absl/strings:cord when we update absl to a newer version. # # NOTE!!!!! (b/150412548 for reference) # DO NOT add dependencies here! # # 1) Read first the purpose of this target and the purpose of the # :xla_compiled_cpu_runtime_{srcs|hrds} targets. # # Key point is: this target represents the sources available under the xla_aot_runtime_srcs # dir in the tensorflow pip package, which enables linking AOT-ed saved modules in customer # projects, including those not using bazel. So adding a dependency here to fix the tests # using this target (xla_compiled_cpu_runtime_standalone) defeats the purpose. # # 2) Most likely you need to split your dependency into 2 file groups and add them # to the :xla_compiled_cpu_runtime_{srcs|hrds} targets accordingly. # # !!!DO NOT EDIT dependencies - please read the above!!! deps = [ "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:dynamic_annotations", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/synchronization", "@ducc//:fft_wrapper", "@eigen_archive//:eigen3", "@llvm-project//mlir:TransformUtils", "@xla//xla/tsl:empty", "//tensorflow/core/framework:numeric_types", "//tensorflow/core/platform:bfloat16", "//tensorflow/core/platform:float8", "//tensorflow/core/platform:stringpiece", # Extra dependencies required for multithreaded runtime objects. "//tensorflow/core/platform:blocking_counter", "//tensorflow/core/platform:logging", "//tensorflow/core/platform:mutex", "@tsl//tsl/platform:blocking_counter", "@tsl//tsl/platform:context", "@tsl//tsl/platform:cord", "@tsl//tsl/platform:logging", "@tsl//tsl/platform:ml_dtypes", "@tsl//tsl/platform:tstring", "@tsl//tsl/platform:types", ] + tf_additional_tensor_coding_deps() + mkl_deps(), alwayslink = 1, ) # copybara:uncomment_begin(google-only) # check_dependencies( # # Please read the whole comment above the `deps` of xla_compiled_cpu_runtime_standalone before # # making any change here. # name = "only_pip_sources_in_xla_lib", # of = [":xla_compiled_cpu_runtime_standalone"], # allowed_match_regexp = "|".join([ # ":xla_compiled_cpu_runtime_standalone", # "//tensorflow/core/framework:numeric_types", # "@xla//xla/tsl/framework:numeric_types", # "@xla//xla/tsl/framework:device_type", # "@xla//xla/tsl/framework:fixedpoint_types", # "@tsl//tsl/platform:bfloat16", # "@tsl//tsl/platform:blocking_counter", # "@xla//xla/tsl/platform:byte_order", # "@tsl//tsl/platform:context", # "@tsl//tsl/platform:cord", # "@tsl//tsl/platform:env_time", # "@tsl//tsl/platform:refcount", # "@tsl//tsl/platform:ml_dtypes", # "@tsl//tsl/platform:logging", # "@tsl//tsl/platform:macros", # "@tsl//tsl/platform:mutex", # "@tsl//tsl/platform:platform", # "@tsl//tsl/platform:stringpiece", # "@tsl//tsl/platform:thread_annotations", # "@tsl//tsl/platform:tstring", # "@tsl//tsl/platform:types", # "@xla//xla/tsl/platform:env_time", # "@xla//xla/tsl/platform:logging", # "@xla//xla/tsl/platform:types", # "@xla//xla/tsl/platform:macros", # "@xla//xla/tsl/platform/default:context", # "@xla//xla/tsl/platform/default:cord", # "@xla//xla/tsl/platform/default:env_time", # "@xla//xla/tsl/platform/default:types", # "@xla//xla/tsl/platform/google:context", # "@xla//xla/tsl/platform/google:cord", # "@xla//xla/tsl/platform/google:env_time", # "@xla//xla/tsl/platform/google:types", # "@xla//xla/tsl/platform/windows:env_time", # "//tensorflow/core/platform:bfloat16", # "//tensorflow/core/platform:blocking_counter", # "//tensorflow/core/platform:byte_order", # "//tensorflow/core/platform:cord", # "//tensorflow/core/platform:env_time", # "//tensorflow/core/platform:float8", # "//tensorflow/core/platform:logging", # "//tensorflow/core/platform:macros", # "//tensorflow/core/platform:mutex", # "//tensorflow/core/platform:platform", # "//tensorflow/core/platform:stringpiece", # "//tensorflow/core/platform:thread_annotations", # "//tensorflow/core/platform:tstring", # "//tensorflow/core/platform:types", # "@xla//xla/tsl:empty", # ]), # dont_match_regexp = "|".join(["//tensorflow"]), # kind = "cc_library", # ) # copybara:uncomment_end cc_library( name = "allocator", srcs = ["allocator.cc"], hdrs = ["allocator.h"], compatible_with = get_compatible_with_portable(), visibility = ["//visibility:public"], deps = [ "@com_google_absl//absl/base:dynamic_annotations", "@com_google_absl//absl/types:span", "@xla//xla/backends/cpu:alignment", "@xla//xla/backends/cpu:buffer_allocation_info", ], ) tf_cc_test( name = "allocator_test", srcs = ["allocator_test.cc"], deps = [ ":allocator", "//tensorflow/core:framework", "//tensorflow/core:test", "//tensorflow/core:test_main", "@xla//xla/backends/cpu:alignment", "@xla//xla/backends/cpu:buffer_allocation_info", ], ) cc_library( name = "xla_compiled_cpu_function", srcs = ["xla_compiled_cpu_function.cc"], hdrs = ["xla_compiled_cpu_function.h"], compatible_with = get_compatible_with_portable(), visibility = ["//visibility:public"], deps = [ # Keep dependencies to a minimum here; this library is used in every AOT # binary produced by tfcompile. ":allocator", "@com_google_absl//absl/log:check", "@com_google_absl//absl/types:span", ":encoded_buffer_allocation_info", "@xla//xla/service:custom_call_status_internal", "@xla//xla/backends/cpu/runtime:rng_state_lib", "@xla//xla/backends/cpu:alignment", "@xla//xla/backends/cpu:buffer_allocation_info", "@xla//xla:executable_run_options", "//tensorflow/core/platform:types", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings:string_view", ], ) cc_library( name = "xla_compiled_cpu_function_thunks", srcs = ["xla_compiled_cpu_function_thunks.cc"], hdrs = ["xla_compiled_cpu_function_thunks.h"], compatible_with = get_compatible_with_portable(), visibility = ["//visibility:public"], deps = [ ":xla_compiled_cpu_function", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/types:span", "@xla//xla:executable_run_options", "@xla//xla/backends/cpu/codegen:aot_compiled_function_library", "@xla//xla/backends/cpu/nanort:nanort_executable", "@xla//xla/backends/cpu/runtime:function_library", "@xla//xla/backends/cpu/runtime/thunk_serdes:all_thunk_serdes", # buildcleaner: keep "@xla//xla/service:executable", # buildcleaner: keep (b/404179184) "@xla//xla/service/cpu:cpu_aot_compilation_result", "@xla//xla/service/cpu:executable_proto_cc", "@xla//xla/tsl/concurrency:async_value", "@xla//xla/tsl/platform:env", ], ) cc_library( name = "xla_compiled_cpu_function_factory", srcs = ["xla_compiled_cpu_function_factory.cc"], hdrs = ["xla_compiled_cpu_function_factory.h"], visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", ], ) cc_library( name = "xla_compiled_cpu_function_thunk_factory_registerer", srcs = ["xla_compiled_cpu_function_thunk_factory_registerer.cc"], visibility = ["//visibility:public"], deps = [ ":xla_compiled_cpu_function_factory", "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function", "//tensorflow/compiler/tf2xla:xla_compiled_cpu_function_thunks", ], alwayslink = 1, ) cc_library( name = "xla_jit_compiled_cpu_function", srcs = ["xla_jit_compiled_cpu_function.cc"], hdrs = ["xla_jit_compiled_cpu_function.h"], visibility = ["//visibility:public"], deps = [ ":encoded_buffer_allocation_info", ":tf2xla", ":tf2xla_proto_cc", ":xla_compiled_cpu_function", ":xla_compiled_cpu_function_thunks", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:errors", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/backends/cpu:buffer_allocation_info", "@xla//xla/backends/cpu:buffer_allocation_info_util", "@xla//xla/backends/cpu/codegen:compiled_function_library", "@xla//xla/backends/cpu/runtime/thunk_serdes:all_thunk_serdes", # buildcleaner: keep "@xla//xla/client:client_library", "@xla//xla/client:executable_build_options", "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/service:platform_util", "@xla//xla/service/cpu:cpu_aot_compilation_result", "@xla//xla/service/cpu:executable_proto_cc", "@xla//xla/stream_executor:platform", ] + if_libtpu( if_false = [ "@xla//xla/service:cpu_plugin", "@xla//xla/service/cpu:cpu_executable", ], if_true = [], ), ) tf_cc_test( name = "graph_compiler_test", srcs = ["graph_compiler_test.cc"], deps = [ ":graph_compiler_util", ":tf2xla_proto_cc", ":xla_compilation_device", ":xla_compiler", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/platform:refcount", "@com_google_googletest//:gtest_main", ], ) cc_library( name = "xla_compiler", srcs = [ "const_analysis.cc", "graph_compiler.cc", "xla_compiler.cc", "xla_cpu_backend.cc", "xla_op_kernel.cc", ] + if_cuda_is_configured([ "xla_gpu_backend.cc", ]) + if_rocm_is_configured([ "xla_gpu_backend.cc", ]), hdrs = [ "const_analysis.h", "graph_compiler.h", "xla_compiler.h", "xla_helpers.h", "xla_op_kernel.h", "xla_op_registry.h", ], copts = tf_copts(), visibility = [":friends"], deps = [ ":common", ":host_compute_metadata_proto_cc", ":layout_util", ":rearrange_function_argument", ":sharding_util", ":side_effect_util", ":tf2xla_util", ":xla_argument", ":xla_compilation_device", ":xla_context", ":xla_expression", ":xla_helpers", ":xla_op_registry", ":xla_resource", "//tensorflow/compiler/jit:common", "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:shape_inference", "//tensorflow/compiler/jit:xla_compile_util", "//tensorflow/compiler/mlir/tensorflow:attribute_utils", "//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy", "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes", "//tensorflow/compiler/mlir/utils:array_container_utils", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "//tensorflow/core/tpu:tpu_defs", "//tensorflow/core/util:overflow", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/memory", "@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:optional", "@com_google_absl//absl/types:span", "@com_google_absl//absl/types:variant", "@tsl//tsl/platform:tensor_float_32_hdr_lib", "@xla//xla:executable_run_options", "@xla//xla:literal", "@xla//xla:protobuf_util", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:util", "@xla//xla:xla_data_proto_cc", "@xla//xla/client", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/core/collectives:clique_id", "@xla//xla/core/collectives:clique_key", "@xla//xla/hlo/builder:value_inference", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util", "@xla//xla/service:computation_placer_hdr", "@xla//xla/service:hlo_proto_cc", "@xla//xla/service/spmd/shardy:constants", "@xla//xla/service/spmd/shardy:utils", "@xla//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import", "@xla//xla/stream_executor:device_memory_allocator", ] + if_libtpu([ ":xla_tpu_backend_registration", ]) + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), alwayslink = 1, ) cc_library( name = "xla_compilation_device", srcs = [ "xla_compilation_device.cc", ], hdrs = [ "xla_compilation_device.h", ], deps = [ ":common", ":frontend_attributes_util", ":sharding_util", ":xla_context", ":xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:session_options", "//tensorflow/core/common_runtime:core_cpu_internal", "@xla//xla/hlo/builder:xla_builder", ], alwayslink = 1, ) cc_library( name = "xla_tpu_backend_registration", srcs = ["xla_tpu_backend.cc"], visibility = ["//visibility:public"], deps = [ ":xla_op_registry", "//tensorflow/core/tpu:tpu_defs", "//tensorflow/core/tpu:tpu_node_device_util", ], alwayslink = 1, ) cc_library( name = "xla_context", srcs = [ "xla_context.cc", ], hdrs = [ "xla_context.h", ], deps = [ ":common", ":xla_expression", ":xla_helpers", ":xla_resource", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/common_runtime:core_cpu_internal", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", "@xla//xla/client:client_library", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder:xla_computation", ], alwayslink = 1, ) cc_library( name = "xla_op_registry", srcs = [ "xla_op_registry.cc", ], hdrs = [ "xla_op_registry.h", ], visibility = [":friends"], deps = [ ":common", ":xla_context", "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:xla_cluster_util", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:session_options", "//tensorflow/core/common_runtime:core_cpu_internal", "//tensorflow/core/common_runtime/next_pluggable_device:next_pluggable_device_factory_hdrs", "//tensorflow/core/platform:stream_executor_no_cuda", "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@xla//xla:util", "@xla//xla/client:client_library", ], alwayslink = 1, ) cc_library( name = "xla_expression", srcs = [ "xla_expression.cc", ], hdrs = [ "xla_expression.h", ], deps = [ ":common", ":xla_resource", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:optional", "@xla//xla/client", "@xla//xla/hlo/builder:value_inference", "@xla//xla/hlo/builder:xla_builder", ], alwayslink = 1, ) cc_library( name = "xla_resource", srcs = [ "xla_resource.cc", ], hdrs = [ "xla_resource.h", ], visibility = [ ":internal", "//learning/deepmind/tensorflow/tpufunc:__pkg__", ], deps = [ ":common", ":sharding_util", ":xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@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:optional", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder:xla_builder", ], alwayslink = 1, ) cc_library( name = "xla_helpers", srcs = [ "xla_helpers.cc", ], hdrs = [ "xla_helpers.h", ], visibility = [":friends"], deps = [ ":common", ":host_compute_metadata_proto_cc", "//tensorflow/compiler/tf2xla/lib:util", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/platform:refcount", "@com_google_absl//absl/container:btree", "@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/time", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@xla//xla:executable_run_options", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla:shape_util", "@xla//xla/core/collectives:clique_id", "@xla//xla/core/collectives:clique_key", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util", "@xla//xla/service:computation_placer_hdr", "@xla//xla/service/gpu:gpu_executable_run_options", "@xla//xla/stream_executor:stream", "@xla//xla/stream_executor:stream_executor_h", ], alwayslink = 1, ) cc_library( name = "xla_argument", srcs = [ "xla_argument.cc", ], hdrs = [ "xla_argument.h", ], deps = [ ":host_compute_metadata_proto_cc", ":xla_argument_proto_cc", ":xla_resource", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@com_google_absl//absl/types:variant", "@llvm-project//llvm:Support", "@xla//xla:shape_util", "@xla//xla/tsl/platform:statusor", ], alwayslink = 1, ) tf_cc_test( name = "xla_argument_test", srcs = ["xla_argument_test.cc"], deps = [ ":xla_argument", ":xla_argument_proto_cc", ":xla_resource", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/framework:tensor_testutil", "@com_google_googletest//:gtest_main", "@xla//xla:shape_util", ], ) tf_proto_library( name = "xla_argument_proto", srcs = ["xla_argument.proto"], protodeps = [ "@xla//xla:xla_data_proto", "//tensorflow/core:protos_all", ], visibility = [ ":friends", "//tensorflow/compiler:__subpackages__", ] + if_google([ "//waymo/ml/deploy:__subpackages__", ]), ) cc_library( name = "layout_util", srcs = [ "layout_util.cc", ], hdrs = [ "layout_util.h", ], visibility = [":friends"], deps = [ ":common", ":xla_argument", ":xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/ir:hlo", ], alwayslink = 1, ) cc_library( name = "common", srcs = [ "literal_util.cc", "shape_util.cc", "type_util.cc", ], hdrs = [ "literal_util.h", "shape_util.h", "type_util.h", ], visibility = [":friends"], deps = [ "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@xla//xla:literal", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", ], ) tf_cc_test( name = "type_util_test", srcs = ["type_util_test.cc"], deps = [ ":common", "//tensorflow/core:framework_types_hdr", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/status:statusor", ], ) cc_library( name = "frontend_attributes_util", srcs = ["frontend_attributes_util.cc"], hdrs = ["frontend_attributes_util.h"], visibility = ["//visibility:public"], deps = [ ":tf2xla_defs", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:optional", "@xla//xla:xla_data_proto_cc", ], ) cc_library( name = "sharding_util", srcs = ["sharding_util.cc"], hdrs = ["sharding_util.h"], visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/mlir/tensorflow:xla_sharding_util", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla/hlo/builder:sharding_builder", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/ir:hlo_sharding", "@xla//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import", ], ) tf_cc_test( name = "sharding_util_test", srcs = ["sharding_util_test.cc"], deps = [ ":sharding_util", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_googletest//:gtest", "@xla//xla:shape_util", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo_sharding", "@xla//xla/tsl/platform:statusor", ], ) # Internal targets below this point. cc_library( name = "tf2xla_defs", hdrs = ["tf2xla_defs.h"], visibility = [":friends"], deps = [ "@com_google_absl//absl/strings", ], ) cc_library( name = "tf2xla_util", srcs = ["tf2xla_util.cc"], hdrs = ["tf2xla_util.h"], visibility = [":friends"], deps = [ ":sharding_util", ":tf2xla_defs", ":tf2xla_proto_cc", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//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/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:xla_data_proto_cc", ], ) tf_cc_test( name = "tf2xla_util_test", srcs = ["tf2xla_util_test.cc"], deps = [ ":sharding_util", ":tf2xla_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:scope", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:math_ops_op_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/status", "@com_google_absl//absl/strings", ], ) tf_cc_test( name = "tf2xla_test", srcs = ["tf2xla_test.cc"], tags = ["no_cuda_on_cpu_tap"], deps = [ ":tf2xla", ":tf2xla_proto_cc", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/status:statusor", "@tsl//tsl/platform:tensor_float_32_hdr_lib", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/service:cpu_plugin", ] + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), ) tf_cc_test( name = "xla_jit_compiled_cpu_function_test", srcs = ["xla_jit_compiled_cpu_function_test.cc"], deps = [ ":tf2xla_proto_cc", ":xla_compiled_cpu_function_thunks", ":xla_jit_compiled_cpu_function", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:statusor", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/client:executable_build_options", "@xla//xla/hlo/testlib:test", "@xla//xla/service:compiler", "@xla//xla/service:platform_util", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor:platform_manager", ], ) tf_cc_test( name = "xla_compiler_test", srcs = [ "xla_compiler_test.cc", "xla_expression_test.cc", ], tags = ["no_cuda_on_cpu_tap"], deps = [ ":common", ":layout_util", ":side_effect_util", ":xla_argument", ":xla_compiler", ":xla_expression", ":xla_resource", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:scope", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "//tensorflow/core/framework:resource_base", "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:types", "@com_google_absl//absl/log", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest", "@tsl//tsl/platform:statusor", "@xla//xla:array", "@xla//xla:debug_options_flags", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla:shape_util", "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", "@xla//xla/client", "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/ir:hlo_sharding", "@xla//xla/hlo/testlib:filecheck", "@xla//xla/service", "@xla//xla/service:cpu_plugin", "@xla//xla/service:hlo_module_util", "@xla//xla/service:hlo_proto_cc", "@xla//xla/service:hlo_proto_util", "@xla//xla/tests:literal_test_util", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], ) tf_cc_test( name = "literal_util_test", srcs = [ "literal_util_test.cc", ], deps = [ ":common", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "//tensorflow/core/platform:types", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc", ], ) tf_cc_test( name = "const_analysis_test", size = "small", srcs = ["const_analysis_test.cc"], deps = [ ":xla_compiler", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:xla_cluster_util", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:ops", "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) cc_library( name = "light_outside_compilation_kernels_for_test", testonly = True, srcs = ["light_outside_compilation_kernels_for_test.cc"], linkstatic = 1, deps = [ ":xla_compiler", "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", "//tensorflow/compiler/tf2xla/kernels:light_outside_compilation", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:types", "@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/stream_executor:device_memory", "@xla//xla/stream_executor:stream", "@xla//xla/tsl/platform:statusor", ], alwayslink = 1, ) tf_gen_op_wrapper_py( name = "test_ops_for_light_outside_compilation", testonly = True, out = "test_ops_for_light_outside_compilation.py", extra_py_deps = [ "//tensorflow/python:pywrap_tfe", "//tensorflow/python/util:dispatch", "//tensorflow/python/util:deprecation", "//tensorflow/python/util:tf_export", ], py_lib_rule = py_library, deps = [ ":light_outside_compilation_kernels_for_test", ], ) cuda_py_strict_test( name = "light_outside_compilation_test", srcs = ["light_outside_compilation_test.py"], tags = [ "no_oss", ], xla_enable_strict_auto_jit = False, deps = [ ":test_ops_for_light_outside_compilation", "//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper", "//tensorflow/python/eager:context", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:ops", "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:nn_ops", "//tensorflow/python/ops:random_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", ] + if_google([":light_outside_compilation_kernels_for_test"]), ) cc_library( name = "functionalize_control_flow_util", srcs = [ "functionalize_control_flow_util.cc", ], hdrs = [ "functionalize_control_flow_util.h", ], deps = [ "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/strings", "@xla//xla:status_macros", ], ) cc_library( name = "functionalize_cond", srcs = [ "functionalize_cond.cc", ], hdrs = [ "functionalize_cond.h", ], deps = [ ":functionalize_control_flow_util", ":tf2xla_defs", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:hash", "@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", "@xla//xla:status_macros", ], ) cc_library( name = "functionalize_control_flow", srcs = [ "functionalize_control_flow.cc", ], hdrs = [ "functionalize_control_flow.h", ], visibility = [":friends"], deps = [ ":functionalize_cond", ":functionalize_control_flow_util", ":functionalize_while", ":tf2xla_util", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:lib", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:optional", "@xla//xla:status_macros", "@xla//xla:union_find", ], ) cc_library( name = "mlir_bridge_pass", srcs = ["mlir_bridge_pass.cc"], hdrs = ["mlir_bridge_pass.h"], visibility = [":internal"], deps = [ ":tf2xla_defs", ":xla_op_registry", "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/mlir:mlir_graph_optimization_pass", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:attribute_utils", "//tensorflow/compiler/mlir/tensorflow:device_util", "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops", "//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy", "//tensorflow/compiler/mlir/tf2xla/api/v1:cluster_tf", "//tensorflow/compiler/mlir/tf2xla/api/v1:tf_dialect_to_executor", "//tensorflow/compiler/mlir/tf2xla/api/v2:cluster_tf", "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_dialect_to_executor", "//tensorflow/compiler/mlir/tf2xla/internal:mlir_bridge_pass_util", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/common_runtime:device_set", "//tensorflow/core/protobuf:for_core_protos_cc", "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", "@xla//xla/tsl/platform:logging", ], alwayslink = 1, ) cc_library( name = "mlir_bridge_pass_registration", srcs = [ "mlir_bridge_pass_registration.cc", ], visibility = [ # We define a new TPU device in TFRT to enable TFRT TPU Runtime. "//learning/brain/tfrt/tf_tpu:__pkg__", ":internal", ], deps = [ ":mlir_bridge_pass", "//tensorflow/compiler/mlir:mlir_graph_optimization_pass", "//tensorflow/compiler/mlir:mlir_graph_optimization_pass_registration", ], alwayslink = 1, ) cc_library( name = "rearrange_function_argument", srcs = [ "rearrange_function_argument.cc", ], hdrs = [ "rearrange_function_argument.h", ], deps = [ ":tf2xla_util", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", "@xla//xla:status_macros", ], ) cc_library( name = "functionalize_control_flow_pass_registration", srcs = [ "functionalize_control_flow_pass_registration.cc", ], visibility = [":friends"], deps = [ ":functionalize_control_flow", ], alwayslink = 1, ) cc_library( name = "functionalize_while", srcs = [ "functionalize_while.cc", ], hdrs = [ "functionalize_while.h", ], deps = [ ":frontend_attributes_util", ":functionalize_cond", ":functionalize_control_flow_util", ":tf2xla_util", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:lib", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@xla//xla:status_macros", "@xla//xla:union_find", ], ) tf_cc_test( name = "functionalize_control_flow_test", srcs = ["functionalize_control_flow_test.cc"], deps = [ ":functionalize_control_flow", ":test_util", ":tf2xla_util", "//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/compiler/tf2xla/cc:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:ops", "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "@xla//xla:status_macros", ], ) tf_cc_test( name = "functionalize_cond_test", srcs = ["functionalize_cond_test.cc"], deps = [ ":functionalize_cond", ":functionalize_control_flow", ":test_util", "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:scope", "//tensorflow/compiler/tf2xla/cc:xla_ops", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:ops", "//tensorflow/core:protos_all_cc", "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "//tensorflow/core/platform:test", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@xla//xla:status_macros", ], ) cc_library( name = "test_util", testonly = 1, srcs = ["test_util.cc"], hdrs = ["test_util.h"], deps = [ "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "@xla//xla:status_macros", ], ) tf_cc_test( name = "xla_op_registry_test", srcs = ["xla_op_registry_test.cc"], deps = [ ":xla_compiler", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log", ], ) cc_library( name = "resource_operation_table", srcs = ["resource_operation_table.cc"], hdrs = ["resource_operation_table.h"], deps = [ "//tensorflow/core:lib", "//tensorflow/core:ops", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) tf_cc_test( name = "resource_operation_table_test", srcs = ["resource_operation_table_test.cc"], deps = [ ":resource_operation_table", ":xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", ], ) cc_library( name = "side_effect_util", srcs = ["side_effect_util.cc"], hdrs = ["side_effect_util.h"], visibility = [":friends"], deps = [ "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/platform:errors", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@xla//xla/tsl/platform:errors", ], ) tf_cuda_cc_test( name = "fused_batchnorm_reserve_space_test", size = "medium", srcs = ["fused_batchnorm_reserve_space_test.cc"], tags = [ "no_cuda_asan", # TODO(b/193450885) ], deps = [ "//tensorflow/cc:cc_ops", "//tensorflow/cc:ops", "//tensorflow/cc:scope", "//tensorflow/compiler/jit", "//tensorflow/compiler/jit:flags", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/strings", "@eigen_archive//:eigen3", ], ) cc_library( name = "mlir_xla_op_kernel", srcs = ["mlir_xla_op_kernel.cc"], hdrs = ["mlir_xla_op_kernel.h"], deps = [ ":xla_compiler", ":xla_expression", "//tensorflow/compiler/jit:xla_compile_util", "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes", "//tensorflow/compiler/mlir/utils:array_container_utils", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/framework:resource_base", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@tsl//tsl/platform:errors", "@tsl//tsl/platform:statusor", "@xla//xla/hlo/builder:xla_builder", ], ) cc_library( name = "resource_util", srcs = ["resource_util.cc"], hdrs = ["resource_util.h"], visibility = [":friends"], deps = [ ":resource_operation_table", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:function_body", "//tensorflow/core/common_runtime:function_utils", "@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/hash", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@xla//xla:status_macros", ], ) tf_cuda_cc_test( name = "resource_util_test", size = "small", srcs = ["resource_util_test.cc"], deps = [ ":resource_util", "//tensorflow/cc:scope", "//tensorflow/core:framework", "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", ], ) filegroup( name = "tf2xla_opset_hdrs", srcs = [ "tf2xla_opset.h", ], visibility = ["//tensorflow/python/util:__pkg__"], ) cc_library( name = "tf2xla_opset", srcs = [ "tf2xla_opset.cc", ], hdrs = ["tf2xla_opset.h"], visibility = [ "//tensorflow/python:__pkg__", "//tensorflow/python/util:__pkg__", ], deps = [ ":tf2xla_util", ":xla_op_registry", "//tensorflow/compiler/jit:xla_device", "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], ) tf_cc_test( name = "tf2xla_opset_test", srcs = [ "tf2xla_opset_test.cc", ], deps = [ ":tf2xla_opset", "//tensorflow/core:test", "//tensorflow/core:test_main", ], )