diff --git a/.bazelrc b/.bazelrc index 9eaa3ebe20276b..f938902fcacca7 100644 --- a/.bazelrc +++ b/.bazelrc @@ -709,6 +709,7 @@ common:rbe_win_base --nobuild_runfile_links common:rbe_win_base --define=override_eigen_strong_inline=true common:rbe_windows_x86_cpu_2022 --config=rbe_win_base --config=windows_x86_cpu_2022 +common:rbe_windows_x86_cpu_2022 --features=-with_wrapper # Override the remote instance and BES instance names to use the new RBE pools common:rbe_windows_x86_cpu_2022 --remote_instance_name=projects/ml-oss-rbe-general/instances/default_instance diff --git a/.github/workflows/pylint-presubmit.yml b/.github/workflows/pylint-presubmit.yml index 3e50cfa713ab3c..0526b66ff4c041 100644 --- a/.github/workflows/pylint-presubmit.yml +++ b/.github/workflows/pylint-presubmit.yml @@ -44,7 +44,7 @@ jobs: - name: Install Python dependencies run: | python -m pip install --upgrade pip - pip install pylint==2.13.9 numpy wheel + pip install pylint==2.13.9 "numpy<2.0.0" wheel - name: Run PyLint on changed files run: | echo "${{ steps.get_file_changes.outputs.files}}" | tr " " "\n" | grep ".py$" | xargs pylint --rcfile=tensorflow/tools/ci_build/pylintrc diff --git a/tensorflow/c/tf_datatype.h b/tensorflow/c/tf_datatype.h index c991fc1f74f2e8..0d1496ccb81ab9 100644 --- a/tensorflow/c/tf_datatype.h +++ b/tensorflow/c/tf_datatype.h @@ -65,7 +65,9 @@ typedef enum TF_DataType { TF_UINT4 = 30, TF_INT2 = 31, TF_UINT2 = 32, - TF_FLOAT4_E2M1FN = 33 // 2 exponent bits, 1 mantissa bit, finite-only + TF_FLOAT4_E2M1FN = 33, // 2 exponent bits, 1 mantissa bit, finite-only + TF_FLOAT8_E8M0FNU = 34 // 8 exponent bits, 0 mantissa bits, finite-only, with + // NaN. } TF_DataType; // TF_DataTypeSize returns the sizeof() for the underlying type corresponding diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD index 93b4bf515cbde8..91e7cf6f300b6f 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -722,6 +722,8 @@ tf_cc_test( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest", + "@xla//xla/tsl/lib/core:status_test_util", + "@xla//xla/tsl/platform:status_matchers", ], ) diff --git a/tensorflow/cc/saved_model/fingerprinting.cc b/tensorflow/cc/saved_model/fingerprinting.cc index b2622a32cb0847..5f184faa5dc2c1 100644 --- a/tensorflow/cc/saved_model/fingerprinting.cc +++ b/tensorflow/cc/saved_model/fingerprinting.cc @@ -174,6 +174,11 @@ absl::StatusOr CreateFingerprintDefPb( SavedModel saved_model; TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), pb_file, &saved_model)); + if (saved_model.meta_graphs_size() == 0) { + return absl::InvalidArgumentError( + "SavedModel (.pb) contains no MetaGraphs."); + } + // Create a copy of `metagraph` which will be used and mutated for fingerprint // computation. FingerprintDef fingerprint_def; diff --git a/tensorflow/cc/saved_model/fingerprinting_test.cc b/tensorflow/cc/saved_model/fingerprinting_test.cc index 135646df74fdf9..bd1758ecafe3b4 100644 --- a/tensorflow/cc/saved_model/fingerprinting_test.cc +++ b/tensorflow/cc/saved_model/fingerprinting_test.cc @@ -17,12 +17,15 @@ limitations under the License. #include +#include #include #include "absl/numeric/int128.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/numbers.h" #include "absl/strings/string_view.h" +#include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/status_matchers.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/platform/env.h" @@ -206,5 +209,19 @@ TEST(FingerprintingTest, TestSingleprint) { const_singleprint); } +TEST(FingerprintingTest, CreateFingerprintDefPbEmptyMetaGraphsReturnsError) { + const std::string model_dir = + io::JoinPath(::testing::TempDir(), "empty_metagraphs_model"); + TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(model_dir)); + const std::string pb_file = io::JoinPath(model_dir, "saved_model.pb"); + SavedModel saved_model; + TF_ASSERT_OK(WriteBinaryProto(Env::Default(), pb_file, saved_model)); + + EXPECT_THAT( + CreateFingerprintDef(model_dir).status(), + tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument, + "SavedModel (.pb) contains no MetaGraphs.")); +} + } // namespace } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index d0877d9af1ac55..add8bc7291af36 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -148,6 +148,7 @@ cc_library( "@xla//xla/client:client_library", "@xla//xla/client:compile_only_client", "@xla//xla/hlo/builder:xla_computation", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/service:compiler", "@xla//xla/service/cpu:cpu_aot_compilation_result", "@xla//xla/service/cpu:cpu_compiler", diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index efc6869f6235de..792ad9bcc5cad5 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -42,6 +42,7 @@ limitations under the License. #include "xla/client/client_library.h" #include "xla/client/compile_only_client.h" #include "xla/hlo/builder/xla_computation.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/service/compiler.h" #include "xla/service/cpu/cpu_aot_compilation_result.h" #include "xla/shape.h" @@ -174,8 +175,10 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config, graph_def, config, &computation, flags.debug_info, flags.debug_info_path_begin_marker)); } else { + TF_ASSIGN_OR_RETURN(auto* pjrt_compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); TF_RETURN_IF_ERROR(ConvertGraphDefToXla(std::move(graph_def), config, - client, &computation)); + pjrt_compiler, &computation)); } if (flags.experimental_quantize && *quantize_xla) { diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a93dc10be415fd..b9410f089772bc 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -156,6 +156,7 @@ cc_library( "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", + "//tensorflow/core/tfrt/common:pjrt_cpu_client_registration", "@com_google_absl//absl/memory", "@xla//xla/stream_executor:platform_manager", ] + if_libtpu( @@ -291,7 +292,6 @@ cc_library( # Public visibility is needed for external TF/XLA backends. visibility = ["//visibility:public"], deps = [ - ":common", ":device_compilation_cluster_signature", ":device_compilation_profiler", ":device_compiler", @@ -304,13 +304,11 @@ cc_library( ":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", @@ -346,7 +344,6 @@ cc_library( "//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", @@ -354,28 +351,21 @@ cc_library( "//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:shape_util", "@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/service:shaped_buffer", + "@xla//xla/stream_executor:platform", "@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", ], @@ -403,11 +393,9 @@ cc_library( deps = [ ":device_compilation_cluster_signature", ":device_compilation_profiler", + ":device_compiler", ":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", @@ -446,6 +434,7 @@ cc_library( "//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/status", "@com_google_absl//absl/types:optional", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", @@ -476,13 +465,12 @@ cc_library( "//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/log", "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:optional", "@xla//xla:parse_flags_from_env", ], ) @@ -496,10 +484,8 @@ cc_library( 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", ], ) @@ -629,7 +615,6 @@ cc_library( ":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", @@ -644,6 +629,7 @@ cc_library( "//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/base", "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -660,18 +646,13 @@ cc_library( "@xla//xla/client:local_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_common", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt:pjrt_executable", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_execute_options", + "@xla//xla/runtime:device_id", "@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/stream_executor:device_address_allocator", "@xla//xla/tsl/framework:device_id_utils", "@xla//xla/tsl/framework:serving_device_selector_policies", ], @@ -875,30 +856,41 @@ cc_library( visibility = [":internal"], deps = [ ":compilability_check_util", + ":device_compilation_profiler", ":device_compiler", + ":pjrt_device_compiler_client", ":variable_info", ":variable_info_util", ":xla_device_no_jit_rewrite_registration", ":xla_launch_util", "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla:xla_resource", "//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", + "//tensorflow/core/platform:statusor", + "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/algorithm:container", + "@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:shape_util", "@xla//xla:status_macros", - "@xla//xla/client:executable_build_options", - "@xla//xla/client:local_client", + "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/translate:portable_api", + "@xla//xla/pjrt:compiled_memory_stats", + "@xla//xla/pjrt:pjrt_client", "@xla//xla/service:hlo_graph_dumper", + "@xla//xla/service:hlo_module_config", + "@xla//xla/service:hlo_proto_util", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor/host:host_platform_id", + "@xla//xla/tsl/platform:errors", + "@xla//xla/tsl/platform:statusor", ], alwayslink = 1, ) @@ -1984,9 +1976,12 @@ tf_cuda_cc_test( ], deps = [ ":device_compilation_cluster_signature", + ":device_compilation_profiler", ":device_compiler", ":device_compiler_client", + ":device_executable_persistor", ":tf_graph_to_hlo_compiler_proto_cc", + ":xla_activity_listener", ":xla_compile_util", ":xla_device_compiler_client", ":xla_gpu_device", @@ -2004,14 +1999,17 @@ tf_cuda_cc_test( "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:env", "//tensorflow/core/platform:errors", + "//tensorflow/core/platform:path", "//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/base", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", @@ -2022,6 +2020,7 @@ tf_cuda_cc_test( "@xla//xla/hlo/builder:xla_computation", "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/lib/strings:proto_serialization", + "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], ) @@ -2075,8 +2074,10 @@ tf_cuda_cc_test( "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", - "@xla//xla/stream_executor:device_memory_allocator", + "@xla//xla/service/cpu:cpu_transfer_manager", + "@xla//xla/stream_executor:device_address_allocator", "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/jit/device_compiler_test.cc b/tensorflow/compiler/jit/device_compiler_test.cc index d42f797526be83..68b6652857da6d 100644 --- a/tensorflow/compiler/jit/device_compiler_test.cc +++ b/tensorflow/compiler/jit/device_compiler_test.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" -#include +#include #include #include #include @@ -26,7 +26,10 @@ limitations under the License. #include #include #include "absl/algorithm/container.h" +#include "absl/base/casts.h" #include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/notification.h" @@ -36,32 +39,37 @@ limitations under the License. #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/math_ops.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" +#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler_client.h" +#include "tensorflow/compiler/jit/device_executable_persistor.h" #include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h" #include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.pb.h" +#include "tensorflow/compiler/jit/xla_activity_listener.h" #include "tensorflow/compiler/jit/xla_compile_util.h" #include "tensorflow/compiler/jit/xla_device_compiler_client.h" #include "tensorflow/compiler/tf2xla/xla_argument.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/client/client_library.h" #include "xla/client/local_client.h" -#include "xla/hlo/builder/xla_computation.h" #include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/strings/proto_serialization.h" +#include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/fake_input.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/resource_base.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/ops_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/errors.h" -#include "tensorflow/core/platform/status.h" -#include "tensorflow/core/platform/status_matchers.h" -#include "tensorflow/core/platform/statusor.h" +#include "tensorflow/core/platform/file_system.h" +#include "tensorflow/core/platform/path.h" +#include "tensorflow/core/platform/refcount.h" #include "tensorflow/core/platform/test.h" #include "tsl/platform/path.h" @@ -186,7 +194,6 @@ class DeviceCompilerTest : public ::testing::Test { XlaCompiler::Options GetDefaultXlaOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_GPU_XLA_JIT); - options.client = xla_device_compiler_->client(); options.flib_def = flib_def_.get(); return options; } @@ -621,7 +628,6 @@ TEST_F(OpsTestBase, CompileSingleOpSuccess) { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_GPU_XLA_JIT); - options.client = GetLocalClient(); options.flib_def = flib_def.get(); std::vector args(1); diff --git a/tensorflow/compiler/jit/device_context_test.cc b/tensorflow/compiler/jit/device_context_test.cc index 33bba30f3db3e1..0b3701db0132bb 100644 --- a/tensorflow/compiler/jit/device_context_test.cc +++ b/tensorflow/compiler/jit/device_context_test.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include @@ -39,10 +38,6 @@ static bool Initialized = [] { class DeviceContextTest : public ::testing::Test { public: void SetDevice(const std::string& device_type) { - auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; - rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type)); - rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type)); - auto device_factory = DeviceFactory::GetFactory(device_type); SessionOptions options; std::vector> devices; diff --git a/tensorflow/compiler/jit/device_executable_persistor_test.cc b/tensorflow/compiler/jit/device_executable_persistor_test.cc index c33dcb9538ffd6..576dc2462894fe 100644 --- a/tensorflow/compiler/jit/device_executable_persistor_test.cc +++ b/tensorflow/compiler/jit/device_executable_persistor_test.cc @@ -122,7 +122,6 @@ class DeviceExecutionPersistorTest : public ::testing::Test { XlaCompiler::Options DefaultXlaOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - options.client = xla_compiler_client_->client(); options.flib_def = flib_def_.get(); return options; } @@ -130,7 +129,9 @@ class DeviceExecutionPersistorTest : public ::testing::Test { XlaCompiler::Options DefaultPjRtOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - options.client = nullptr; + if (pjrt_compiler_client_ != nullptr) { + options.client = pjrt_compiler_client_->client(); + } options.flib_def = flib_def_.get(); return options; } diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index a0a0d45736f1e8..a65f1ae7c7edf8 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -15,18 +15,22 @@ limitations under the License. #include "tensorflow/compiler/jit/flags.h" +#include +#include #include -#include // NOLINT #include +#include #include #include "absl/base/call_once.h" +#include "absl/log/log.h" #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" #include "absl/strings/strip.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dump_graph.h" #include "xla/parse_flags_from_env.h" -#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/tpu/kernels/sparse_core_xla_flags_defaults.h" #include "tensorflow/core/util/command_line_flags.h" @@ -270,6 +274,12 @@ void AllocateAndParseFlags() { ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_ = true; ops_flags->tf_xla_use_device_api.enabled_for_all_ = false; ops_flags->tf_xla_use_device_api.enabled_for_gpu_ = true; + ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaLaunch( + DeviceType(DEVICE_GPU)); + ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaCompileOnDemand( + DeviceType(DEVICE_GPU)); + ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaCompileAndRun( + DeviceType(DEVICE_GPU)); call_module_flags = new XlaCallModuleFlags; // The `enable_mlir_bridge` flag allows the user to explicitly request that @@ -414,6 +424,11 @@ void AllocateAndParseFlags() { AppendMarkForCompilationPassFlagsInternal(flag_list); xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list); + if (!ops_flags->tf_xla_use_device_api.enabled_for_gpu_) { + LOG(WARNING) + << "tf_xla_enable_device_api_for_gpu is ignored; assuming true"; + } + mlir_flags = new MlirCommonFlags; if (!enable_mlir_bridge_is_explicit) { mlir_flags->tf_mlir_enable_mlir_bridge = diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index 96154b892ae5b0..a2a1b810444dc4 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -22,9 +22,7 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" -#include "absl/types/optional.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/platform/types.h" #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/util/command_line_flags.h" @@ -90,7 +88,7 @@ struct MarkForCompilationPassFlags { bool tf_xla_disable_resource_variable_safety_checks_for_debugging; // If true names of clustered operations will be computed deterministically - // so that they remain stable from run to run of auto clusteing. + // so that they remain stable from run to run of auto clustering. bool tf_xla_deterministic_cluster_names; // If non-empty, JIT-compiled executables are saved to and loaded from the @@ -109,7 +107,7 @@ struct MarkForCompilationPassFlags { // for debugging. Defaults to false. bool tf_xla_disable_strict_signature_checks; - // Specifies the persistance cache prefix. Default is "xla_compile_cache" + // Specifies the persistence cache prefix. Default is "xla_compile_cache" std::string tf_xla_persistent_cache_prefix; }; @@ -163,7 +161,9 @@ struct XlaOpsCommonFlags { } bool IsEnabledInXlaLaunchForDevice(const DeviceType& device_type) const { - if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; + if (device_type.type_string() == "GPU") { + return true; + } return enabled_for_all_ || (enabled_for_xla_launch_ && xla_launch_allowed_devices_.contains(device_type.type_string())); @@ -180,7 +180,9 @@ struct XlaOpsCommonFlags { bool IsEnabledInXlaCompileOnDemandForDevice( const DeviceType& device_type) const { - if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; + if (device_type.type_string() == "GPU") { + return true; + } return enabled_for_all_ || (enabled_for_compile_on_demand_ && xla_compile_on_demand_allowed_devices_.contains( @@ -198,14 +200,14 @@ struct XlaOpsCommonFlags { bool IsEnabledInXlaCompileAndRunForDevice( const DeviceType& device_type) const { - if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; + if (device_type.type_string() == "GPU") { + return true; + } return enabled_for_all_ || (enabled_for_compile_and_run_ && xla_compile_and_run_allowed_devices_.contains( device_type.type_string())); } - bool IsEnabledForGpu() const { return enabled_for_gpu_; } - // If true, uses Device API (PjRt) for single device compilation and // execution of functions marked for JIT compilation i.e. jit_compile=True. // Defaults to false. @@ -277,7 +279,7 @@ struct BuildXlaOpsPassFlags { // TensorCore sequencing will be used. bool tf_xla_disable_full_embedding_pipelining; - // Whether to enable automatical embedding pipelining when summary ops are + // Whether to enable automatic embedding pipelining when summary ops are // detected in the graph. bool tf_xla_disable_full_embedding_pipelining_with_summaries; diff --git a/tensorflow/compiler/jit/get_compiler_ir.cc b/tensorflow/compiler/jit/get_compiler_ir.cc index e50869f2870e5c..88635fb9aa122f 100644 --- a/tensorflow/compiler/jit/get_compiler_ir.cc +++ b/tensorflow/compiler/jit/get_compiler_ir.cc @@ -19,31 +19,40 @@ limitations under the License. #include #include #include -#include #include #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "absl/strings/ascii.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tensorflow/compiler/jit/compilability_check_util.h" +#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" +#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" #include "tensorflow/compiler/jit/xla_compiler_options_util.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/executable_build_options.h" -#include "xla/client/local_client.h" +#include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/translate/portable_api.h" +#include "xla/pjrt/compiled_memory_stats.h" +#include "xla/pjrt/pjrt_client.h" #include "xla/service/hlo_graph_dumper.h" +#include "xla/service/hlo_module_config.h" +#include "xla/service/hlo_proto_util.h" +#include "xla/shape.h" #include "xla/status_macros.h" #include "xla/stream_executor/host/host_platform_id.h" #include "xla/stream_executor/platform.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/common_runtime/eager/tensor_handle.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" @@ -54,49 +63,27 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/refcount.h" -#include "tensorflow/core/platform/statusor.h" -#include "tsl/platform/errors.h" -#include "tsl/platform/statusor.h" +#include "tensorflow/core/tpu/tpu_defs.h" namespace tensorflow { -static absl::StatusOr> BuildExecutable( - xla::LocalClient* local_client, - const XlaCompiler::CompilationResult& result, - const XlaCompiler::Options& options, - const bool xla_embed_ir_in_executable = false) { - std::vector argument_layouts( - result.xla_input_shapes.size()); - for (int i = 0, end = result.xla_input_shapes.size(); i < end; ++i) { - argument_layouts[i] = &result.xla_input_shapes[i]; - } - xla::ExecutableBuildOptions build_options; - if (result.collective_info) { - build_options.set_num_replicas(result.collective_info->group_size); - } - build_options.set_device_ordinal( - options.device_ordinal != -1 ? options.device_ordinal - : local_client->default_device_ordinal()); - build_options.set_result_layout(result.xla_output_shape); - build_options.set_device_allocator(options.device_allocator.get()); - build_options.set_alias_passthrough_params(options.alias_passthrough_params); - build_options.mutable_debug_options()->set_xla_detailed_logging( - options.detailed_logging); - // If the embed_ir_in_executable is set, hlo_proto will be dumped in - // executable. The hlo_proto contains HLO modules and buffer assignment. - build_options.mutable_debug_options()->set_xla_embed_ir_in_executable( - xla_embed_ir_in_executable); +static absl::StatusOr> +BuildExecutable(xla::PjRtClient* pjrt_client, + const XlaCompiler::CompilationResult& result, + const XlaCompiler::Options& options, + const bool xla_embed_ir_in_executable = false) { + xla::CompileOptions compile_options = GetPjRtCompileOptions(options, result); + compile_options.executable_build_options.mutable_debug_options() + ->set_xla_embed_ir_in_executable(xla_embed_ir_in_executable); TF_ASSIGN_OR_RETURN( - std::vector> executables, - local_client->Compile(*result.computation, argument_layouts, - build_options)); - TF_RET_CHECK(executables.size() == 1); - return std::move(executables[0]); + std::unique_ptr executable, + pjrt_client->CompileAndLoad(*result.computation, compile_options)); + return executable; } static absl::StatusOr BuildHLOString( IrExportStage stage, const XlaCompiler::CompilationResult& result, - xla::LocalClient* local_client, const XlaCompiler::Options& options) { + xla::PjRtClient* pjrt_client, const XlaCompiler::Options& options) { switch (stage) { case IrExportStage::STABLEHLO: case IrExportStage::STABLEHLO_SERIALIZED: @@ -135,33 +122,41 @@ static absl::StatusOr BuildHLOString( return new_module->ToString(opts); } case IrExportStage::OPTIMIZED_HLO: - case IrExportStage::OPTIMIZED_HLO_SERIALIZED: { - TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - BuildExecutable(local_client, result, options)); - xla::Executable* new_executable = executable->executable(); + case IrExportStage::OPTIMIZED_HLO_SERIALIZED: + case IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED: + case IrExportStage::OPTIMIZED_HLO_DOT: { + TF_RET_CHECK(pjrt_client != nullptr); + bool xla_embed_ir_in_executable = + (stage == IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED); + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + BuildExecutable(pjrt_client, result, options, + xla_embed_ir_in_executable)); + TF_ASSIGN_OR_RETURN(std::vector> modules, + executable->GetHloModules()); + TF_RET_CHECK(!modules.empty()); if (stage == IrExportStage::OPTIMIZED_HLO_SERIALIZED) { - return new_executable->module().ToProto().SerializeAsString(); + return modules[0]->ToProto().SerializeAsString(); + } else if (stage == IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED) { + xla::HloProto hlo_proto = xla::MakeHloProto(*modules[0]); + absl::StatusOr memory_stats = + executable->GetCompiledMemoryStats(); + if (memory_stats.ok() && + !memory_stats->serialized_buffer_assignment.empty()) { + hlo_proto.mutable_buffer_assignment()->ParseFromString( + memory_stats->serialized_buffer_assignment); + } + return hlo_proto.SerializeAsString(); + } else if (stage == IrExportStage::OPTIMIZED_HLO_DOT) { + absl::StatusOr graph = xla::RenderGraph( + *modules[0]->entry_computation(), "Visualization", + /*debug_options=*/{}, xla::RenderedGraphFormat::kDot, + /*hlo_render_options=*/{}); + TF_RETURN_IF_ERROR(graph.status()); + return *graph; } else { - return new_executable->module().ToString(); + return modules[0]->ToString(); } } - case IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED: { - TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - BuildExecutable(local_client, result, options, - /*xla_embed_ir_in_executable=*/true)); - return executable->executable()->hlo_proto()->SerializeAsString(); - } - case IrExportStage::OPTIMIZED_HLO_DOT: { - TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - BuildExecutable(local_client, result, options)); - absl::StatusOr graph = xla::RenderGraph( - *executable->executable()->module().entry_computation(), - "Visualization", - /*debug_options=*/{}, xla::RenderedGraphFormat::kDot, - /*hlo_render_options=*/{}); - TF_RETURN_IF_ERROR(graph.status()); - return *graph; - } } } @@ -318,7 +313,7 @@ absl::StatusOr> PrepareXlaCompilerArgs( absl::StatusOr CompileAndBuildHLOString( IrExportStage stage, const XlaCompiler::Options& options, - xla::LocalClient* local_client, const NameAttrList& function, + xla::PjRtClient* pjrt_client, const NameAttrList& function, const std::vector& args) { XlaCompiler::CompileOptions compile_options; compile_options.always_return_tuple = false; @@ -329,7 +324,7 @@ absl::StatusOr CompileAndBuildHLOString( TF_RETURN_IF_ERROR( compiler.CompileFunction(compile_options, function, args, &result)); - return BuildHLOString(stage, result, local_client, options); + return BuildHLOString(stage, result, pjrt_client, options); } /** @@ -352,8 +347,8 @@ absl::StatusOr GetCompilerIr( absl::Span input_arg_shape_and_dtype, absl::Span input_handles, CompilerArgSource compiler_arg_source) { - using XlaDeviceCompiler = - DeviceCompiler; + using PjRtDeviceCompiler = + DeviceCompiler; se::Stream* stream = nullptr; if (const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info = @@ -375,33 +370,32 @@ absl::StatusOr GetCompilerIr( compiler_arg_source)); XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(dev); - auto compilation_device_type = platform_info.device_type(); - if (platform_info.device_type() != DEVICE_TPU) { + DeviceType compilation_device_type = platform_info.device_type(); + if (platform_info.device_type() == DEVICE_TPU) { + compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); + } else { TF_ASSIGN_OR_RETURN(compilation_device_type, GetCompilationDeviceType(platform_info.device_type())); } - XlaDeviceCompiler* xla_device_compiler; - TF_RETURN_IF_ERROR(dev->resource_manager()->LookupOrCreate( - dev->resource_manager()->default_container(), "xla_device_compiler", - &xla_device_compiler, [&](XlaDeviceCompiler** xla_device_compiler) { - return BuildXlaDeviceCompiler(dev, flr, platform_info, - compilation_device_type, - xla_device_compiler); - })); - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + PjRtDeviceCompiler* pjrt_device_compiler = nullptr; + DeviceCompilationProfiler* profiler = nullptr; + TF_RETURN_IF_ERROR(GetOrCreatePjRtDeviceCompilerAndProfiler( + platform_info, dev->resource_manager(), flr, &pjrt_device_compiler, + &profiler)); + core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); + core::ScopedUnref profiler_ref(profiler); XlaCompiler::Options options; if (platform_info.device_type() == DEVICE_TPU && stream == nullptr) { - options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr); + options = GenerateCompilerOptionsForTfrtTpu(compilation_device_type, *flr); } else { - options = GenerateCompilerOptions(*xla_device_compiler, *flr, dev, stream, - platform_info, - /*has_ref_vars=*/false); + options = GenerateCompilerOptionsForPjRt(*flr, dev, platform_info, + pjrt_device_compiler); } - return CompileAndBuildHLOString(stage, options, xla_device_compiler->client(), - function, args); + return CompileAndBuildHLOString( + stage, options, pjrt_device_compiler->client(), function, args); } absl::StatusOr GetCompilerIr( @@ -411,9 +405,6 @@ absl::StatusOr GetCompilerIr( absl::Span input_arg_shape_and_dtype, absl::Span input_handles, CompilerArgSource compiler_arg_source) { - using XlaDeviceCompiler = - DeviceCompiler; - TF_RETURN_IF_ERROR( ValidateGetCompilerIrTfrtTpu(platform_name, /*stream=*/nullptr, stage)); @@ -445,27 +436,16 @@ absl::StatusOr GetCompilerIr( /*pjrt_device_metadata=*/nullptr, /*device_allocator=*/nullptr); DeviceType compilation_device_type = platform_info.device_type(); - if (platform_info.device_type() != DEVICE_TPU) { + if (platform_info.device_type() == DEVICE_TPU) { + compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); + } else { TF_ASSIGN_OR_RETURN(compilation_device_type, GetCompilationDeviceType(platform_info.device_type())); } - XlaDeviceCompiler* xla_device_compiler; - const std::string xla_device_compiler_name = absl::StrCat( - absl::AsciiStrToLower(platform_name), "_xla_device_compiler"); - TF_RETURN_IF_ERROR( - context->HostCPU()->resource_manager()->LookupOrCreate( - context->HostCPU()->resource_manager()->default_container(), - xla_device_compiler_name, &xla_device_compiler, - [&](XlaDeviceCompiler** xla_device_compiler) { - return BuildXlaDeviceCompiler(/*dev=*/nullptr, flr, platform_info, - compilation_device_type, - xla_device_compiler); - })); - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); XlaCompiler::Options options; if (platform_info.device_type() == DEVICE_TPU) { - options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr); + options = GenerateCompilerOptionsForTfrtTpu(compilation_device_type, *flr); } else { options.device_type = compilation_device_type; options.flib_def = flr->GetFunctionLibraryDefinition(); @@ -475,7 +455,7 @@ absl::StatusOr GetCompilerIr( options.alias_passthrough_params = !platform_info.is_on_xla_device(); } - return CompileAndBuildHLOString(stage, options, /*local_client=*/nullptr, + return CompileAndBuildHLOString(stage, options, /*pjrt_client=*/nullptr, function, args); } diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 322daf4594c42b..4b8e481d718b2b 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include #include @@ -38,7 +37,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" -#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/jit/flags.h" @@ -48,49 +46,32 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_activity.pb.h" #include "tensorflow/compiler/jit/xla_activity_listener.h" #include "tensorflow/compiler/jit/xla_compile_util.h" -#include "tensorflow/compiler/jit/xla_compiler_options_util.h" -#include "tensorflow/compiler/jit/xla_host_recv_device_context.h" -#include "tensorflow/compiler/jit/xla_host_send_device_context.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" -#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "xla/client/local_client.h" -#include "xla/executable_run_options.h" -#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/service/executable.h" -#include "xla/service/gpu/gpu_executable_run_options.h" -#include "xla/stream_executor/host/host_platform_id.h" -#include "xla/tsl/concurrency/async_value_ref.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/macros.h" +#include "xla/tsl/platform/statusor.h" #include "xla/tsl/protobuf/error_codes.pb.h" #include "tensorflow/core/framework/allocator.h" -#include "tensorflow/core/framework/control_flow.h" #include "tensorflow/core/framework/device.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" -#include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/monitoring/counter.h" #include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/refcount.h" -#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/status.h" #include "tensorflow/core/platform/threadpool.h" #include "tensorflow/core/platform/tstring.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/device_name_utils.h" -#include "tensorflow/core/util/stream_executor_util.h" #include "tsl/platform/thread_annotations.h" -#include "tsl/profiler/lib/traceme.h" // OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that // in error case, it returns RET instead of void. @@ -106,8 +87,6 @@ limitations under the License. namespace tensorflow { namespace { -using XlaDeviceCompiler = - DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; @@ -201,172 +180,11 @@ class ExecutableClosureStore { void operator=(const ExecutableClosureStore&) = delete; }; -using XlaExecutableClosure = - ExecutableClosure; -using XlaExecutableClosureStore = - ExecutableClosureStore; using PjRtExecutableClosure = ExecutableClosure; using PjRtExecutableClosureStore = ExecutableClosureStore; -se::Stream* GetStream(OpKernelContext* ctx) { - return ctx->op_device_context() ? ctx->op_device_context()->stream() - : nullptr; -} - -XlaComputationLaunchContext GetLaunchContext( - const XlaPlatformInfo& platform_info, OpKernelContext* ctx, - xla::LocalClient* client, - stream_executor::DeviceAddressAllocator* allocator) { - se::Stream* stream = GetStream(ctx); - int device_ordinal = stream ? stream->parent()->device_ordinal() - : client->default_device_ordinal(); - XlaComputationLaunchContext launch_context( - client, allocator, device_ordinal, - /*allocate_xla_tensors=*/platform_info.is_on_xla_device(), - /*use_multiple_streams=*/platform_info.UseMultipleStreams()); - return launch_context; -} - -absl::Status GetTaskName(const absl::string_view device_name, - std::string* task_name) { - std::string ignored; - if (!DeviceNameUtils::SplitDeviceName(device_name, task_name, &ignored)) { - return absl::InvalidArgumentError( - absl::StrCat("Unable to parse device name: ", device_name)); - } - - return absl::OkStatus(); -} - -// Provide SendDeviceMemoryFunction for XLA host callbacks. This callback -// handles transferring from device to host. -xla::SendDeviceMemoryFunction GetSendDeviceMemoryFunction( - OpKernelContext* ctx, const std::string& program_key) { - return - [ctx, program_key]( - int64_t channel_id, se::Stream* stream, const xla::Shape& shape, - const stream_executor::DeviceAddressBase& device_memory_base, - const absl::flat_hash_map& frontend_attrs) - -> absl::StatusOr>> { - auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous"); - - // Generate the Rendezvous key. - const std::string& rendezvous_key_base = - absl::StrCat(program_key, iter->second); - - const std::string& src_device = ctx->device()->name(); - std::string task_prefix; - TF_RETURN_IF_ERROR(GetTaskName(src_device, &task_prefix)); - const std::string dst_device = - absl::StrCat(task_prefix, "/device:CPU:0"); - const std::string& rendezvous_key = - Rendezvous::CreateKey(src_device, /*src_incarnation=*/1, dst_device, - rendezvous_key_base, FrameAndIter(0, 0)); - VLOG(2) << "Rendezvous Key for receiving at host: " << rendezvous_key; - - RendezvousInterface::ParsedKey parsed_key; - TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key)); - - TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent()); - tsl::AsyncValueRef> done_event = - tsl::MakeConstructedAsyncValueRef>( - std::move(event)); - - Rendezvous::Args args; - // Rendezvous::Args owns the device context pointer. - args.device_context = new XlaHostRecvDeviceContext( - stream, device_memory_base, shape, done_event); - - Tensor host_tensor; - TF_RETURN_IF_ERROR( - ctx->rendezvous()->Send(parsed_key, args, host_tensor, false)); - - return std::move(done_event); - }; -} - -// Provide RecvDeviceMemoryFunction for XLA host callbacks. This callback -// handles transferring from host to device. -xla::RecvDeviceMemoryFunction GetRecvDeviceMemoryFunction( - OpKernelContext* ctx, const std::string& program_key) { - return - [ctx, program_key]( - int64_t channel_id, se::Stream* stream, const xla::Shape& shape, - stream_executor::DeviceAddressBase* device_memory_base, - const absl::flat_hash_map& frontend_attrs) - -> absl::StatusOr>> { - auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous"); - - // Generate the Rendezvous key. - const std::string& rendezvous_key_base = - absl::StrCat(program_key, iter->second); - - const std::string& dst_device = ctx->device()->name(); - std::string task_prefix; - TF_RETURN_IF_ERROR(GetTaskName(dst_device, &task_prefix)); - const std::string src_device = - absl::StrCat(task_prefix, "/device:CPU:0"); - const std::string& rendezvous_key = - Rendezvous::CreateKey(src_device, /*src_incarnation=*/1, dst_device, - rendezvous_key_base, FrameAndIter(0, 0)); - VLOG(2) << "Rendezvous Key for sending from host: " << rendezvous_key; - - RendezvousInterface::ParsedKey parsed_key; - TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key)); - - TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent()); - tsl::AsyncValueRef> done_event = - tsl::MakeConstructedAsyncValueRef>( - std::move(event)); - - Rendezvous::Args args; - // Rendezvous::Args owns the device context pointer. - args.device_context = new XlaHostSendDeviceContext( - stream, device_memory_base, shape, done_event); - - Tensor device_tensor; - bool is_dead; - TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv( - parsed_key, args, &device_tensor, /*is_dead=*/&is_dead)); - - return std::move(done_event); - }; -} - -absl::StatusOr RunExecutable( - const XlaPlatformInfo& platform_info, - const XlaComputationLaunchContext& launch_context, - std::vector execution_inputs, - xla::ExecutableRunOptions run_options, xla::LocalExecutable* executable, - OpKernelContext* ctx, stream_executor::DeviceAddressAllocator* allocator) { - VLOG(2) << "Executing Xla Computation."; - Env* env = Env::Default(); - auto start_time = env->NowMicros(); - - se::Stream* stream = GetStream(ctx); - run_options.set_stream(GetStream(ctx)); - run_options.set_allocator(allocator); - run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); - run_options.set_rng_seed(GetXLARandomSeed()); - - absl::StatusOr execution_output; - bool run_synchronous = - !stream || platform_info.platform_id() == se::host::kHostPlatformId; - if (run_synchronous) { - execution_output = - executable->Run(std::move(execution_inputs), run_options); - } else { - execution_output = - executable->RunAsync(std::move(execution_inputs), run_options); - } - - auto elapsed = env->NowMicros() - start_time; - VLOG(2) << "Elapsed time for Xla Executable Run: " << elapsed << "us"; - return execution_output; -} - absl::StatusOr< std::pair, ResourceVarsSnapshot>> GetXlaCompilerArgsAndSnapshotVariables( @@ -391,59 +209,6 @@ GetXlaCompilerArgsAndSnapshotVariables( return result; } -absl::Status CompileToLocalExecutable( - OpKernelContext* ctx, const NameAttrList& function, bool has_ref_vars, - const XlaPlatformInfo& platform_info, - const std::vector& args, - DeviceCompileMode compile_mode, bool may_alias_resource_update, - xla::LocalClient** client, - const XlaCompiler::CompilationResult** compilation_result, - xla::LocalExecutable** executable) { - // We store information about the JIT-compiled XLA computation - // in the ResourceMgr. - ResourceMgr* rm = ctx->resource_manager(); - if (!rm) { - return absl::InternalError("No resource manager."); - } - - TF_ASSIGN_OR_RETURN(DeviceType compilation_device_type, - GetCompilationDeviceType(platform_info.device_type())); - - XlaDeviceCompiler* xla_device_compiler; - TF_RETURN_IF_ERROR(rm->LookupOrCreate( - rm->default_container(), "xla_device_compiler", &xla_device_compiler, - [&](XlaDeviceCompiler** xla_device_compiler) { - return BuildXlaDeviceCompiler(ctx->device(), ctx->function_library(), - platform_info, compilation_device_type, - xla_device_compiler); - })); - DeviceCompilationProfiler* profiler; - TF_RETURN_IF_ERROR(rm->LookupOrCreate( - rm->default_container(), "device_compilation_profiler", &profiler, - [](DeviceCompilationProfiler** profiler) { - *profiler = new DeviceCompilationProfiler(); - return absl::OkStatus(); - })); - // Hold the reference to the XLA device compiler and profiler during - // evaluation. (We could probably free them sooner because the ResourceMgr - // will retain references, but this is more obviously correct.) - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); - core::ScopedUnref profiler_ref(profiler); - - *client = static_cast(xla_device_compiler->client()); - - XlaCompiler::Options options = GenerateCompilerOptions( - *xla_device_compiler, *ctx->function_library(), ctx->device(), - GetStream(ctx), platform_info, has_ref_vars); - - XlaCompiler::CompileOptions compile_options = - GenerateCompileOptions(has_ref_vars, may_alias_resource_update); - - return xla_device_compiler->CompileIfNeeded( - options, function, args, compile_options, compile_mode, profiler, - compilation_result, executable); -} - absl::Status GetUpdatedVariables( const OpKernelContext* ctx, absl::Span inputs, absl::Span variable_indices, @@ -517,9 +282,6 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { std::vector xla_compiler_args; const XlaCompiler::CompilationResult* compilation_result; - xla::LocalClient* client; // Not owned. - xla::LocalExecutable* executable; // Not owned. - xla::PjRtClient* pjrt_client; // Not owned. xla::PjRtLoadedExecutable* pjrt_executable; // Not owned. @@ -548,66 +310,26 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { xla_compiler_args = std::move(status_or_xla_compiler_args.value()); } - bool use_pjrt = GetXlaOpsCommonFlags() - ->tf_xla_use_device_api.IsEnabledInXlaLaunchForDevice( - platform_info_.device_type()); - if (use_pjrt) { - VLOG(2) << "Compiling using PJRT"; - absl::Status status = CompileToPjRtLoadedExecutable( - *ctx, platform_info_, function_, xla_compiler_args, - DeviceCompileMode::kStrict, has_ref_vars_, - /*may_alias_resource_update=*/true, &compilation_result, &pjrt_client, - &pjrt_executable); - OP_REQUIRES_OK_ASYNC(ctx, status, done); - - VLOG(2) << "Compiled using PJRT: " << status; - VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); - VLOG(2) << "compilation_result != nullptr: " - << (compilation_result != nullptr); - VLOG(2) << "Executing using PJRT."; - - auto run_pjrt_cluster = [ctx, pjrt_client, pjrt_executable, - compilation_result, done, inputs, - resources = resources_]() { - // Separate scope so that VariableInfo locks are released before done() is - // called. - { - std::vector variable_infos; - OP_REQUIRES_OK_ASYNC( - ctx, - GetUpdatedVariables(ctx, inputs, resources, *compilation_result, - &variable_infos), - done); - OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)), - done); - OP_REQUIRES_OK_ASYNC( - ctx, - RunPjRtExecutable(inputs, variable_infos, *compilation_result, - pjrt_client, pjrt_executable, ctx), - done); - } - VLOG(2) << "Done executing with PJRT."; - done(); - }; - - RunInThreadPoolIfCollectivesPresent(*compilation_result, run_pjrt_cluster); - return; - } - - absl::Status status = CompileToLocalExecutable( - ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_, - xla_compiler_args, DeviceCompileMode::kStrict, - /*may_alias_resource_update=*/true, &client, &compilation_result, - &executable); + VLOG(2) << "Compiling using PJRT"; + absl::Status status = CompileToPjRtLoadedExecutable( + *ctx, platform_info_, function_, xla_compiler_args, + DeviceCompileMode::kStrict, has_ref_vars_, + /*may_alias_resource_update=*/true, &compilation_result, &pjrt_client, + &pjrt_executable); OP_REQUIRES_OK_ASYNC(ctx, status, done); - // Continuation of the execution, may be run in a different thread. - auto run_xla_cluster = [ctx, client, executable, compilation_result, done, - inputs, resources = resources_]() { - // Separate scope so that VariableInfo locks are released before done is + VLOG(2) << "Compiled using PJRT: " << status; + VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); + VLOG(2) << "compilation_result != nullptr: " + << (compilation_result != nullptr); + VLOG(2) << "Executing using PJRT."; + + auto run_pjrt_cluster = [ctx, pjrt_client, pjrt_executable, + compilation_result, done, inputs, + resources = resources_]() { + // Separate scope so that VariableInfo locks are released before done() is // called. { - auto platform_info = XlaPlatformInfoFromDevice(ctx->device()); std::vector variable_infos; OP_REQUIRES_OK_ASYNC( ctx, @@ -616,61 +338,17 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { done); OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)), done); - absl::flat_hash_map resource_var_ptrs; - for (int i = 0; i < resources.size(); i++) { - resource_var_ptrs[resources[i]] = variable_infos[i].var()->tensor(); - } - - std::shared_ptr allocator = - GetAllocator(ctx->device(), GetStream(ctx), platform_info); - XlaComputationLaunchContext launch_context = - GetLaunchContext(platform_info, ctx, client, allocator.get()); - - const xla::HloInputOutputAliasConfig& input_output_alias = - executable->executable()->module().input_output_alias_config(); - absl::StatusOr> execution_inputs = - launch_context.PopulateInputs( - ctx, compilation_result, resource_var_ptrs, - /*missing_ctx_input_prefix=*/0, input_output_alias); - OP_REQUIRES_OK_ASYNC(ctx, execution_inputs.status(), done); - - xla::gpu::GpuExecutableRunOptions gpu_options; - xla::DeviceAssignment device_assignment; - xla::ExecutableRunOptions run_options; - if (compilation_result->collective_info.has_value()) { - OP_REQUIRES_OK_ASYNC(ctx, - ResolveDeviceAssignment( - ctx, *compilation_result->collective_info, - run_options, device_assignment, gpu_options), - done); - } - - // Hardcode run id to always be zero: TF distributed strategy - // differentiates between subsequent runs using dependency edges. This - // is safe, as only TF dist-strat can produce distributed ops, and we - // can rely on TF dist-strat invariants. - xla::RunId run_id(0); - run_options.set_run_id(run_id); - - absl::StatusOr execution_output = RunExecutable( - platform_info, launch_context, std::move(*execution_inputs), - run_options, executable, ctx, allocator.get()); - OP_REQUIRES_ASYNC(ctx, execution_output.ok(), execution_output.status(), - done); - OP_REQUIRES_OK_ASYNC( ctx, - launch_context.PopulateOutputs( - ctx, compilation_result, execution_output->ConsumeResult(), - /*missing_ctx_input_prefix=*/0, absl::MakeSpan(variable_infos), - input_output_alias, resource_var_ptrs), + RunPjRtExecutable(inputs, variable_infos, *compilation_result, + pjrt_client, pjrt_executable, ctx), done); - VLOG(1) << "Done"; } + VLOG(2) << "Done executing with PJRT."; done(); }; - RunInThreadPoolIfCollectivesPresent(*compilation_result, run_xla_cluster); + RunInThreadPoolIfCollectivesPresent(*compilation_result, run_pjrt_cluster); } namespace { @@ -762,8 +440,6 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaCompileOp " << def().name() << (must_compile_ ? "(must-compile)" : ""); const XlaCompiler::CompilationResult* kernel = nullptr; - xla::LocalClient* client = nullptr; - xla::LocalExecutable* executable = nullptr; xla::PjRtClient* pjrt_client = nullptr; xla::PjRtLoadedExecutable* pjrt_executable = nullptr; ResourceVarsSnapshot variables_snapshot; @@ -783,14 +459,9 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { : DeviceCompileMode::kLazy; }(); - bool use_pjrt = - GetXlaOpsCommonFlags() - ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( - platform_info_.device_type()); - if (GetXlaOpsCommonFlags()->tf_xla_always_defer_compilation || cannot_compile_cluster) { - executable = nullptr; + pjrt_executable = nullptr; } else { auto args_and_variables_snapshot = GetXlaCompilerArgsAndSnapshotVariables( resources_, constants_, inputs, ctx); @@ -801,19 +472,12 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { // Do not alias resource updates as locking variables in XlaCompile and // unlocking them in XlaRun may lead to deadlocks. - absl::Status status; - if (use_pjrt) { - VLOG(2) << "Using PJRT for compilation. Function name: " - << function_.name(); - status = CompileToPjRtLoadedExecutable( - *ctx, platform_info_, function_, args, compile_mode, has_ref_vars_, - /*may_alias_resource_update=*/false, &kernel, &pjrt_client, - &pjrt_executable); - } else { - status = CompileToLocalExecutable( - ctx, function_, has_ref_vars_, platform_info_, args, compile_mode, - /*may_alias_resource_update=*/false, &client, &kernel, &executable); - } + VLOG(2) << "Using PJRT for compilation. Function name: " + << function_.name(); + absl::Status status = CompileToPjRtLoadedExecutable( + *ctx, platform_info_, function_, args, compile_mode, has_ref_vars_, + /*may_alias_resource_update=*/false, &kernel, &pjrt_client, + &pjrt_executable); if (compile_mode != DeviceCompileMode::kLazy || status.code() != error::UNIMPLEMENTED) { OP_REQUIRES_OK(ctx, status); @@ -826,7 +490,6 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { BroadcastOptimizationRemark( XlaOptimizationRemark::UNIMPLEMENTED_OPERATION, status.ToString()) .IgnoreError(); - executable = nullptr; pjrt_executable = nullptr; mutex_lock guard(cannot_compile_cluster_mu_); cannot_compile_cluster_ = true; @@ -839,7 +502,7 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { Allocator* cpu_allocator = ctx->device()->GetAllocator(host_alloc_attrs); // Async compilation returns nullptr executable without an error. - if (!executable && !pjrt_executable) { + if (!pjrt_executable) { DCHECK(!must_compile_); Tensor compilation_key(cpu_allocator, DT_STRING, TensorShape({})); Tensor compilation_successful(cpu_allocator, DT_BOOL, TensorShape({})); @@ -854,21 +517,12 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { // hit. This is because we at least need new snapshots of the resource // variables. Tensor compilation_key(cpu_allocator, DT_STRING, TensorShape({})); - if (use_pjrt) { - PjRtExecutableClosureStore::KeyT key = - PjRtExecutableClosureStore::Global()->Produce(PjRtExecutableClosure( - pjrt_client, pjrt_executable, kernel, std::move(variables_snapshot), - constants_.size())); - compilation_key.flat()(0) = key; - VLOG(2) << "Compiled with PJRT. compilation_key: " << key; - } else { - XlaExecutableClosureStore::KeyT key = - XlaExecutableClosureStore::Global()->Produce(XlaExecutableClosure( - client, executable, kernel, std::move(variables_snapshot), - constants_.size())); - compilation_key.flat()(0) = key; - VLOG(2) << "Compiled with XLA. compilation_key: " << key; - } + PjRtExecutableClosureStore::KeyT key = + PjRtExecutableClosureStore::Global()->Produce(PjRtExecutableClosure( + pjrt_client, pjrt_executable, kernel, std::move(variables_snapshot), + constants_.size())); + compilation_key.flat()(0) = key; + VLOG(2) << "Compiled with PJRT. compilation_key: " << key; Tensor compilation_successful(cpu_allocator, DT_BOOL, TensorShape({})); compilation_successful.flat()(0) = true; @@ -877,123 +531,40 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { ctx->set_output(1, compilation_successful); } -XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) - : OpKernel(ctx), platform_info_(XlaPlatformInfoFromDevice(ctx->device())) {} +XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} void XlaRunOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaRunOp " << def().name(); Tensor key_tensor = ctx->input(ctx->num_inputs() - 1); - bool use_pjrt = - GetXlaOpsCommonFlags() - ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( - platform_info_.device_type()); - - if (use_pjrt) { - const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat()(0); - PjRtExecutableClosure closure = - PjRtExecutableClosureStore::Global()->Consume(key); - - // Fetch inputs from the OpKernelContext. Inputs are the same as the ones - // for XlaCompile, except that the must-be-constant inputs that appear in - // the beginning are stripped off and the closure key is appended as the - // last input. So the inputs look like: input tensors, resource variables, - // closure key tensor. - std::vector inputs = InputsFromContext(ctx); - absl::flat_hash_map variable_snapshots; - for (const auto& [variable_index, variable_tensor] : - closure.resource_var_snapshots()) { - variable_snapshots.emplace(variable_index, variable_tensor.has_value() - ? &variable_tensor.value() - : nullptr); - } - - { - absl::StatusOr> updated_variables = - GatherVariableInfo(ctx, *closure.compilation_result(), - closure.num_constant_args()); - OP_REQUIRES_OK(ctx, updated_variables.status()); - OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables))); - OP_REQUIRES_OK( - ctx, RunPjRtExecutable(closure.num_constant_args(), inputs, - variable_snapshots, *updated_variables, - *closure.compilation_result(), - closure.client(), closure.executable(), ctx)); - } + const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat()(0); + PjRtExecutableClosure closure = + PjRtExecutableClosureStore::Global()->Consume(key); - OP_REQUIRES_OK(ctx, absl::OkStatus()); - return; - } - - const XlaExecutableClosureStore::KeyT& key = key_tensor.flat()(0); - - XlaExecutableClosure closure = - XlaExecutableClosureStore::Global()->Consume(key); - std::shared_ptr allocator = - GetAllocator(ctx->device(), GetStream(ctx), platform_info_); - XlaComputationLaunchContext launch_context = - GetLaunchContext(platform_info_, ctx, closure.client(), allocator.get()); - - // We're missing the must-be-constant inputs, tell `PopulateInputs` - // about this. We don't actually need these inputs because they've - // already been baked into the compiled kernel. - const xla::HloInputOutputAliasConfig& input_output_alias = - closure.executable()->executable()->module().input_output_alias_config(); - absl::StatusOr> execution_inputs; - absl::flat_hash_map snapshot_ptrs; - { - tsl::profiler::TraceMe hlo_module_activity( - [&] { - return absl::StrCat( - "Populate Inputs (", - closure.compilation_result()->xla_input_shapes.size(), ")"); - }, - tsl::profiler::TraceMeLevel::kInfo); - - for (const auto& [variable_index, variable_tensor] : - closure.resource_var_snapshots()) { - snapshot_ptrs.emplace(variable_index, variable_tensor.has_value() - ? &variable_tensor.value() - : nullptr); - } - execution_inputs = launch_context.PopulateInputs( - ctx, closure.compilation_result(), snapshot_ptrs, - /*missing_ctx_input_prefix=*/closure.num_constant_args(), - input_output_alias); - OP_REQUIRES_OK(ctx, execution_inputs.status()); + // Fetch inputs from the OpKernelContext. Inputs are the same as the ones + // for XlaCompile, except that the must-be-constant inputs that appear in + // the beginning are stripped off and the closure key is appended as the + // last input. So the inputs look like: input tensors, resource variables, + // closure key tensor. + std::vector inputs = InputsFromContext(ctx); + absl::flat_hash_map variable_snapshots; + for (const auto& [variable_index, variable_tensor] : + closure.resource_var_snapshots()) { + variable_snapshots.emplace(variable_index, variable_tensor.has_value() + ? &variable_tensor.value() + : nullptr); } - xla::ExecutableRunOptions run_options; - - // Host callbacks used for HLO send/recv. - xla::SendDeviceMemoryFunction send_function = - GetSendDeviceMemoryFunction(ctx, key); - run_options.set_send_device_memory_function(&send_function); - xla::RecvDeviceMemoryFunction recv_function = - GetRecvDeviceMemoryFunction(ctx, key); - run_options.set_recv_device_memory_function(&recv_function); - - absl::StatusOr execution_output = RunExecutable( - platform_info_, launch_context, std::move(*execution_inputs), run_options, - closure.executable(), ctx, allocator.get()); - OP_REQUIRES(ctx, execution_output.ok(), execution_output.status()); - - tsl::profiler::TraceMe hlo_module_activity( - [&] { - return absl::StrCat("Populate Outputs (", ctx->num_outputs(), ")"); - }, - tsl::profiler::TraceMeLevel::kInfo); - - absl::StatusOr> variable_infos = GatherVariableInfo( - ctx, *closure.compilation_result(), closure.num_constant_args()); - OP_REQUIRES_OK(ctx, variable_infos.status()); - OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*variable_infos))); + absl::StatusOr> updated_variables = + GatherVariableInfo(ctx, *closure.compilation_result(), + closure.num_constant_args()); + OP_REQUIRES_OK(ctx, updated_variables.status()); + OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables))); OP_REQUIRES_OK( ctx, - launch_context.PopulateOutputs( - ctx, closure.compilation_result(), execution_output->ConsumeResult(), - /*missing_ctx_input_prefix=*/closure.num_constant_args(), - absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs)); + RunPjRtExecutable(closure.num_constant_args(), inputs, variable_snapshots, + *updated_variables, *closure.compilation_result(), + closure.client(), closure.executable(), ctx)); } XlaMergeOp::XlaMergeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} diff --git a/tensorflow/compiler/jit/kernels/xla_ops.h b/tensorflow/compiler/jit/kernels/xla_ops.h index 1e42fe5ff38741..96ca28cdac5f55 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.h +++ b/tensorflow/compiler/jit/kernels/xla_ops.h @@ -16,25 +16,17 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ #define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ -#include #include -#include "tensorflow/compiler/jit/device_compiler.h" -#include "tensorflow/compiler/jit/xla_device.h" -#include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" -#include "xla/stream_executor/integrations/tf_allocator_adapter.h" -#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/attr_value.pb.h" -#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/util/stream_executor_util.h" +#include "tensorflow/core/platform/mutex.h" +#include "tsl/platform/thread_annotations.h" namespace tensorflow { - // XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp. // The only difference is that it does not require arguments to follow // the "constants, then regular args, then resources" order. @@ -76,8 +68,8 @@ class XlaLocalLaunchBase : public AsyncOpKernel { // Once all inputs are present, and their shapes are known, the op can // use a 'DeviceCompiler' to compile and execute code which is specific // to the shapes of input Tensors. -// XlaLocalLaunchOp uses xla::LocalClient::Compile() and -// xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device +// XlaLocalLaunchOp uses xla::PjRtClient and +// xla::PjRtLoadedExecutable, and passes arguments into/out of XLA in device // memory. class XlaLocalLaunchOp : public XlaLocalLaunchBase { public: @@ -125,9 +117,6 @@ class XlaRunOp : public OpKernel { explicit XlaRunOp(OpKernelConstruction* ctx); void Compute(OpKernelContext* ctx) override; - - private: - const XlaPlatformInfo platform_info_; }; class XlaMergeOp : public OpKernel { diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc index f1a648f15592cf..c69b9ce20350f8 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc @@ -17,68 +17,40 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h" -#include -#include -#include -#include #include -#include "absl/container/flat_hash_map.h" #include "absl/log/check.h" #include "absl/log/log.h" -#include "absl/memory/memory.h" +#include "absl/status/status.h" #include "absl/types/span.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" #include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" -#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" -#include "tensorflow/compiler/jit/xla_compile_util.h" #include "tensorflow/compiler/jit/xla_compiler_options_util.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" -#include "tensorflow/compiler/tf2xla/const_analysis.h" -#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/local_client.h" -#include "xla/executable_run_options.h" -#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/tf_pjrt_client.h" -#include "xla/service/executable.h" -#include "xla/service/gpu/gpu_executable_run_options.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/refcount.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/statusor.h" -#include "tensorflow/core/tfrt/common/pjrt_util.h" -#include "tsl/platform/errors.h" namespace tensorflow { namespace { -using XlaDeviceCompiler = - DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; -XlaCompiler::CompileOptions GetCompileOptions(bool for_pjrt = false) { +XlaCompiler::CompileOptions GetCompileOptions() { XlaCompiler::CompileOptions compile_options; compile_options.is_entry_computation = true; - // Optimization: where possible, have the computation return a naked array - // rather than a one-element tuple. - compile_options.always_return_tuple = false; - if (for_pjrt) { - compile_options.use_tuple_arg = false; - compile_options.always_return_tuple = true; - } - + compile_options.always_return_tuple = true; + compile_options.use_tuple_arg = false; return compile_options; } @@ -111,70 +83,6 @@ XlaCompileOnDemandOp::XlaCompileOnDemandOp(OpKernelConstruction* ctx) function_(GetDeviceCompilerFunction(ctx->def())), canonical_function_(Canonicalize(function_)) {} -absl::Status XlaCompileOnDemandOp::Run( - const ResourceVarsSnapshot& variable_args, - const XlaCompiler::CompilationResult* result, - const XlaDeviceCompiler* xla_device_compiler, - xla::LocalExecutable* executable, OpKernelContext* ctx) { - xla::LocalClient* client = - static_cast(xla_device_compiler->client()); - - se::Stream* stream = - ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; - std::shared_ptr allocator_ptr = - GetAllocator(ctx->device(), stream, platform_info_); - stream_executor::DeviceAddressAllocator* allocator = allocator_ptr.get(); - XlaComputationLaunchContext launch_context( - client, allocator, client->default_device_ordinal(), - /*allocate_xla_tensors=*/platform_info_.xla_device_metadata() != nullptr, - platform_info_.xla_device_metadata() - ? platform_info_.xla_device_metadata()->UseMultipleStreams() - : false); - - absl::flat_hash_map snapshot_ptrs; - for (auto& p : variable_args) { - snapshot_ptrs.emplace(p.first, - p.second.has_value() ? &p.second.value() : nullptr); - } - - const xla::HloInputOutputAliasConfig& input_output_alias = - executable->executable()->module().input_output_alias_config(); - absl::StatusOr> execution_inputs = - launch_context.PopulateInputs(ctx, result, snapshot_ptrs, - /*missing_ctx_input_prefix=*/0, - input_output_alias); - TF_RETURN_IF_ERROR(execution_inputs.status()); - - VLOG(2) << "Executing computation: " << name(); - xla::ExecutableRunOptions run_options; - xla::gpu::GpuExecutableRunOptions gpu_options; - xla::DeviceAssignment device_assignment; - if (result->collective_info) { - TF_RETURN_IF_ERROR(ResolveDeviceAssignment(ctx, *result->collective_info, - run_options, device_assignment, - gpu_options)); - } - - run_options.set_stream(stream); - run_options.set_allocator(allocator); - run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); - run_options.set_rng_seed(GetXLARandomSeed()); - - absl::StatusOr run_result = - executable->Run(std::move(execution_inputs).value(), run_options); - TF_RETURN_IF_ERROR(run_result.status()); - xla::ExecutionOutput execution_output = std::move(run_result).value(); - absl::StatusOr> variable_infos = - GatherVariableInfo(ctx, *result, 0); - TF_RETURN_IF_ERROR(variable_infos.status()); - TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(*variable_infos))); - TF_RETURN_IF_ERROR(launch_context.PopulateOutputs( - ctx, result, execution_output.ConsumeResult(), - /*missing_ctx_input_prefix=*/0, absl::MakeSpan(*variable_infos), - input_output_alias, snapshot_ptrs)); - return absl::OkStatus(); -} - absl::Status XlaCompileOnDemandOp::Compile( const std::vector& args, OpKernelContext* ctx, PjRtDeviceCompiler** pjrt_device_compiler, @@ -190,60 +98,17 @@ absl::Status XlaCompileOnDemandOp::Compile( platform_info_, *pjrt_device_compiler); // No detailed logging for on demand op. options.detailed_logging = false; - XlaCompiler::CompileOptions compile_options = GetCompileOptions(true); - - return (*pjrt_device_compiler) - ->CompileSingleOpIfNeeded(options, function_, canonical_function_, args, - compile_options, ctx, *profiler, result, - executable); -} - -absl::Status XlaCompileOnDemandOp::Compile( - const std::vector& args, OpKernelContext* ctx, - XlaDeviceCompiler** xla_device_compiler, - DeviceCompilationProfiler** profiler, - const XlaCompiler::CompilationResult** result, - xla::LocalExecutable** executable) { - // We store information about the JIT-compiled XLA computation - // in the ResourceMgr. - ResourceMgr* rm = ctx->resource_manager(); - CHECK(rm); - - TF_ASSIGN_OR_RETURN(DeviceType compilation_device_type, - GetCompilationDeviceType(platform_info_.device_type())); - - TF_RETURN_IF_ERROR(rm->LookupOrCreate( - rm->default_container(), "xla_device_compiler", xla_device_compiler, - [&](XlaDeviceCompiler** xla_device_compiler) { - return BuildXlaDeviceCompiler(ctx->device(), ctx->function_library(), - platform_info_, compilation_device_type, - xla_device_compiler); - })); - - TF_RETURN_IF_ERROR(rm->LookupOrCreate( - rm->default_container(), "device_compilation_profiler", profiler, - [](DeviceCompilationProfiler** profiler) { - *profiler = new DeviceCompilationProfiler(); - return absl::OkStatus(); - })); - - XlaCompiler::Options options = GenerateCompilerOptions( - **xla_device_compiler, *ctx->function_library(), ctx->device(), - ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr, - platform_info_, /*has_ref_vars=*/true); - // No detailed logging from on demand op. - options.detailed_logging = false; XlaCompiler::CompileOptions compile_options = GetCompileOptions(); - return (*xla_device_compiler) + return (*pjrt_device_compiler) ->CompileSingleOpIfNeeded(options, function_, canonical_function_, args, compile_options, ctx, *profiler, result, executable); } void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) { - const XlaCompiler::CompilationResult* result; - DeviceCompilationProfiler* profiler; + const XlaCompiler::CompilationResult* result = nullptr; + DeviceCompilationProfiler* profiler = nullptr; OP_REQUIRES(ctx, ctx->function_library(), absl::InternalError("Function library missing")); @@ -255,70 +120,35 @@ void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) { std::vector variable_indices = GetResourceVariableIndicesFromContext(ctx); - bool use_pjrt = - GetXlaOpsCommonFlags() - ->tf_xla_use_device_api.IsEnabledInXlaCompileOnDemandForDevice( - platform_info_.device_type()); - if (use_pjrt) { - std::vector variables; - std::vector args; - // Lock variables for the whole duration of compile + execute. - OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( - *ctx, inputs, *constant_indices_or, - variable_indices, &variables, &args)); - - PjRtDeviceCompiler* pjrt_device_compiler = nullptr; - xla::PjRtLoadedExecutable* pjrt_executable = nullptr; - absl::Status status = Compile(args, ctx, &pjrt_device_compiler, &profiler, - &result, &pjrt_executable); - // Hold the reference to the XLA device compiler and profiler during - // evaluation. (We could probably free them sooner because the ResourceMgr - // will retain references, but this is more obviously correct.) - // We must also ensure these pointers are freed even if compilation fails. - core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); - core::ScopedUnref profiler_ref(profiler); - OP_REQUIRES_OK(ctx, status); - - VLOG(2) << "Compiled op with PJRT: " << ctx->status(); - VLOG(2) << "result != nullptr: " << (result != nullptr); - VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); - VLOG(2) << "Executing with PJRT ..."; - - OP_REQUIRES_OK(ctx, RunPjRtExecutable(inputs, variables, *result, - pjrt_device_compiler->client(), - pjrt_executable, ctx)); - - VLOG(2) << "Completed executing with PJRT!"; - } else { - ResourceVarsSnapshot variable_args; - std::vector args; - // Lock variables only for generating XlaCompiler::Arguments and then - // release them. - { - std::vector variables; - OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( - *ctx, inputs, *constant_indices_or, - variable_indices, &variables, &args)); - OP_REQUIRES_OK(ctx, SnapshotResourceVariables(ctx, variable_indices, - variables, &variable_args)); - } - - XlaDeviceCompiler* xla_device_compiler = nullptr; - xla::LocalExecutable* executable = nullptr; - absl::Status status = Compile(args, ctx, &xla_device_compiler, &profiler, - &result, &executable); - // Hold the reference to the XLA device compiler and profiler during - // evaluation. (We could probably free them sooner because the ResourceMgr - // will retain references, but this is more obviously correct.) - // We must also ensure these pointers are freed even if compilation fails. - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); - core::ScopedUnref profiler_ref(profiler); - OP_REQUIRES_OK(ctx, status); - - // Locks are acquired again when populating the `ctx` outputs. - OP_REQUIRES_OK( - ctx, Run(variable_args, result, xla_device_compiler, executable, ctx)); - } + std::vector variables; + std::vector args; + // Lock variables for the whole duration of compile + execute. + OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( + *ctx, inputs, *constant_indices_or, variable_indices, + &variables, &args)); + + PjRtDeviceCompiler* pjrt_device_compiler = nullptr; + xla::PjRtLoadedExecutable* pjrt_executable = nullptr; + absl::Status status = Compile(args, ctx, &pjrt_device_compiler, &profiler, + &result, &pjrt_executable); + // Hold the reference to the XLA device compiler and profiler during + // evaluation. (We could probably free them sooner because the ResourceMgr + // will retain references, but this is more obviously correct.) + // We must also ensure these pointers are freed even if compilation fails. + core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); + core::ScopedUnref profiler_ref(profiler); + OP_REQUIRES_OK(ctx, status); + + VLOG(2) << "Compiled op with PJRT: " << ctx->status(); + VLOG(2) << "result != nullptr: " << (result != nullptr); + VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); + VLOG(2) << "Executing with PJRT ..."; + + OP_REQUIRES_OK(ctx, RunPjRtExecutable(inputs, variables, *result, + pjrt_device_compiler->client(), + pjrt_executable, ctx)); + + VLOG(2) << "Completed executing with PJRT!"; } } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.h b/tensorflow/compiler/jit/xla_compile_on_demand_op.h index 113972706271e7..e368de9f5e2236 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.h +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.h @@ -21,20 +21,16 @@ limitations under the License. #include +#include "absl/status/status.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" #include "tensorflow/compiler/jit/device_compilation_profiler.h" -#include "tensorflow/compiler/jit/variable_info.h" -#include "tensorflow/compiler/jit/variable_info_util.h" -#include "tensorflow/compiler/jit/xla_launch_util.h" +#include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/lib/core/status.h" namespace tensorflow { @@ -47,14 +43,6 @@ class XlaCompileOnDemandOp : public OpKernel { void Compute(OpKernelContext* ctx) override; private: - absl::Status Compile(const std::vector& args, - OpKernelContext* ctx, - DeviceCompiler** - xla_device_compiler, - DeviceCompilationProfiler** profiler, - const XlaCompiler::CompilationResult** result, - xla::LocalExecutable** executable); - absl::Status Compile(const std::vector& args, OpKernelContext* ctx, DeviceCompiler* - xla_device_compiler, - xla::LocalExecutable* executable, OpKernelContext* ctx); - const XlaPlatformInfo platform_info_; // Canonicalized function to compile derived from the Op attributes. diff --git a/tensorflow/compiler/jit/xla_compile_util.cc b/tensorflow/compiler/jit/xla_compile_util.cc index c9c80d244b35b9..021d7f55911ad6 100644 --- a/tensorflow/compiler/jit/xla_compile_util.cc +++ b/tensorflow/compiler/jit/xla_compile_util.cc @@ -78,12 +78,6 @@ absl::StatusOr> CreateSingleOpGraph( return graph; } -bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type) { - const auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; - return rollout_config.IsEnabledInXlaLaunchForDevice(device_type) || - rollout_config.IsEnabledInXlaCompileOnDemandForDevice(device_type) || - rollout_config.IsEnabledInXlaCompileAndRunForDevice(device_type); -} std::string GetPjRtDeviceCompilerResourceName(const DeviceType& device_type) { return absl::StrCat(kPjRtDeviceCompilerResourceName, "_", diff --git a/tensorflow/compiler/jit/xla_compile_util.h b/tensorflow/compiler/jit/xla_compile_util.h index d722ba8e784e61..81dc51169a91cf 100644 --- a/tensorflow/compiler/jit/xla_compile_util.h +++ b/tensorflow/compiler/jit/xla_compile_util.h @@ -45,9 +45,6 @@ absl::StatusOr> CreateSingleOpGraph( const NodeDef& node_def, absl::Span args, absl::Span result_types); -// Checks if single device compilation and execution with PJRT is enabled for -// `device_type` in either the XlaLaunch op or the XlaCompileOnDemand op. -bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type); // Gets the resource name of the PjRt DeviceCompiler for `device_type`. std::string GetPjRtDeviceCompilerResourceName(const DeviceType& device_type); diff --git a/tensorflow/compiler/jit/xla_compile_util_test.cc b/tensorflow/compiler/jit/xla_compile_util_test.cc index 7e55498ec42d33..fa529c603e8fa9 100644 --- a/tensorflow/compiler/jit/xla_compile_util_test.cc +++ b/tensorflow/compiler/jit/xla_compile_util_test.cc @@ -75,49 +75,6 @@ TEST_F(OpsTestBase, CreateSingleOpGraph) { EXPECT_EQ(retval_input_node->name(), "identity_op"); } -TEST(XlaCompileUtilTest, PjRtXlaLaunchFlagTest) { - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on, but no device is allowlisted. - auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; - rollout_config.enabled_for_xla_launch_ = true; - - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on, some device is allowlisted, but the requested one isn't. - rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_GPU)); - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on and the requested device is allowlisted. - rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_CPU)); - EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // The requested device is allowlisted, but the flag is turned off. - rollout_config.enabled_for_xla_launch_ = false; - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); -} - -TEST(XlaCompileUtilTest, PjRtXlaCompileOnDemandFlagTest) { - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on, but no device is allowlisted. - auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; - rollout_config.enabled_for_compile_on_demand_ = true; - - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on, some device is allowlisted, but the requested one isn't. - rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(DEVICE_GPU)); - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // Flag is turned on and the requested device is allowlisted. - rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(DEVICE_CPU)); - EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); - - // The requested device is allowlisted, but the flag is turned off. - rollout_config.enabled_for_compile_on_demand_ = false; - EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); -} TEST(XlaCompileUtilTest, PjRtDeviceCompilerResourceName) { EXPECT_EQ(GetPjRtDeviceCompilerResourceName(DeviceType(DEVICE_TPU)), diff --git a/tensorflow/compiler/jit/xla_compiler_options_util.cc b/tensorflow/compiler/jit/xla_compiler_options_util.cc index e17815cf4e91f4..31c5f80053deb3 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util.cc +++ b/tensorflow/compiler/jit/xla_compiler_options_util.cc @@ -20,19 +20,15 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/stream_executor/host/host_platform_id.h" -#include "xla/stream_executor/stream.h" #include "xla/tsl/framework/device_id_utils.h" +#include "xla/tsl/framework/device_type.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/types.h" namespace tensorflow { namespace { -using XlaDeviceCompiler = - DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; @@ -53,41 +49,12 @@ inline void LogOptions(const XlaCompiler::Options& options) { } } // namespace -XlaCompiler::Options GenerateCompilerOptions( - const XlaDeviceCompiler& xla_device_compiler, - const FunctionLibraryRuntime& function_library, DeviceBase* device, - se::Stream* stream, const XlaPlatformInfo& platform_info, - bool has_ref_vars) { - XlaCompiler::Options options; - options.client = static_cast(xla_device_compiler.client()); - if (stream != nullptr) { - options.device_ordinal = stream->parent()->device_ordinal(); - } - options.device_type = xla_device_compiler.device_type(); - options.flib_def = function_library.GetFunctionLibraryDefinition(); - options.graph_def_version = function_library.graph_def_version(); - options.allow_cpu_custom_calls = - (platform_info.platform_id() == se::host::kHostPlatformId); - options.device_allocator = GetAllocator(device, stream, platform_info); - if (platform_info.xla_device_metadata()) { - options.shape_determination_fns = - platform_info.xla_device_metadata()->default_shape_determination_fns(); - } - // If reference variables are not present in the graph, we can safely alias - // passthrough parameters without performing a copy. - options.alias_passthrough_params = - !has_ref_vars && !platform_info.is_on_xla_device(); - - LogOptions(options); - return options; -} - XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu( - const XlaDeviceCompiler& xla_device_compiler, + const tsl::DeviceType& device_type, const FunctionLibraryRuntime& function_library) { XlaCompiler::Options options; // TODO(b/238830423): consider device_ordinal and shape_determination_fns. - options.device_type = xla_device_compiler.device_type(); + options.device_type = device_type; options.flib_def = function_library.GetFunctionLibraryDefinition(); options.graph_def_version = function_library.graph_def_version(); options.allow_cpu_custom_calls = false; @@ -137,6 +104,9 @@ XlaCompiler::Options GenerateCompilerOptionsForPjRt( } else if (pjrt_device_compiler != nullptr) { options.device_type = pjrt_device_compiler->device_type(); } + if (pjrt_device_compiler != nullptr) { + options.client = pjrt_device_compiler->client(); + } // TODO(b/255826209): Confirm below options are correctly set after testing. options.allow_cpu_custom_calls = false; options.alias_passthrough_params = false; diff --git a/tensorflow/compiler/jit/xla_compiler_options_util.h b/tensorflow/compiler/jit/xla_compiler_options_util.h index 8c6fb8e0eaf2ab..96ca45311d9756 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util.h +++ b/tensorflow/compiler/jit/xla_compiler_options_util.h @@ -19,26 +19,16 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/stream_executor/stream.h" +#include "xla/tsl/framework/device_type.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" namespace tensorflow { -// Returns created options for the XLA compiler. -XlaCompiler::Options GenerateCompilerOptions( - const DeviceCompiler& - xla_device_compiler, - const FunctionLibraryRuntime& function_library, DeviceBase* device, - se::Stream* stream, const XlaPlatformInfo& platform_info, - bool has_ref_vars); - // Returns created options for XLA compiler when TFRT-TPU is used. XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu( - const DeviceCompiler& - xla_device_compiler, + const tsl::DeviceType& device_type, const FunctionLibraryRuntime& function_library); // Returns created options for XLA compiler when PjRt (Device API) is used for diff --git a/tensorflow/compiler/jit/xla_compiler_options_util_test.cc b/tensorflow/compiler/jit/xla_compiler_options_util_test.cc index 0ea5d3222204fb..539acab667b198 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util_test.cc +++ b/tensorflow/compiler/jit/xla_compiler_options_util_test.cc @@ -28,19 +28,15 @@ limitations under the License. #include "tensorflow/compiler/jit/pjrt_device_compiler_client.h" #include "tensorflow/compiler/jit/test_util.h" #include "tensorflow/compiler/jit/xla_device.h" -#include "tensorflow/compiler/jit/xla_device_compiler_client.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_argument.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "xla/client/client_library.h" -#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" #include "xla/shape.h" -#include "xla/stream_executor/device_memory_allocator.h" -#include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/platform.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" @@ -54,25 +50,11 @@ limitations under the License. namespace tensorflow { namespace { -using XlaDeviceCompiler = - DeviceCompiler; -using XlaDeviceExecutablePersistor = - DeviceExecutablePersistor; using PjRtDeviceCompiler = DeviceCompiler; using PjRtDeviceExecutablePersistor = DeviceExecutablePersistor; -XlaDeviceCompiler* CreateXlaDeviceCompiler(DeviceType device_type, - xla::LocalClient* local_client) { - auto persistor = std::make_unique( - XlaDeviceExecutablePersistor::Config(), device_type); - auto compiler_client = - std::make_unique(local_client); - return new XlaDeviceCompiler(std::move(persistor), - std::move(compiler_client)); -} - PjRtDeviceCompiler* CreatePjRtDeviceCompiler(DeviceType device_type, xla::PjRtClient* pjrt_client) { auto persistor = std::make_unique( @@ -229,103 +211,13 @@ TEST_F(XlaCompilerOptionsTest, PjRtOptionsNonXlaDevice) { tensorflow::XlaLayoutPreference::kNoPreference); } -TEST_F(XlaCompilerOptionsTest, XlaOptions) { - device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU}); - Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU); - - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); - DeviceType device_type = DeviceType(DEVICE_XLA_GPU); - DeviceType compilation_device_type = DeviceType(DEVICE_GPU_XLA_JIT); - - auto xla_device_compiler = - CreateXlaDeviceCompiler(compilation_device_type, client); - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); - - se::Platform::Id platform_id = se::host::kHostPlatformId; - auto xla_device_metadata = CreateXlaDeviceMetadata(compilation_device_type); - std::shared_ptr custom_allocator; - XlaPlatformInfo platform_info( - device_type, platform_id, xla_device_metadata.get(), - /*pjrt_device_metadata=*/nullptr, custom_allocator); - - XlaCompiler::Options options = - GenerateCompilerOptions(*xla_device_compiler, *device_setup_.flr(), - device, nullptr, platform_info, false); - - EXPECT_EQ(options.device_type, compilation_device_type); - EXPECT_NE(options.flib_def, nullptr); - EXPECT_EQ(options.graph_def_version, TF_GRAPH_DEF_VERSION); - EXPECT_TRUE(options.allow_cpu_custom_calls); - EXPECT_NE(options.device_allocator, nullptr); - EXPECT_FALSE(options.alias_passthrough_params); - // Check if options have the supplied shape determination functions set. - TF_ASSERT_OK_AND_ASSIGN( - auto shape, options.shape_determination_fns.shape_representation_fn( - TensorShape(), DT_FLOAT, false, - tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout)); - EXPECT_EQ(shape, xla::Shape()); - EXPECT_EQ(options.shape_determination_fns.layout_preference_fn( - TensorShape(), DT_FLOAT, std::nullopt), - tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout); -} - -TEST_F(XlaCompilerOptionsTest, XlaOptionsHasRefVarsNoXlaDeviceMetadata) { - device_setup_.AddDevicesAndSetUp({DEVICE_CPU}); - Device* device = device_setup_.GetDevice(DEVICE_CPU); - - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); - DeviceType device_type = DeviceType(DEVICE_CPU); - DeviceType compilation_device_type = DeviceType(DEVICE_CPU_XLA_JIT); - - auto xla_device_compiler = - CreateXlaDeviceCompiler(compilation_device_type, client); - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); - - se::Platform::Id platform_id = se::host::kHostPlatformId; - std::shared_ptr custom_allocator; - XlaPlatformInfo platform_info( - device_type, platform_id, /*xla_device_metadata=*/nullptr, - /*pjrt_device_metadata=*/nullptr, custom_allocator); - - XlaCompiler::Options options = - GenerateCompilerOptions(*xla_device_compiler, *device_setup_.flr(), - device, nullptr, platform_info, false); - - EXPECT_EQ(options.device_type, compilation_device_type); - EXPECT_NE(options.flib_def, nullptr); - EXPECT_EQ(options.graph_def_version, TF_GRAPH_DEF_VERSION); - EXPECT_TRUE(options.allow_cpu_custom_calls); - EXPECT_NE(options.device_allocator, nullptr); - EXPECT_TRUE(options.alias_passthrough_params); - // Check whether options have default shape determination functions set. - TF_ASSERT_OK_AND_ASSIGN( - auto shape, options.shape_determination_fns.shape_representation_fn( - TensorShape(), DT_FLOAT, false, - tensorflow::XlaLayoutPreference::kNoPreference)); - xla::ShapeProto shape_proto; - shape_proto.set_element_type(xla::PrimitiveType::F32); - shape_proto.mutable_layout(); - TF_ASSERT_OK_AND_ASSIGN(auto expected_shape, - xla::Shape::FromProto(shape_proto)); - EXPECT_EQ(shape, expected_shape); - EXPECT_EQ(options.shape_determination_fns.layout_preference_fn( - TensorShape(), DT_FLOAT, std::nullopt), - tensorflow::XlaLayoutPreference::kNoPreference); -} - TEST_F(XlaCompilerOptionsTest, TfRtTpuOptions) { device_setup_.AddDevicesAndSetUp({DEVICE_TPU_NODE}); - // Just use the default local client for testing purposes. - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); DeviceType compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); - auto xla_device_compiler = - CreateXlaDeviceCompiler(compilation_device_type, client); - core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); - XlaCompiler::Options options = GenerateCompilerOptionsForTfrtTpu( - *xla_device_compiler, *device_setup_.flr()); + compilation_device_type, *device_setup_.flr()); EXPECT_EQ(options.device_type, compilation_device_type); EXPECT_NE(options.flib_def, nullptr); diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc index 0c8b36e49d1ef4..45488c44577684 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -17,48 +17,48 @@ limitations under the License. #include +#include +#include #include #include #include +#include #include "absl/base/call_once.h" -#include "absl/memory/memory.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/match.h" -#include "tensorflow/compiler/jit/defs.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" #include "tensorflow/compiler/jit/pjrt_device_context.h" #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h" -#include "tensorflow/compiler/jit/xla_compile_util.h" -#include "tensorflow/compiler/jit/xla_device_context.h" +#include "tensorflow/compiler/jit/xla_tensor.h" +#include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/client/client_library.h" -#include "xla/service/stream_pool.h" -#include "tensorflow/core/common_runtime/device.h" -#include "tensorflow/core/common_runtime/device_factory.h" +#include "xla/client/local_client.h" +#include "xla/service/shaped_buffer.h" +#include "xla/shape.h" +#include "xla/stream_executor/platform.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/common_runtime/dma_helper.h" -#include "tensorflow/core/common_runtime/function.h" -#include "tensorflow/core/common_runtime/graph_constructor.h" -#include "tensorflow/core/common_runtime/renamed_device.h" +#include "tensorflow/core/common_runtime/local_device.h" #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/kernel_def.pb.h" -#include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/lib/core/notification.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/stream_executor_no_cuda.h" -#include "tensorflow/core/profiler/lib/traceme.h" +#include "tensorflow/core/graph/types.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/threadpool.h" #include "tensorflow/core/public/session_options.h" -#include "tensorflow/core/public/version.h" #include "tensorflow/core/tfrt/common/async_value_tensor.h" -#include "tensorflow/core/util/device_name_utils.h" -#include "tensorflow/core/util/dump_graph.h" -#include "tensorflow/core/util/stream_executor_util.h" namespace tensorflow { @@ -77,59 +77,6 @@ absl::Status DefaultPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) { return absl::OkStatus(); } -// Caches a XlaDeviceAllocator per pair. A -// XlaDeviceAllocator is created on demand and is associated with a -// XlaDevice. It outlives the device itself (for instance, the buffer -// backing a tensor holds a pointer to the allocator for book-keeping, -// and this buffer can outlast the device). -class XlaDeviceAllocatorState { - public: - // Creates or returns a cached XlaDeviceAllocator for a given - // backend and device_ordinal. - static XlaDeviceAllocator* GetOrCreateXlaDeviceAllocator( - const xla::Backend* backend, int device_ordinal); - - private: - // Returns the singleton instance of XlaDeviceAllocatorState. - static XlaDeviceAllocatorState& Singleton(); - XlaDeviceAllocatorState(); - ~XlaDeviceAllocatorState(); - - mutex allocator_mutex_; // Guards the singleton allocator state. - std::unordered_map, - std::unique_ptr, - hash>> - allocators_ TF_GUARDED_BY(allocator_mutex_); - - XlaDeviceAllocatorState(const XlaDeviceAllocatorState&) = delete; - void operator=(const XlaDeviceAllocatorState&) = delete; -}; - -/* static */ XlaDeviceAllocatorState& XlaDeviceAllocatorState::Singleton() { - static auto a = new XlaDeviceAllocatorState; - return *a; -} - -XlaDeviceAllocatorState::XlaDeviceAllocatorState() = default; -XlaDeviceAllocatorState::~XlaDeviceAllocatorState() = default; - -XlaDeviceAllocator* XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( - const xla::Backend* backend, int device_ordinal) { - XlaDeviceAllocatorState& state = Singleton(); - mutex_lock lock(state.allocator_mutex_); - - auto it = state.allocators_.find({backend, device_ordinal}); - if (it != state.allocators_.end()) { - return it->second.get(); - } - - std::unique_ptr alloc = - std::make_unique( - backend->stream_executors()[device_ordinal]); - XlaDeviceAllocator* alloc_ptr = alloc.get(); - state.allocators_[{backend, device_ordinal}] = std::move(alloc); - return alloc_ptr; -} namespace { @@ -194,11 +141,6 @@ const DeviceType& XlaDevice::Metadata::jit_device_type() const { return GetMetadataFromDevice(ctx->device(), metadata); } -/* static */ mutex XlaDevice::global_mu_(LINKER_INITIALIZED); -/* static */ std::vector>* - XlaDevice::global_compute_streams_ = - new std::vector>; - XlaDevice::XlaDevice(const SessionOptions& session_options, const Options& options) : LocalDevice(session_options, @@ -217,27 +159,15 @@ XlaDevice::XlaDevice(const SessionOptions& session_options, platform_(options.platform), intra_op_parallelism_threads_( session_options.config.intra_op_parallelism_threads()), - use_multiple_streams_(options.use_multiple_streams), shape_determination_fns_(options.shape_determination_fns), - allowed_devices_(options.allowed_devices), - use_global_compute_stream_(options.use_global_compute_stream) { + allowed_devices_(options.allowed_devices) { if (options.shape_determination_fns.empty()) { LOG(ERROR) << "shape_representation_fns must be non-empty."; } VLOG(1) << "Created XLA device " << options.compilation_device_name << " " << options.device_ordinal << " " << this; - VLOG(1) << "XlaDevice options: use_multiple_streams: " - << options.use_multiple_streams << " use_global_compute_stream: " - << options.use_global_compute_stream; thread_pool_.reset(new thread::ThreadPool(session_options.env, "xla_device", /*num_threads=*/1)); - - // We have multiple device to device streams to allow for some concurrency - // between transfers. The particular value of '4' is chosen fairly - // arbitrarily. It may be necessary to make this tunable via - // XlaDevice::Options. - static constexpr int kNumDeviceToDeviceStreams = 4; - device_to_device_streams_.resize(kNumDeviceToDeviceStreams); } XlaDevice::~XlaDevice() { @@ -248,19 +178,6 @@ XlaDevice::~XlaDevice() { } } -absl::StatusOr XlaDevice::GetOrCreateClient() const { - // We lazily create the client because the platform commits to the - // details of the host hardware when the client is created, so we - // don't want to do it until we get a chance to hook the platform up - // to a simulator. - - xla::LocalClientOptions options; - options.set_platform(platform_) - .set_allowed_devices(allowed_devices_) - .set_intra_op_parallelism_threads(intra_op_parallelism_threads_); - return xla::ClientLibrary::GetOrCreateLocalClient(options); -} - Allocator* XlaDevice::GetAllocator(AllocatorAttributes attr) { mutex_lock lock(mu_); return GetAllocatorLocked(attr); @@ -272,17 +189,9 @@ Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) { } if (xla_allocator_ == nullptr) { - if (UsePjRtForSingleDeviceCompilation(device_name_)) { - VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator"; - pjrt_allocator_ = std::make_unique(); - xla_allocator_ = pjrt_allocator_.get(); - } else { - // TODO(b/78468222): This can fail, at least when the backend is GPU and - // there is no GPU on the host. - xla::Backend* backend = GetOrCreateClient().value()->mutable_backend(); - xla_allocator_ = XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( - backend, device_ordinal_); - } + VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator"; + pjrt_allocator_ = std::make_unique(); + xla_allocator_ = pjrt_allocator_.get(); } return xla_allocator_; } @@ -292,136 +201,24 @@ absl::Status XlaDevice::EnsureDeviceContextOk() { return GetDeviceContextLocked().status(); } -absl::Status XlaDevice::EnsureStreamOkLocked( - xla::Backend* backend, const std::string& name, - std::shared_ptr* stream, bool* stream_was_changed) { - if (!(*stream) || !(*stream)->ok()) { - xla::StreamPool::Ptr ptr; - TF_ASSIGN_OR_RETURN(ptr, backend->BorrowStream(device_ordinal_)); - *stream = std::shared_ptr(std::move(ptr)); - VLOG(1) << "XlaDevice " << this << " new " << name - << " stream=" << (*stream); - *stream_was_changed = true; - } - return absl::OkStatus(); -} - absl::StatusOr> XlaDevice::GetDeviceContextLocked() { - if (UsePjRtForSingleDeviceCompilation(device_name_)) { - if (device_contexts_.empty()) { - for (const auto& iter : shape_determination_fns_) { - auto device_context = new PjRtDeviceContext(iter); - VLOG(1) << "XlaDevice " << this << " new PjRtDeviceContext " - << device_context; - device_contexts_.emplace_back(device_context); - } - if (use_accelerator_device_info_) { - auto accelerator_device_info = - std::make_unique(); - accelerator_device_info->default_context = device_contexts_.at(0); - set_tensorflow_accelerator_device_info(accelerator_device_info.get()); - accelerator_device_info_ = std::move(accelerator_device_info); - VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " - << accelerator_device_info_.get(); - } - } - - return device_contexts_; - } - - TF_ASSIGN_OR_RETURN(xla::LocalClient * client, GetOrCreateClient()); - xla::Backend* backend = client->mutable_backend(); - - // Ensure all our streams are valid, borrowing new streams if necessary. - bool need_new_device_context = device_contexts_.empty(); - if (use_global_compute_stream_) { - mutex_lock lock(global_mu_); - if (global_compute_streams_->size() <= device_ordinal_) { - global_compute_streams_->resize(device_ordinal_ + 1, nullptr); + if (device_contexts_.empty()) { + for (const auto& iter : shape_determination_fns_) { + auto device_context = new PjRtDeviceContext(iter); + VLOG(1) << "XlaDevice " << this << " new PjRtDeviceContext " + << device_context; + device_contexts_.emplace_back(device_context); } - - auto& global_stream = global_compute_streams_->at(device_ordinal_); - if (global_stream != nullptr && global_stream->ok()) { - stream_ = global_stream; - } else { - // Directly create the stream here instead of borrowing from the stream - // pool to avoid potential lifetime issues. - TF_ASSIGN_OR_RETURN( - stream_, - backend->stream_executors()[device_ordinal_]->CreateStream()); - TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, - &need_new_device_context)); - (*global_compute_streams_)[device_ordinal_] = stream_; - } - } else { - TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, - &need_new_device_context)); - } - - std::shared_ptr host_to_device_stream; - std::shared_ptr device_to_host_stream; - std::vector> device_to_device_streams; - if (use_multiple_streams_) { - TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream", - &host_to_device_stream_, - &need_new_device_context)); - for (std::shared_ptr& stream : device_to_device_streams_) { - TF_RETURN_IF_ERROR( - EnsureStreamOkLocked(backend, "device_to_device_stream", &stream, - &need_new_device_context)); + if (use_accelerator_device_info_) { + auto accelerator_device_info = + std::make_unique(); + accelerator_device_info->default_context = device_contexts_.at(0); + set_tensorflow_accelerator_device_info(accelerator_device_info.get()); + accelerator_device_info_ = std::move(accelerator_device_info); + VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " + << accelerator_device_info_.get(); } - host_to_device_stream = host_to_device_stream_; - device_to_device_streams = device_to_device_streams_; - // The data transfer requests from device to host could arrive out of order, - // so a single stream would cause deadlock. For this case, - // xla_device_context would borrow a stream for each transfer request. - device_to_host_stream = nullptr; - } else { - host_to_device_stream = stream_; - device_to_host_stream = stream_; - device_to_device_streams = {stream_}; - } - - if (!need_new_device_context) { - return device_contexts_; - } - - // At this point we know we need a new device context. - // Call GetAllocator for the side-effect of ensuring the allocator is created. - GetAllocatorLocked({}); - for (const auto& iter : device_contexts_) { - iter->Unref(); - } - // The XlaDeviceContext keeps a reference count to the streams, and the - // XlaDeviceContext remains live for the duration of a Executor run. This - // ensures that the streams remain live for the duration of a run, even if - // an error is encountered and the streams are replaced with new ones. - for (const auto& iter : shape_determination_fns_) { - auto device_context = new XlaDeviceContext( - stream_, host_to_device_stream, device_to_host_stream, - device_to_device_streams, client, iter, thread_pool_.get()); - VLOG(1) << "XlaDevice " << this << " new XlaDeviceContext " - << device_context; - device_contexts_.emplace_back(device_context); - } - - // Create and set a new AcceleratorDeviceInfo, if necessary. - // - // TODO(b/78232898): This isn't thread-safe; there is a race between the call - // to set_tensorflow_accelerator_device_info() with ops that call the getter - // tensorflow_accelerator_device_info(). This isn't trivially fixed by adding - // locking to those methods; see the bug for details. Our only saving grace at - // the moment is that this race doesn't seem to occur in practice. - if (use_accelerator_device_info_) { - auto accelerator_device_info = - std::make_unique(); - accelerator_device_info->stream = stream_.get(); - accelerator_device_info->default_context = device_contexts_.at(0); - set_tensorflow_accelerator_device_info(accelerator_device_info.get()); - accelerator_device_info_ = std::move(accelerator_device_info); - VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " - << accelerator_device_info_.get(); } return device_contexts_; @@ -483,21 +280,6 @@ void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, absl::Status XlaDevice::Sync() { VLOG(1) << "XlaDevice::Sync"; - tsl::profiler::TraceMe activity("XlaDevice::Sync", - tsl::profiler::TraceMeLevel::kInfo); - std::shared_ptr stream; - { - mutex_lock lock(mu_); - stream = stream_; - } - if (!stream) return absl::OkStatus(); - - absl::Status status = stream->BlockHostUntilDone(); - TF_RETURN_IF_ERROR(status); - if (!stream->ok()) { - return absl::InternalError("XlaDevice::Sync() failed."); - } - VLOG(1) << "XlaDevice::Sync completed"; return absl::OkStatus(); } @@ -565,24 +347,7 @@ absl::Status XlaDevice::HandleDeviceError() { return absl::OkStatus(); } -absl::Status XlaDevice::RefreshStatus() { - std::shared_ptr stream; - { - mutex_lock lock(mu_); - stream = stream_; - } - if (!stream) { - return absl::OkStatus(); - } - absl::Status status = stream->RefreshStatus(); - if (!status.ok()) { - // Ignore errors from HandleDeviceError, since by definition the status is - // already non-ok, so there's nothing extra to report if HandleDeviceError - // itself returns an error. - HandleDeviceError().IgnoreError(); - } - return status; -} +absl::Status XlaDevice::RefreshStatus() { return absl::OkStatus(); } XlaDeviceOpRegistrations* RegisterXlaDeviceKernels( const char* device, const char* jit_device, diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index bedbf87a3cd2e8..30396631d4a5e7 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -205,15 +205,8 @@ class XlaDevice : public LocalDevice { absl::Status RefreshStatus() override TF_LOCKS_EXCLUDED(mu_); - private: - absl::StatusOr GetOrCreateClient() const; Allocator* GetAllocatorLocked(AllocatorAttributes attr) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_); - absl::Status EnsureStreamOkLocked(xla::Backend* backend, - const std::string& name, - std::shared_ptr* stream, - bool* stream_was_changed) - TF_EXCLUSIVE_LOCKS_REQUIRED(mu_); // Return a vector of device context, ordered by the sequence in the given // shape_representation_fns. @@ -240,24 +233,6 @@ class XlaDevice : public LocalDevice { Allocator* xla_allocator_ TF_GUARDED_BY(mu_) = nullptr; // Not owned. std::unique_ptr pjrt_allocator_ TF_GUARDED_BY(mu_); - // Stream associated with this device. Operations enqueued on this - // stream are executed on the device. Operations include data - // copying back and forth between CPU and the device, and - // computations enqueued by XLA. - std::shared_ptr stream_ TF_GUARDED_BY(mu_); - // If false, only stream_ is valid and all computation and transfers use - // stream_. If true, computation is performed by stream_ and transfers are - // performed by host_to_device/device_to_device stream or borrowing a stream - // for each device to host transfer. - const bool use_multiple_streams_; - // If use_multiple_streams_, host to device transfers are performed using this - // stream. - std::shared_ptr host_to_device_stream_ TF_GUARDED_BY(mu_); - // If use_multiple_streams_, transfers between different devices are performed - // using these streams. - std::vector> device_to_device_streams_ - TF_GUARDED_BY(mu_); - // See comments in options. std::vector shape_determination_fns_; @@ -288,15 +263,6 @@ class XlaDevice : public LocalDevice { // platform will have resources allocated. For GPUs this will be // filled from visible_gpu_devices list from session configuration. std::optional> allowed_devices_; - - const bool use_global_compute_stream_; - - // A static vector with device_ordinal as its index, describing the global - // compute streams used in each XLA device. It is only used if - // `use_global_compute_stream` in `XlaDevice::Options` is set to true. - static mutex global_mu_; - static std::vector>* global_compute_streams_ - TF_GUARDED_BY(global_mu_); }; // Builds OpKernel registrations on 'device' for the JIT operators diff --git a/tensorflow/compiler/jit/xla_gpu_device.cc b/tensorflow/compiler/jit/xla_gpu_device.cc index 95d3db26a322d1..fadf7903aaa2e0 100644 --- a/tensorflow/compiler/jit/xla_gpu_device.cc +++ b/tensorflow/compiler/jit/xla_gpu_device.cc @@ -158,7 +158,7 @@ REGISTER_LOCAL_DEVICE_FACTORY(DEVICE_XLA_GPU, XlaGpuDeviceFactory); // Kernel registrations -constexpr std::array kAllXlaGpuTypes = {{DT_UINT8, +constexpr std::array kAllXlaGpuTypes = {{DT_UINT8, DT_QUINT8, DT_UINT16, DT_INT8, @@ -179,6 +179,7 @@ constexpr std::array kAllXlaGpuTypes = {{DT_UINT8, DT_FLOAT8_E4M3FNUZ, DT_FLOAT8_E4M3B11FNUZ, DT_FLOAT8_E5M2FNUZ, + DT_FLOAT8_E8M0FNU, DT_INT4, DT_UINT4}}; diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 12c9f7d884ad66..223d3c5e2b9bd4 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_launch_util.h" -#include #include #include #include @@ -24,6 +23,7 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/casts.h" #include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -38,40 +38,29 @@ limitations under the License. #include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" -#include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" -#include "xla/client/local_client.h" #include "xla/future.h" -#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_common.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" -#include "xla/service/executable.h" -#include "xla/service/maybe_owning_device_memory.h" -#include "xla/service/shaped_buffer.h" -#include "xla/service/transfer_manager.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_execute_options.h" +#include "xla/runtime/device_id.h" #include "xla/shape.h" -#include "xla/shape_util.h" #include "xla/status_macros.h" -#include "xla/stream_executor/device_memory.h" -#include "xla/stream_executor/device_memory_allocator.h" -#include "xla/stream_executor/event.h" -#include "xla/stream_executor/host/host_platform_id.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/device_address_allocator.h" #include "xla/tsl/framework/device_id_utils.h" #include "xla/tsl/framework/serving_device_selector_policies.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/util.h" +#include "tensorflow/core/common_runtime/composite_device.h" #include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" -#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_handle.h" @@ -79,28 +68,12 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/refcount.h" -#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/tfrt/common/async_value_tensor.h" -#include "tsl/platform/casts.h" namespace tensorflow { namespace { -using xla::ScopedShapedBuffer; -using xla::ShapedBuffer; - -// Fetch the platform Id from device. -se::Platform::Id XlaPlatformInfoFromDevice(DeviceBase* device_base) { - auto device = static_cast(device_base); - se::Platform::Id platform_id = nullptr; - if (device->device_type() == DEVICE_CPU) { - platform_id = se::host::kHostPlatformId; - } - - return platform_id; -} absl::flat_hash_map CreateVariableLookup( const std::vector& variables) { @@ -136,177 +109,6 @@ absl::StatusOr> GetConstantInputIndicesFromContext( return constant_input_indices; } -XlaComputationLaunchContext::XlaComputationLaunchContext( - xla::LocalClient* client, - stream_executor::DeviceAddressAllocator* xla_allocator, int device_ordinal, - bool allocate_xla_tensors, bool use_multiple_streams) - : client_(client), - xla_allocator_(xla_allocator), - allocate_xla_tensors_(allocate_xla_tensors), - use_multiple_streams_(use_multiple_streams), - device_ordinal_(device_ordinal) { - if (use_multiple_streams_) { - CHECK(allocate_xla_tensors_) << "To use multiple streams correctly we must " - "be allocating XLA tensors!"; - } -} - -// Fills in `execution_input` with `buffer` for `index`. -static void PopulateExecutionInputBuffer( - xla::ExecutionInput& execution_input, const xla::ShapeIndex& index, - stream_executor::DeviceAddressBase buffer, bool donate_buffer, - int device_ordinal, stream_executor::DeviceAddressAllocator* allocator) { - xla::MaybeOwningDeviceAddress* in_buffer = - execution_input.MutableBuffer(index); - if (donate_buffer) { - // Here we pass ownership of the buffer to execution_input without releasing - // ownership from the caller of PopulateExecutionInputBuffer. If execution - // succeeds, we'll take back that duplicate ownership in - // GetOrCreateTensorForOutput. If execution fails, the ExecutionInput will - // release that duplicate ownership automatically. - *in_buffer = stream_executor::ScopedDeviceAddress( - buffer, device_ordinal, allocator); - } else { - *in_buffer = buffer; - } -} - -absl::StatusOr> -XlaComputationLaunchContext::PopulateInputs( - OpKernelContext* ctx, - const XlaCompiler::CompilationResult* compilation_result, - const absl::flat_hash_map& resource_vars, - int missing_ctx_input_prefix, - const xla::HloInputOutputAliasConfig& input_output_alias) { - std::vector arguments; - arguments.reserve(compilation_result->xla_input_shapes.size()); - - xla::ShapeIndex root_index = {}; - - for (int i = 0; i < compilation_result->xla_input_shapes.size(); ++i) { - int arg_num = compilation_result->input_mapping[i]; - CHECK_GE(arg_num, missing_ctx_input_prefix); - const xla::Shape& device_shape = compilation_result->xla_input_shapes[i]; - - auto resource_var_it = resource_vars.find(arg_num); - bool is_resource_variable = resource_var_it != resource_vars.end(); - bool is_updated_resource_variable = - is_resource_variable && - absl::c_any_of(compilation_result->resource_updates, - [&](const XlaCompiler::ResourceUpdate& update) { - // XlaCompiler records `arg_num` (instead of kernel - // parameters) in `resource_updates`. - return update.input_index == arg_num && - update.modified; - }); - - const Tensor* t = is_resource_variable - ? resource_var_it->second - : &(ctx->input(arg_num - missing_ctx_input_prefix)); - CHECK(t); - bool donate_buffer = t->RefCountIsOne() && is_updated_resource_variable && - input_output_alias.ParameterHasAlias(i, root_index); - VLOG(3) << "Processing input: " << i - << "; is_resource_variable=" << is_resource_variable - << "; is_updated_resource_variable=" << is_updated_resource_variable - << "; donate_buffer=" << donate_buffer; - - if (use_multiple_streams_) { - CHECK(ctx->op_device_context() && ctx->op_device_context()->stream()) - << "Must have a stream available when using XLA tensors!"; - XlaTensor* xla_tensor = XlaTensor::FromTensor(t); - CHECK(xla_tensor); - xla_tensor->WaitForDefinitionEventOnStream( - ctx->op_device_context()->stream()); - } - - arguments.emplace_back(&device_shape); - xla::ExecutionInput& execution_input = arguments.back(); - stream_executor::DeviceAddressBase dmem = - XlaTensor::DeviceMemoryFromTensor(*t); - PopulateExecutionInputBuffer(execution_input, root_index, dmem, - donate_buffer, device_ordinal_, - xla_allocator_); - } - return std::move(arguments); -} - -// Construct the tensor for the given type and buffer. -static Tensor MakeTensor(DataType dtype, const TensorShape& shape, - stream_executor::DeviceAddressBase buffer, - Allocator* allocator) { - size_t expected_size = shape.num_elements() * DataTypeSize(dtype); - auto* tensor_buffer = new XlaTensorBuffer(buffer.opaque(), expected_size, - buffer.size(), allocator); - Tensor t(dtype, shape, tensor_buffer); - tensor_buffer->Unref(); - return t; -} - -// Get aliased tensor from output, or make a new one for the corresponding -// output operation. Transfers ownership of the buffer from output to the -// returned tensor. -static absl::StatusOr GetOrCreateTensorForOutput( - xla::ScopedShapedBuffer& output, int output_num, OpKernelContext* ctx, - int missing_ctx_input_prefix, - const xla::HloInputOutputAliasConfig& input_output_alias, - absl::Span input_mapping, - const absl::flat_hash_map& resource_vars_snapshots, - DataType output_dtype, const TensorShape& output_shape, - Allocator* output_allocator, bool allocate_xla_tensors, se::Stream* stream, - bool use_multiple_streams, std::shared_ptr definition_event) { - xla::ShapeIndex output_index = input_output_alias.shape().IsTuple() - ? xla::ShapeIndex({output_num}) - : xla::ShapeIndex({}); - CHECK(input_output_alias.shape().IsTuple() || output_num == 0); - if (std::optional alias = - input_output_alias.GetAliasedParameter(output_index)) { - VLOG(3) << "Found alias: " << alias->ToString(); - int tf_param = - input_mapping[alias->parameter_number] - missing_ctx_input_prefix; - const Tensor input_tensor = - ctx->input(tf_param).dtype() != DT_RESOURCE - ? ctx->input(tf_param) - : *resource_vars_snapshots.at(missing_ctx_input_prefix + tf_param); - stream_executor::DeviceAddressBase input_buffer = - XlaTensor::DeviceMemoryFromTensor(input_tensor); - stream_executor::DeviceAddressBase output_buffer = - output.buffer({output_num}); - if (input_buffer.opaque() == output_buffer.opaque()) { - // In the case of a donated buffer, both input_tensor and output think - // they have ownership of the buffer (see comment in - // PopulateExecutionInputBuffer). Release ownership from output to avoid - // double free. - output.set_buffer(stream_executor::ScopedDeviceAddress(), - {output_num}); - return input_tensor; - } - } - - if (allocate_xla_tensors) { - Tensor output_tensor; - TF_RETURN_IF_ERROR( - ctx->allocate_temp(output_dtype, output_shape, &output_tensor)); - if (output_tensor.TotalBytes() > 0) { - XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor); - TF_RET_CHECK(xla_tensor); - xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num})); - if (use_multiple_streams) { - xla_tensor->ResetDefinitionEvent(definition_event, stream); - } - } - return output_tensor; - } - - stream_executor::DeviceAddressBase output_buffer = - output.buffer({output_num}); - Tensor output_tensor = - MakeTensor(output_dtype, output_shape, output_buffer, output_allocator); - output.set_buffer(stream_executor::ScopedDeviceAddress(), - {output_num}); - return output_tensor; -} - // Sets output `output_num` for `ctx` provided it is known at a compile time. absl::Status SetOutputForConstant( OpKernelContext* ctx, bool requires_copy_to_device, @@ -386,157 +188,6 @@ absl::StatusOr> GatherVariableInfo( return std::move(out); } -absl::Status XlaComputationLaunchContext::PopulateOutputs( - OpKernelContext* ctx, - const XlaCompiler::CompilationResult* compilation_result, - ScopedShapedBuffer output, int missing_ctx_input_prefix, - absl::Span variable_infos, - const xla::HloInputOutputAliasConfig& input_output_alias, - const absl::flat_hash_map& resource_vars) { - se::Stream* stream = - ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; - Allocator* allocator = ctx->device()->GetAllocator({}); - - // Computation output should always be a tuple. - VLOG(2) << "Result tuple shape: " << output.on_host_shape().ToString(); - VLOG(2) << "Result tuple shape (on device): " - << output.on_device_shape().ToString(); - CHECK_EQ(ctx->num_outputs(), compilation_result->outputs.size()); - - // If the on-host-shape isn't a tuple, create a new single-element tuple - // buffer with a nullptr root index table. This allows the code below to treat - // output as a tuple unconditionally. - if (!output.on_host_shape().IsTuple()) { - ShapedBuffer nontuple_buffer = output.release(); - ShapedBuffer buffer( - xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_device_shape()}), - output.device_ordinal()); - buffer.buffers().CopySubtreeFrom(nontuple_buffer.buffers(), - /*src_index=*/{}, - /*dst_index=*/{0}); - output = ScopedShapedBuffer(std::move(buffer), output.memory_allocator()); - } - - std::shared_ptr definition_event; - if (use_multiple_streams_ && stream) { - TF_ASSIGN_OR_RETURN(definition_event, stream->parent()->CreateEvent()); - TF_RETURN_IF_ERROR(stream->RecordEvent(definition_event.get())); - } - - for (const XlaOutputDescription& descr : compilation_result->outputs) { - if (descr.type == DT_VARIANT) { - return xla::Unimplemented( - "Support for TensorList crossing the XLA/TF boundary " - "is not implemented"); - } - } - - std::vector output_tensor_shapes; - output_tensor_shapes.reserve(ctx->num_outputs()); - if (output.on_host_shape().is_dynamic()) { - const se::Platform* platform = nullptr; - if (stream != nullptr) { - platform = stream->parent()->GetPlatform(); - } else { - // Stream is not set for the host platform. - TF_ASSIGN_OR_RETURN(platform, - se::PlatformManager::PlatformWithId( - XlaPlatformInfoFromDevice(ctx->device()))); - } - TF_ASSIGN_OR_RETURN(auto transfer_manager, - xla::TransferManager::GetForPlatform(platform)); - - xla::Shape output_device_shape = output.on_device_shape(); - TF_RETURN_IF_ERROR(transfer_manager->ReadDynamicShapes( - stream, &output, &output_device_shape)); - - output.set_shapes(output_device_shape, output_device_shape); - for (int i = 0; i < ctx->num_outputs(); ++i) { - const xla::Shape& subshape = - xla::ShapeUtil::GetSubshape(output_device_shape, {i}); - TensorShape shape; - TF_RETURN_IF_ERROR(XLAShapeToTensorShape(subshape, &shape)); - output_tensor_shapes.push_back(shape); - } - } else { - for (int i = 0; i < ctx->num_outputs(); ++i) { - output_tensor_shapes.push_back(compilation_result->outputs[i].shape); - } - } - - // Copy XLA results to the OpOutputList. - int output_num = 0; - for (int i = 0, end = ctx->num_outputs(); i < end; ++i) { - const TensorShape& shape = output_tensor_shapes[i]; - const DataType& type = compilation_result->outputs[i].type; - VLOG(2) << "Populating output for retval " << i << " shape " - << shape.DebugString() << " type " << DataTypeString(type); - - if (compilation_result->outputs[i].is_constant) { - TF_RETURN_IF_ERROR(SetOutputForConstant( - ctx, /*requires_copy_to_device=*/stream != nullptr, - compilation_result, i)); - } else if (type == DT_RESOURCE) { - int input_index = - compilation_result->outputs[i].input_index - missing_ctx_input_prefix; - TF_RET_CHECK(input_index >= 0 && input_index < ctx->num_inputs()) - << "Invalid input for outputs " << i << ": " << input_index; - ctx->set_output(i, ctx->input(input_index)); - } else { - TF_ASSIGN_OR_RETURN( - Tensor output_tensor, - GetOrCreateTensorForOutput( - output, output_num, ctx, missing_ctx_input_prefix, - input_output_alias, compilation_result->input_mapping, - resource_vars, ctx->expected_output_dtype(i), shape, allocator, - allocate_xla_tensors_, stream, use_multiple_streams_, - definition_event)); - ctx->set_output(i, output_tensor); - ++output_num; - } - } - - // input_index -> index into variable_infos. - absl::flat_hash_map variable_info_lookup; - for (int i = 0; i < variable_infos.size(); i++) { - variable_info_lookup.emplace(variable_infos[i].index(), i); - } - - // Apply variable updates, if any. - for (int i = 0, end = compilation_result->resource_updates.size(); i < end; - ++i) { - const XlaCompiler::ResourceUpdate& write = - compilation_result->resource_updates[i]; - int actual_input_index = write.input_index - missing_ctx_input_prefix; - CHECK_GE(actual_input_index, 0); - CHECK_LT(actual_input_index, ctx->num_inputs()); - Var* var = variable_infos[variable_info_lookup[actual_input_index]].var(); - CHECK(var); - - VLOG(2) << "Updating variable #" << i - << " at input index: " << actual_input_index << " with shape " - << write.shape.DebugString() << "; variable tensor has shape: " - << var->tensor()->shape().DebugString(); - - if (var->is_initialized && var->tensor()->dtype() != write.type) { - return absl::InternalError("Mismatched type in variable write"); - } - - TF_ASSIGN_OR_RETURN( - Tensor output_tensor, - GetOrCreateTensorForOutput(output, output_num, ctx, - missing_ctx_input_prefix, input_output_alias, - compilation_result->input_mapping, - resource_vars, write.type, write.shape, - allocator, allocate_xla_tensors_, stream, - use_multiple_streams_, definition_event)); - var->is_initialized |= write.modified; - *var->tensor() = output_tensor; - ++output_num; - } - return absl::OkStatus(); -} - absl::StatusOr> XlaComputationLaunchContext::BuildXlaCompilerArguments( absl::Span must_be_constant_idxs, @@ -815,7 +466,8 @@ absl::Status PopulateCtxOutputsFromPjRtExecutableOutputs( xla::ExecuteOptions GetPjRtExecuteOptions( const DeviceType& device_type, - absl::flat_hash_set non_donatable_input_indices) { + absl::flat_hash_set non_donatable_input_indices, + xla::ExecuteContext* execute_context) { xla::ExecuteOptions options; // Hardcode run id to always be one: TF distributed strategy // differentiates between subsequent runs using dependency edges. This @@ -834,6 +486,7 @@ xla::ExecuteOptions GetPjRtExecuteOptions( // tensorflow/compiler/xla/pjrt/pjrt_c_api_client.cc;l=923-927;rcl=519286815 options.use_major_to_minor_data_layout_for_callbacks = true; options.non_donatable_input_indices = std::move(non_donatable_input_indices); + options.context = execute_context; return options; } @@ -867,14 +520,17 @@ absl::Status RunPjRtExecutable( OpKernelContext* ctx) { const CompositeDevice::AcceleratorDeviceInfo* accelerator_device_info = ctx->device()->tensorflow_accelerator_device_info(); - const bool use_pjrt_tensor_buffer = - (accelerator_device_info == nullptr) - ? true - : accelerator_device_info->use_pjrt_tensor_buffer; - const DeviceType& device_type = GetDeviceType(ctx); - const int pjrt_device_id = - tsl::GetDeviceIdFromDeviceParsedName(ctx->device()->parsed_name()); + const bool use_pjrt_tensor_buffer = + device_type == DEVICE_GPU || accelerator_device_info == nullptr || + accelerator_device_info->use_pjrt_tensor_buffer; + int pjrt_device_id = 0; + if (device_type == DEVICE_CPU || device_type == DEVICE_XLA_CPU) { + pjrt_device_id = 0; + } else { + pjrt_device_id = + tsl::GetDeviceIdFromDeviceParsedName(ctx->device()->parsed_name()); + } TF_ASSIGN_OR_RETURN( xla::PjRtDevice * device, pjrt_client->LookupAddressableDevice(xla::LocalDeviceId(pjrt_device_id))); @@ -898,12 +554,17 @@ absl::Status RunPjRtExecutable( executable->FingerprintExecutable()); device_selector_resource->selector()->Enqueue(pjrt_device_id, fingerprint); } + const Eigen::ThreadPoolDevice* intra_op_thread_pool = nullptr; + if (ctx != nullptr && ctx->device() != nullptr && + ctx->device()->has_eigen_cpu_device()) { + intra_op_thread_pool = ctx->device()->eigen_cpu_device(); + } TF_ASSIGN_OR_RETURN( std::vector> execute_outputs, RunPjRtExecutable(num_missing_prefix_ctx_inputs, inputs, variable_snapshots, updated_variables, device_type, use_pjrt_tensor_buffer, compilation_result, device, - pjrt_client, executable)); + pjrt_client, executable, intra_op_thread_pool)); if (device_selector_resource != nullptr) { device_selector_resource->selector()->Completed(pjrt_device_id, /*had_error=*/false); @@ -922,7 +583,8 @@ absl::StatusOr>> RunPjRtExecutable( const DeviceType& device_type, bool use_pjrt_tensor_buffer, const XlaCompiler::CompilationResult& compilation_result, xla::PjRtDevice* device, xla::PjRtClient* pjrt_client, - xla::PjRtLoadedExecutable* executable) { + xla::PjRtLoadedExecutable* executable, + const Eigen::ThreadPoolDevice* intra_op_thread_pool) { std::vector executable_args; executable_args.reserve(compilation_result.input_mapping.size()); std::vector> owned_executable_args; @@ -935,32 +597,37 @@ absl::StatusOr>> RunPjRtExecutable( std::vector> execute_outputs; std::optional> future; + + std::unique_ptr cpu_execute_context; + if (pjrt_client->platform_id() == xla::CpuId() && + intra_op_thread_pool != nullptr) { + cpu_execute_context = std::make_unique(); + cpu_execute_context->set_intra_op_thread_pool(intra_op_thread_pool); + } + xla::ExecuteOptions execute_options = + GetPjRtExecuteOptions(device_type, std::move(non_donatable_input_indices), + cpu_execute_context.get()); + if (executable->num_replicas() != 1 || executable->num_partitions() != 1) { - TF_ASSIGN_OR_RETURN( - execute_outputs, - executable->ExecuteSharded( - executable_args, device, - GetPjRtExecuteOptions(device_type, - std::move(non_donatable_input_indices)), - future)); + TF_ASSIGN_OR_RETURN(execute_outputs, + executable->ExecuteSharded(executable_args, device, + execute_options, future)); } else { - TF_ASSIGN_OR_RETURN( - execute_outputs, - executable->ExecutePortable( - executable_args, device, - GetPjRtExecuteOptions(device_type, - std::move(non_donatable_input_indices)), - future)); + TF_ASSIGN_OR_RETURN(execute_outputs, + executable->ExecutePortable(executable_args, device, + execute_options, future)); } - // We need to ensure the PjRtBuffers owned by `owned_executable_args` live - // until execution is complete. We do this by capturing - // `owned_executable_args` by move, so it is owned by the lambda - // that is executed when the future returned by ExecutePortable/ExecuteSharded - // is ready i.e. when the execution is complete. - if (!owned_executable_args.empty() && future.has_value()) { - future->OnReady([owned_executable_args = - std::move(owned_executable_args)](absl::Status s) {}); + // We need to ensure the PjRtBuffers owned by `owned_executable_args` and + // the `cpu_execute_context` live until execution is complete. We do this by + // capturing them by move, so they are owned by the lambda that is executed + // when the future returned by ExecutePortable/ExecuteSharded is ready + // i.e. when the execution is complete. + if ((!owned_executable_args.empty() || cpu_execute_context != nullptr) && + future.has_value()) { + future->OnReady([owned_executable_args = std::move(owned_executable_args), + cpu_execute_context = + std::move(cpu_execute_context)](absl::Status s) {}); } return execute_outputs; diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 57ed70bb565fa8..14ab1968e0b947 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -18,24 +18,35 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ #define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ -#include +#include +#include #include -#include #include #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" #include "tensorflow/compiler/jit/variable_info.h" -#include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "xla/client/local_client.h" +#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/service/maybe_owning_device_memory.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/service/executable.h" #include "xla/service/shaped_buffer.h" -#include "xla/stream_executor/device_memory_allocator.h" +#include "xla/stream_executor/device_address_allocator.h" #include "tensorflow/core/framework/allocation_description.pb.h" +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/device.h" +#include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/framework/types.h" + +namespace Eigen { +struct ThreadPoolDevice; +} // namespace Eigen namespace tensorflow { @@ -108,7 +119,8 @@ absl::Status PopulateCtxOutputsFromPjRtExecutableOutputs( // Returns the options used for executing a PjRtLoadedExecutable. xla::ExecuteOptions GetPjRtExecuteOptions( const DeviceType& device_type, - absl::flat_hash_set non_donatable_input_indices); + absl::flat_hash_set non_donatable_input_indices, + xla::ExecuteContext* execute_context = nullptr); // Returns the device ordinal from the parsed name of the device. int GetDeviceOrdinal(const DeviceBase* device); @@ -152,25 +164,13 @@ absl::StatusOr>> RunPjRtExecutable( const DeviceType& device_type, bool use_pjrt_tensor_buffer, const XlaCompiler::CompilationResult& compilation_result, xla::PjRtDevice* device, xla::PjRtClient* pjrt_client, - xla::PjRtLoadedExecutable* executable); + xla::PjRtLoadedExecutable* executable, + const Eigen::ThreadPoolDevice* intra_op_thread_pool = nullptr); // Helper class to perform the marshalling of TensorFlow inputs and outputs to // ShapedBuffers suitable for passing to an XLA computation. class XlaComputationLaunchContext { public: - // Create a new launch context. 'allocate_xla_tensors' is true if allocated - // output tensors and variables are always XlaTensors. If false they are - // assumed to be "normal" device pointers. - // If 'use_multiple_streams' is true, tensors may be defined and used on - // multiple streams and so se::Events must be defined and waited for. If - // 'use_multiple_streams' is true, 'allocate_xla_tensors' must also be true - // because we track inter-stream dependencies through events inside XlaTensor - // objects. - XlaComputationLaunchContext( - xla::LocalClient* client, - stream_executor::DeviceAddressAllocator* xla_allocator, - int device_ordinal, bool allocate_xla_tensors, bool use_multiple_streams); - // Builds a XlaCompiler::Argument vector from the arguments to an XlaLaunch // op. // Precondition: variables in `variable_args` are locked. @@ -179,45 +179,6 @@ class XlaComputationLaunchContext { absl::Span inputs, absl::Span variable_args, Device* device); - - // Add all inputs within `ctx` as XLA arguments (returned by arguments()). - // `variables` is a map from TensorFlow argument number to resource variable. - // - // Assumes that the first `missing_ctx_input_prefix` inputs to the kernel are - // missing and adjusts input indices accordingly. All elements in kernel's - // input_mapping must be greater than or equal to `missing_ctx_input_prefix` - // (in other words, no inputs actually required by the kernel can be missing). - absl::StatusOr> PopulateInputs( - OpKernelContext* ctx, - const XlaCompiler::CompilationResult* compilation_result, - const absl::flat_hash_map& resource_vars, - int missing_ctx_input_prefix, - const xla::HloInputOutputAliasConfig& input_output_alias); - - // Given the XLA output in `output`, populate all outputs of `ctx`. Also - // writes out the resource variable updates. - // - // Updates to all resource variables are written in a single atomic operation. - // This models *->Write dependencies between resource variable operations. - // See jit/resource_operation_safety_analysis for details. - // - // - // Assumes that the first `missing_ctx_input_prefix` inputs to the - // compilation_result are missing and adjusts input indices accordingly. - absl::Status PopulateOutputs( - OpKernelContext* ctx, - const XlaCompiler::CompilationResult* compilation_result, - xla::ScopedShapedBuffer output, int missing_ctx_input_prefix, - absl::Span variable_infos, - const xla::HloInputOutputAliasConfig& input_output_alias, - const absl::flat_hash_map& resource_vars); - - private: - xla::LocalClient* client_; - stream_executor::DeviceAddressAllocator* xla_allocator_; - bool allocate_xla_tensors_; - bool use_multiple_streams_; - int device_ordinal_; }; // A simple TensorBuffer implementation that allows us to create Tensors that diff --git a/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc b/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc index e3f32f8403379a..dd7c4f6090ba6a 100644 --- a/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc +++ b/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc @@ -130,7 +130,6 @@ class PjRtExecutionUtilGpuTest : public OpsTestBase { delete tensor; } tensors_.clear(); - device_context_->Unref(); core::ScopedUnref device_compiler_ref(device_compiler_); core::ScopedUnref profiler_ref(profiler_); } @@ -224,8 +223,8 @@ TEST_F(PjRtExecutionUtilGpuTest, PreparePjRtExecutableArguments) { std::vector input_mapping{3, 4}; std::vector variables; - auto pjrt_device = - pjrt_client_->LookupAddressableDevice(device_->parsed_name().id); + auto pjrt_device = pjrt_client_->LookupAddressableDevice( + xla::LocalDeviceId(device_->parsed_name().id)); TF_EXPECT_OK(pjrt_device.status()); std::vector exec_args; diff --git a/tensorflow/compiler/jit/xla_launch_util_test.cc b/tensorflow/compiler/jit/xla_launch_util_test.cc index 22ce12448ce347..d42a0a5f362019 100644 --- a/tensorflow/compiler/jit/xla_launch_util_test.cc +++ b/tensorflow/compiler/jit/xla_launch_util_test.cc @@ -96,8 +96,6 @@ class PjRtExecutionUtilTest : public OpsTestBase { // Add and setup the XLA_CPU device. auto device_type = DeviceType(DEVICE_XLA_CPU); - rollout_config.AllowForDeviceInXlaLaunch(device_type); - rollout_config.AllowForDeviceInXlaCompileOnDemand(device_type); auto jit_device_type = DeviceType(DEVICE_CPU_XLA_JIT); auto device = diff --git a/tensorflow/compiler/mlir/lite/schema/schema.fbs b/tensorflow/compiler/mlir/lite/schema/schema.fbs index faeffadb2831eb..b51dce02b63040 100644 --- a/tensorflow/compiler/mlir/lite/schema/schema.fbs +++ b/tensorflow/compiler/mlir/lite/schema/schema.fbs @@ -88,10 +88,52 @@ table BlockwiseQuantization { // index to the zero point tensor. If zero_points is -1, the zero point is // assumed to be 0, following the convention of optional tensors in tflite. zero_points: int; - // The block size of the tensor. + + // Superseded by block_shape, which is the general form. Only expresses + // blocking along the last dimension: block_size = n is equivalent to + // block_shape = [1, ..., 1, n]. Retained so that models written before + // block_shape existed remain readable. block_size: int; + + // Note: This field is not supported yet. + // The N-D block shape, one entry per dimension of the quantized tensor. + // block_shape[d] is the extent of a single block along dimension d, so the + // element at index i is dequantized with the parameters at block coordinate + // floor(i / block_shape), and the scale and zero point tensors have shape + // ceil(tensor_shape / block_shape). + // + // Along each axis d there are ceil(tensor_shape[d] / block_shape[d]) + // distinct parameters. The two extremes are: + // * block_shape[d] == tensor_shape[d]: a single block spans the axis, so + // one parameter is shared across all of d. + // * block_shape[d] == 1: the finest granularity, a distinct parameter for + // every index along d. + // + // block_shape always has one entry per dimension of the tensor, so every + // existing granularity is expressible. For a tensor of shape S: + // * per-tensor: block_shape[d] = S[d] for every d. + // * per-channel with quantized_dimension = c: block_shape[c] = 1, and + // block_shape[d] = S[d] for every other d. + // * MultiAxisQuantization with quantized_dimensions = C and block_size 0: + // block_shape[c] = 1 for every c in C, and block_shape[d] = S[d] for + // every other d. + // * MultiAxisQuantization on a 3D tensor [E, N, K] with + // quantized_dimensions = [0, 1] and block_size = b on the innermost + // dimension: block_shape = [1, 1, b]. + // + // A scale or zero point tensor may also have extent 1 on an axis where + // ceil(tensor_shape / block_shape) is greater than 1, in which case it + // broadcasts to the number of blocks on that axis. + // + // If block_shape is non-empty, block_size is ignored. + block_shape:[int]; } +// Parameters for multi-axis quantization. +// +// Note: Superseded by BlockwiseQuantization with block_shape, which provides +// a unified N-D representation covering per-channel, multi-axis, and blockwise +// schemes. table MultiAxisQuantization { // index to the scale tensor, the tensor can be found in tensors array in // subgraph. @@ -103,6 +145,8 @@ table MultiAxisQuantization { // The block size of the tensor if it's block-wise quantized. // For per-channel, 0. + // + // Superseded by BlockwiseQuantization.block_shape. block_size: int; // The quantized dimensions, same as the quantization_dimension field in diff --git a/tensorflow/compiler/mlir/lite/schema/schema_generated.h b/tensorflow/compiler/mlir/lite/schema/schema_generated.h index e5c3bd4f98d378..cb9f9a479665f9 100755 --- a/tensorflow/compiler/mlir/lite/schema/schema_generated.h +++ b/tensorflow/compiler/mlir/lite/schema/schema_generated.h @@ -5235,6 +5235,7 @@ struct BlockwiseQuantizationT : public ::flatbuffers::NativeTable { int32_t scales = 0; int32_t zero_points = 0; int32_t block_size = 0; + std::vector block_shape{}; }; struct BlockwiseQuantization FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table { @@ -5243,7 +5244,8 @@ struct BlockwiseQuantization FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE { VT_SCALES = 4, VT_ZERO_POINTS = 6, - VT_BLOCK_SIZE = 8 + VT_BLOCK_SIZE = 8, + VT_BLOCK_SHAPE = 10 }; int32_t scales() const { return GetField(VT_SCALES, 0); @@ -5254,12 +5256,17 @@ struct BlockwiseQuantization FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta int32_t block_size() const { return GetField(VT_BLOCK_SIZE, 0); } + const ::flatbuffers::Vector *block_shape() const { + return GetPointer *>(VT_BLOCK_SHAPE); + } template bool Verify(::flatbuffers::VerifierTemplate &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_SCALES, 4) && VerifyField(verifier, VT_ZERO_POINTS, 4) && VerifyField(verifier, VT_BLOCK_SIZE, 4) && + VerifyOffset(verifier, VT_BLOCK_SHAPE) && + verifier.VerifyVector(block_shape()) && verifier.EndTable(); } BlockwiseQuantizationT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const; @@ -5280,6 +5287,9 @@ struct BlockwiseQuantizationBuilder { void add_block_size(int32_t block_size) { fbb_.AddElement(BlockwiseQuantization::VT_BLOCK_SIZE, block_size, 0); } + void add_block_shape(::flatbuffers::Offset<::flatbuffers::Vector> block_shape) { + fbb_.AddOffset(BlockwiseQuantization::VT_BLOCK_SHAPE, block_shape); + } explicit BlockwiseQuantizationBuilder(::flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -5295,14 +5305,31 @@ inline ::flatbuffers::Offset CreateBlockwiseQuantization( ::flatbuffers::FlatBufferBuilder &_fbb, int32_t scales = 0, int32_t zero_points = 0, - int32_t block_size = 0) { + int32_t block_size = 0, + ::flatbuffers::Offset<::flatbuffers::Vector> block_shape = 0) { BlockwiseQuantizationBuilder builder_(_fbb); + builder_.add_block_shape(block_shape); builder_.add_block_size(block_size); builder_.add_zero_points(zero_points); builder_.add_scales(scales); return builder_.Finish(); } +inline ::flatbuffers::Offset CreateBlockwiseQuantizationDirect( + ::flatbuffers::FlatBufferBuilder &_fbb, + int32_t scales = 0, + int32_t zero_points = 0, + int32_t block_size = 0, + const std::vector *block_shape = nullptr) { + auto block_shape__ = block_shape ? _fbb.CreateVector(*block_shape) : 0; + return tflite::CreateBlockwiseQuantization( + _fbb, + scales, + zero_points, + block_size, + block_shape__); +} + ::flatbuffers::Offset CreateBlockwiseQuantization(::flatbuffers::FlatBufferBuilder &_fbb, const BlockwiseQuantizationT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr); struct MultiAxisQuantizationT : public ::flatbuffers::NativeTable { @@ -17552,6 +17579,7 @@ inline void BlockwiseQuantization::UnPackTo(BlockwiseQuantizationT *_o, const :: { auto _e = scales(); _o->scales = _e; } { auto _e = zero_points(); _o->zero_points = _e; } { auto _e = block_size(); _o->block_size = _e; } + { auto _e = block_shape(); if (_e) { _o->block_shape.resize(_e->size()); for (::flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->block_shape[_i] = _e->Get(_i); } } else { _o->block_shape.resize(0); } } } inline ::flatbuffers::Offset CreateBlockwiseQuantization(::flatbuffers::FlatBufferBuilder &_fbb, const BlockwiseQuantizationT *_o, const ::flatbuffers::rehasher_function_t *_rehasher) { @@ -17565,11 +17593,13 @@ inline ::flatbuffers::Offset BlockwiseQuantization::Pack( auto _scales = _o->scales; auto _zero_points = _o->zero_points; auto _block_size = _o->block_size; + auto _block_shape = _o->block_shape.size() ? _fbb.CreateVector(_o->block_shape) : 0; return tflite::CreateBlockwiseQuantization( _fbb, _scales, _zero_points, - _block_size); + _block_size, + _block_shape); } inline MultiAxisQuantizationT *MultiAxisQuantization::UnPack(const ::flatbuffers::resolver_function_t *_resolver) const { diff --git a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir index d06ce02b70d299..8c15335815bf4b 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir @@ -275,4 +275,71 @@ func.func @BatchMatmulResourceConstantToFullyConnected(%arg0: tensor<16x1024xf32 // CHECK-NOT: "tfl.batch_matmul" %0 = "tfl.batch_matmul"(%arg0, %weight) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> func.return %0 : tensor<16x128xf32> +} + +// The cases below pin `IsRhsFoldableIntoFullyConnected`, which both directions +// of the batch_matmul/fully_connected rewrite share. They come in pairs: an rhs +// shape that batch_matmul -> fully_connected accepts, and the same shape in the +// fully_connected -> batch_matmul position, which must be refused. If the two +// ever disagree the pair rewrites each other forever, so the refusal half is +// load bearing rather than incidental. + +// A dequantize folds into the filter whether or not its input is a constant. +// CHECK-LABEL: Batchmatmul2FullyconnectedDequantizedNonConstY +func.func @Batchmatmul2FullyconnectedDequantizedNonConstY(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128x!quant.uniform>) -> tensor<16x128xf32> { + %0 = "tfl.dequantize"(%arg1) : (tensor<1024x128x!quant.uniform>) -> tensor<1024x128xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulDequantizedNonConst +func.func @NotFuseTransposeFCRhsToBatchMatmulDequantizedNonConst(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128x!quant.uniform>, %arg2: none) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %0 = "tfl.dequantize"(%arg1) : (tensor<1024x128x!quant.uniform>) -> tensor<1024x128xf32> + %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %2 = "tfl.fully_connected"(%arg0, %1, %arg2) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32> + func.return %2 : tensor<16x128xf32> +} + +// A transpose only rearranges the constant, so the rhs is still foldable. +// CHECK-LABEL: Batchmatmul2FullyconnectedTransposeChainConstY +func.func @Batchmatmul2FullyconnectedTransposeChainConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %weight = arith.constant dense_resource<__elided__> : tensor<128x1024xf32> + %0 = "tfl.transpose"(%weight, %cst) : (tensor<128x1024xf32>, tensor<2xi32>) -> tensor<1024x128xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulTransposeChainConstant +func.func @NotFuseTransposeFCRhsToBatchMatmulTransposeChainConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %weight = arith.constant dense_resource<__elided__> : tensor<128x1024xf32> + %0 = "tfl.transpose"(%weight, %cst) : (tensor<128x1024xf32>, tensor<2xi32>) -> tensor<1024x128xf32> + %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32> + func.return %2 : tensor<16x128xf32> +} + +// The chain is walked to the constant rather than only one op deep, so a +// reshape stacked under a transpose is still recognized. +// CHECK-LABEL: Batchmatmul2FullyconnectedReshapeTransposeChainConstY +func.func @Batchmatmul2FullyconnectedReshapeTransposeChainConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %cst_shape = arith.constant dense<[128, 1024]> : tensor<2xi32> + %weight = arith.constant dense_resource<__elided__> : tensor<64x2048xf32> + %0 = "tfl.reshape"(%weight, %cst_shape) : (tensor<64x2048xf32>, tensor<2xi32>) -> tensor<128x1024xf32> + %1 = "tfl.transpose"(%0, %cst) : (tensor<128x1024xf32>, tensor<2xi32>) -> tensor<1024x128xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %2 = "tfl.batch_matmul"(%arg0, %1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %2 : tensor<16x128xf32> } \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul.td b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul.td index 8fef699049f7a5..9b02acd26618de 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul.td +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul.td @@ -23,7 +23,16 @@ include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td" include "tensorflow/compiler/mlir/lite/utils/utils.td" include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td" -def NotFromFoldableChain : Constraint>; +// Mirrors `IsRhsFoldableIntoFullyConnected` in optimize_batch_matmul_pass.cc, +// which is the single gate shared with +// `ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs`. That pattern turns +// a batch_matmul over such a value into a fully_connected, so a pattern going +// the other way has to refuse the same values or the two rewrite each other +// forever. +def IsRhsFoldableIntoFullyConnectedPred + : CPred<"IsRhsFoldableIntoFullyConnected($0)">; +def NotFoldableIntoFullyConnected + : Constraint>; def IsResultRankEqualTo : Constraint($0.getType().front()).getRank() == " @@ -72,4 +81,4 @@ def FuseTransposeFCRhsToBatchMatmul : Pat< //coverage. This pattern will now work for all the transpose->fc patterns //where the fc has not done imlicit reshape of the output (IsResultRankEqualTo $output, $lhs), - (NotFromFoldableChain $rhs)]>; + (NotFoldableIntoFullyConnected $rhs)]>; diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index 9bcda8d49a1902..f90ba3d574d314 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -44,17 +44,19 @@ namespace mlir { namespace TFL { namespace { -// Checks whether the producer of `value` is part of a chain that can be folded -// into a constant. This includes DequantizeOp, or chains involving ReshapeOp -// and SplitOp originating from a DequantizeOp. -bool NotFromFoldableChain(mlir::Value value) { +// Checks whether `value` is, or can be folded into, a constant. This covers a +// plain constant, the dequantize wrappers a quantized constant weight hides +// behind, and chains of ops that only rearrange one of those. +bool IsFromFoldableChain(mlir::Value value) { mlir::Operation* defining_op = value.getDefiningOp(); while (defining_op) { - if (mlir::isa(defining_op) || + // `QConstOp` carries its value in an attribute but is neither ConstantLike + // nor foldable, so it has to be named rather than left to `m_Constant`. + if (mlir::isa(defining_op) || defining_op->hasTrait() || matchPattern(defining_op, mlir::m_Constant())) { - return false; + return true; } // Look through ops that don't change the constant nature. @@ -69,7 +71,35 @@ bool NotFromFoldableChain(mlir::Value value) { break; } } - return true; + return false; +} + +// Checks whether a batch_matmul whose rhs is `value` can be rewritten into a +// fully_connected, i.e. whether `value` can become that op's filter. +// +// This is the single gate shared by both directions of the +// batch_matmul/fully_connected rewrite, and the two have to agree on it: +// `ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs` converts a matmul +// over such a value into a fully_connected, while the +// `NotFoldableIntoFullyConnected` constraint in `optimize_batch_matmul.td` +// stops `FuseTransposeFCRhsToBatchMatmul` converting it straight back. +bool IsRhsFoldableIntoFullyConnected(mlir::Value value) { + auto ranked_type = mlir::dyn_cast(value.getType()); + if (!ranked_type || ranked_type.getRank() != 2) { + return false; + } + + if (IsFromFoldableChain(value)) { + return true; + } + + // A signed int4/int8 quantized rhs becomes the filter as-is, so unlike the + // case above it does not have to be constant. + auto quantized_type = + mlir::dyn_cast(ranked_type.getElementType()); + return quantized_type != nullptr && quantized_type.isSigned() && + (quantized_type.getStorageTypeIntegralWidth() == 4 || + quantized_type.getStorageTypeIntegralWidth() == 8); } // Converts batch_matmul operation to fully_connected if rhs is a @@ -79,51 +109,11 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(TFL::BatchMatMulOp bmm_op, PatternRewriter& rewriter) const override { - bool is_int_quantized_rank_2_value = false; - if (auto rhs_type = mlir::dyn_cast( - getElementTypeOrSelf(bmm_op.getY().getType()))) { - int64_t rhs_type_rank = bmm_op.getY().getType().getRank(); - bool rhs_i4_or_i8 = rhs_type.getStorageTypeIntegralWidth() == 4 || - rhs_type.getStorageTypeIntegralWidth() == 8; - if (rhs_i4_or_i8 && rhs_type.isSigned() && rhs_type_rank == 2) { - is_int_quantized_rank_2_value = true; - } - } - - ElementsAttr constant = nullptr; - Value rhs = bmm_op.getY(); - // If there is a reshape, look through it. - if (auto reshape = rhs.getDefiningOp()) { - rhs = reshape.getInput(); - } - - ElementsAttr elements_constant; - if (matchPattern(rhs, m_Constant(&elements_constant))) { - constant = elements_constant; - } else if (auto dq = rhs.getDefiningOp()) { - Value q_input = dq.getInput(); - if (matchPattern(q_input, m_Constant(&elements_constant))) { - constant = elements_constant; - } else if (auto q = q_input.getDefiningOp()) { - if (matchPattern(q.getInput(), m_Constant(&elements_constant))) { - constant = elements_constant; - } - } else if (auto pseudo_q = q_input.getDefiningOp()) { - constant = pseudo_q.getValue(); - } else if (auto const_op = q_input.getDefiningOp()) { - constant = const_op.getValue(); - } - } - - const bool is_rank_2_constant = - constant && - mlir::cast(bmm_op.getY().getType()).getRank() == 2; - - if (!is_rank_2_constant && !is_int_quantized_rank_2_value) { + if (!IsRhsFoldableIntoFullyConnected(bmm_op.getY())) { return rewriter.notifyMatchFailure( bmm_op, - "rhs is neither a constant with rank 2 nor int4 quantized nor a " - "dequantized value."); + "rhs is neither a rank 2 value foldable into a constant nor a rank 2 " + "signed int4/int8 quantized value."); } // Create a tfl.transpose op that performs ZX transpose on `input`. diff --git a/tensorflow/compiler/mlir/quantization/common/uniform_quantized_types_test.cc b/tensorflow/compiler/mlir/quantization/common/uniform_quantized_types_test.cc index fc57cac4276af3..ca8d14021a1b37 100644 --- a/tensorflow/compiler/mlir/quantization/common/uniform_quantized_types_test.cc +++ b/tensorflow/compiler/mlir/quantization/common/uniform_quantized_types_test.cc @@ -609,11 +609,10 @@ TEST_F(IsOpFullyQuantizedTest, TrueIfOpFullyQuantized) { auto func_op = module_op->lookupSymbol("fully_quantized_add"); ASSERT_THAT(func_op, NotNull()); - auto add_op_itr = func_op.getBody().op_begin(); - ASSERT_THAT(add_op_itr, - Ne(func_op.getBody().op_end())); + auto add_ops = func_op.getBody().getOps(); + ASSERT_THAT(add_ops.begin(), Ne(add_ops.end())); - EXPECT_TRUE(IsOpFullyQuantized(*add_op_itr)); + EXPECT_TRUE(IsOpFullyQuantized(*add_ops.begin())); } TEST_F(IsOpFullyQuantizedTest, FalseIfOpNotQuantized) { @@ -630,11 +629,10 @@ TEST_F(IsOpFullyQuantizedTest, FalseIfOpNotQuantized) { auto func_op = module_op->lookupSymbol("not_quantized_add"); ASSERT_THAT(func_op, NotNull()); - auto add_op_itr = func_op.getBody().op_begin(); - ASSERT_THAT(add_op_itr, - Ne(func_op.getBody().op_end())); + auto add_ops = func_op.getBody().getOps(); + ASSERT_THAT(add_ops.begin(), Ne(add_ops.end())); - EXPECT_FALSE(IsOpFullyQuantized(*add_op_itr)); + EXPECT_FALSE(IsOpFullyQuantized(*add_ops.begin())); } TEST_F(IsOpFullyQuantizedTest, FalseIfOpPartiallyQuantized) { @@ -651,13 +649,11 @@ TEST_F(IsOpFullyQuantizedTest, FalseIfOpPartiallyQuantized) { auto func_op = module_op->lookupSymbol("quantize"); ASSERT_THAT(func_op, NotNull()); - auto uniform_quantize_op_itr = - func_op.getBody().op_begin(); - ASSERT_THAT( - uniform_quantize_op_itr, - Ne(func_op.getBody().op_end())); + auto uniform_quantize_ops = + func_op.getBody().getOps(); + ASSERT_THAT(uniform_quantize_ops.begin(), Ne(uniform_quantize_ops.end())); - EXPECT_FALSE(IsOpFullyQuantized(*uniform_quantize_op_itr)); + EXPECT_FALSE(IsOpFullyQuantized(*uniform_quantize_ops.begin())); } using IsOpNotQuantizedTest = QuantizationTestBase; @@ -676,11 +672,10 @@ TEST_F(IsOpNotQuantizedTest, TrueIfOpNotQuantized) { auto func_op = module_op->lookupSymbol("not_quantized_add"); ASSERT_THAT(func_op, NotNull()); - auto add_op_itr = func_op.getBody().op_begin(); - ASSERT_THAT(add_op_itr, - Ne(func_op.getBody().op_end())); + auto add_ops = func_op.getBody().getOps(); + ASSERT_THAT(add_ops.begin(), Ne(add_ops.end())); - EXPECT_TRUE(IsOpNotQuantized(*add_op_itr)); + EXPECT_TRUE(IsOpNotQuantized(*add_ops.begin())); } TEST_F(IsOpNotQuantizedTest, FalseIfOpQuantized) { @@ -697,11 +692,10 @@ TEST_F(IsOpNotQuantizedTest, FalseIfOpQuantized) { auto func_op = module_op->lookupSymbol("quantized_add"); ASSERT_THAT(func_op, NotNull()); - auto add_op_itr = func_op.getBody().op_begin(); - ASSERT_THAT(add_op_itr, - Ne(func_op.getBody().op_end())); + auto add_ops = func_op.getBody().getOps(); + ASSERT_THAT(add_ops.begin(), Ne(add_ops.end())); - EXPECT_FALSE(IsOpNotQuantized(*add_op_itr)); + EXPECT_FALSE(IsOpNotQuantized(*add_ops.begin())); } TEST_F(IsOpNotQuantizedTest, FalseIfOpPartiallyQuantized) { @@ -718,15 +712,13 @@ TEST_F(IsOpNotQuantizedTest, FalseIfOpPartiallyQuantized) { auto func_op = module_op->lookupSymbol("quantize"); ASSERT_THAT(func_op, NotNull()); - auto uniform_quantize_op_itr = - func_op.getBody().op_begin(); - ASSERT_THAT( - uniform_quantize_op_itr, - Ne(func_op.getBody().op_end())); + auto uniform_quantize_ops = + func_op.getBody().getOps(); + ASSERT_THAT(uniform_quantize_ops.begin(), Ne(uniform_quantize_ops.end())); // `uniform_quantize` is considered partially quantized because its output is // a quantized tensor whereas its input is not quantized. - EXPECT_FALSE(IsOpNotQuantized(*uniform_quantize_op_itr)); + EXPECT_FALSE(IsOpNotQuantized(*uniform_quantize_ops.begin())); } using UniformQuantizedTypeTest = QuantizationTestBase; diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD index 420737dc98007b..5f0f6ed959ed2d 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD @@ -404,9 +404,7 @@ tf_cc_test( "@llvm-project//llvm:Support", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/client:client_library", - "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", + "@xla//xla/pjrt:pjrt_compiler", ], ) diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc index 2e9d7a8794bb6f..74ed3547fce3dd 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc @@ -21,10 +21,8 @@ limitations under the License. #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h" -#include "xla/client/client_library.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" @@ -46,10 +44,8 @@ class LegalizeTFQuantTest : public Test { mlir_to_hlo_args.rollout_state = tensorflow::ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED; mlir_to_hlo_args.mlir_module = mlir_module_string; - tensorflow::se::Platform* platform = - tensorflow::se::PlatformManager::PlatformWithName("Host").value(); - auto client = - xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); + auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); + TF_ASSERT_OK(compiler.status()); tensorflow::tpu::TPUCompileMetadataProto metadata_proto; // Set up an arg per arg_shape with the specified type. for (int i = 0; i < arg_shapes.size(); ++i) { @@ -69,7 +65,7 @@ class LegalizeTFQuantTest : public Test { mlir_to_hlo_args, metadata_proto, use_tuple_args, /*device_type=*/"XLA_TPU_JIT", custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, client) + &arg_core_mapping, &per_core_arg_shapes, *compiler) .status()); } }; diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type_test.cc b/tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type_test.cc index 4d95f799029225..a54e2ce832ef9c 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type_test.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type_test.cc @@ -40,6 +40,7 @@ TEST(IsLargeFloatTypeTest, scalars) { EXPECT_FALSE(IsLargeFloatType(Float8E4M3B11FNUZType::get(context.get()))); EXPECT_FALSE(IsLargeFloatType(Float8E5M2FNUZType::get(context.get()))); EXPECT_FALSE(IsLargeFloatType(Float8E5M2Type::get(context.get()))); + EXPECT_FALSE(IsLargeFloatType(Float8E8M0FNUType::get(context.get()))); EXPECT_FALSE(IsLargeFloatType(Float16Type::get(context.get()))); EXPECT_FALSE(IsLargeFloatType(BFloat16Type::get(context.get()))); EXPECT_TRUE(IsLargeFloatType(Float32Type::get(context.get()))); @@ -66,6 +67,8 @@ TEST(IsLargeFloatTypeTest, tensors) { RankedTensorType::get({2, 2}, Float8E5M2FNUZType::get(context.get())))); EXPECT_FALSE(IsLargeFloatType( RankedTensorType::get({2, 2}, Float8E5M2Type::get(context.get())))); + EXPECT_FALSE(IsLargeFloatType( + RankedTensorType::get({2, 2}, Float8E8M0FNUType::get(context.get())))); EXPECT_FALSE(IsLargeFloatType( RankedTensorType::get({2, 2}, BFloat16Type::get(context.get())))); EXPECT_TRUE(IsLargeFloatType( @@ -96,6 +99,8 @@ TEST(ToBfloat16TypeTest, scalars) { Float8E5M2FNUZType::get(context.get())); EXPECT_EQ(ToBfloat16Type(Float8E5M2Type::get(context.get())), Float8E5M2Type::get(context.get())); + EXPECT_EQ(ToBfloat16Type(Float8E8M0FNUType::get(context.get())), + Float8E8M0FNUType::get(context.get())); EXPECT_EQ(ToBfloat16Type(Float16Type::get(context.get())), Float16Type::get(context.get())); EXPECT_EQ(ToBfloat16Type(BFloat16Type::get(context.get())), @@ -137,6 +142,10 @@ TEST(ToBfloat16TypeTest, tensors) { EXPECT_EQ(ToBfloat16Type(RankedTensorType::get( {2, 2}, Float8E5M2Type::get(context.get()))), RankedTensorType::get({2, 2}, Float8E5M2Type::get(context.get()))); + EXPECT_EQ( + ToBfloat16Type( + RankedTensorType::get({2, 2}, Float8E8M0FNUType::get(context.get()))), + RankedTensorType::get({2, 2}, Float8E8M0FNUType::get(context.get()))); EXPECT_EQ(ToBfloat16Type( RankedTensorType::get({2, 2}, Float16Type::get(context.get()))), RankedTensorType::get({2, 2}, Float16Type::get(context.get()))); diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h index 64b5d2e141f13d..2d96141aec7549 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h @@ -55,10 +55,9 @@ LogicalResult VerifyTypeRangesAreCompatible(Operation *op, // Fold Arithmetic Op if one of the operands is a constant known to be an // Identity (e.g. X+0, X*1, etc...). For commutative operations fold if // known identity value is either lhs or rhs. -template < - typename OpT, - typename std::enable_if::value>::type * = nullptr> +template ::value>* = nullptr> OpFoldResult IdentityArithmeticOpFolder(OpT arithmetic_op, ArrayRef operands) { auto lhs_type = mlir::cast(arithmetic_op.getX().getType()); @@ -86,7 +85,7 @@ OpFoldResult IdentityArithmeticOpFolder(OpT arithmetic_op, // Check that we have a constant operand on one side (candidate for identity). const bool is_commutative = - (std::is_same::value || std::is_same::value); + (std::is_same_v || std::is_same_v); auto lhs_attr = mlir::dyn_cast_or_null(operands[0]); auto rhs_attr = mlir::dyn_cast_or_null(operands[1]); if (!rhs_attr && !(is_commutative && lhs_attr)) return {}; @@ -94,8 +93,8 @@ OpFoldResult IdentityArithmeticOpFolder(OpT arithmetic_op, // Mul and Div ops have identity value one while AddV2 and SubOp have identity // value zero. const int identity = - (std::is_same::value || std::is_same::value || - std::is_same::value) + (std::is_same_v || std::is_same_v || + std::is_same_v) ? 1 : 0; diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td index 29161518550279..53b78049f72704 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td @@ -15516,7 +15516,7 @@ e.g. Max(segment_ids) should be equal to `num_segments` - 1 for a 1-d segment_id With inconsistent num_segments, the op still runs. only difference is, the output takes the size of num_segments irrespective of size of segment_ids and data. for num_segments less than expected output size, the last elements are ignored -for num_segments more than the expected output size, last elements are assigned +for num_segments more than the expected output size, last elements are assigned smallest possible value for the specific numeric type. For example: @@ -15686,7 +15686,7 @@ e.g. Max(segment_ids) should be equal to `num_segments` - 1 for a 1-d segment_id With inconsistent num_segments, the op still runs. only difference is, the output takes the size of num_segments irrespective of size of segment_ids and data. for num_segments less than expected output size, the last elements are ignored -for num_segments more than the expected output size, last elements are assigned +for num_segments more than the expected output size, last elements are assigned the largest possible value for the specific numeric type. For example: @@ -15792,7 +15792,7 @@ The only difference with SegmentProd is the additional input `num_segments`. This helps in evaluating the output shape in compile time. `num_segments` should be consistent with segment_ids. e.g. Max(segment_ids) - 1 should be equal to `num_segments` for a 1-d segment_ids -With inconsistent num_segments, the op still runs. only difference is, +With inconsistent num_segments, the op still runs. only difference is, the output takes the size of num_segments irrespective of size of segment_ids and data. for num_segments less than expected output size, the last elements are ignored for num_segments more than the expected output size, last elements are assigned 1. @@ -21456,6 +21456,9 @@ values. These keys should be in byte format.}]>:$keys, Arg, [{A list of float tensors containing prediction values.}]>:$predictions_list, Arg:$step, Arg:$timestamp_usec, + Arg:$alternate_keys, StrArrayAttr:$prediction_names, BoolAttr:$training, diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td index 057f540817fe4f..3cf992052b2885 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td @@ -332,6 +332,7 @@ def TF_Float8E5M2Ref : TF_TensorFlowType<"Float8E5M2Ref", "float8e5m2ref">; def TF_Float8E4M3FNUZRef : TF_TensorFlowType<"Float8E4M3FNUZRef", "float8e4m3fnuzref">; def TF_Float8E4M3B11FNUZRef : TF_TensorFlowType<"Float8E4M3B11FNUZRef", "float8e4m3b11fnuzref">; def TF_Float8E5M2FNUZRef : TF_TensorFlowType<"Float8E5M2FNUZRef", "float8e5m2fnuzref">; +def TF_Float8E8M0FNURef : TF_TensorFlowType<"Float8E8M0FNURef", "float8e8m0fnuref">; // Complex reference types def TF_Complex64Ref : TF_TensorFlowType<"Complex64Ref", "complex64ref">; @@ -457,6 +458,7 @@ def TF_Float8E5M2 : AnyTypeOf<[F8E5M2, TF_Float8E5M2Ref], "float8e5m2">; def TF_Float8E4M3FNUZ : AnyTypeOf<[F8E4M3FNUZ, TF_Float8E4M3FNUZRef], "float8e4m3fnuz">; def TF_Float8E4M3B11FNUZ : AnyTypeOf<[F8E4M3B11FNUZ, TF_Float8E4M3B11FNUZRef], "float8e4m3b11fnuz">; def TF_Float8E5M2FNUZ : AnyTypeOf<[F8E5M2FNUZ, TF_Float8E5M2FNUZRef], "float8e5m2fnuz">; +def TF_Float8E8M0FNU : AnyTypeOf<[F8E8M0FNU, TF_Float8E8M0FNURef], "float8e8m0fnu">; def TF_F32OrF64 : AnyTypeOf<[TF_Float32, TF_Float64], "32/64-bit float">; @@ -477,6 +479,7 @@ def TF_Float8E5M2Tensor : TensorOf<[TF_Float8E5M2]>; def TF_Float8E4M3FNUZTensor : TensorOf<[TF_Float8E4M3FNUZ]>; def TF_Float8E4M3B11FNUZTensor : TensorOf<[TF_Float8E4M3B11FNUZ]>; def TF_Float8E5M2FNUZTensor : TensorOf<[TF_Float8E5M2FNUZ]>; +def TF_Float8E8M0FNUTensor : TensorOf<[TF_Float8E8M0FNU]>; //===----------------------------------------------------------------------===// // Complex types (including corresponding reference types) diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc index 78d0ba7d52f182..27e0cf2c5f9fcb 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc @@ -306,8 +306,8 @@ void BatchFunctionOp::eraseArguments(const BitVector& erase_indices) { //===----------------------------------------------------------------------===// template ::value>::type* = nullptr> + std::enable_if_t::value>* = nullptr> static LogicalResult Verify(OpT op) { if (!HasRankAtLeast(op.getX(), 2)) { return op.emitOpError("requires lhs operand to have rank at least two"); @@ -1142,8 +1142,9 @@ std::optional CollectiveGatherV2Op::GetResourceInstanceStr() { // ConcatOp and ConcatV2Op //===----------------------------------------------------------------------===// -template ::value>::type* = nullptr> +template ::value>* = + nullptr> static LogicalResult Verify(OpT op) { // TODO(hinsu): Convert variadic length attributes to derived attributes. Operation::operand_range values = op.getValues(); @@ -1519,9 +1520,8 @@ void ConcatV2Op::getCanonicalizationPatterns(RewritePatternSet& results, //===----------------------------------------------------------------------===// template ::value>::type* = - nullptr> + std::enable_if_t::value>* = nullptr> static LogicalResult Verify(OpT op) { if (!IsOfRankOrUnranked(op.getAxis(), 0)) return op.emitOpError("requires scalar axis operand"); @@ -1773,8 +1773,9 @@ static LogicalResult VerifyConvOpAttributes( // Verifies that, // * Number of input channels is divisible by the number of filter input // channels -template ::value>::type* = nullptr> +template ::value>* = + nullptr> static LogicalResult Verify(OpT op) { int num_spatial_dims = std::is_same() ? 2 : 3; int num_dims = 2 + num_spatial_dims; @@ -1891,8 +1892,8 @@ LogicalResult Conv2DOp::UpdateDataFormat(StringRef data_format) { // Verifies the inferred return type of the given operation. template ::value>::type* = nullptr> + std::enable_if_t::value>* = nullptr> static LogicalResult inferConvReturnTypeComponents( std::optional location, OpT op, ArrayRef explicit_padding, diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc index f739af998ab676..dae6526df2a1e9 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc @@ -719,7 +719,7 @@ int GetLengthOfRange(FloatOrInt start, FloatOrInt limit, FloatOrInt delta) { // Refer to the implementation in // tensorflow/lite/kernels/range.cc. FloatOrInt diff = limit - start; - if (std::is_integral::value) { + if (std::is_integral_v) { return ((std::abs(diff) + std::abs(delta) - 1) / std::abs(delta)); } return std::ceil(std::abs(diff / delta)); diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_types.def b/tensorflow/compiler/mlir/tensorflow/ir/tf_types.def index ddbd29815ec841..e8706d86789e89 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_types.def +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_types.def @@ -73,6 +73,7 @@ HANDLE_TF_REF_TYPE(Float8E5M2Ref, FLOAT8_E5M2_REF, "float8e5m2ref") HANDLE_TF_REF_TYPE(Float8E4M3FNUZRef, FLOAT8_E4M3FNUZ_REF, "float8e4m3fnuzref") HANDLE_TF_REF_TYPE(Float8E4M3B11FNUZRef, FLOAT8_E4M3B11FNUZ_REF, "float8e4m3b11fnuzref") HANDLE_TF_REF_TYPE(Float8E5M2FNUZRef, FLOAT8_E5M2FNUZ_REF, "float8e5m2fnuzref") +HANDLE_TF_REF_TYPE(Float8E8M0FNURef, FLOAT8_E8M0FNU_REF, "float8e8m0fnuref") #ifndef HANDLE_LAST_TF_TYPE #define HANDLE_LAST_TF_TYPE(class, enumerant, name) \ diff --git a/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.cc b/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.cc index 48f1370e8af210..974feb613764ad 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.cc @@ -280,6 +280,7 @@ absl::StatusOr ConvertTensor(const Tensor& input_tensor, case DT_FLOAT8_E4M3FNUZ: case DT_FLOAT8_E4M3B11FNUZ: case DT_FLOAT8_E5M2FNUZ: + case DT_FLOAT8_E8M0FNU: return ConvertTensorOfCustomFloatType(input_tensor, type); case DT_INT4: return ConvertSubByteTensor(input_tensor, type); @@ -723,6 +724,10 @@ absl::Status ConvertToTensorProto(const ElementsAttr attr, TF_RETURN_IF_ERROR(ConvertFloat8ElementsAttr( attr, output->mutable_float8_val())); break; + case DT_FLOAT8_E8M0FNU: + TF_RETURN_IF_ERROR(ConvertFloat8ElementsAttr( + attr, output->mutable_float8_val())); + break; case tensorflow::DT_INT4: TF_RETURN_IF_ERROR(ConvertIntElementsAttr( attr, output->mutable_int_val(), output->mutable_tensor_content())); diff --git a/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor_test.cc b/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor_test.cc index 9a7d1de2a7a6bc..79393cdcfe6491 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor_test.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/convert_tensor_test.cc @@ -158,6 +158,9 @@ TEST_F(ConvertTensorTest, Simple) { ASSERT_NO_FATAL_FAILURE(VerifyConversion( {tsl::float8_e5m2fnuz{1.0}, tsl::float8_e5m2fnuz{-1.0}}, DT_FLOAT8_E5M2FNUZ, mlir::Float8E5M2FNUZType::get(&context))); + ASSERT_NO_FATAL_FAILURE(VerifyConversion( + {tsl::float8_e8m0fnu{1.0}, tsl::float8_e8m0fnu{2.0}}, DT_FLOAT8_E8M0FNU, + mlir::Float8E8M0FNUType::get(&context))); ASSERT_NO_FATAL_FAILURE(VerifyConversion( {static_cast(1), static_cast(-1)}, DT_INT4, diff --git a/tensorflow/compiler/mlir/tensorflow/utils/convert_type.cc b/tensorflow/compiler/mlir/tensorflow/utils/convert_type.cc index e1f93278404153..e22f7c823e703c 100644 --- a/tensorflow/compiler/mlir/tensorflow/utils/convert_type.cc +++ b/tensorflow/compiler/mlir/tensorflow/utils/convert_type.cc @@ -97,6 +97,9 @@ absl::Status ConvertDataType(DataType dtype, Builder builder, Type* type) { case tensorflow::DT_FLOAT8_E5M2FNUZ: *type = builder.getType(); return absl::OkStatus(); + case tensorflow::DT_FLOAT8_E8M0FNU: + *type = builder.getType(); + return absl::OkStatus(); case DT_INT4: // build a **signless** integer type. *type = builder.getIntegerType(4); @@ -151,6 +154,9 @@ absl::Status ConvertScalarTypeToDataType(Type type, DataType* dtype) { } else if (llvm::isa(type)) { *dtype = DT_FLOAT8_E5M2FNUZ; return absl::OkStatus(); + } else if (llvm::isa(type)) { + *dtype = DT_FLOAT8_E8M0FNU; + return absl::OkStatus(); } else if (auto itype = mlir::dyn_cast(type)) { switch (itype.getWidth()) { case 1: diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD index 996123eb6308dc..b67e1381abfab4 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD @@ -170,9 +170,9 @@ cc_library( "@llvm-project//mlir:Support", "@xla//xla:shape_util", "@xla//xla:status_macros", - "@xla//xla/client:compile_only_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/mlir_hlo:hlo_dialect_registration", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/proto:compile_options_proto_cc", "@xla//xla/service:hlo_proto_cc", ], @@ -213,11 +213,9 @@ tf_cc_test( "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/client:client_library", "@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/tsl/lib/core:status_test_util", "@xla//xla/tsl/lib/monitoring:test_utils", "@xla//xla/tsl/platform:errors", diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc index 82c13f5fd71cd5..c751b4e6a4a02f 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h" -#include #include #include #include @@ -49,9 +48,9 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/compile_only_client.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/mlir_hlo/mhlo/IR/register.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -218,13 +217,13 @@ absl::Status CompileTFFunctionWithoutMlir( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client, + xla::PjRtCompiler* compiler, XlaCompiler::CompilationResult* compilation_result) { absl::Status comp_status = CompileTFFunctionToHlo( *function_computation.flib_def, function_computation.graph_def_version, shape_determination_funcs, arg_shapes, device_type, function_computation.guaranteed_constants, *function_computation.function, - metadata, client, arg_core_mapping, per_core_arg_shapes, use_tuple_args, + metadata, compiler, arg_core_mapping, per_core_arg_shapes, use_tuple_args, compilation_result); if (comp_status.ok()) { phase2_bridge_compilation_status->GetCell(kOldBridgeNoMlirSuccess) @@ -246,7 +245,7 @@ absl::Status CompileMLIRTFFunction( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client, + xla::PjRtCompiler* compiler, XlaCompiler::CompilationResult* compilation_result) { if (!mlir::SetTPUInfeedLayout(mlir_module)) return absl::InternalError("Failed to set layouts attribute"); @@ -284,7 +283,7 @@ absl::Status CompileMLIRTFFunction( TF_RETURN_IF_ERROR(CompileTFFunctionToHlo( *flib_def, versions.producer(), shape_determination_funcs, arg_shapes, - device_type, consts, func, metadata, client, arg_core_mapping, + device_type, consts, func, metadata, compiler, arg_core_mapping, per_core_arg_shapes, use_tuple_args, compilation_result)); if (compilation_result != nullptr && @@ -308,13 +307,13 @@ absl::Status CompileMLIRTFFunction( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client, + xla::PjRtCompiler* compiler, XlaCompiler::CompilationResult* compilation_result) { if (mlir_computation.mlir_module_op.has_value()) { return CompileMLIRTFFunction( mlir_computation.mlir_module_op.value(), metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, client, compilation_result); + per_core_arg_shapes, compiler, compilation_result); } mlir::DialectRegistry registry; @@ -326,10 +325,10 @@ absl::Status CompileMLIRTFFunction( TF_RETURN_IF_ERROR(DeserializeMlirModule(mlir_computation.mlir_module, &context, &mlir_module)); - return CompileMLIRTFFunction(*mlir_module, metadata, use_tuple_args, - shape_determination_funcs, arg_shapes, - device_type, arg_core_mapping, - per_core_arg_shapes, client, compilation_result); + return CompileMLIRTFFunction( + *mlir_module, metadata, use_tuple_args, shape_determination_funcs, + arg_shapes, device_type, arg_core_mapping, per_core_arg_shapes, compiler, + compilation_result); } } // namespace @@ -343,7 +342,7 @@ absl::Status CompileTensorflowGraphToHlo( tsl::DeviceType device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client, + xla::PjRtCompiler* compiler, XlaCompiler::CompilationResult* compilation_result) { LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the " "old (non-MLIR) tf2xla bridge"; @@ -360,14 +359,14 @@ absl::Status CompileTensorflowGraphToHlo( TF_RETURN_IF_ERROR(CompileMLIRTFFunction( std::get<0>(computation), metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, client, compilation_result)); + per_core_arg_shapes, compiler, compilation_result)); } else { FunctionToHloArgs function_computation = std::get<1>(computation); TF_RETURN_IF_ERROR(CompileTFFunctionWithoutMlir( function_computation, metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, client, compilation_result)); + per_core_arg_shapes, compiler, compilation_result)); } phase2_bridge_compilation_time->GetCell(kBridgePhase2Config) diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h index 44d4bd83e0e31c..d70cbad1fdd4c8 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h @@ -23,7 +23,7 @@ limitations under the License. #include "absl/types/variant.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/compile_only_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/proto/compile_options.pb.h" #include "xla/shape.h" #include "tensorflow/core/framework/tensor_shape.h" @@ -46,7 +46,7 @@ absl::Status CompileTensorflowGraphToHlo( tsl::DeviceType device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client, + xla::PjRtCompiler* compiler, XlaCompiler::CompilationResult* compilation_result); } // namespace v1 diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc index 7098fb36a7073e..aea90d2b38fd4d 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc @@ -35,12 +35,10 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/client_library.h" #include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h" #include "xla/mlir_hlo/mhlo/IR/register.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/monitoring/test_utils.h" @@ -75,7 +73,6 @@ static constexpr char kCompilationTimeStreamzName[] = static constexpr char kCompilationStatusStreamzName[] = "/tensorflow/core/tf2xla/api/v1/phase2_compilation_status"; -static constexpr char kPlatformName[] = "Host"; constexpr char kEntryFuncName[] = "main"; static constexpr char kMlirModuleStr[] = R"( @@ -99,10 +96,8 @@ absl::StatusOr CompileWithComputation( bool use_serialized_mlir) { XlaCompilationResult compilation_result; - se::Platform* platform = - se::PlatformManager::PlatformWithName(kPlatformName).value(); - auto client = - xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); + TF_ASSIGN_OR_RETURN(auto* compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); bool use_tuple_args = true; std::vector arg_core_mapping; @@ -135,7 +130,7 @@ absl::StatusOr CompileWithComputation( tensorflow::tf2xla::v1::CompileTensorflowGraphToHlo( computation, metadata_proto, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType("XLA_TPU_JIT"), &arg_core_mapping, - &per_core_arg_shapes, client, &compilation_result); + &per_core_arg_shapes, compiler, &compilation_result); if (!compilation_status.ok()) return compilation_status; diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD index 343ab53fbe32d9..5631953b4b6737 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD @@ -56,13 +56,12 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:variant", "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_proto_cc", - "@xla//xla/client:compile_only_client", "@xla//xla/hlo/ir:hlo", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/proto:compile_options_proto_cc", "@xla//xla/tsl/platform:statusor", ], @@ -91,9 +90,7 @@ tf_cc_test( "@com_google_googletest//:gtest", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/client:client_library", - "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/tsl/lib/monitoring:test_utils", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc index c391db870fac9f..2f05202cc3355e 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc @@ -35,8 +35,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/reproducer.pb.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/compile_only_client.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/lib/monitoring/sampler.h" @@ -158,7 +158,7 @@ absl::StatusOr LegalizeMlirToHlo( const std::vector& arg_shapes, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client) { + xla::PjRtCompiler* compiler) { CompilationTimer timer; auto record_time = llvm::scope_exit([&timer] { phase2_bridge_compilation_time->GetCell(kFullBridge) @@ -174,7 +174,7 @@ absl::StatusOr LegalizeMlirToHlo( TF_RETURN_IF_ERROR(tf2xla::v1::CompileTensorflowGraphToHlo( computation, metadata, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping, - per_core_arg_shapes, client, compilation_result.get())); + per_core_arg_shapes, compiler, compilation_result.get())); DumpHloCompilationResult("legalize_tf_fallback.hlo", compilation_result.get()) @@ -185,7 +185,7 @@ absl::StatusOr LegalizeMlirToHlo( auto combined_bridge_status = internal::LegalizeTfToHlo( std::get<0>(computation), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, - per_core_arg_shapes, custom_legalization_passes, client, + per_core_arg_shapes, custom_legalization_passes, compiler, compilation_result.get()); if (combined_bridge_status.ok()) { diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h index 4a33f9760326ca..fc8d4cb8bfae9f 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h @@ -21,14 +21,14 @@ limitations under the License. #include #include "absl/status/status.h" -#include "absl/types/variant.h" -#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h" +#include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/compile_only_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/proto/compile_options.pb.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/shape.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -48,7 +48,7 @@ namespace v2 { // arg_shapes - The shapes of the args. // arg_core_mapping - Which args go on which cores. // per_core_arg_shapes - For each core, the shapes for each argument. -// client - The Xla Compilation client. +// compiler - The Xla PjRtCompiler. absl::StatusOr LegalizeMlirToHlo( const std::variant& computation, const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args, @@ -58,7 +58,7 @@ absl::StatusOr LegalizeMlirToHlo( const std::vector& arg_shapes, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::CompileOnlyClient* client); + xla::PjRtCompiler* compiler = nullptr); }; // namespace v2 }; // namespace tf2xla diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc index 428bf6d0290dcb..7c69394204d2d0 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc @@ -28,10 +28,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h" #include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/client_library.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/monitoring/test_utils.h" #include "xla/tsl/platform/statusor.h" @@ -254,10 +252,8 @@ TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) { /*graph_def_version=*/0, {&guaranteed_constants}}; - se::Platform* cpu_platform = - se::PlatformManager::PlatformWithName("Host").value(); - auto client = - xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform).value(); + auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); + TF_ASSERT_OK(compiler.status()); std::vector arg_shapes; TPUCompileMetadataProto metadata_proto; @@ -272,7 +268,7 @@ TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) { /*device_type=*/"XLA_CPU_JIT", custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, client); + &arg_core_mapping, &per_core_arg_shapes, *compiler); EXPECT_FALSE(compile_result.ok()); } diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD index 2fe5b204522837..6213d9ad025f90 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD @@ -36,7 +36,6 @@ cc_library( "//tensorflow/compiler/tf2xla:xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:lib", - "//tensorflow/core:test", "//tensorflow/core/protobuf:for_core_protos_cc", "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc", "//tensorflow/core/tpu/kernels:tpu_compile_op_support", @@ -44,9 +43,8 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/client:client_library", - "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", ], ) diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc index 17d2490cd8c37e..8470889ba76f30 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc @@ -24,10 +24,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h" #include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/client/client_library.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/protobuf/config.pb.h" @@ -51,10 +49,8 @@ absl::StatusOr CompileMlirModule( mlir_to_hlo_args.rollout_state = rollout_state; mlir_to_hlo_args.mlir_module = mlir_module_str; - TF_ASSIGN_OR_RETURN(se::Platform * platform, - se::PlatformManager::PlatformWithName("Host")); - TF_ASSIGN_OR_RETURN( - auto client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform)); + TF_ASSIGN_OR_RETURN(auto* compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); std::vector arg_shapes; TPUCompileMetadataProto metadata_proto; @@ -72,7 +68,7 @@ absl::StatusOr CompileMlirModule( return LegalizeMlirToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args, device_type, custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, client); + &arg_core_mapping, &per_core_arg_shapes, compiler); } } // namespace testing diff --git a/tensorflow/compiler/mlir/tf2xla/internal/BUILD b/tensorflow/compiler/mlir/tf2xla/internal/BUILD index 04ecb10a60da8d..ce091086781921 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/internal/BUILD @@ -185,8 +185,8 @@ cc_library( "@llvm-project//mlir:Pass", "@stablehlo//:register", "@xla//xla:shape_util", - "@xla//xla/client:compile_only_client", "@xla//xla/mlir_hlo:hlo_dialect_registration", + "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], @@ -206,7 +206,6 @@ tf_cc_test( "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test_main", - "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/protobuf:for_core_protos_cc", "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc", "//tensorflow/core/tpu/kernels:tpu_compile_op_support", @@ -216,11 +215,11 @@ tf_cc_test( "@llvm-project//mlir:Pass", "@stablehlo//:register", "@xla//xla:shape_util", - "@xla//xla/client:client_library", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/tsl/lib/monitoring:cell_reader", "@xla//xla/tsl/platform:errors", + "@xla//xla/tsl/platform:statusor", ], ) diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc index fcfbc0997cae05..5ad143b8104fd6 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc @@ -35,8 +35,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/compile_only_client.h" #include "xla/mlir_hlo/mhlo/IR/register.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/platform/errors.h" @@ -72,7 +72,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) { + xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result) { if (!mlir::SetTPUInfeedLayout(mlir_module_op)) return absl::AbortedError("Failed to set layouts attribute"); @@ -93,7 +93,7 @@ absl::StatusOr LegalizeTfToHlo( absl::Status old_bridge_status = v1::CompileTensorflowGraphToHlo( mlir_to_hlo_args, metadata, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping, - per_core_arg_shapes, client, compilation_result); + per_core_arg_shapes, compiler, compilation_result); TF_RETURN_IF_ERROR(CheckAndIncrementCounter( old_bridge_status, MlirBridgeSecondPhaseMetric::kMlirCombinedOldFailure)); @@ -114,7 +114,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) { + xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result) { LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the " "Combined MLIR Tf2Xla Bridge."; @@ -122,7 +122,7 @@ absl::StatusOr LegalizeTfToHlo( return LegalizeTfToHlo(computation.mlir_module_op.value(), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, per_core_arg_shapes, - custom_legalization_passes, client, + custom_legalization_passes, compiler, compilation_result); } @@ -142,7 +142,7 @@ absl::StatusOr LegalizeTfToHlo( return LegalizeTfToHlo(mlir_module.get(), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, per_core_arg_shapes, - custom_legalization_passes, client, + custom_legalization_passes, compiler, compilation_result); } diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h index 5c6e0f12f02b08..c28b395ed69725 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h @@ -20,7 +20,7 @@ limitations under the License. #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/compile_only_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -40,7 +40,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result); + xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result); }; // namespace internal }; // namespace tf2xla diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc index 324a7f8d346a6f..77323a4ed42e2d 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc @@ -33,15 +33,14 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/client_library.h" #include "xla/mlir_hlo/mhlo/IR/register.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" +#include "xla/tsl/lib/monitoring/cell_reader.h" #include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" -#include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -50,10 +49,10 @@ namespace tensorflow { namespace tf2xla { namespace internal { -using ::tensorflow::monitoring::testing::CellReader; using tpu::MlirToHloArgs; using tpu::ShardingAndIndex; using tpu::TPUCompileMetadataProto; +using tsl::monitoring::testing::CellReader; static constexpr char kMlirLegalizeCount[] = "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count"; @@ -102,10 +101,8 @@ absl::StatusOr CompileMlirModule( mlir_to_hlo_args.mlir_module_op = mlir_module.get(); } - se::Platform* platform = - se::PlatformManager::PlatformWithName("Host").value(); - auto client = - xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); + TF_ASSIGN_OR_RETURN(auto* compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); std::vector arg_shapes = {{1}}; TPUCompileMetadataProto metadata_proto; @@ -123,7 +120,7 @@ absl::StatusOr CompileMlirModule( /*device_type=*/"XLA_TPU_JIT", /*shape_determination_fns=*/{}, arg_shapes, &arg_core_mapping, &per_core_arg_shapes, - custom_legalization_passes, client, + custom_legalization_passes, compiler, compilation_result.get()); } diff --git a/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD b/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD index 6ef93d44495b82..e217bebb7a98e1 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD @@ -28,8 +28,8 @@ cc_library( ], deps = [ "//tensorflow/compiler/mlir/tensorflow", - "//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags", "//tensorflow/compiler/mlir/tensorflow/translate/tools:file_tf_mlir_translate", + "//tensorflow/compiler/mlir/tf2xla/api/v2:mlir_roundtrip_flags", "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph", "//tensorflow/compiler/mlir/tools:translate_cl_options", "//tensorflow/compiler/tf2xla:xla_compiler", @@ -46,13 +46,8 @@ cc_library( "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", "@tsl//tsl/platform:protobuf", - "@xla//xla/client:client_library", - "@xla//xla/client:compile_only_client", - "@xla//xla/service/cpu:cpu_compiler", - "@xla//xla/service/cpu:cpu_transfer_manager", - "@xla//xla/stream_executor:platform_manager", - "@xla//xla/stream_executor/host:host_platform", - "@xla//xla/stream_executor/host:host_platform_id", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", "@xla//xla/tsl/platform:status", ], alwayslink = 1, diff --git a/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc b/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc index 8edb01f90063f6..53b69438337437 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc +++ b/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc @@ -27,16 +27,13 @@ limitations under the License. #include "mlir/Support/LogicalResult.h" // from @llvm-project #include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project #include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h" -#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h" #include "tensorflow/compiler/mlir/tensorflow/translate/tools/file_tf_mlir_translate.h" +#include "tensorflow/compiler/mlir/tf2xla/api/v2/mlir_roundtrip_flags.h" #include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h" #include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "xla/client/client_library.h" -#include "xla/client/compile_only_client.h" -#include "xla/stream_executor/host/host_platform_id.h" -#include "xla/stream_executor/platform_manager.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/tsl/platform/status.h" #include "tensorflow/core/common_runtime/graph_constructor.h" #include "tensorflow/core/framework/function.h" @@ -58,10 +55,10 @@ static constexpr char kMlirToGraphCompilationCheckName[] = static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT"; static Status CompileGraph(tensorflow::Graph* graph, - xla::CompileOnlyClient* client) { - if (!graph || !client) { + xla::PjRtCompiler* compiler) { + if (!graph || !compiler) { return Status(absl::StatusCode::kInvalidArgument, - "Invalid graph or client"); + "Invalid graph or compiler"); } tensorflow::FunctionDefLibrary flib; @@ -70,9 +67,9 @@ static Status CompileGraph(tensorflow::Graph* graph, tensorflow::XlaCompiler::Options options; options.device_type = tensorflow::DeviceType(kArbitraryDeviceName); - options.client = client; + options.compiler = compiler; options.flib_def = flib_def.get(); - tensorflow::XlaCompiler compiler(options); + tensorflow::XlaCompiler xla_compiler(options); std::unique_ptr graph_copy = std::make_unique(tensorflow::OpRegistry::Global()); @@ -80,9 +77,9 @@ static Status CompileGraph(tensorflow::Graph* graph, tensorflow::XlaCompiler::CompileOptions compile_options; tensorflow::XlaCompiler::CompilationResult result; - return compiler.CompileGraph(compile_options, - kMlirToGraphCompilationCheckName, - std::move(graph_copy), {}, &result); + return xla_compiler.CompileGraph(compile_options, + kMlirToGraphCompilationCheckName, + std::move(graph_copy), {}, &result); } static mlir::OwningOpRef GraphdefToMlirTranslateFunction( @@ -123,19 +120,15 @@ static mlir::LogicalResult MlirToGraphTranslateFunction( return mlir::failure(); } - // Use Host platform, which should always exist, to make sure graphs compile. - auto platform = stream_executor::PlatformManager::PlatformWithId( - stream_executor::host::kHostPlatformId); - if (!platform.ok()) { + auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); + if (!compiler.ok()) { return mlir::failure(); } - auto client = - xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform.value()); tensorflow::XlaOpRegistry::RegisterCompilationKernels(); // Verify that the resulting graph can compile. - if (client.ok() && !CompileGraph(graph.get(), client.value()).ok()) { + if (!CompileGraph(graph.get(), compiler.value()).ok()) { return mlir::failure(); } diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD index 83c59a84218e31..7ae680f60668c2 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD @@ -162,6 +162,7 @@ cc_library( "//tensorflow/core/tpu/kernels/xla:host_compute_ops", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -172,15 +173,15 @@ cc_library( "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@tsl//tsl/platform:fingerprint", + "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", - "@xla//xla/client:client_library", "@xla//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo", "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", "@xla//xla/python/ifrt", "@xla//xla/service:device_assignment", "@xla//xla/service:hlo_proto_cc", - "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/platform:errors", ], ) diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc index b1b8aa73d877fd..5a810d4d7c5245 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/log/vlog_is_on.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" @@ -51,14 +52,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/client/client_library.h" #include "xla/hlo/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/python/ifrt/client.h" #include "xla/python/ifrt/layout.h" #include "xla/service/device_assignment.h" #include "xla/shape.h" -#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/strings/proto_serialization.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" @@ -71,6 +70,7 @@ limitations under the License. #include "tensorflow/core/protobuf/tpu/topology.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" #include "tsl/platform/fingerprint.h" +#include "tsl/platform/protobuf.h" namespace tensorflow { namespace ifrt_serving { @@ -244,12 +244,8 @@ absl::StatusOr CompileTfToHlo(const Tf2HloArg& arg) { } VLOG(1) << "device_type: " << device_type; - TF_ASSIGN_OR_RETURN( - auto* platform, - stream_executor::PlatformManager::PlatformWithName("Host")); - TF_ASSIGN_OR_RETURN( - auto* client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform)); - + TF_ASSIGN_OR_RETURN(auto* compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); std::vector arg_shapes; arg_shapes.reserve(arg.input_dtypes_and_shapes.size()); @@ -286,7 +282,7 @@ absl::StatusOr CompileTfToHlo(const Tf2HloArg& arg) { tensorflow::XlaShapeLayoutHelpers::ShapeDeterminationFns( tensorflow::UseNoPreferenceLayoutFn(), arg.shape_representation_fn), - arg_shapes, &arg_core_mapping, &per_core_arg_shapes, client)); + arg_shapes, &arg_core_mapping, &per_core_arg_shapes, compiler)); for (auto arg_shapes_iter = per_core_arg_shapes.begin() + 1; arg_shapes_iter != per_core_arg_shapes.end(); ++arg_shapes_iter) { diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 56b3fc99746d7d..7291da1ad16f9a 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -110,10 +110,6 @@ py_test( ], ) -#LINT.IfChange(combined_tests) -# If you add a new tf_xla_py_strict_test please either add the test file to one of the combined test -# targets that matches in all tags and other settings or add a new combined test target. - tf_xla_combined_py_test( name = "combined_ops_test_a", size = "medium", @@ -276,9 +272,7 @@ tf_xla_combined_py_test( # go/keep-sorted end ], ) -#LINT.ThenChange(:individual_tests) -#LINT.IfChange(individual_tests) tf_xla_py_strict_test( name = "adadelta_test", size = "medium", @@ -2589,13 +2583,12 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core:protos_all_py", + "//tensorflow/core/protobuf:for_core_protos_py_proto", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", - "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2620,13 +2613,12 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core:protos_all_py", + "//tensorflow/core/protobuf:for_core_protos_py_proto", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", - "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2650,13 +2642,12 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core:protos_all_py", + "//tensorflow/core/protobuf:for_core_protos_py_proto", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", - "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2747,23 +2738,21 @@ tf_xla_py_strict_test( ) # copybara:uncomment_begin(google-only) -# tf_xla_py_strict_test( -# name = "xla_dump_to_sponge_test", -# size = "medium", -# srcs = ["xla_dump_to_sponge_test.py"], -# tags = [ -# "optonly", -# ], -# deps = [ -# ":xla_test", -# "//third_party/py/numpy", -# "//tensorflow/python/ops:array_ops", -# "//tensorflow/python/ops:math_ops", -# "//tensorflow/python/platform:test", -# ], -# ) -# copybara:uncomment_end -#LINT.ThenChange(:combined_tests) +tf_xla_py_strict_test( + name = "xla_dump_to_sponge_test", + size = "medium", + srcs = ["xla_dump_to_sponge_test.py"], + tags = [ + "optonly", + ], + deps = [ + ":xla_test", + "//tensorflow/python/ops:array_ops", + "//tensorflow/python/ops:math_ops", + "//tensorflow/python/platform:test", + "//third_party/py/numpy", + ], +) cuda_py_strict_test( name = "xla_device_gpu_test", diff --git a/tensorflow/compiler/tests/const_test.py b/tensorflow/compiler/tests/const_test.py index 423c92f2abb015..dffd2c303d5f9b 100644 --- a/tensorflow/compiler/tests/const_test.py +++ b/tensorflow/compiler/tests/const_test.py @@ -51,6 +51,7 @@ def testConst(self): dtypes.float8_e4m3fnuz, dtypes.float8_e4m3b11fnuz, dtypes.float8_e5m2fnuz, + dtypes.float8_e8m0fnu, } for dtype in types: with self.subTest(dtype=dtype): @@ -61,6 +62,7 @@ def testConst(self): dtypes.uint16, dtypes.uint32, dtypes.uint64, + dtypes.float8_e8m0fnu, ]: values = [0., 1., dtype.min, dtype.max] else: diff --git a/tensorflow/compiler/tests/xla_call_module_test.py b/tensorflow/compiler/tests/xla_call_module_test.py index de4a3ab8018cfc..4e6f56172c9bf7 100644 --- a/tensorflow/compiler/tests/xla_call_module_test.py +++ b/tensorflow/compiler/tests/xla_call_module_test.py @@ -1590,7 +1590,8 @@ def f(x): errors.InvalidArgumentError, 'Cannot deserialize computation: .+', ): - f(x) + with self.test_scope(): + f(x) if __name__ == '__main__': diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 59283a2bbe0d59..513459692a6370 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -186,10 +186,8 @@ cc_library( 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", @@ -201,9 +199,13 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:status", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", - "@xla//xla/client", "@xla//xla/hlo/builder:xla_computation", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", ], ) @@ -241,7 +243,6 @@ cc_library( "@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", ], @@ -557,29 +558,37 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:errors", + "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", + "@xla//xla:executable_run_options", "@xla//xla:shape_util", + "@xla//xla:status_macros", "@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/pjrt:common_pjrt_client", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt:pjrt_executable", + "@xla//xla/pjrt/cpu:cpu_client", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", + "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", + "@xla//xla/service:buffer_assignment", "@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 = [], @@ -662,8 +671,12 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -688,12 +701,16 @@ cc_library( "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/service:device_assignment", "@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", + "@xla//xla/tsl/platform:logging", ] + if_libtpu([ ":xla_tpu_backend_registration", ]) + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), @@ -816,11 +833,18 @@ cc_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:optional", - "@xla//xla/client", + "@com_google_absl//absl/strings", + "@xla//xla:literal", + "@xla//xla:shape_util", "@xla//xla/hlo/builder:value_inference", "@xla//xla/hlo/builder:xla_builder", + "@xla//xla/hlo/builder:xla_computation", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt:pjrt_executable", ], alwayslink = 1, ) @@ -1171,14 +1195,20 @@ tf_cc_test( "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", "@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", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt:pjrt_executable", + "@xla//xla/pjrt/cpu:cpu_client", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", ] + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), ) @@ -1190,20 +1220,22 @@ tf_cc_test( ":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", + "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/client:executable_build_options", "@xla//xla/hlo/testlib:test", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", "@xla//xla/service:compiler", - "@xla//xla/service:platform_util", "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_id", "@xla//xla/stream_executor:platform_manager", + "@xla//xla/tsl/platform:statusor", ], ) @@ -1237,34 +1269,31 @@ tf_cc_test( "//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/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/ir:hlo_sharding", "@xla//xla/hlo/testlib:filecheck", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt:pjrt_executable", + "@xla//xla/pjrt/cpu:cpu_client", + "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", + "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@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", diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index 4d3b87a177a25a..a4fb37e1dcf55c 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -15,42 +15,43 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/graph_compiler.h" -#include -#include +#include +#include +#include +#include #include #include +#include "absl/container/inlined_vector.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" -#include "tensorflow/compiler/tf2xla/literal_util.h" -#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" -#include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/tf2xla/xla_expression.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" -#include "xla/client/client_library.h" +#include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/hlo/builder/xla_builder.h" -#include "tensorflow/core/common_runtime/device.h" -#include "tensorflow/core/common_runtime/executor.h" -#include "tensorflow/core/common_runtime/function.h" -#include "tensorflow/core/common_runtime/graph_constructor.h" -#include "tensorflow/core/common_runtime/graph_optimizer.h" +#include "xla/status_macros.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/logging.h" +#include "xla/tsl/platform/statusor.h" +#include "tensorflow/core/common_runtime/function_utils.h" +#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/graph/algorithm.h" -#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/graph/graph_node_util.h" #include "tensorflow/core/graph/validate.h" -#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" -#include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/monitoring/counter.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/public/version.h" -#include "tensorflow/core/util/dump_graph.h" namespace tensorflow { diff --git a/tensorflow/compiler/tf2xla/kernels/unravel_index_op.cc b/tensorflow/compiler/tf2xla/kernels/unravel_index_op.cc index 06560a59c60ed2..8e3b1c72d2b448 100644 --- a/tensorflow/compiler/tf2xla/kernels/unravel_index_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unravel_index_op.cc @@ -61,14 +61,15 @@ class UnravelIndexOp : public XlaOpKernel { OP_REQUIRES(ctx, TensorShapeUtils::IsVector(indices_shape) || TensorShapeUtils::IsScalar(indices_shape), - errors::InvalidArgument( + absl::InvalidArgumentError(absl::StrCat( "The indices can only be scalar or vector, got \"", - indices_shape.DebugString(), "\"")); + indices_shape.DebugString(), "\""))); const TensorShape dims_shape = ctx->InputShape(1); OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dims_shape), - errors::InvalidArgument("The dims can only be 1-D, got \"", - dims_shape.DebugString(), "\"")); + absl::InvalidArgumentError( + absl::StrCat("The dims can only be 1-D, got \"", + dims_shape.DebugString(), "\""))); std::vector dims; OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &dims)); @@ -77,11 +78,11 @@ class UnravelIndexOp : public XlaOpKernel { int64_t product = 1; for (int64_t i = 0; i < rank; ++i) { OP_REQUIRES(ctx, dims[i] > 0, - errors::InvalidArgument( + absl::InvalidArgumentError(absl::StrCat( "Input dims cannot be zero or negative, got dim = ", - dims[i], " at index ", i)); + dims[i], " at index ", i))); OP_REQUIRES(ctx, product <= std::numeric_limits::max() / dims[i], - errors::InvalidArgument( + absl::InvalidArgumentError( "Input dims product is causing integer overflow")); product *= dims[i]; } @@ -116,9 +117,10 @@ class UnravelIndexOp : public XlaOpKernel { // [rank] stride/dims constants along a new trailing indices // dimension via BroadcastInDim, then combine elementwise. const int64_t num_indices = indices_shape.num_elements(); - OP_REQUIRES(ctx, num_indices > 0, - errors::InvalidArgument("received empty tensor indices: ", - indices_shape.DebugString())); + OP_REQUIRES( + ctx, num_indices > 0, + absl::InvalidArgumentError(absl::StrCat( + "received empty tensor indices: ", indices_shape.DebugString()))); xla::XlaOp indices_bcast = xla::Broadcast(indices, {rank}); xla::XlaOp trailing_bcast = xla::BroadcastInDim(trailing_const, {rank, num_indices}, {0}); diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc index 4634d2500bd458..4abaa081653f8f 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -159,7 +159,7 @@ absl::Status ConvertLoopInvariantsToConst( const NameAttrList& cond_name_attr, std::vector* args, std::vector* compile_time_const_arg_indices, - int* num_compile_time_const_args, xla::Client* client) { + int* num_compile_time_const_args) { std::vector loop_invariants(ctx->num_inputs()); GetLoopInvariants(ctx, body_name_attr, &loop_invariants); @@ -381,7 +381,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES_OK(ctx, ConvertLoopInvariantsToConst( ctx, body_name_attr_, cond_name_attr_, &arguments, &compile_time_const_arg_indices, - &num_compile_time_const_args, compiler->client())); + &num_compile_time_const_args)); } VLOG(1) << "Compiling body"; diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc index 49c0eaad39b1d6..8acadcc451dcb9 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.cc +++ b/tensorflow/compiler/tf2xla/tf2xla.cc @@ -15,38 +15,31 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/tf2xla.h" -#include +#include #include #include -#include #include #include +#include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" #include "tensorflow/compiler/aot/aot_only_var_handle_op.h" #include "tensorflow/compiler/tf2xla/graph_compiler_util.h" -#include "tensorflow/compiler/tf2xla/shape_util.h" -#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_computation.h" -#include "tensorflow/core/common_runtime/function.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/tsl/platform/errors.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/node_def.pb.h" -#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/versions.pb.h" -#include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" -#include "tensorflow/core/graph/node_builder.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/platform/errors.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/dump_graph.h" namespace tensorflow { @@ -56,8 +49,9 @@ namespace { // graph symbolically, with each op building up the XLA HLO. absl::Status ConvertGraphToXla(std::unique_ptr graph, const tf2xla::Config& config, - xla::Client* client, - xla::XlaComputation* computation) { + xla::PjRtCompiler* compiler, + xla::XlaComputation* computation, + xla::PjRtClient* client) { XlaOpRegistry::RegisterCompilationKernels(); for (Node* node : graph->nodes()) { node->set_assigned_device_name( @@ -69,19 +63,20 @@ absl::Status ConvertGraphToXla(std::unique_ptr graph, PopulateXlaArgs(config, &xla_args); // Compile the graph into an XLA computation. XlaCompiler::Options compiler_options; + compiler_options.compiler = compiler; compiler_options.client = client; compiler_options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); compiler_options.flib_def = &graph->flib_def(); compiler_options.graph_def_version = graph->versions().producer(); compiler_options.allow_cpu_custom_calls = true; - XlaCompiler compiler(compiler_options); + XlaCompiler xla_compiler(compiler_options); XlaCompiler::CompilationResult result; XlaCompiler::CompileOptions options; options.alias_resource_update = true; - TF_RETURN_IF_ERROR(compiler.CompileGraph( + TF_RETURN_IF_ERROR(xla_compiler.CompileGraph( options, "tfcompile", std::move(graph), xla_args, &result)); *computation = std::move(*result.computation); @@ -159,13 +154,14 @@ absl::Status ConvertVarHandlesToAotVarHandles(GraphDef* graph_def) { absl::Status ConvertGraphDefToXla(GraphDef graph_def, const tf2xla::Config& config, - xla::Client* client, - xla::XlaComputation* computation) { + xla::PjRtCompiler* compiler, + xla::XlaComputation* computation, + xla::PjRtClient* client) { std::unique_ptr graph; TF_RETURN_IF_ERROR(ConvertVarHandlesToAotVarHandles(&graph_def)); TF_RETURN_IF_ERROR(InitGraph(graph_def, config, &graph)); - TF_RETURN_IF_ERROR( - ConvertGraphToXla(std::move(graph), config, client, computation)); + TF_RETURN_IF_ERROR(ConvertGraphToXla(std::move(graph), config, compiler, + computation, client)); return absl::OkStatus(); } diff --git a/tensorflow/compiler/tf2xla/tf2xla.h b/tensorflow/compiler/tf2xla/tf2xla.h index 095ad49afc6b66..d38b563cc7a2ac 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.h +++ b/tensorflow/compiler/tf2xla/tf2xla.h @@ -18,11 +18,15 @@ limitations under the License. #include "absl/strings/string_view.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" -#include "xla/client/client.h" #include "xla/hlo/builder/xla_computation.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/platform/status.h" +namespace xla { +class PjRtClient; +class PjRtCompiler; +} // namespace xla + namespace tensorflow { // Converts a tensorflow::GraphDef into an xla::XlaComputation. The given @@ -30,12 +34,14 @@ namespace tensorflow { // fetches. Each feed is a positional input argument for the generated // computation, while each fetch is a positional output argument. // -// The computation is built in the context of the given `client`, which may -// subsequently be used to compile or execute the computation. +// The computation is built in the context of the given `compiler` (and optional +// `client`), which may subsequently be used to compile or execute the +// computation. absl::Status ConvertGraphDefToXla(GraphDef graph_def, const tf2xla::Config& config, - xla::Client* client, - xla::XlaComputation* computation); + xla::PjRtCompiler* compiler, + xla::XlaComputation* computation, + xla::PjRtClient* client = nullptr); // Similar to ConvertGraphDefToXla, but uses MLIR and handle debug information. // diff --git a/tensorflow/compiler/tf2xla/tf2xla_test.cc b/tensorflow/compiler/tf2xla/tf2xla_test.cc index 72bd28f2b47a8c..b609deb78d3eb5 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_test.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_test.cc @@ -15,24 +15,26 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/tf2xla.h" +#include #include +#include "absl/status/status.h" +#include "absl/strings/string_view.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" -#include "xla/client/client_library.h" -#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_computation.h" -#include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/pjrt/cpu/cpu_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" +#include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_shape.pb.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/statusor.h" -#include "tensorflow/core/platform/stringpiece.h" #include "tensorflow/core/platform/test.h" #include "tsl/platform/tensor_float_32_utils.h" @@ -109,35 +111,46 @@ tf2xla::Config SumConfig() { return config; } +// NOLINTBEGIN + TEST(ConvertGraphDefToXla, Sum) { GraphDef graph_def = SumGraph(); tf2xla::Config config = SumConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto client, + xla::GetPjRtCpuClient(xla::CpuClientOptions())); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation, + client.get())); // Set up arguments. auto x_literal = xla::LiteralUtil::CreateR0(10); auto y_literal = xla::LiteralUtil::CreateR0(32); - auto x_global_or = client->TransferToServer(x_literal); - auto y_global_or = client->TransferToServer(y_literal); - TF_EXPECT_OK(x_global_or.status()); - TF_EXPECT_OK(y_global_or.status()); - std::unique_ptr x_global = std::move(x_global_or.value()); - std::unique_ptr y_global = std::move(y_global_or.value()); + TF_ASSERT_OK_AND_ASSIGN( + auto x_buf, + client->BufferFromHostLiteral(x_literal, client->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto y_buf, + client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); // Execute and check result. - auto result_or = - client->ExecuteAndTransfer(computation, {x_global.get(), y_global.get()}); - TF_EXPECT_OK(result_or.status()); - xla::Literal result = std::move(result_or.value()); - EXPECT_EQ("(\ns32[] 42\n)", result.ToString()); + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + client->CompileAndLoad(computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto results, + loaded_exec->Execute({{x_buf.get(), y_buf.get()}}, execute_options)); + TF_ASSERT_OK_AND_ASSIGN(auto result, results[0][0]->ToLiteralSync()); + EXPECT_EQ(42, result->Get({})); config.mutable_feed(0)->mutable_id()->set_output_index( 123); /* invalid output_index */ - EXPECT_TRUE(absl::IsInvalidArgument( - ConvertGraphDefToXla(graph_def, config, client, &computation))); + EXPECT_TRUE(absl::IsInvalidArgument(ConvertGraphDefToXla( + graph_def, config, compiler, &computation, client.get()))); } GraphDef EinsumGraph(DataType dtype = DT_FLOAT) { @@ -180,9 +193,10 @@ TEST(ConvertGraphDefToXla, EinsumIsConvertedToDotWithDefaultPrecision) { GraphDef graph_def = EinsumGraph(); tf2xla::Config config = EinsumConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -209,9 +223,10 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = EinsumGraph(); tf2xla::Config config = EinsumConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -238,9 +253,10 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = EinsumGraph(DT_BFLOAT16); tf2xla::Config config = EinsumConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -300,9 +316,10 @@ TEST(ConvertGraphDefToXla, Conv2DIsConvertedToConvolutionWithDefaultPrecision) { GraphDef graph_def = Conv2DGraph(); tf2xla::Config config = Conv2DConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); int num_convolutions = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -329,9 +346,10 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = Conv2DGraph(); tf2xla::Config config = Conv2DConfig(); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); int num_convolutions = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -362,31 +380,42 @@ TEST(ConvertGraphDefToXla, SumWithUnusedArgument) { (*unused->mutable_attr())["dtype"] = TypeAttrValue(DT_INT32); config.add_feed()->mutable_id()->set_node_name("unused"); - xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(auto client, + xla::GetPjRtCpuClient(xla::CpuClientOptions())); + TF_ASSERT_OK_AND_ASSIGN(auto compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation, + client.get())); // Set up arguments. auto x_literal = xla::LiteralUtil::CreateR0(10); auto y_literal = xla::LiteralUtil::CreateR0(32); - auto x_global_or = client->TransferToServer(x_literal); - auto y_global_or = client->TransferToServer(y_literal); - auto unused_global_or = client->TransferToServer(y_literal); - TF_EXPECT_OK(x_global_or.status()); - TF_EXPECT_OK(y_global_or.status()); - TF_EXPECT_OK(unused_global_or.status()); - std::unique_ptr x_global = std::move(x_global_or.value()); - std::unique_ptr y_global = std::move(y_global_or.value()); - std::unique_ptr unused_global = - std::move(unused_global_or.value()); + TF_ASSERT_OK_AND_ASSIGN( + auto x_buf, + client->BufferFromHostLiteral(x_literal, client->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto y_buf, + client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto unused_buf, + client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); // Execute and check result. - auto result_or = client->ExecuteAndTransfer( - computation, {x_global.get(), y_global.get(), unused_global.get()}); - TF_EXPECT_OK(result_or.status()); - xla::Literal result = std::move(result_or.value()); - EXPECT_EQ("(\ns32[] 42\n)", result.ToString()); + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + client->CompileAndLoad(computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto results, + loaded_exec->Execute({{x_buf.get(), y_buf.get(), unused_buf.get()}}, + execute_options)); + TF_ASSERT_OK_AND_ASSIGN(auto result, results[0][0]->ToLiteralSync()); + EXPECT_EQ(42, result->Get({})); } +// NOLINTEND + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/type_util.cc b/tensorflow/compiler/tf2xla/type_util.cc index f0c72ec4a91bab..f9d2f631dc0e1c 100644 --- a/tensorflow/compiler/tf2xla/type_util.cc +++ b/tensorflow/compiler/tf2xla/type_util.cc @@ -87,6 +87,9 @@ absl::Status DataTypeToPrimitiveType(DataType data_type, case tensorflow::DT_FLOAT8_E5M2FNUZ: *type = xla::F8E5M2FNUZ; return absl::OkStatus(); + case tensorflow::DT_FLOAT8_E8M0FNU: + *type = xla::F8E8M0FNU; + return absl::OkStatus(); case tensorflow::DT_FLOAT4_E2M1FN: *type = xla::F4E2M1FN; return absl::OkStatus(); @@ -125,6 +128,7 @@ absl::StatusOr EncodePrimitiveTypeAsDataType( {xla::F8E4M3FNUZ, DT_FLOAT8_E4M3FNUZ}, {xla::F8E4M3B11FNUZ, DT_FLOAT8_E4M3B11FNUZ}, {xla::F8E5M2FNUZ, DT_FLOAT8_E5M2FNUZ}, + {xla::F8E8M0FNU, DT_FLOAT8_E8M0FNU}, {xla::F4E2M1FN, DT_FLOAT4_E2M1FN}, {xla::BF16, DT_BFLOAT16}, {xla::F16, DT_HALF}, diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index b9abd5006a958a..446381a14ede76 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -46,6 +46,8 @@ limitations under the License. #include "xla/hlo/builder/xla_builder.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/hlo/ir/hlo_sharding.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/status_macros.h" #include "xla/stream_executor/device_memory_allocator.h" @@ -185,7 +187,8 @@ class XlaCompiler { // -1 indicates the default device should be used. int device_ordinal = -1; - xla::Client* client = nullptr; + xla::PjRtCompiler* compiler = nullptr; + xla::PjRtClient* client = nullptr; // Function library in which to find function definitions. Must be non-null. const FunctionLibraryDefinition* flib_def = nullptr; @@ -342,7 +345,8 @@ class XlaCompiler { const std::string& host_compute_name, xla::XlaOp handle); const Options& options() const { return options_; } - xla::Client* client() const { return options_.client; } + xla::PjRtCompiler* compiler() const { return options_.compiler; } + xla::PjRtClient* client() const { return options_.client; } FunctionLibraryRuntime* flib_runtime() const { return flib_runtime_; } void PushNodeTokenMapping(); diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 32dfd44f5fba1f..c6982f86eab099 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -37,12 +38,9 @@ limitations under the License. #include "tensorflow/cc/ops/data_flow_ops.h" #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/functional_ops.h" -#include "tensorflow/cc/ops/list_ops.h" #include "tensorflow/cc/ops/math_ops.h" #include "tensorflow/cc/ops/resource_variable_ops.h" -#include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/tf2xla/layout_util.h" -#include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" @@ -51,30 +49,31 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/array.h" -#include "xla/client/client.h" -#include "xla/client/client_library.h" -#include "xla/client/local_client.h" #include "xla/debug_options_flags.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/hlo/builder/xla_computation.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/testlib/filecheck.h" #include "xla/layout_util.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/pjrt/cpu/cpu_client.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" +#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_module_util.h" -#include "xla/service/hlo_proto_util.h" -#include "xla/service/service.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/status_macros.h" #include "xla/tests/literal_test_util.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" -#include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/graph_constructor.h" #include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/fake_input.h" @@ -98,19 +97,17 @@ limitations under the License. #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/node_builder.h" #include "tensorflow/core/kernels/ops_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/test.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/public/version.h" -#include "tsl/platform/statusor.h" namespace tensorflow { class XlaCompilerTest : public ::testing::Test { protected: void SetUp() override { - client_ = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(pjrt_client_, + xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); + TF_ASSERT_OK_AND_ASSIGN(pjrt_compiler_, + xla::GetDefaultPjRtCompiler(xla::CpuName())); XlaOpRegistry::RegisterCompilationKernels(); @@ -121,7 +118,8 @@ class XlaCompilerTest : public ::testing::Test { XlaCompiler::Options DefaultOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - options.client = client_; + options.compiler = pjrt_compiler_; + options.client = pjrt_client_.get(); options.flib_def = flib_def_.get(); return options; } @@ -130,7 +128,8 @@ class XlaCompilerTest : public ::testing::Test { return compiler->local_flib_def_.get(); } - xla::Client* client_; + std::unique_ptr pjrt_client_; + xla::PjRtCompiler* pjrt_compiler_; std::unique_ptr flib_def_; }; @@ -237,7 +236,11 @@ TEST_F(XlaCompilerTest, EmptyReturnValues) { std::move(graph), /*args=*/{}, &result)); - TF_ASSERT_OK(client_->Execute(*result.computation, {}).status()); + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK(loaded_exec->Execute({{}}, execute_options).status()); } // Tests compilation and execution of a graph that adds two tensors. @@ -270,20 +273,28 @@ TEST_F(XlaCompilerTest, Simple) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client_ - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN( + auto param0_buffer, + pjrt_client_->BufferFromHostLiteral(param0_literal, + pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto param1_buffer, + pjrt_client_->BufferFromHostLiteral(param1_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR1({4, 143}); - xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal)); } absl::StatusOr> LoadModuleFromHloProto( @@ -368,18 +379,27 @@ TEST_F(XlaCompilerTest, OutOfOrderGraph) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client_ - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); - - EXPECT_TRUE(xla::LiteralTestUtil::Equal(param0_literal, actual_literal)); + TF_ASSERT_OK_AND_ASSIGN( + auto param0_buffer, + pjrt_client_->BufferFromHostLiteral(param0_literal, + pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto param1_buffer, + pjrt_client_->BufferFromHostLiteral(param1_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); + + EXPECT_TRUE(xla::LiteralTestUtil::Equal(param0_literal, *actual_literal)); } // Tests that the compiler can correctly propagate the layout assigned by @@ -420,7 +440,7 @@ TEST_F(XlaCompilerTest, HonorShapeRepresentationFnForUnwrittenResource) { args, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithDenseLayout(xla::S32, {2, 3}, {0, 1}); - // Check that the return shapes are correctly tranposed. + // Check that the return shapes are correctly transposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed})); } @@ -522,7 +542,9 @@ TEST_F(XlaCompilerTest, HonorShapeRepresentationFnForRetVal) { // Check that the return shapes are correctly tranposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed, transposed})); - EXPECT_EQ(result.computation->GetProgramShape().value().result(), + TF_ASSERT_OK_AND_ASSIGN(auto program_shape, + result.computation->GetProgramShape()); + EXPECT_EQ(program_shape.result(), xla::ShapeUtil::MakeTupleShape({transposed, transposed})); } @@ -562,7 +584,7 @@ TEST_F(XlaCompilerTest, TransposeVariables) { std::move(graph), args, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithDenseLayout(xla::S32, {2, 3}, {1, 0}); - // Check that the return shapes are correctly tranposed. + // Check that the return shapes are correctly transposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed, transposed})); } @@ -582,7 +604,7 @@ TEST_F(XlaCompilerTest, UnrankedFakeParam) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "compile", std::move(graph), {}, &result)); - // Check that the return shapes are correctly tranposed. + // Check that the return shapes are correctly transposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape( {xla::ShapeUtil::MakeShape(xla::S32, {0})})); @@ -711,18 +733,28 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - - std::unique_ptr actual = - client_->Execute(*result.computation, {param0_data.get()}).value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN( + auto param0_buffer, + pjrt_client_->BufferFromHostLiteral(param0_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, pjrt_client_->CompileAndLoad(*result.computation, + xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get()}}, execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, + actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, + actual[0][1]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR0(7); xla::Literal expected1 = xla::LiteralUtil::CreateR1({-7, -42}); - xla::Literal expected = - xla::LiteralUtil::MakeTuple({&expected0, &expected1}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); } } @@ -937,25 +969,58 @@ TEST_F(XlaCompilerTest, CanPassTensorArraysToAndFromComputation) { update.tensor_array_gradients_accessed); // Tests that the generated computation works. + // PjRtCpuClient does not support nested tuples, so wrap the computation to + // flatten the input and output tuples. + xla::XlaBuilder b("flatten_wrapper"); + xla::XlaOp p0 = + xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::S32, {2}), "base"); + xla::XlaOp p1 = + xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::S32, {2}), "grad2"); + xla::XlaOp call = + xla::Call(&b, *result.computation, {xla::Tuple(&b, {p0, p1})}); + xla::XlaOp read_op = xla::GetTupleElement(call, 0); + xla::XlaOp resource = xla::GetTupleElement(call, 1); + xla::XlaOp out_base = xla::GetTupleElement(resource, 0); + xla::XlaOp out_grad1 = xla::GetTupleElement(resource, 1); + xla::XlaOp out_grad2 = xla::GetTupleElement(resource, 2); + xla::Tuple(&b, {read_op, out_base, out_grad1, out_grad2}); + TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation flattened_comp, b.Build()); + xla::Literal input_base = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal input_grad2 = xla::LiteralUtil::CreateR1({-3, 101}); - xla::Literal input = xla::LiteralUtil::MakeTuple({&input_base, &input_grad2}); - std::unique_ptr param0_data = - client_->TransferToServer(input).value(); - - std::unique_ptr actual = - client_->Execute(*result.computation, {param0_data.get()}).value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN(auto input_base_buffer, + pjrt_client_->BufferFromHostLiteral( + input_base, pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN(auto input_grad2_buffer, + pjrt_client_->BufferFromHostLiteral( + input_grad2, pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(flattened_comp, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, loaded_exec->Execute( + {{input_base_buffer.get(), input_grad2_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_output_read, + actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_output_base, + actual[0][1]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_output_grad1, + actual[0][2]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_output_grad2, + actual[0][3]->ToLiteralSync()); xla::Literal output_read = xla::LiteralUtil::CreateR0(42); xla::Literal output_base = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal output_grad1 = xla::LiteralUtil::CreateR1({0, 1}); xla::Literal output_grad2 = xla::LiteralUtil::CreateR1({-3, 101}); - xla::Literal output_resource = - xla::LiteralUtil::MakeTuple({&output_base, &output_grad1, &output_grad2}); - xla::Literal expected_literal = - xla::LiteralUtil::MakeTuple({&output_read, &output_resource}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_read, *actual_output_read)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_base, *actual_output_base)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_grad1, *actual_output_grad1)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_grad2, *actual_output_grad2)); } // Tests compilation and execution of a graph that adds two tensors. @@ -1055,8 +1120,8 @@ FunctionDef FillFn() { TEST_F(XlaCompilerTest, FunctionCallWithConstants) { // Certain operations in a function, "Fill" for example, requires the // operator's argument to be a compile-time constant instead of a parameter. - // This testcase tests if XlaCompiler can handle such operators inside - // function calls. + // This case tests if XlaCompiler can handle such operators inside function + // calls. XlaCompiler compiler(DefaultOptions()); FunctionDefLibrary flib; @@ -1187,25 +1252,32 @@ TEST_F(XlaCompilerTest, SliceWithDynamicBegins) { } void RunAndCheckVariablesComputation( - xla::Client* client, const XlaCompiler::CompilationResult& result) { + xla::PjRtClient* client, const XlaCompiler::CompilationResult& result) { xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - std::unique_ptr param0_data = - client->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN(auto param0_buffer, + client->BufferFromHostLiteral( + param0_literal, client->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN(auto param1_buffer, + client->BufferFromHostLiteral( + param1_literal, client->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + client->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR1({5, 144}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({4, 143}); - xla::Literal expected_literal = - xla::LiteralUtil::MakeTuple({&expected0, &expected1}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); } // Tests a simple graph that reads and writes a variable. @@ -1242,7 +1314,7 @@ TEST_F(XlaCompilerTest, Variables) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", std::move(graph), args, &result)); - RunAndCheckVariablesComputation(client_, result); + RunAndCheckVariablesComputation(pjrt_client_.get(), result); } TEST_F(XlaCompilerTest, ResultLayoutSingle) { @@ -1357,15 +1429,20 @@ TEST_F(XlaCompilerTest, ReturnResourceHandleOnly) { // Tests that the generated computation works. xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client_->Execute(*result.computation, {param1_data.get()}).value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); - - xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + TF_ASSERT_OK_AND_ASSIGN( + auto param1_buffer, + pjrt_client_->BufferFromHostLiteral(param1_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param1_buffer.get()}}, execute_options)); + + EXPECT_EQ(actual[0].size(), 0); } TEST_F(XlaCompilerTest, ReturnResourceHandle) { @@ -1403,7 +1480,7 @@ TEST_F(XlaCompilerTest, ReturnResourceHandle) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", std::move(graph), args, &result)); - RunAndCheckVariablesComputation(client_, result); + RunAndCheckVariablesComputation(pjrt_client_.get(), result); } absl::StatusOr> BuildTestGraph() { @@ -1458,44 +1535,52 @@ TEST_F(XlaCompilerTest, VariableRepresentationShapeFunction) { TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), args, &result)); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, - client_->GetComputationShape(*result.computation)); + TF_ASSERT_OK_AND_ASSIGN(xla::ProgramShape program_shape, + result.computation->GetProgramShape()); - ASSERT_EQ(program_shape->parameters_size(), 2); + ASSERT_EQ(program_shape.parameters_size(), 2); EXPECT_TRUE( - xla::ShapeUtil::Compatible(program_shape->parameters(0), + xla::ShapeUtil::Compatible(program_shape.parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {2, 2}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape->parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape.parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape->result(), - xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {2, 2}), - xla::ShapeUtil::MakeShape(xla::S32, {4})}))); + program_shape.result(), xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::S32, {2, 2}), + xla::ShapeUtil::MakeShape(xla::S32, {4})}))); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR2({{4, 55}, {1, -3}}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({22, 11, 33, 404}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client_ - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN( + auto param0_buffer, + pjrt_client_->BufferFromHostLiteral(param0_literal, + pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto param1_buffer, + pjrt_client_->BufferFromHostLiteral(param1_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR2({{27, 67}, {35, 402}}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({26, 66, 34, 401}); - xla::Literal expected_literal = - xla::LiteralUtil::MakeTuple({&expected0, &expected1}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); } TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { @@ -1532,43 +1617,51 @@ TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), args, &result)); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, - client_->GetComputationShape(*result.computation)); + TF_ASSERT_OK_AND_ASSIGN(xla::ProgramShape program_shape, + result.computation->GetProgramShape()); - ASSERT_EQ(program_shape->parameters_size(), 2); + ASSERT_EQ(program_shape.parameters_size(), 2); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape->parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape.parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape->parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape.parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape->result(), - xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {4}), - xla::ShapeUtil::MakeShape(xla::S32, {4})}))); + program_shape.result(), xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::S32, {4}), + xla::ShapeUtil::MakeShape(xla::S32, {4})}))); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({4, 55, 1, -3}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({22, 11, 33, 404}); - std::unique_ptr param0_data = - client_->TransferToServer(param0_literal).value(); - std::unique_ptr param1_data = - client_->TransferToServer(param1_literal).value(); - - std::unique_ptr actual = - client_ - ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN( + auto param0_buffer, + pjrt_client_->BufferFromHostLiteral(param0_literal, + pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN( + auto param1_buffer, + pjrt_client_->BufferFromHostLiteral(param1_literal, + pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR1({27, 67, 35, 402}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({26, 66, 34, 401}); - xla::Literal expected_literal = - xla::LiteralUtil::MakeTuple({&expected0, &expected1}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); } // Tests a graph which has a function with an invalid op. @@ -1751,9 +1844,9 @@ TEST_F(XlaCompilerTest, OpsWithTensorListInput) { { Scope scope = Scope::NewRootScope().ExitOnError(); std::unique_ptr graph(new Graph(OpRegistry::Global())); - ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); + (void)ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); auto result = ops::Const(scope, {true}, {}); - ops::_Retval(scope.WithOpName("ret"), result, 0); + (void)ops::_Retval(scope.WithOpName("ret"), result, 0); TF_ASSERT_OK(scope.ToGraph(graph.get())); FunctionDef fdef; TF_ASSERT_OK(GraphToFunctionDef(*graph, "cond", &fdef)); @@ -1764,7 +1857,7 @@ TEST_F(XlaCompilerTest, OpsWithTensorListInput) { Scope scope = Scope::NewRootScope().ExitOnError(); std::unique_ptr graph(new Graph(OpRegistry::Global())); auto arg = ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); - ops::_Retval(scope.WithOpName("ret"), arg, 0); + (void)ops::_Retval(scope.WithOpName("ret"), arg, 0); TF_ASSERT_OK(scope.ToGraph(graph.get())); FunctionDef fdef; TF_ASSERT_OK(GraphToFunctionDef(*graph, "body", &fdef)); @@ -1906,26 +1999,35 @@ TEST_F(XlaCompilerTest, WhileWithResources) { xla::Literal literal0 = xla::LiteralUtil::CreateR0(0); xla::Literal literal1 = xla::LiteralUtil::CreateR0(2); xla::Literal literal2 = xla::LiteralUtil::CreateR0(1); - std::unique_ptr data0 = - client_->TransferToServer(literal0).value(); - std::unique_ptr data1 = - client_->TransferToServer(literal1).value(); - std::unique_ptr data2 = - client_->TransferToServer(literal2).value(); - - std::unique_ptr actual = - client_ - ->Execute(*result.computation, - {data0.get(), data1.get(), data2.get()}) - .value(); - xla::Literal actual_literal = client_->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN(auto buffer0, + pjrt_client_->BufferFromHostLiteral( + literal0, pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN(auto buffer1, + pjrt_client_->BufferFromHostLiteral( + literal1, pjrt_client_->memory_spaces()[0])); + TF_ASSERT_OK_AND_ASSIGN(auto buffer2, + pjrt_client_->BufferFromHostLiteral( + literal2, pjrt_client_->memory_spaces()[0])); + + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN( + auto actual, + loaded_exec->Execute({{buffer0.get(), buffer1.get(), buffer2.get()}}, + execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal2, actual[0][2]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR0(10); xla::Literal expected1 = xla::LiteralUtil::CreateR0(2); xla::Literal expected2 = xla::LiteralUtil::CreateR0(1); - xla::Literal expected_literal = - xla::LiteralUtil::MakeTuple({&expected0, &expected1, &expected2}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected2, *actual_literal2)); } TEST_F(XlaCompilerTest, SetShardingForReturnedTuple) { @@ -2128,7 +2230,6 @@ TEST_F(OpsTestBase, CompileSingleOp) { XlaCompiler::SingleOpCompileArgument single_op_arg(*context_); - xla::Client* client = xla::ClientLibrary::LocalClientOrDie(); XlaOpRegistry::RegisterCompilationKernels(); FunctionDefLibrary flib; std::unique_ptr flib_def( @@ -2136,7 +2237,13 @@ TEST_F(OpsTestBase, CompileSingleOp) { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - options.client = client; + auto pjrt_client = xla::GetPjRtCpuClient(xla::CpuClientOptions()); + TF_ASSERT_OK(pjrt_client.status()); + options.client = pjrt_client->get(); + auto pjrt_compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); + if (pjrt_compiler.ok()) { + options.compiler = *pjrt_compiler; + } options.flib_def = flib_def.get(); XlaCompiler compiler(options); @@ -2153,13 +2260,18 @@ TEST_F(OpsTestBase, CompileSingleOp) { single_op_arg, args, &result)); // Tests that the generated computation works. - std::unique_ptr actual = - client->Execute(*result.computation, {}).value(); - xla::Literal actual_literal = client->Transfer(*actual).value(); + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + (*pjrt_client) + ->CompileAndLoad(*result.computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN(auto actual, + loaded_exec->Execute({{}}, execute_options)); + + TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); xla::Literal expected0 = xla::LiteralUtil::CreateR2({{6.9, 4.2}}); - xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal)); } TEST_F(XlaCompilerTest, SetShardingForParametersAndReturnValues) { diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc index d26a53e2163e25..a9baa94c9a8ed2 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.cc +++ b/tensorflow/compiler/tf2xla/xla_expression.cc @@ -15,11 +15,36 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_expression.h" +#include +#include +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/hlo/builder/value_inference.h" +#include "xla/hlo/builder/xla_builder.h" +#include "xla/hlo/builder/xla_computation.h" +#include "xla/layout.h" +#include "xla/layout_util.h" +#include "xla/literal.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/shape.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/macros.h" +#include "xla/tsl/platform/statusor.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" -#include "tensorflow/core/lib/core/errors.h" namespace tensorflow { @@ -148,7 +173,7 @@ absl::StatusOr XlaExpression::ResolveDynamism() const { } absl::StatusOr> XlaExpression::ResolveConstant( - xla::Client* client, bool dynamic_dimension_is_minus_one, + xla::PjRtClient* client, bool dynamic_dimension_is_minus_one, xla::ValueInferenceMode mode) const { switch (kind()) { case Kind::kConstant: @@ -191,17 +216,28 @@ absl::StatusOr> XlaExpression::ResolveConstant( return {std::nullopt}; } - if (!client) + if (!client) { return absl::InvalidArgumentError("client is required to resolve constant"); + } TF_ASSIGN_OR_RETURN(xla::XlaComputation constant_graph, handle().builder()->BuildConstantSubGraph( handle(), dynamic_dimension_is_minus_one)); - TF_ASSIGN_OR_RETURN(xla::Literal literal, - client->ComputeConstant(constant_graph, &layout)); + TF_ASSIGN_OR_RETURN( + std::unique_ptr loaded_exec, + client->CompileAndLoad(constant_graph, xla::CompileOptions())); + xla::ExecuteOptions execute_options; + TF_ASSIGN_OR_RETURN(auto results, + loaded_exec->Execute({{}}, execute_options)); + if (results.empty() || results[0].empty() || results[0][0] == nullptr) { + return absl::InternalError("Constant computation produced no output"); + } + TF_ASSIGN_OR_RETURN(std::shared_ptr literal, + results[0][0]->ToLiteralSync()); Tensor tensor; - TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype(), &tensor)); + TF_RETURN_IF_ERROR( + LiteralToHostTensor(literal->Relayout(layout), dtype(), &tensor)); return {tensor}; } diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h index ed0041fc9942a0..a00cd0682e01d7 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.h +++ b/tensorflow/compiler/tf2xla/xla_expression.h @@ -16,14 +16,19 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ -#include "absl/types/optional.h" +#include +#include +#include + +#include "absl/status/statusor.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" -#include "xla/client/client.h" #include "xla/hlo/builder/value_inference.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/shape.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/statusor.h" +#include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { @@ -126,7 +131,7 @@ class XlaExpression { // optional if it cannot be resolved. Returns an error if passed a resource // expression. absl::StatusOr> ResolveConstant( - xla::Client* client, bool dynamic_dimension_is_minus_one = false, + xla::PjRtClient* client, bool dynamic_dimension_is_minus_one = false, xla::ValueInferenceMode mode = xla::ValueInferenceMode::kValue) const; // ResolveDynamism computes where a value inside this op is dynamic or can be diff --git a/tensorflow/compiler/tf2xla/xla_expression_test.cc b/tensorflow/compiler/tf2xla/xla_expression_test.cc index 797002476aeb1c..c989e5484a8c8a 100644 --- a/tensorflow/compiler/tf2xla/xla_expression_test.cc +++ b/tensorflow/compiler/tf2xla/xla_expression_test.cc @@ -15,19 +15,29 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_expression.h" +#include #include +#include +#include +#include -#include "absl/memory/memory.h" +#include "absl/log/log.h" +#include "absl/status/statusor.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" -#include "xla/client/client_library.h" -#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/hlo/builder/xla_computation.h" #include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" +#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/shape_util.h" -#include "xla/status_macros.h" #include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/statusor.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { @@ -36,7 +46,8 @@ namespace { class XlaExpressionTest : public ::testing::Test { protected: void SetUp() override { - client_ = xla::ClientLibrary::LocalClientOrDie(); + TF_ASSERT_OK_AND_ASSIGN(client_, + xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); builder_ = std::make_unique("acomputation"); constant_ = test::AsScalar(42); op_ = xla::ConstantR0(builder_.get(), 7); @@ -50,7 +61,7 @@ class XlaExpressionTest : public ::testing::Test { /*tensor_array_multiple_writes_aggregate=*/false); } - xla::Client* client_; + std::unique_ptr client_; std::unique_ptr builder_; Tensor constant_; xla::XlaOp op_; @@ -86,10 +97,16 @@ TEST_F(XlaExpressionTest, AsXlaOp) { XlaExpression::Constant(constant_).AsXlaOp(builder_.get()); TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation computation, builder_->BuildConstantSubGraph(const_as_op)); - TF_ASSERT_OK_AND_ASSIGN(xla::Literal value, - client_->ComputeConstant(computation)); + ASSERT_NE(client_, nullptr); + TF_ASSERT_OK_AND_ASSIGN( + auto loaded_exec, + client_->CompileAndLoad(computation, xla::CompileOptions{})); + xla::ExecuteOptions execute_options; + TF_ASSERT_OK_AND_ASSIGN(auto results, + loaded_exec->Execute({{}}, execute_options)); + TF_ASSERT_OK_AND_ASSIGN(auto value, results[0][0]->ToLiteralSync()); EXPECT_TRUE(xla::LiteralTestUtil::Equal( - xla::LiteralUtil::CreateR0(42), value)); + xla::LiteralUtil::CreateR0(42), *value)); } TEST_F(XlaExpressionTest, GetShape) { @@ -110,27 +127,27 @@ TEST_F(XlaExpressionTest, GetShape) { } TEST_F(XlaExpressionTest, ResolveConstant) { - EXPECT_FALSE(XlaExpression().ResolveConstant(client_).ok()); - EXPECT_FALSE(XlaExpression::Invalid().ResolveConstant(client_).ok()); + EXPECT_FALSE(XlaExpression().ResolveConstant(client_.get()).ok()); + EXPECT_FALSE(XlaExpression::Invalid().ResolveConstant(client_.get()).ok()); EXPECT_FALSE(XlaExpression::Resource(resource_.get()) - .ResolveConstant(client_) + .ResolveConstant(client_.get()) ->has_value()); TF_ASSERT_OK_AND_ASSIGN( std::optional op_constant, - XlaExpression::XlaOp(op_, DT_INT32).ResolveConstant(client_)); + XlaExpression::XlaOp(op_, DT_INT32).ResolveConstant(client_.get())); ASSERT_TRUE(op_constant.has_value()); test::ExpectTensorEqual(test::AsScalar(7), *op_constant); TF_ASSERT_OK_AND_ASSIGN(std::optional op_nonconstant, XlaExpression::XlaOp(non_constant_op_, DT_FLOAT) - .ResolveConstant(client_)); + .ResolveConstant(client_.get())); EXPECT_FALSE(op_nonconstant.has_value()); TF_ASSERT_OK_AND_ASSIGN( std::optional constant_constant, - XlaExpression::Constant(constant_).ResolveConstant(client_)); + XlaExpression::Constant(constant_).ResolveConstant(client_.get())); ASSERT_TRUE(constant_constant.has_value()); test::ExpectTensorEqual(constant_, *constant_constant); } @@ -138,11 +155,11 @@ TEST_F(XlaExpressionTest, ResolveConstant) { TEST_F(XlaExpressionTest, ResolveConstantOnResource) { XlaExpression constant_resource = XlaExpression::ConstantResource(constant_, resource_.get()); - EXPECT_TRUE(constant_resource.ResolveConstant(client_).ok()); + EXPECT_TRUE(constant_resource.ResolveConstant(client_.get()).ok()); EXPECT_TRUE(resource_->SetZeroValue(builder_.get()).ok()); LOG(ERROR) << "Resource is overwritten: " << resource_->IsOverwritten(); absl::StatusOr> resolved_constant = - constant_resource.ResolveConstant(client_); + constant_resource.ResolveConstant(client_.get()); EXPECT_TRUE(resolved_constant.ok()); EXPECT_FALSE(resolved_constant->has_value()); } diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc index de0e7bf45eae79..40c2247a79858c 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc @@ -16,12 +16,16 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h" #include +#include #include #include +#include #include #include +#include "absl/base/casts.h" #include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/tf2xla.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" @@ -29,28 +33,28 @@ limitations under the License. #include "xla/backends/cpu/buffer_allocation_info.h" #include "xla/backends/cpu/buffer_allocation_info_util.h" #include "xla/backends/cpu/codegen/compiled_function_library.h" -#include "xla/client/client_library.h" #include "xla/client/executable_build_options.h" -#include "xla/client/local_client.h" +#include "xla/executable_run_options.h" #include "xla/hlo/builder/xla_computation.h" +#include "xla/pjrt/common_pjrt_client.h" +#include "xla/pjrt/cpu/cpu_client.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" +#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" +#include "xla/service/buffer_assignment.h" #include "xla/service/cpu/cpu_aot_compilation_result.h" #include "xla/service/cpu/cpu_executable.h" -#include "xla/service/platform_util.h" -#include "xla/shape_util.h" -#include "xla/stream_executor/platform.h" +#include "xla/shape.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/graph.pb.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/errors.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/types.h" -#include "tsl/platform/casts.h" -#include "tsl/platform/statusor.h" namespace tensorflow { namespace { -constexpr char kHostPlatform[] = "Host"; // Returns the index of the result in the temp buffers. absl::StatusOr ComputeResultIndex( @@ -107,19 +111,21 @@ void CollectNames(const T& entries, std::vector* nonempty_names, XlaJitCompiledCpuFunction::Compile( const GraphDef& graph_def, const tf2xla::Config& config, const xla::ExecutableBuildOptions& build_options) { + TF_ASSIGN_OR_RETURN(xla::PjRtCompiler * compiler, + xla::GetDefaultPjRtCompiler(xla::CpuName())); + + TF_ASSIGN_OR_RETURN(std::unique_ptr client, + xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); + // Convert the graph_def into an xla::XlaComputation. - TF_ASSIGN_OR_RETURN(se::Platform * platform, - xla::PlatformUtil::GetPlatform(kHostPlatform)); - TF_ASSIGN_OR_RETURN(xla::LocalClient * client, - xla::ClientLibrary::GetOrCreateLocalClient(platform)); xla::XlaComputation computation; - TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla(graph_def, config, client, - &computation)); + TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla( + graph_def, config, compiler, &computation, client.get())); // Get and verify the program shape. - TF_ASSIGN_OR_RETURN(std::unique_ptr program_shape, - client->GetComputationShape(computation)); - if (program_shape->result().element_type() != xla::TUPLE) { + TF_ASSIGN_OR_RETURN(xla::ProgramShape program_shape, + computation.GetProgramShape()); + if (program_shape.result().element_type() != xla::TUPLE) { // The XlaCompiler we use to build the xla computation always generates a // tuple result, and XlaCompiledCpuFunction relies on this for simpler // calling semantics. @@ -128,24 +134,35 @@ XlaJitCompiledCpuFunction::Compile( } // The parameter names are currently meaningless, and redundant with the rest // of our metadata, so clear them out to avoid confusion and save space. - program_shape->clear_parameter_names(); + program_shape.clear_parameter_names(); // Compute arg shapes, needed to compile the executable. - std::vector arg_shapes; - arg_shapes.reserve(program_shape->parameters_size()); - for (int i = 0; i < program_shape->parameters_size(); ++i) { - arg_shapes.push_back(&program_shape->parameters(i)); + std::vector arg_shapes; + arg_shapes.reserve(program_shape.parameters_size()); + for (int i = 0; i < program_shape.parameters_size(); ++i) { + arg_shapes.push_back(program_shape.parameters(i)); } - // Compile the executable. The static_cast to the CpuExecutable subclass is - // necessary since the raw function and buffer assignments are only available - // there. - TF_ASSIGN_OR_RETURN(auto executables, - client->Compile(computation, arg_shapes, build_options)); - TF_RET_CHECK(executables.size() == 1); - std::unique_ptr executable = std::move(executables[0]); + xla::CompileOptions compile_options; + compile_options.executable_build_options = build_options; + compile_options.argument_layouts = std::move(arg_shapes); + + // Compile the executable. + TF_ASSIGN_OR_RETURN(const xla::PjRtTopologyDescription* topology, + client->GetTopologyDescription()); + TF_ASSIGN_OR_RETURN(std::unique_ptr pjrt_executable, + compiler->Compile(std::move(compile_options), computation, + *topology, client.get())); + + auto* pjrt_cpu_executable = + absl::down_cast(pjrt_executable.get()); + if (!pjrt_cpu_executable) { + return absl::InternalError( + "Expected PjRtCpuExecutable from CPU compilation."); + } xla::cpu::CpuExecutable* cpu_executable = - static_cast(executable->executable()); + static_cast( + pjrt_cpu_executable->cpu_executable().get()); const xla::BufferAssignment& buffer_assignment = cpu_executable->buffer_assignment(); @@ -154,10 +171,6 @@ XlaJitCompiledCpuFunction::Compile( xla::cpu::CreateBufferAllocationInfos(cpu_executable->module(), buffer_assignment); - std::vector buffer_allocation_infos = - xla::cpu::CreateBufferAllocationInfos(cpu_executable->module(), - buffer_assignment); - std::vector arg_index_table = xla::cpu::CreateArgIndexTable(buffer_infos); std::vector result_index_table = @@ -216,12 +229,13 @@ XlaJitCompiledCpuFunction::Compile( compiled_function_library->GetTypelessSymbolsMap(); } - jit->executable_ = std::move(executable); + jit->pjrt_client_ = std::move(client); + jit->pjrt_executable_ = std::move(pjrt_executable); jit->buffer_infos_ = std::move(buffer_infos); jit->arg_index_table_ = std::move(arg_index_table); jit->result_index_table_ = std::move(result_index_table); jit->program_shape_ = - std::make_unique(program_shape->ToProto()); + std::make_unique(program_shape.ToProto()); XlaCompiledCpuFunction::set_static_data_compilation_result_proto( &jit->static_data_, jit->compilation_result_proto_.get()); XlaCompiledCpuFunction::set_static_data_function_library_symbol_map( @@ -269,4 +283,13 @@ XlaJitCompiledCpuFunction::Compile( return std::move(jit_unique_ptr); } +const Eigen::ThreadPoolDevice* +XlaJitCompiledCpuFunction::eigen_intra_op_thread_pool_device() const { + auto* common_client = + absl::down_cast(pjrt_client_.get()); + auto* raw_client = + absl::down_cast(common_client->raw_client()); + return raw_client->eigen_intraop_device(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h index 6f61f472a2fd5a..6f175c74d6af4c 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h @@ -16,20 +16,23 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_JIT_COMPILED_CPU_FUNCTION_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_JIT_COMPILED_CPU_FUNCTION_H_ +#include #include #include #include #include "absl/container/flat_hash_map.h" #include "absl/log/check.h" -#include "tensorflow/compiler/tf2xla/encoded_buffer_allocation_info.h" +#include "absl/status/statusor.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" -#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function_thunks.h" +#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h" #include "xla/backends/cpu/buffer_allocation_info.h" -#include "xla/client/local_client.h" +#include "xla/client/executable_build_options.h" +#include "xla/executable_run_options.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_executable.h" #include "xla/service/cpu/executable.pb.h" #include "tensorflow/core/framework/graph.pb.h" -#include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -62,16 +65,21 @@ class XlaJitCompiledCpuFunction { return static_data_; } - const xla::LocalExecutable& LocalExecutable() const { - CHECK(executable_); // Crash ok - return *executable_; + const xla::PjRtExecutable& pjrt_executable() const { + CHECK(pjrt_executable_ != nullptr); // Crash ok + return *pjrt_executable_; } + const Eigen::ThreadPoolDevice* eigen_intra_op_thread_pool_device() const; + private: XlaJitCompiledCpuFunction() : compilation_result_proto_(nullptr) {} + // The client holds the thread pool used for execution. + std::unique_ptr pjrt_client_; + // The executable holds the underlying function. - std::unique_ptr executable_; + std::unique_ptr pjrt_executable_; // The compilation result proto. std::unique_ptr compilation_result_proto_; diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc index 3864e4fdeca467..05fd83d967cc30 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h" +#include #include #include @@ -25,25 +26,24 @@ limitations under the License. #include "xla/client/executable_build_options.h" #include "xla/hlo/testlib/test.h" #include "xla/service/compiler.h" -#include "xla/service/platform_util.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_id.h" #include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/platform/test.h" -#include "tensorflow/core/platform/types.h" +#include "tsl/platform/protobuf.h" namespace tensorflow { namespace { -using ::testing::HasSubstr; - PLATFORM_DEFINE_ID(kFakePlatformId, FakePlatform); AttrValue TypeAttrValue(DataType type) { @@ -347,9 +347,6 @@ TEST(XlaJitCompiledCpuFunction, CanCompileWithAdditionalPlatform) { return std::unique_ptr(nullptr); }); - EXPECT_THAT(xla::PlatformUtil::GetDefaultPlatform().status().message(), - HasSubstr(kFakePlatformId->ToName())); - GraphDef graph_def = SumGraph(); tf2xla::Config config = SumConfig(); TF_ASSERT_OK_AND_ASSIGN( diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index 26461ddaa0c0a4..02720052b9c6d1 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.h +++ b/tensorflow/compiler/tf2xla/xla_op_registry.h @@ -66,7 +66,7 @@ constexpr std::array kNumericTypes = { DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64, DT_COMPLEX128, DT_BFLOAT16}}; -constexpr std::array kCpuAllTypes = {{DT_UINT8, +constexpr std::array kCpuAllTypes = {{DT_UINT8, DT_QUINT8, DT_UINT16, DT_UINT32, @@ -89,10 +89,11 @@ constexpr std::array kCpuAllTypes = {{DT_UINT8, DT_FLOAT8_E4M3FNUZ, DT_FLOAT8_E4M3B11FNUZ, DT_FLOAT8_E5M2FNUZ, + DT_FLOAT8_E8M0FNU, DT_INT4, DT_UINT4}}; -constexpr std::array kGpuAllTypes = {{DT_UINT8, +constexpr std::array kGpuAllTypes = {{DT_UINT8, DT_QUINT8, DT_UINT16, DT_UINT32, @@ -115,6 +116,7 @@ constexpr std::array kGpuAllTypes = {{DT_UINT8, DT_FLOAT8_E4M3FNUZ, DT_FLOAT8_E4M3B11FNUZ, DT_FLOAT8_E5M2FNUZ, + DT_FLOAT8_E8M0FNU, DT_INT4, DT_UINT4}}; diff --git a/tensorflow/core/api_def/base_api/api_def_DecodeJXL.pbtxt b/tensorflow/core/api_def/base_api/api_def_DecodeJXL.pbtxt index 8b1e434ef43899..aaf7eecc03d255 100644 --- a/tensorflow/core/api_def/base_api/api_def_DecodeJXL.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_DecodeJXL.pbtxt @@ -18,7 +18,13 @@ END Number of color channels for the decoded image. END } - summary: "Decode a JPEG XL encoded image to a uint8 tensor." + attr { + name: "dtype" + description: <= 3 with shape `batch_dims + [height, width, channels]`. +END + } + out_arg { + name: "contents" + description: <