From 3a62a9d6df2b2689b2037fe58e92e402895c3e82 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Fri, 27 Sep 2024 14:20:40 -0700 Subject: [PATCH] Fork `xla::ExecuteOptions` into `xla::ifrt::ExecuteOptions` Only the fields that are currently being used by any IFRT implementation are copied to the fork. Two fields that are currently being used as a passthrough in PjRt-IFRT but are removed are: * `untuple_result`: Since IFRT does not have XLA tuple types, all PjRt-IFRT invocations must have already been setting this to true. So this CL changes PjRt-IFRT to unconditionally set `untuple_result` to true when invoking PjRt and gets rid of this field from `xla.ifrt.ExecuteOptions`. * `use_major_to_minor_data_layout_for_callbacks`: The meaning of this field is very specific to PjRt. Since this field is set to true in every IFRT invocation that cares about this field, this CL instead changes PjRt-IFRT to internally always set this field to true and avoid exposing this to IFRT. In order to not break the IFRT Proxy's version compatibility, the forked `ExecuteOptionsProto` uses the same field tags as the original proto. PiperOrigin-RevId: 679727285 --- xla/python/ifrt/BUILD | 7 +++++ xla/python/ifrt/executable.cc | 25 +++++++++++++++++ xla/python/ifrt/executable.h | 27 +++++++++++++++++-- xla/python/ifrt/execute_options.proto | 26 ++++++++++++++++++ xla/python/ifrt_proxy/common/BUILD | 1 + .../ifrt_proxy/common/ifrt_service.proto | 4 +-- xla/python/pjrt_ifrt/pjrt_executable.cc | 10 ++++--- xla/python/py_executable.cc | 4 +-- xla/python/py_executable.h | 4 +-- 9 files changed, 95 insertions(+), 13 deletions(-) create mode 100644 xla/python/ifrt/execute_options.proto diff --git a/xla/python/ifrt/BUILD b/xla/python/ifrt/BUILD index 27118ddf854c3..78cb7d039bcd0 100644 --- a/xla/python/ifrt/BUILD +++ b/xla/python/ifrt/BUILD @@ -87,6 +87,7 @@ cc_library( ":attribute_map", ":device_proto_cc", ":dtype_proto_cc", + ":execute_options_proto_cc", ":remap_plan_proto_cc", ":serdes", ":shape_proto_cc", @@ -109,6 +110,7 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/container:node_hash_set", "@com_google_absl//absl/functional:function_ref", @@ -165,6 +167,11 @@ xla_cc_test( ], ) +tf_proto_library( + name = "execute_options_proto", + srcs = ["execute_options.proto"], +) + xla_cc_test( name = "future_test", size = "small", diff --git a/xla/python/ifrt/executable.cc b/xla/python/ifrt/executable.cc index 77cabe7f6a938..31f5aee05ebae 100644 --- a/xla/python/ifrt/executable.cc +++ b/xla/python/ifrt/executable.cc @@ -15,11 +15,36 @@ limitations under the License. #include "xla/python/ifrt/executable.h" +#include "absl/status/statusor.h" +#include "xla/python/ifrt/execute_options.pb.h" + namespace xla { namespace ifrt { char Executable::ID = 0; char LoadedExecutable::ID = 0; +absl::StatusOr ExecuteOptions::ToProto() const { + ExecuteOptionsProto proto; + + proto.set_launch_id(launch_id); + proto.mutable_non_donatable_input_indices()->Add( + non_donatable_input_indices.begin(), non_donatable_input_indices.end()); + + return proto; +} + +absl::StatusOr ExecuteOptions::FromProto( + const xla::ifrt::ExecuteOptionsProto& proto) { + ExecuteOptions options; + + options.launch_id = proto.launch_id(); + options.non_donatable_input_indices.insert( + proto.non_donatable_input_indices().begin(), + proto.non_donatable_input_indices().end()); + + return options; +} + } // namespace ifrt } // namespace xla diff --git a/xla/python/ifrt/executable.h b/xla/python/ifrt/executable.h index 6b642bd5d178d..adcf1e6ebe940 100644 --- a/xla/python/ifrt/executable.h +++ b/xla/python/ifrt/executable.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" @@ -33,6 +34,7 @@ limitations under the License. #include "xla/python/ifrt/attribute_map.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" +#include "xla/python/ifrt/execute_options.pb.h" #include "xla/python/ifrt/future.h" #include "xla/tsl/concurrency/ref_count.h" @@ -104,6 +106,28 @@ class Executable : public llvm::RTTIExtends { static char ID; // NOLINT }; +struct ExecuteOptions { + // If non-zero, identifies this execution as part of a potentially + // multi-device launch. This can be used to detect scheduling errors, e.g. if + // multi-host programs are launched in different orders on different hosts, + // the launch IDs may be used by the runtime to detect the mismatch. + int32_t launch_id = 0; + + // A set of indices denoting the input arrays that should not be donated. An + // input array may be non-donable, for example, if it is referenced more than + // once. Since such runtime information is not available at compile time, the + // compiler might mark the input as `may-alias`, which could lead IFRT to + // donate the input array when it should not. By defining this set of indices, + // a higher-level IFRT caller can instruct IFRT client not to donate specific + // input arrays. + absl::flat_hash_set non_donatable_input_indices; + + absl::StatusOr ToProto() const; + + static absl::StatusOr FromProto( + const ExecuteOptionsProto& proto); +}; + // Wraps a computation that has been fully compiled and loaded for execution. class LoadedExecutable : public llvm::RTTIExtends { @@ -176,8 +200,7 @@ class LoadedExecutable // `LoadedExecutable` methods. - // Short-term alias. - using ExecuteOptions = ::xla::ExecuteOptions; + using ExecuteOptions = xla::ifrt::ExecuteOptions; // Result from an execution. struct ExecuteResult { diff --git a/xla/python/ifrt/execute_options.proto b/xla/python/ifrt/execute_options.proto new file mode 100644 index 0000000000000..6cc20c0996f4f --- /dev/null +++ b/xla/python/ifrt/execute_options.proto @@ -0,0 +1,26 @@ +/* Copyright 2024 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +syntax = "proto3"; + +package xla.ifrt; + +message ExecuteOptionsProto { + bool untuple_result = 2; + int32 launch_id = 3; + repeated int32 non_donatable_input_indices = 7; + + reserved 1, 4 to 6, 8; +} diff --git a/xla/python/ifrt_proxy/common/BUILD b/xla/python/ifrt_proxy/common/BUILD index aa971c34b7155..9d6b3bbecb4d9 100644 --- a/xla/python/ifrt_proxy/common/BUILD +++ b/xla/python/ifrt_proxy/common/BUILD @@ -72,6 +72,7 @@ tf_proto_library( "//xla/pjrt:execute_options_proto", "//xla/python/ifrt:attribute_map_proto", "//xla/python/ifrt:dtype_proto", + "//xla/python/ifrt:execute_options_proto", "//xla/python/ifrt:remap_plan_proto", "//xla/python/ifrt:serdes_proto", "//xla/python/ifrt:shape_proto", diff --git a/xla/python/ifrt_proxy/common/ifrt_service.proto b/xla/python/ifrt_proxy/common/ifrt_service.proto index 4a342c253af9c..3f17ee69abbd5 100644 --- a/xla/python/ifrt_proxy/common/ifrt_service.proto +++ b/xla/python/ifrt_proxy/common/ifrt_service.proto @@ -17,9 +17,9 @@ syntax = "proto3"; package xla.ifrt.proxy; import "google/protobuf/any.proto"; -import "xla/pjrt/execute_options.proto"; import "xla/python/ifrt/attribute_map.proto"; import "xla/python/ifrt/dtype.proto"; +import "xla/python/ifrt/execute_options.proto"; import "xla/python/ifrt/remap_plan.proto"; import "xla/python/ifrt/serdes.proto"; import "xla/python/ifrt/shape.proto"; @@ -428,7 +428,7 @@ message LoadedExecutableMetadataResponse { message LoadedExecutableExecuteRequest { fixed64 loaded_executable_handle = 1; repeated fixed64 args_handles = 2; - xla.ExecuteOptionsProto execute_options = 3; + xla.ifrt.ExecuteOptionsProto execute_options = 3; repeated int32 device_ids = 4; } message LoadedExecutableExecuteResponse { diff --git a/xla/python/pjrt_ifrt/pjrt_executable.cc b/xla/python/pjrt_ifrt/pjrt_executable.cc index 19c441cdea448..7225a3a40db8f 100644 --- a/xla/python/pjrt_ifrt/pjrt_executable.cc +++ b/xla/python/pjrt_ifrt/pjrt_executable.cc @@ -542,7 +542,11 @@ PjRtLoadedExecutable::Execute( const bool returned_future_supported = pjrt_loaded_executable_->IsReturnedFutureSupported(); - auto opts = options; + xla::ExecuteOptions opts; + opts.untuple_result = true; + opts.launch_id = options.launch_id; + opts.use_major_to_minor_data_layout_for_callbacks = true; + opts.non_donatable_input_indices = options.non_donatable_input_indices; if (!all_loaded_host_callbacks_->empty() && !returned_future_supported) { return Internal( @@ -565,9 +569,7 @@ PjRtLoadedExecutable::Execute( contexts.push_back(CreateHostCallbackStateAndAppendSendRecvCallbacks( host_send_recv_callback->host_callback(), /*host_memory_for_device_manager=*/nullptr, send_callbacks, - recv_callbacks, - /*use_major_to_minor_data_layout_for_callbacks=*/ - options.use_major_to_minor_data_layout_for_callbacks)); + recv_callbacks, opts.use_major_to_minor_data_layout_for_callbacks)); } } opts.send_callbacks = host_callback_states->send_callbacks; diff --git a/xla/python/py_executable.cc b/xla/python/py_executable.cc index b7395ad779305..7e68c944f9438 100644 --- a/xla/python/py_executable.cc +++ b/xla/python/py_executable.cc @@ -93,13 +93,11 @@ PyLoadedExecutable::PyLoadedExecutable( if (next_) { next_->prev_ = this; } - options_.untuple_result = true; if (fingerprint_) { options_.launch_id = tsl::Fingerprint32(*fingerprint_); VLOG(1) << "Fingerprint for executable " << ifrt_loaded_executable_->name() << ": " << *fingerprint_; } - options_.use_major_to_minor_data_layout_for_callbacks = true; } PyLoadedExecutable::~PyLoadedExecutable() { @@ -203,7 +201,7 @@ void PopulateExecuteShardedResults( template > absl::StatusOr ExecuteShardedOnLocalDevicesInternal( - const ExecuteOptions& options, const nb_class_ptr& client, + const ifrt::ExecuteOptions& options, const nb_class_ptr& client, ifrt::LoadedExecutable* ifrt_loaded_executable, absl::Span args, std::optional>>& returned_futures, bool attach_status_to_results) { diff --git a/xla/python/py_executable.h b/xla/python/py_executable.h index ed34ce99ef1a8..e032ee7b4acdd 100644 --- a/xla/python/py_executable.h +++ b/xla/python/py_executable.h @@ -227,7 +227,7 @@ class PyLoadedExecutable { return exec->shared_ptr_pjrt_loaded_executable(); } - const ExecuteOptions& options() const { return options_; } + const ifrt::ExecuteOptions& options() const { return options_; } const std::optional& fingerprint() const { return fingerprint_; } // Keep `obj` alive as long as PyLoadedExecutable. @@ -246,7 +246,7 @@ class PyLoadedExecutable { std::optional fingerprint_; // The options to pass to `executable_.Execute`. - ExecuteOptions options_; + ifrt::ExecuteOptions options_; // Python objects to keep alive as requested by user. std::vector keepalives_;