From 2c93352577068415c4b55ce201b34e2cd572d49b Mon Sep 17 00:00:00 2001 From: Sannidhya Chauhan Date: Fri, 28 Aug 2026 04:34:17 -0700 Subject: [PATCH 01/11] Allow ProfilerController to consume data in the kCollectData state. PiperOrigin-RevId: 972507321 --- .../xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc | 2 +- .../xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc index 41f1cc1a27344e..490c96fffe792a 100644 --- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc +++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_collection.cc @@ -70,7 +70,7 @@ absl::StatusOr ProfilerCollection::Consume() { data_vector.push_back(std::move(result->data)); total_estimated_size_bytes += result->estimated_size_bytes; } else if (absl::IsUnimplemented(result.status())) { - LOG(WARNING) << "Profiler consume not implemented: " << result.status(); + VLOG(1) << "Profiler consume not implemented: " << result.status(); data_vector.push_back(std::any()); } else { LOG(ERROR) << "Profiler consume failed: " << result.status(); diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc index 2bf12431880c33..fd311ccd70ea2c 100644 --- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc +++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_controller.cc @@ -88,7 +88,7 @@ absl::Status ProfilerController::CollectData( } absl::StatusOr ProfilerController::Consume() { - if (state_ != ProfilerState::kStart && state_ != ProfilerState::kStop) { + if (state_ == ProfilerState::kInit) { return absl::AbortedError("Consume called in the wrong order."); } if (!status_.ok()) { From 3b6d61bd25b9684406becae4912c451331b93174 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 04:55:17 -0700 Subject: [PATCH 02/11] [XLA:GPU] Add DegenerateDimensionRewriter to GPU pipeline. This pass removes unnecessary size-1 from the module. Degenerate dimensions are generally no-op, but it add unnecessary reshapes/bitcasts to the graph that can sometime prevent better fusion and tiling decision or cause problem with emitter pipelines, like Triton. Reverts 327ef3d139c4b26e76eefd36bd542240ad8d8107 PiperOrigin-RevId: 972515964 --- third_party/xla/xla/service/gpu/BUILD | 1 - third_party/xla/xla/service/gpu/gpu_compiler.cc | 9 --------- 2 files changed, 10 deletions(-) diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index bbe83ba91d3b30..d4de5b0fd9462a 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1971,7 +1971,6 @@ cc_library( "//xla/hlo/transforms/simplifiers:broadcast_canonicalizer", "//xla/hlo/transforms/simplifiers:conditional_canonicalizer", "//xla/hlo/transforms/simplifiers:convert_mover", - "//xla/hlo/transforms/simplifiers:degenerate_dimension_rewriter", "//xla/hlo/transforms/simplifiers:dot_merger", "//xla/hlo/transforms/simplifiers:dynamic_dimension_simplifier", "//xla/hlo/transforms/simplifiers:flatten_call_graph", diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index a5408c0f2f3587..15258447cc2958 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -209,7 +209,6 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/broadcast_canonicalizer.h" #include "xla/hlo/transforms/simplifiers/conditional_canonicalizer.h" #include "xla/hlo/transforms/simplifiers/convert_mover.h" -#include "xla/hlo/transforms/simplifiers/degenerate_dimension_rewriter.h" #include "xla/hlo/transforms/simplifiers/dot_merger.h" #include "xla/hlo/transforms/simplifiers/dynamic_dimension_simplifier.h" #include "xla/hlo/transforms/simplifiers/flatten_call_graph.h" @@ -978,14 +977,6 @@ absl::Status RunOptimizationPasses( pipeline.AddPass(); pipeline.AddPass( gpu_target_config.device_description.gpu_compute_capability()); - - // It's important to run AlgebraicSimplifier after - // DegenerateDimensionRewriter before ReshapeMover. - // DegenerateDimensionRewriter introduces reshape to remove size-1 dims from - // ops like iota and broadcast, and algebraic simplifier has patterns to - // fold reshape(iota) and reshape(broadcast). If we run ReshapeMover first, - // it will move these reshapes down the graph, and prevent the folding. - pipeline.AddPass(); pipeline.AddPass(layout_insensitive_algsimp_opts, gpu_version); pipeline.AddPass(); From 9dd68c038c3cf18e58456fd8f96178db4e82eea8 Mon Sep 17 00:00:00 2001 From: Levon Ter-Grigoryan Date: Fri, 28 Aug 2026 05:01:08 -0700 Subject: [PATCH 03/11] [XLA:GPU] Use symmetric memory and peer address lookup for collective kernel thunks. Create a symmetric memory for scratch buffers and use it to get a peer address of memory chunk instead of running separate rendzevous. PiperOrigin-RevId: 972518088 --- .../xla/xla/backends/gpu/runtime/BUILD | 6 +- .../gpu/runtime/collective_kernel_api.cc | 36 ----------- .../gpu/runtime/collective_kernel_api.h | 9 --- .../gpu/runtime/collective_kernel_thunk.cc | 64 +++++++++++++++++-- .../gpu/runtime/collective_kernel_thunk.h | 3 + .../runtime/collective_kernel_thunk_test.cc | 20 ++++-- 6 files changed, 82 insertions(+), 56 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index ad0fc565c7b22c..714dadf77ac100 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1961,8 +1961,9 @@ cc_library( hdrs = ["collective_kernel_thunk.h"], deps = [ ":all_reduce", - ":collective_kernel_api", + ":collective_cliques", ":collective_kernel_thunk_proto_cc", + ":collective_memory", ":collective_params", ":collective_thunk", ":collective_thunk_proto_cc", @@ -1974,7 +1975,9 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/backends/gpu/collectives:gpu_clique_key", + "//xla/backends/gpu/collectives:gpu_communicator", "//xla/core/collectives:rank_id", + "//xla/core/collectives:symmetric_memory", "//xla/runtime:buffer_use", "//xla/runtime:device_id", "//xla/service:buffer_assignment", @@ -1992,6 +1995,7 @@ cc_library( "//xla/stream_executor/gpu:all_reduce_kernel", "//xla/stream_executor/gpu:collective_kernel_metadata", "//xla/tsl/util:safe_reinterpret_cast", + "//xla/tsl/util:tied_ref", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc index dce95619458e2c..89204ae5635988 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.cc @@ -148,41 +148,5 @@ size_t GetMultiGpuBarrierSignalBufferSize() { size_t GetMultiGpuBarrierSignalValueSize() { return sizeof(uint32_t); } -absl::StatusOr> CollectParamToPeers( - const GpuCliqueKey& clique_key, RankId rank, - stream_executor::Stream* stream, - std::vector parameters) { - std::vector param_to_peers_ptrs; - - size_t num_parameters = parameters.size(); - // Exchange device parameters with all ranks in the clique. - ABSL_ASSIGN_OR_RETURN( - auto device_parameters, - GpuCliqueRendezvous::Join(clique_key, rank, std::move(parameters))); - - // Collect pointers to device buffers from all participating ranks. - param_to_peers_ptrs.reserve(num_parameters * clique_key.num_devices()); - - absl::flat_hash_map> - peer_to_parameters(clique_key.num_devices()); - - using DeviceParameters = std::vector; - - for (auto peer = RankId(0); peer < RankId(clique_key.num_devices()); ++peer) { - ABSL_ASSIGN_OR_RETURN(const DeviceParameters& peer_parameters, - device_parameters->at(peer)); - peer_to_parameters[peer.value()] = std::move(peer_parameters); - } - - for (int parameter = 0; parameter < num_parameters; ++parameter) { - for (int peer = 0; peer < clique_key.num_devices(); ++peer) { - param_to_peers_ptrs.push_back( - peer_to_parameters[peer][parameter].opaque()); - } - } - - return param_to_peers_ptrs; -} - } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h index ff4bc5101d7593..cd601d10610dd9 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_api.h @@ -21,8 +21,6 @@ limitations under the License. #include #include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "xla/backends/gpu/collectives/gpu_clique_key.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/symmetric_memory.h" #include "xla/stream_executor/device_address.h" @@ -57,13 +55,6 @@ size_t GetMultiGpuBarrierSignalBufferSize(); // Returns the size of the barrier signal value in bytes. size_t GetMultiGpuBarrierSignalValueSize(); -// Collect the pointers to the parameters at the peer devices. -// The size of the returned vector is num_parameters * num_devices. -absl::StatusOr> CollectParamToPeers( - const GpuCliqueKey& clique_key, RankId rank, - stream_executor::Stream* stream, - std::vector parameters); - } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_KERNEL_API_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc index 5294bcb0bc783c..e48d12ca03c5ef 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc @@ -35,15 +35,18 @@ limitations under the License.*/ #include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/gpu_clique_key.h" +#include "xla/backends/gpu/collectives/gpu_communicator.h" #include "xla/backends/gpu/runtime/all_reduce.h" -#include "xla/backends/gpu/runtime/collective_kernel_api.h" +#include "xla/backends/gpu/runtime/collective_cliques.h" #include "xla/backends/gpu/runtime/collective_kernel_thunk.pb.h" +#include "xla/backends/gpu/runtime/collective_memory.h" #include "xla/backends/gpu/runtime/collective_params.h" #include "xla/backends/gpu/runtime/collective_thunk.h" #include "xla/backends/gpu/runtime/collective_thunk.pb.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/core/collectives/rank_id.h" +#include "xla/core/collectives/symmetric_memory.h" #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" @@ -62,6 +65,7 @@ limitations under the License.*/ #include "xla/stream_executor/stream.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/util/safe_reinterpret_cast.h" +#include "xla/tsl/util/tied_ref.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -455,8 +459,37 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { const size_t num_parameters = parameters.size(); const size_t param_to_peers_ptrs_size_bytes = num_parameters * clique_key.num_devices() * sizeof(uint64_t); + TF_RET_CHECK(params.collective_params != nullptr) + << "Collective params must not be null in " + "CollectiveKernelThunk::Initialize"; + TF_RET_CHECK(params.collective_cliques != nullptr) + << "Collective cliques must not be null in " + "CollectiveKernelThunk::Initialize"; + TF_RET_CHECK(params.collective_memory != nullptr) + << "Collective memory must not be null in " + "CollectiveKernelThunk::Initialize"; + + ABSL_ASSIGN_OR_RETURN(GpuCommunicator * comm, + params.collective_cliques->GetComm(clique_key, *rank)); + + if (memory_state->scratch_symmetric_memories.empty()) { + memory_state->scratch_symmetric_memories.reserve( + memory_state->scratch_allocations.size()); + for (size_t i = 0; i < memory_state->scratch_allocations.size(); ++i) { + se::DeviceAddressBase addr = + memory_state->scratch_allocations[i].address(); + ABSL_ASSIGN_OR_RETURN(std::unique_ptr symmetric_memory, + comm->CreateSymmetricMemory(addr)); + ABSL_ASSIGN_OR_RETURN(tsl::TiedRef tied_symmetric_memory, + params.collective_cliques->Tie( + clique_key, std::move(symmetric_memory))); + memory_state->scratch_symmetric_memories.push_back( + std::move(tied_symmetric_memory)); + } + } + std::vector multimem_addresses; - if (RequiresMultimem(kernel_spec_) && params.collective_memory != nullptr) { + if (RequiresMultimem(kernel_spec_)) { multimem_addresses.resize(num_parameters, nullptr); for (size_t i = 0; i < num_parameters; ++i) { auto [mmem, offset] = params.collective_memory->FindSymmetricMemory( @@ -469,9 +502,30 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { } } } - ABSL_ASSIGN_OR_RETURN(std::vector param_to_peers_ptrs, - CollectParamToPeers(clique_key, state->rank, params.stream, - std::move(parameters))); + + static constexpr auto is_multimem_buffer = + [](const IoBufferSpec& spec) -> bool { return spec.requires_multimem; }; + int32_t scratch_buffers_index = + absl::c_count_if(kernel_spec_.input_buffer_specs, is_multimem_buffer) + + absl::c_count_if(kernel_spec_.output_buffer_specs, is_multimem_buffer); + std::vector param_to_peers_ptrs(num_parameters * + clique_key.num_devices()); + for (size_t i = scratch_buffers_index; i < num_parameters; ++i) { + const size_t scratch_index = i - scratch_buffers_index; + auto sym_mem = + memory_state->scratch_symmetric_memories[scratch_index].Lock(); + TF_RET_CHECK(sym_mem != nullptr) + << "Symmetric memory for scratch buffer " << scratch_index + << " is no longer valid"; + for (int device_rank = 0; device_rank < clique_key.num_devices(); + ++device_rank) { + ABSL_ASSIGN_OR_RETURN(se::DeviceAddressBase peer_address, + sym_mem->peer_addr(RankId(device_rank))); + const size_t parameter_offset = i * clique_key.num_devices(); + param_to_peers_ptrs[parameter_offset + device_rank] = + peer_address.opaque(); + } + } const size_t multimem_size_bytes = multimem_addresses.size() * sizeof(void*); state->metadata = params.executor->Allocate( diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h index 0316d660cf4b81..760b3643accba4 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h @@ -37,6 +37,7 @@ limitations under the License.*/ #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/backends/gpu/runtime/traced_command.h" #include "xla/core/collectives/rank_id.h" +#include "xla/core/collectives/symmetric_memory.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/stream_executor/device_address.h" @@ -44,6 +45,7 @@ limitations under the License.*/ #include "xla/stream_executor/gpu/all_reduce_kernel.h" #include "xla/stream_executor/kernel.h" #include "xla/stream_executor/stream.h" +#include "xla/tsl/util/tied_ref.h" namespace xla::gpu { @@ -132,6 +134,7 @@ class CollectiveKernelThunk : public TracedCommand { // Per-executor scratch memory. struct StreamMemory { std::vector scratch_allocations; + std::vector> scratch_symmetric_memories; }; // Per-executor state that needs to be synchronized for access. diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index ec65e9093d7f75..ef1f736b916900 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -372,6 +372,7 @@ absl::StatusOr RunCollectiveKernelThunk( initialize_params.stream = stream.get(); initialize_params.buffer_allocations = &buffer_allocations; initialize_params.collective_params = &collective_params; + initialize_params.collective_cliques = &collective_cliques; initialize_params.src = {kKernelSource}; initialize_params.collective_memory = &collective_memory; @@ -477,7 +478,7 @@ TEST(CollectiveKernelThunkTest, MultiprocessTest) { /*is_multimem_enabled=*/false, /*use_ptx=*/true); EXPECT_THAT(RunCollectiveKernelThunkOnDevices(metadata, /*emulate_multiprocess=*/true), - StatusIs(absl::StatusCode::kInvalidArgument)); + StatusIs(absl::StatusCode::kNotFound)); } TEST(CollectiveKernelThunkTest, BufferUses) { @@ -615,19 +616,29 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) { &allocations1}; ASSERT_OK(collective_kernel_thunk->Prepare(prepare_params)); + CollectiveMemoryCache collective_memory_cache; + ASSERT_OK_AND_ASSIGN( + CollectiveCliques collective_cliques, + AcquireCollectiveCliques(collective_params, clique_requests)); + ASSERT_OK_AND_ASSIGN( + CollectiveMemory collective_memory, + AcquireCollectiveMemory(collective_params, collective_cliques, + memory_requests, collective_memory_cache)); + Thunk::InitializeParams initialize_params; initialize_params.executor = executor; initialize_params.stream = stream.get(); initialize_params.buffer_allocations = &allocations1; initialize_params.collective_params = &collective_params; + initialize_params.collective_cliques = &collective_cliques; initialize_params.src.text = kKernelSource; + initialize_params.collective_memory = &collective_memory; ASSERT_OK(collective_kernel_thunk->Initialize(initialize_params)); ASSERT_OK(stream->BlockHostUntilDone()); Thunk::ExecuteParams params1 = Thunk::ExecuteParams::Create( run_options, allocations1, stream.get(), trace_stream.get(), - &collective_params, /*collective_cliques=*/nullptr, - /*collective_memory=*/nullptr); + &collective_params, &collective_cliques, &collective_memory); CommandStateManager state; Command::RecordParams record_params = {state}; @@ -648,8 +659,7 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) { BufferAllocations updated_allocations({src2, dst2}, 0, nullptr); Thunk::ExecuteParams params2 = Thunk::ExecuteParams::Create( run_options, updated_allocations, stream.get(), trace_stream.get(), - &collective_params, /*collective_cliques=*/nullptr, - /*collective_memory=*/nullptr); + &collective_params, &collective_cliques, &collective_memory); std::vector updated_allocs = {0, 1}; Command::RecordParams update_record_params = {state, std::move(updated_allocs)}; From 86017959633faf04f70ecfc407914e3e5ac5b4be Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Fri, 28 Aug 2026 05:03:17 -0700 Subject: [PATCH 04/11] PR #47701: [xla] Extract a separate collective_rendezvous library Imported from GitHub PR https://github.com/openxla/xla/pull/47701 **NFC:** splitting functionality into separate library Extract `collective_rendezvous` from `collective_ops_utils` as it's really a runtime support and node used for collective ops analysis/optimization. Copybara import of the project: -- 6b91ae42a6453e1adcac7e7835a55179db1b9dd5 by Eugene Zhulenev : [xla] Extract a separate collective_rendezvous library Merging this change closes #47701 PiperOrigin-RevId: 972518994 --- .../xla/xla/backends/cpu/collectives/BUILD | 7 +- .../cpu/collectives/cpu_collectives.cc | 2 +- .../cpu/collectives/cpu_collectives.h | 2 +- .../cpu/collectives/gloo_collectives_test.cc | 2 +- .../collectives/in_process_communicator.cc | 3 +- .../cpu/collectives/in_process_communicator.h | 2 +- .../xla/xla/backends/cpu/runtime/BUILD | 13 ++- .../backends/cpu/runtime/all_gather_thunk.cc | 2 +- .../backends/cpu/runtime/all_reduce_thunk.cc | 2 +- .../backends/cpu/runtime/all_reduce_thunk.h | 2 +- .../backends/cpu/runtime/all_to_all_thunk.cc | 2 +- .../cpu/runtime/collective_permute_thunk.cc | 2 +- .../backends/cpu/runtime/collective_thunk.h | 2 +- .../cpu/runtime/reduce_scatter_thunk.cc | 2 +- .../cpu/runtime/reduce_scatter_thunk.h | 2 +- .../xla/xla/megascale/c_api_client/BUILD | 2 +- .../c_api_client/megascale_c_api_client.cc | 2 +- third_party/xla/xla/pjrt/c/BUILD | 2 +- .../pjrt/c/pjrt_c_api_collectives_internal.cc | 2 +- third_party/xla/xla/service/BUILD | 15 ++- .../xla/xla/service/collective_ops_utils.h | 85 ----------------- .../xla/xla/service/collective_rendezvous.cc | 58 ++++++++++++ .../xla/xla/service/collective_rendezvous.h | 91 +++++++++++++++++++ 23 files changed, 192 insertions(+), 112 deletions(-) create mode 100644 third_party/xla/xla/service/collective_rendezvous.cc create mode 100644 third_party/xla/xla/service/collective_rendezvous.h diff --git a/third_party/xla/xla/backends/cpu/collectives/BUILD b/third_party/xla/xla/backends/cpu/collectives/BUILD index 639f786f171147..e9882d7ba95928 100644 --- a/third_party/xla/xla/backends/cpu/collectives/BUILD +++ b/third_party/xla/xla/backends/cpu/collectives/BUILD @@ -120,7 +120,7 @@ cc_library( "//xla/core/collectives:collectives_registry", "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "@com_google_absl//absl/base", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -163,7 +163,8 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", - "//xla/service:collective_ops_utils", + "//xla/core/collectives:reduction_kind", + "//xla/service:collective_rendezvous", "//xla/service:rendezvous", "//xla/stream_executor:device_address", "//xla/tsl/lib/math:math_util", @@ -249,7 +250,7 @@ xla_cc_test( "//xla/pjrt/distributed:in_memory_key_value_store", "//xla/pjrt/distributed:key_value_store_interface", "//xla/runtime:device_id", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/stream_executor:device_address", "//xla/tsl/concurrency:async_value", "//xla/tsl/lib/core:status_test_util", diff --git a/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.cc b/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.cc index c5acdaacd9ccbd..3f563785e70ced 100644 --- a/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.cc +++ b/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.cc @@ -23,7 +23,7 @@ limitations under the License. #include "xla/core/collectives/collectives.h" #include "xla/core/collectives/collectives_registry.h" #include "xla/core/collectives/communicator.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/util.h" #include "tsl/platform/casts.h" diff --git a/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.h b/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.h index 4e864c139f890b..582c97ad87553f 100644 --- a/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.h +++ b/third_party/xla/xla/backends/cpu/collectives/cpu_collectives.h @@ -27,7 +27,7 @@ limitations under the License. #include "xla/core/collectives/collectives.h" #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/util.h" #include "xla/xla_data.pb.h" diff --git a/third_party/xla/xla/backends/cpu/collectives/gloo_collectives_test.cc b/third_party/xla/xla/backends/cpu/collectives/gloo_collectives_test.cc index 518d5f71b128a9..dc5f327a0f2855 100644 --- a/third_party/xla/xla/backends/cpu/collectives/gloo_collectives_test.cc +++ b/third_party/xla/xla/backends/cpu/collectives/gloo_collectives_test.cc @@ -37,7 +37,7 @@ limitations under the License. #include "xla/pjrt/distributed/in_memory_key_value_store.h" #include "xla/pjrt/distributed/key_value_store_interface.h" #include "xla/runtime/device_id.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/stream_executor/device_address.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/env.h" diff --git a/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.cc b/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.cc index 058b59ff544cdb..152f73a9bbe865 100644 --- a/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.cc +++ b/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.cc @@ -38,10 +38,11 @@ limitations under the License. #include "absl/types/span.h" #include "xla/backends/cpu/collectives/cpu_collectives.h" #include "xla/core/collectives/rank_id.h" +#include "xla/core/collectives/reduction_kind.h" #include "xla/debug_options_flags.h" #include "xla/future.h" #include "xla/primitive_util.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/service/rendezvous.h" #include "xla/stream_executor/device_address.h" #include "xla/tsl/lib/math/math_util.h" diff --git a/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.h b/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.h index 55691f1a0860f6..20561962e66b88 100644 --- a/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.h +++ b/third_party/xla/xla/backends/cpu/collectives/in_process_communicator.h @@ -26,8 +26,8 @@ limitations under the License. #include "absl/types/span.h" #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" +#include "xla/core/collectives/reduction_kind.h" #include "xla/future.h" -#include "xla/service/collective_ops_utils.h" #include "xla/stream_executor/device_address.h" #include "xla/util.h" #include "xla/xla_data.pb.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/BUILD b/third_party/xla/xla/backends/cpu/runtime/BUILD index d325b7e528965c..fc20ecd3f5fe5a 100644 --- a/third_party/xla/xla/backends/cpu/runtime/BUILD +++ b/third_party/xla/xla/backends/cpu/runtime/BUILD @@ -386,7 +386,7 @@ cc_library( "//xla:shape_util", "//xla/backends/cpu/collectives:cpu_collectives", "//xla/core/collectives:communicator", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:inlined_vector", @@ -513,8 +513,9 @@ cc_library( "//xla:util", "//xla/backends/cpu/collectives:cpu_collectives", "//xla/core/collectives:communicator", + "//xla/core/collectives:reduction_kind", "//xla/service:buffer_assignment", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:inlined_vector", @@ -540,7 +541,7 @@ cc_library( "//xla/backends/cpu/collectives:cpu_collectives", "//xla/core/collectives:communicator", "//xla/service:buffer_assignment", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:logging", "//xla/tsl/platform:statusor", @@ -565,7 +566,8 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/backends/cpu/collectives:cpu_collectives", "//xla/core/collectives:communicator", - "//xla/service:collective_ops_utils", + "//xla/core/collectives:reduction_kind", + "//xla/service:collective_rendezvous", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:inlined_vector", @@ -592,7 +594,7 @@ cc_library( "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", "//xla/service:buffer_assignment", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/service:computation_placer", "//xla/tsl/concurrency:async_value", "//xla/tsl/platform:statusor", @@ -630,6 +632,7 @@ cc_library( "//xla/runtime:resource_use", "//xla/service:buffer_assignment", "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/service:computation_placer", "//xla/service:hlo_proto_cc", "//xla/stream_executor:device_address", diff --git a/third_party/xla/xla/backends/cpu/runtime/all_gather_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/all_gather_thunk.cc index de24beeeb8063d..7291a51f47aa06 100644 --- a/third_party/xla/xla/backends/cpu/runtime/all_gather_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/all_gather_thunk.cc @@ -31,7 +31,7 @@ limitations under the License. #include "xla/backends/cpu/runtime/thunk.h" #include "xla/core/collectives/communicator.h" #include "xla/future.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.cc index dedc406c18021b..7852ea5b0e1c73 100644 --- a/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.cc @@ -35,7 +35,7 @@ limitations under the License. #include "xla/future.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.h b/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.h index a639b818539640..7a1a96d94fc06d 100644 --- a/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.h +++ b/third_party/xla/xla/backends/cpu/runtime/all_reduce_thunk.h @@ -20,8 +20,8 @@ limitations under the License. #include "absl/status/statusor.h" #include "xla/backends/cpu/runtime/collective_thunk.h" +#include "xla/core/collectives/reduction_kind.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/collective_ops_utils.h" #include "xla/tsl/concurrency/async_value_ref.h" namespace xla::cpu { diff --git a/third_party/xla/xla/backends/cpu/runtime/all_to_all_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/all_to_all_thunk.cc index d17125c61dd4b2..3a96a22abe8f4d 100644 --- a/third_party/xla/xla/backends/cpu/runtime/all_to_all_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/all_to_all_thunk.cc @@ -29,7 +29,7 @@ limitations under the License. #include "xla/core/collectives/communicator.h" #include "xla/future.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/collective_permute_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/collective_permute_thunk.cc index 129a8b11a6bea1..7d44b6951897cf 100644 --- a/third_party/xla/xla/backends/cpu/runtime/collective_permute_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/collective_permute_thunk.cc @@ -38,7 +38,7 @@ limitations under the License. #include "xla/core/collectives/rank_id.h" #include "xla/future.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/service/computation_placer.h" #include "xla/shape.h" #include "xla/shape_util.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/collective_thunk.h b/third_party/xla/xla/backends/cpu/runtime/collective_thunk.h index 0ba4c66178133f..4ea2476b8a4d9c 100644 --- a/third_party/xla/xla/backends/cpu/runtime/collective_thunk.h +++ b/third_party/xla/xla/backends/cpu/runtime/collective_thunk.h @@ -35,7 +35,7 @@ limitations under the License. #include "xla/runtime/device_id.h" #include "xla/runtime/resource_use.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/shape.h" #include "xla/stream_executor/device_address.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.cc index 11adc9f67ce522..30843abbb70dfc 100644 --- a/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.cc +++ b/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.cc @@ -32,7 +32,7 @@ limitations under the License. #include "xla/core/collectives/communicator.h" #include "xla/future.h" #include "xla/primitive_util.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/concurrency/async_value_ref.h" diff --git a/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.h b/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.h index d0efd32e677be6..3919a5bade889f 100644 --- a/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.h +++ b/third_party/xla/xla/backends/cpu/runtime/reduce_scatter_thunk.h @@ -20,7 +20,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "xla/backends/cpu/runtime/collective_thunk.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/core/collectives/reduction_kind.h" #include "xla/tsl/concurrency/async_value_ref.h" #include "xla/xla_data.pb.h" diff --git a/third_party/xla/xla/megascale/c_api_client/BUILD b/third_party/xla/xla/megascale/c_api_client/BUILD index 6aac0b5d4f63db..1a0d5c3b9c3266 100644 --- a/third_party/xla/xla/megascale/c_api_client/BUILD +++ b/third_party/xla/xla/megascale/c_api_client/BUILD @@ -58,7 +58,7 @@ cc_library( "//xla/pjrt/c_api_client:pjrt_c_api_client", "//xla/pjrt/c_api_client:pjrt_c_api_multi_slice_config", "//xla/pjrt/plugin:plugin_names", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/stream_executor:device_address", "//xla/tsl/platform:logging", "//xla/tsl/platform:macros", diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc index deed08ecf60bdc..8a26c9c477274a 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc @@ -64,7 +64,7 @@ limitations under the License. #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/pjrt/plugin/plugin_names.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/stream_executor/device_address.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/macros.h" diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD index 6c545ce6fd25e3..c0d66b6772925d 100644 --- a/third_party/xla/xla/pjrt/c/BUILD +++ b/third_party/xla/xla/pjrt/c/BUILD @@ -172,7 +172,7 @@ cc_library( "//xla/core/collectives:rank_id", "//xla/core/collectives:reduction_kind", "//xla/runtime:device_id", - "//xla/service:collective_ops_utils", + "//xla/service:collective_rendezvous", "//xla/tsl/concurrency:future", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_collectives_internal.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_collectives_internal.cc index d87e2d3dce831e..c3530c1b4168b4 100644 --- a/third_party/xla/xla/pjrt/c/pjrt_c_api_collectives_internal.cc +++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_collectives_internal.cc @@ -41,7 +41,7 @@ limitations under the License. #include "xla/pjrt/c/pjrt_c_api_helpers.h" #include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h" #include "xla/runtime/device_id.h" -#include "xla/service/collective_ops_utils.h" +#include "xla/service/collective_rendezvous.h" #include "xla/tsl/concurrency/future.h" typedef struct PJRT_Collectives_Communicator { diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index ec981df50f3178..443fdfb96c9a24 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -5816,6 +5816,19 @@ xla_cc_test( ], ) +cc_library( + name = "collective_rendezvous", + srcs = ["collective_rendezvous.cc"], + hdrs = ["collective_rendezvous.h"], + deps = [ + "//xla:executable_run_options", + "//xla/runtime:device_id", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + ], +) + cc_library( name = "collective_ops_utils", srcs = ["collective_ops_utils.cc"], @@ -5826,7 +5839,6 @@ cc_library( ":hlo_module_config", ":pattern_matcher", ":source_target_pairs", - "//xla:executable_run_options", "//xla:literal", "//xla:literal_util", "//xla:shape_util", @@ -5847,7 +5859,6 @@ cc_library( "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@jsoncpp_git//:jsoncpp", diff --git a/third_party/xla/xla/service/collective_ops_utils.h b/third_party/xla/xla/service/collective_ops_utils.h index c339c4f7a36c9d..72b2aa73505866 100644 --- a/third_party/xla/xla/service/collective_ops_utils.h +++ b/third_party/xla/xla/service/collective_ops_utils.h @@ -25,12 +25,9 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/statusor.h" -#include "absl/strings/str_format.h" -#include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/core/collectives/reduction_kind.h" -#include "xla/executable_run_options.h" #include "xla/hlo/ir/collective_op_group_mode.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" @@ -294,88 +291,6 @@ void CopyCollectiveGroupKey(const HloInstruction& source, // Removes the collective_group_key attribute. void ClearCollectiveGroupKey(HloInstruction& instruction); -//===----------------------------------------------------------------------===// -// Collective rendezvous. -//===----------------------------------------------------------------------===// - -// Key that identifies a particular Rendezvous object in our global hashtable. -// This determines which calls to ExecuteOnStream communicate with each other. -// The rules are as follows. -// -// * Only ops with the same RunId can communicate with each other. (This is the -// whole purpose of RunId). -// -// * Only ops with the same set of participating replicas can communicate with -// each other. This is how we separate out different replica groups (e.g. a -// single AllReduce HLO might do two reductions, between say GPUs {0,2} and -// {1,3}). -// -// * Only ops with the same opcode can communicate with each other. At the -// moment we only support kAllReduce, so we don't check for this explicitly. -// -// * For cross-module all-reduces (i.e. instr->channel_id().has_value()), -// only ops with the same value for channel_id() can communicate with each -// other. -// -// * For cross-replica (i.e. same-module) all-reduces (i.e. -// !channel_id().has_value()), only ops from the same module (as -// identified by its unique_id()) can communicate with each other. -// -struct RendezvousKey { - enum CollectiveOpKind { - kCrossModule, - kCrossReplica, - }; - - explicit RendezvousKey(const RunId& run_id, - std::vector global_devices, - int num_local_participants, - CollectiveOpKind collective_op_kind, int64_t op_id) - : run_id(run_id), - global_devices(std::move(global_devices)), - num_local_participants(num_local_participants), - collective_op_kind(collective_op_kind), - op_id(op_id) {} - - template - friend H AbslHashValue(H h, const RendezvousKey& k) { - return H::combine(std::move(h), k.run_id, k.global_devices, - k.num_local_participants, k.collective_op_kind, k.op_id); - } - friend bool operator==(const RendezvousKey& a, const RendezvousKey& b) { - return a.run_id == b.run_id && a.global_devices == b.global_devices && - a.num_local_participants == b.num_local_participants && - a.collective_op_kind == b.collective_op_kind && // - a.op_id == b.op_id; - } - friend bool operator!=(const RendezvousKey& a, const RendezvousKey& b) { - return !(a == b); - } - - absl::string_view CollectiveOpKindString() const { - switch (collective_op_kind) { - case kCrossModule: - return "cross_module"; - case kCrossReplica: - return "cross_replica"; - } - } - - std::string ToString() const { - return absl::StrFormat( - "RendezvousKey{run_id=%s, global_devices=[%s], " - "num_local_participants=%d, collective_op_kind=%s, op_id=%d}", - run_id.ToString(), absl::StrJoin(global_devices, ", "), - num_local_participants, CollectiveOpKindString(), op_id); - } - - RunId run_id; - std::vector global_devices; - int num_local_participants; - CollectiveOpKind collective_op_kind; - int64_t op_id; -}; - //===----------------------------------------------------------------------===// // Collective execution utilities. //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/service/collective_rendezvous.cc b/third_party/xla/xla/service/collective_rendezvous.cc new file mode 100644 index 00000000000000..7fd8644bbb1655 --- /dev/null +++ b/third_party/xla/xla/service/collective_rendezvous.cc @@ -0,0 +1,58 @@ +/* Copyright 2026 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. +==============================================================================*/ + +#include "xla/service/collective_rendezvous.h" + +#include +#include +#include +#include + +#include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" +#include "absl/strings/string_view.h" +#include "xla/executable_run_options.h" +#include "xla/runtime/device_id.h" + +namespace xla { + +RendezvousKey::RendezvousKey(const RunId& run_id, + std::vector global_devices, + int num_local_participants, + CollectiveOpKind collective_op_kind, int64_t op_id) + : run_id(run_id), + global_devices(std::move(global_devices)), + num_local_participants(num_local_participants), + collective_op_kind(collective_op_kind), + op_id(op_id) {} + +absl::string_view RendezvousKey::CollectiveOpKindString() const { + switch (collective_op_kind) { + case kCrossModule: + return "cross_module"; + case kCrossReplica: + return "cross_replica"; + } +} + +std::string RendezvousKey::ToString() const { + return absl::StrFormat( + "RendezvousKey{run_id=%s, global_devices=[%s], " + "num_local_participants=%d, collective_op_kind=%s, op_id=%d}", + run_id.ToString(), absl::StrJoin(global_devices, ", "), + num_local_participants, CollectiveOpKindString(), op_id); +} + +} // namespace xla diff --git a/third_party/xla/xla/service/collective_rendezvous.h b/third_party/xla/xla/service/collective_rendezvous.h new file mode 100644 index 00000000000000..1df703992964f5 --- /dev/null +++ b/third_party/xla/xla/service/collective_rendezvous.h @@ -0,0 +1,91 @@ +/* Copyright 2026 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. +==============================================================================*/ + +#ifndef XLA_SERVICE_COLLECTIVE_RENDEZVOUS_H_ +#define XLA_SERVICE_COLLECTIVE_RENDEZVOUS_H_ + +#include +#include +#include +#include + +#include "absl/strings/string_view.h" +#include "xla/executable_run_options.h" +#include "xla/runtime/device_id.h" + +namespace xla { + +// Key that identifies a particular Rendezvous object in our global hashtable. +// This determines which calls to ExecuteOnStream communicate with each other. +// The rules are as follows. +// +// * Only ops with the same RunId can communicate with each other. (This is the +// whole purpose of RunId). +// +// * Only ops with the same set of participating replicas can communicate with +// each other. This is how we separate out different replica groups (e.g. a +// single AllReduce HLO might do two reductions, between say GPUs {0,2} and +// {1,3}). +// +// * Only ops with the same opcode can communicate with each other. At the +// moment we only support kAllReduce, so we don't check for this explicitly. +// +// * For cross-module all-reduces (i.e. instr->channel_id().has_value()), +// only ops with the same value for channel_id() can communicate with each +// other. +// +// * For cross-replica (i.e. same-module) all-reduces (i.e. +// !channel_id().has_value()), only ops from the same module (as +// identified by its unique_id()) can communicate with each other. +// +struct RendezvousKey { + enum CollectiveOpKind { + kCrossModule, + kCrossReplica, + }; + + explicit RendezvousKey(const RunId& run_id, + std::vector global_devices, + int num_local_participants, + CollectiveOpKind collective_op_kind, int64_t op_id); + + template + friend H AbslHashValue(H h, const RendezvousKey& k) { + return H::combine(std::move(h), k.run_id, k.global_devices, + k.num_local_participants, k.collective_op_kind, k.op_id); + } + friend bool operator==(const RendezvousKey& a, const RendezvousKey& b) { + return a.run_id == b.run_id && a.global_devices == b.global_devices && + a.num_local_participants == b.num_local_participants && + a.collective_op_kind == b.collective_op_kind && // + a.op_id == b.op_id; + } + friend bool operator!=(const RendezvousKey& a, const RendezvousKey& b) { + return !(a == b); + } + + absl::string_view CollectiveOpKindString() const; + std::string ToString() const; + + RunId run_id; + std::vector global_devices; + int num_local_participants; + CollectiveOpKind collective_op_kind; + int64_t op_id; +}; + +} // namespace xla + +#endif // XLA_SERVICE_COLLECTIVE_RENDEZVOUS_H_ From 5b4ac04eca39a13071ad10a2e0f9f7a5133a371e Mon Sep 17 00:00:00 2001 From: James Spooner Date: Fri, 28 Aug 2026 05:05:14 -0700 Subject: [PATCH 05/11] Fix schedule-aware collective CSE to advance candidates when distance threshold is exceeded. Iterate by reference in the earlier collectives loop so reassigning earlier_coll updates the vector entry instead of being a local no-op. PiperOrigin-RevId: 972519816 --- .../spmd/schedule_aware_collective_ops_cse.cc | 2 +- .../schedule_aware_collective_ops_cse_test.cc | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse.cc b/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse.cc index 993655772ced77..6b92f9b2673207 100644 --- a/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse.cc +++ b/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse.cc @@ -136,7 +136,7 @@ absl::StatusOr RunOnComputation(HloComputation* comp, bool for_replicas, coll->operand(0))]; bool found = false; int64_t coll_height = height[coll]; - for (HloInstruction* earlier_coll : earlier_colls) { + for (HloInstruction*& earlier_coll : earlier_colls) { if (!ShapeUtil::Equal(earlier_coll->shape(), coll->shape())) { continue; } diff --git a/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse_test.cc b/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse_test.cc index d8affb65247160..5cc9426e73f0d6 100644 --- a/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse_test.cc +++ b/third_party/xla/xla/service/spmd/schedule_aware_collective_ops_cse_test.cc @@ -325,6 +325,43 @@ ENTRY entry { EXPECT_TRUE(RunFileCheck(module->ToString(), hlo_string).value()); } +TEST_F(CollectiveOpsCseTest, MultipleCollectivesWithDistantPredecessor) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + param0 = s32[1,8]{1,0} parameter(0) + ag0 = s32[2,8]{1,0} all-gather(param0), replica_groups={{0,1}}, dimensions={0}, + channel_id=0, use_global_device_ids=true + chain0 = s32[2,8]{1,0} negate(ag0) + chain1 = s32[2,8]{1,0} negate(chain0) + chain2 = s32[2,8]{1,0} negate(chain1) + chain3 = s32[2,8]{1,0} negate(chain2) + chain4 = s32[2,8]{1,0} negate(chain3) + chain5 = s32[2,8]{1,0} negate(chain4) + ag1 = s32[2,8]{1,0} all-gather(param0), replica_groups={{0,1}}, dimensions={0}, + channel_id=0, use_global_device_ids=true + ag2 = s32[2,8]{1,0} all-gather(param0), replica_groups={{0,1}}, dimensions={0}, + channel_id=0, use_global_device_ids=true + add1 = s32[2,8]{1,0} add(chain5, ag1) + add2 = s32[2,8]{1,0} add(chain5, ag2) + ROOT tuple = (s32[2,8]{1,0}, s32[2,8]{1,0}) tuple(add1, add2) +})"; + // Run with distance_threshold = 3 + auto module_status = RunPass(hlo_string, /*distance_threshold=*/3); + EXPECT_TRUE(module_status.status().ok()); + auto module = std::move(module_status).value(); + HloInstruction* tuple = module->entry_computation()->root_instruction(); + EXPECT_EQ(tuple->opcode(), HloOpcode::kTuple); + HloInstruction* add1 = tuple->mutable_operand(0); + HloInstruction* add2 = tuple->mutable_operand(1); + EXPECT_EQ(add1->opcode(), HloOpcode::kAdd); + EXPECT_EQ(add2->opcode(), HloOpcode::kAdd); + // ag1 and ag2 are close to each other, so ag2 should have been replaced with + // ag1! + EXPECT_EQ(add1->operand(1), add2->operand(1)); +} + } // namespace } // namespace spmd } // namespace xla From e7e48554f5c663eb2c1d8c7a1d3b580e97bb896f Mon Sep 17 00:00:00 2001 From: Emilio Cota Date: Fri, 28 Aug 2026 08:49:04 -0700 Subject: [PATCH 06/11] [re-land][xla:cpu] add full msan support Msan support is currently incomplete. For example, we don't add the appropriate msan pass to LLVM's pipeline; instead, we manually call __msan_unpoison here and there. In AOT (tfcompile) the situation is even worse, since we only add these annotations to some of the emitted code. This change adds complete support for msan (and msan-track-origins) to XLA:CPU in both JIT and AOT modes. In particular: - Adds a single boolean config flag to determine whether msan should be enabled in compiled code. This is then wired up through the compiler, removing several #ifdef MEMORY_SANITIZER (or equivalent) checks that we had scattered around the code base. (OK, this is two knobs because we also support track-origins.) - Adds the msan pass to LLVM as documented in https://github.com/google/sanitizers/wiki/MemorySanitizerJIT This requires a small TLS emulation library, which (1) in JIT gives access to host symbols from JIT'ed code, and (2) in AOT (tfcompile) must be linked in the final binary so that the compiled code can reliably access msan symbols. What's missing from this CL: 1. Removal of now obsolete annotations 2. Setting of the appropriate flags in tfcompile to enable the new path in AOT. These two changes will land once users have had time to recompile their compilers (in case they embed tfcompile into them), about 1-2 weeks. Note: the msan support implemented in this CL supersedes cl/959445928, i.e. https://github.com/openxla/xla/commit/70f2607611. That means the xla_backend_extra option added there for msan on AOT will be removed when landing the two missing pieces from this CL. Reverts ef7547c6576081827969042b60dfee5be28b02d7 PiperOrigin-RevId: 972614006 --- tensorflow/compiler/aot/compile.cc | 6 + tensorflow/compiler/aot/flags.cc | 6 + tensorflow/compiler/aot/flags.h | 2 + tensorflow/compiler/aot/tfcompile.bzl | 1 + tensorflow/compiler/tf2xla/BUILD | 1 + .../xla/xla/backends/cpu/codegen/BUILD | 7 + .../codegen/builtin_definition_generator.cc | 51 ++++++- .../codegen/builtin_definition_generator.h | 8 ++ .../xla/backends/cpu/codegen/ir_compiler.cc | 133 ++++++++++++++++-- .../xla/backends/cpu/codegen/ir_compiler.h | 27 +++- .../backends/cpu/codegen/ir_compiler_test.cc | 47 ++++++- .../backends/cpu/codegen/jit_compiler_test.cc | 33 ++--- third_party/xla/xla/backends/cpu/nanort/BUILD | 1 + .../backends/cpu/nanort/nanort_client_test.cc | 53 +++++++ .../xla/xla/backends/cpu/runtime/BUILD | 16 +++ .../backends/cpu/runtime/msan_emulated_tls.cc | 80 +++++++++++ .../backends/cpu/runtime/msan_emulated_tls.h | 47 +++++++ third_party/xla/xla/service/compiler.h | 17 +++ third_party/xla/xla/service/cpu/BUILD | 1 + .../xla/xla/service/cpu/cpu_compiler.cc | 5 + third_party/xla/xla/service/cpu/cpu_runtime.h | 7 +- .../cpu/restricted/cpu_aot_compiler_test.cc | 10 ++ 22 files changed, 511 insertions(+), 48 deletions(-) create mode 100644 third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc create mode 100644 third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index efc6869f6235de..d1fbed6a8783ed 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -206,6 +206,12 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config, flags.sanitize_abilists_dataflow, ',', absl::SkipEmpty())); } + if (flags.sanitize_memory || flags.sanitize_memory_track_origins > 0) { + aot_opts.set_sanitize_memory(true); + aot_opts.set_sanitize_memory_track_origins( + flags.sanitize_memory_track_origins); + } + TF_RETURN_IF_ERROR( ConfigureKernelNamingConvention(aot_opts, computation, flags.cpp_class)); diff --git a/tensorflow/compiler/aot/flags.cc b/tensorflow/compiler/aot/flags.cc index 567426f53c7631..6555d22d05c325 100644 --- a/tensorflow/compiler/aot/flags.cc +++ b/tensorflow/compiler/aot/flags.cc @@ -86,6 +86,12 @@ void AppendMainFlags(std::vector* flag_list, MainFlags* flags) { "Enable DataFlow Sanitizer pass."}, {"sanitize_abilists_dataflow", &flags->sanitize_abilists_dataflow, "Comma separated list of ABIList file paths."}, + {"sanitize_memory", &flags->sanitize_memory, + "Enable Memory Sanitizer pass."}, + {"sanitize_memory_track_origins", &flags->sanitize_memory_track_origins, + "Controls MSan track origins level (0=disabled, 1=without store " + "history, 2=with store history). Setting to >0 implies " + "--sanitize_memory."}, {"gen_name_to_index", &flags->gen_name_to_index, "Generate name-to-index data for Lookup{Arg,Result}Index methods."}, {"gen_program_shape", &flags->gen_program_shape, diff --git a/tensorflow/compiler/aot/flags.h b/tensorflow/compiler/aot/flags.h index 5d0f93f7d67b88..f71fb80e0973f4 100644 --- a/tensorflow/compiler/aot/flags.h +++ b/tensorflow/compiler/aot/flags.h @@ -48,6 +48,8 @@ struct MainFlags { // Sanitizer pass options bool sanitize_dataflow = false; std::string sanitize_abilists_dataflow; + bool sanitize_memory = false; + int32_t sanitize_memory_track_origins = 0; // C++ codegen options bool gen_name_to_index = false; diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 911c917350d775..83c41ba27c2d20 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -358,6 +358,7 @@ def _tf_library( "@xla//xla/backends/cpu/runtime:sort_lib", "@xla//xla/backends/cpu/runtime:topk_lib", "@xla//xla/backends/cpu/runtime:convolution_lib", + "@xla//xla/backends/cpu/runtime:msan_emulated_tls", "@xla//xla/service/cpu:runtime_matmul", "@xla//xla/service/cpu:runtime_single_threaded_matmul", "@eigen_archive//:eigen3", diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 92786fc8f30ed7..f111426835174d 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -485,6 +485,7 @@ cc_library( "@com_google_absl//absl/types:span", ":encoded_buffer_allocation_info", "@xla//xla/service:custom_call_status_internal", + "@xla//xla/backends/cpu/runtime:msan_emulated_tls", "@xla//xla/backends/cpu/runtime:rng_state_lib", "@xla//xla/backends/cpu:alignment", "@xla//xla/backends/cpu:buffer_allocation_info", diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD index 40f90b5b5e18df..188a9299c67e8f 100644 --- a/third_party/xla/xla/backends/cpu/codegen/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/BUILD @@ -48,8 +48,12 @@ cc_library( deps = [ ":builtin_fp16", ":builtin_pow", + "//xla/backends/cpu/runtime:msan_emulated_tls", + "//xla/service/cpu:cpu_runtime", + "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@llvm-project//llvm:Core", "@llvm-project//llvm:OrcJIT", "@llvm-project//llvm:OrcShared", @@ -124,18 +128,21 @@ cc_library( "//xla:util", "//xla:xla_proto_cc", "//xla/backends/cpu:target_machine_options", + "//xla/backends/cpu/runtime:msan_emulated_tls", "//xla/codegen:intrinsic_lib", "//xla/codegen/intrinsic", "//xla/codegen/intrinsic:intrinsic_compiler_lib", "//xla/service:hlo_module_config", "//xla/service/cpu:backend_config_proto_cc", "//xla/service/cpu:cpu_options", + "//xla/service/cpu:cpu_runtime", "//xla/service/cpu:executable_proto_cc", "//xla/service/llvm_ir:llvm_util", "//xla/tools:llvm_targets", # fixdeps: keep "//xla/tsl/platform:logging", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", + "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/log", diff --git a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc index 97e056df2882c2..858223ff375ba1 100644 --- a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc +++ b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.cc @@ -28,19 +28,24 @@ limitations under the License. #include #include +#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" +#include "absl/log/log.h" // IWYU pragma: keep #include "llvm/ADT/StringRef.h" #include "llvm/ExecutionEngine/JITSymbol.h" #include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h" #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/CoreContainers.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" // IWYU pragma: keep (msan) #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h" #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h" #include "llvm/IR/DataLayout.h" #include "llvm/Support/Error.h" #include "xla/backends/cpu/codegen/builtin_fp16.h" #include "xla/backends/cpu/codegen/builtin_pow.h" +#include "xla/backends/cpu/runtime/msan_emulated_tls.h" +#include "xla/service/cpu/cpu_runtime.h" namespace xla::cpu { @@ -267,9 +272,8 @@ static Registry CreateRegistry() { #endif -#ifdef MEMORY_SANITIZER - registry["__msan_unpoison"] = SymbolDef(__msan_unpoison); -#endif + registry[runtime::kMsanEmutlsGetAddressBridgeSymbolName] = + SymbolDef(__xla_cpu_runtime_emutls_get_address); return registry; } @@ -280,22 +284,55 @@ static Registry CreateRegistry() { BuiltinDefinitionGenerator::BuiltinDefinitionGenerator( llvm::DataLayout data_layout) - : data_layout_(std::move(data_layout)) {} + : data_layout_(std::move(data_layout)) { +#ifdef ABSL_HAVE_MEMORY_SANITIZER + // Resolve MSan runtime functions (e.g. __msan_warning*) from the current + // process via dlsym. This is more future-proof than explicitly intercepting + // __msan_* functions; these functions do change between LLVM versions. + auto is_msan_symbol = [](const llvm::orc::SymbolStringPtr& name) { + return (*name).starts_with("__msan_"); + }; + auto generator = + llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( + data_layout_.getGlobalPrefix(), is_msan_symbol); + if (generator) { + process_generator_ = std::move(*generator); + } else { + LOG(WARNING) << "Failed to initialize dynamic library generator for MSan: " + << llvm::toString(generator.takeError()); + } +#endif +} llvm::Error BuiltinDefinitionGenerator::tryToGenerate( - llvm::orc::LookupState&, llvm::orc::LookupKind kind, - llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags, + llvm::orc::LookupState& ls, llvm::orc::LookupKind kind, + llvm::orc::JITDylib& jit_dylib, llvm::orc::JITDylibLookupFlags flags, const llvm::orc::SymbolLookupSet& names) { llvm::orc::SymbolMap symbols; symbols.reserve(names.size()); +#ifdef ABSL_HAVE_MEMORY_SANITIZER + llvm::orc::SymbolLookupSet msan_names; +#endif - for (const auto& [name, flags] : names) { + for (const auto& [name, name_flags] : names) { if (auto symbol = ResolveBuiltinSymbol(data_layout_, *name)) { symbols[name] = *symbol; +#ifdef ABSL_HAVE_MEMORY_SANITIZER + } else if ((*name).starts_with("__msan_")) { + msan_names.add(name, name_flags); +#endif } } cantFail(jit_dylib.define(llvm::orc::absoluteSymbols(std::move(symbols)))); + +#ifdef ABSL_HAVE_MEMORY_SANITIZER + if (!msan_names.empty() && process_generator_) { + return process_generator_->tryToGenerate(ls, kind, jit_dylib, flags, + msan_names); + } +#endif + return llvm::Error::success(); } diff --git a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h index 7689e7b13e425c..23ea4174cc1553 100644 --- a/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h +++ b/third_party/xla/xla/backends/cpu/codegen/builtin_definition_generator.h @@ -16,7 +16,12 @@ limitations under the License. #ifndef XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_ #define XLA_BACKENDS_CPU_CODEGEN_BUILTIN_DEFINITION_GENERATOR_H_ +#include +#include + +#include "absl/base/config.h" // IWYU pragma: keep #include "llvm/ExecutionEngine/Orc/Core.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" // IWYU pragma: keep #include "llvm/IR/DataLayout.h" #include "llvm/Support/Error.h" @@ -44,6 +49,9 @@ class BuiltinDefinitionGenerator : public llvm::orc::DefinitionGenerator { private: llvm::DataLayout data_layout_; +#ifdef ABSL_HAVE_MEMORY_SANITIZER + std::unique_ptr process_generator_; +#endif }; } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc index bda6062f77be13..0a451a6807f33a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/backends/cpu/codegen/ir_compiler.h" #include +#include #include #include #include @@ -24,6 +25,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/call_once.h" +#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/nullability.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -38,17 +40,25 @@ limitations under the License. #include "llvm-c/Target.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Analysis/CGSCCPassManager.h" +#include "llvm/Analysis/GlobalsModRef.h" #include "llvm/Analysis/LoopAnalysisManager.h" #include "llvm/Analysis/RuntimeLibcallInfo.h" #include "llvm/Analysis/TargetLibraryInfo.h" #include "llvm/ExecutionEngine/ExecutionEngine.h" #include "llvm/ExecutionEngine/Orc/Mangling.h" +#include "llvm/IR/Attributes.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/DataLayout.h" +#include "llvm/IR/DerivedTypes.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/LLVMContext.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/Metadata.h" #include "llvm/IR/PassManager.h" #include "llvm/IR/Verifier.h" #include "llvm/MC/MCContext.h" #include "llvm/Object/ObjectFile.h" +#include "llvm/Pass.h" #include "llvm/Passes/OptimizationLevel.h" #include "llvm/Passes/PassBuilder.h" #include "llvm/Passes/StandardInstrumentations.h" @@ -65,14 +75,17 @@ limitations under the License. #include "llvm/TargetParser/Triple.h" #include "llvm/Transforms/IPO/AlwaysInliner.h" #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h" +#include "llvm/Transforms/Instrumentation/MemorySanitizer.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" #include "xla/backends/cpu/codegen/polynomial_approximations.h" +#include "xla/backends/cpu/runtime/msan_emulated_tls.h" #include "xla/backends/cpu/target_machine_options.h" #include "xla/codegen/intrinsic/intrinsic.h" #include "xla/codegen/intrinsic/intrinsic_compiler_lib.h" #include "xla/codegen/intrinsic_lib.h" #include "xla/service/cpu/backend_config.pb.h" #include "xla/service/cpu/cpu_options.h" +#include "xla/service/cpu/cpu_runtime.h" #include "xla/service/hlo_module_config.h" #include "xla/service/llvm_ir/llvm_util.h" #include "xla/tsl/platform/logging.h" @@ -81,6 +94,8 @@ limitations under the License. namespace xla::cpu { +static constexpr char kNoSanitizeMemoryAttr[] = "no_sanitize_memory"; + namespace internal { static absl::once_flag targets_init; @@ -236,9 +251,9 @@ std::unique_ptr IrCompiler::Create( llvm::TargetOptions target_options, Options options, CompilationHooks hooks) { TargetMachineBuilder target_machine_builder = - IrCompiler::InferTargetMachineBuilder(std::move(target_options), - options.opt_level, - options.target_machine_options); + IrCompiler::InferTargetMachineBuilder( + std::move(target_options), options.opt_level, + options.target_machine_options, options.msan_enabled); return std::make_unique(target_machine_builder, std::move(options), std::move(hooks)); @@ -254,14 +269,19 @@ IrCompiler::IrCompiler(TargetMachineBuilder target_machine_builder, absl::StatusOr> IrCompiler::InferTargetMachine( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options) { + const TargetMachineOptions& target_machine_options, bool msan_enabled) { auto attrs_vec = target_machine_options.GetTargetMachineFeaturesVector(); llvm::SmallVector attrs(attrs_vec.begin(), attrs_vec.end()); + llvm::TargetOptions effective_target_options = target_options; + if (msan_enabled) { + effective_target_options.EmulatedTLS = true; + } + absl::call_once(internal::targets_init, &internal::InitializeTargets); std::unique_ptr target_machine( llvm::EngineBuilder() - .setTargetOptions(target_options) + .setTargetOptions(effective_target_options) .setOptLevel(opt_level) .selectTarget( /*TargetTriple=*/llvm::Triple(target_machine_options.triple()), @@ -279,10 +299,10 @@ IrCompiler::InferTargetMachine( IrCompiler::TargetMachineBuilder IrCompiler::InferTargetMachineBuilder( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options) { - return [target_options, opt_level, target_machine_options] { - return InferTargetMachine(target_options, opt_level, - target_machine_options); + const TargetMachineOptions& target_machine_options, bool msan_enabled) { + return [target_options, opt_level, target_machine_options, msan_enabled] { + return InferTargetMachine(target_options, opt_level, target_machine_options, + msan_enabled); }; } @@ -332,6 +352,10 @@ llvm::Expected> IrCompiler::operator()( } } + if (options_.msan_enabled) { + InjectMsanEmulatedTls(module); + } + std::unique_ptr mc_memory_buffer = EmitMachineCode(module, target_machine->get()); @@ -426,12 +450,45 @@ llvm::Error IrCompiler::RunIrPasses(llvm::Module& module, pb.registerLoopAnalyses(lam); pb.crossRegisterProxies(lam, fam, cgam, mam); - llvm::ModulePassManager pm; - - if (options_.dfsan_enabled) { - pm.addPass(llvm::DataFlowSanitizerPass(options_.dfsan_abi_list_files)); + if (options_.msan_enabled) { + for (auto& function : module) { + if (!function.isDeclaration() && + !function.hasFnAttribute(kNoSanitizeMemoryAttr)) { + function.addFnAttr(llvm::Attribute::SanitizeMemory); + } + } } + pb.registerOptimizerLastEPCallback([&](llvm::ModulePassManager& mpm, + llvm::OptimizationLevel level, + llvm::ThinOrFullLTOPhase) { + if (options_.dfsan_enabled) { + mpm.addPass(llvm::DataFlowSanitizerPass(options_.dfsan_abi_list_files)); + } + + if (options_.msan_enabled) { + llvm::MemorySanitizerOptions msan_options( + options_.msan_track_origins, /*Recover=*/false, /*Kernel=*/false, + // Set eager checks to true. This is important to avoid msan flakes on + // KernelThunk's call frame pointer argument with AOT kernels: eager + // checks + nonnull/noundef annotations on the pointer make the AOT + // kernel never read the pointer's msan shadow memory. Without either, + // the AOT kernel would read shadow memory, but the host would not + // have written to it when compiled with eager checks enabled, which + // is the default in Clang. + // Note that if the host is built without eager checks, things still + // work: the host will write the call frame pointer's shadow memory, + // and the AOT kernel won't read it. For pointers flowing in the + // opposite direction (i.e. AOT -> host, such as the return pointer + // from a kernel) we do not mark them as noundef/nonnull, so we + // always write their corresponding shadow memory. + /*EagerChecks=*/true); + mpm.addPass(llvm::MemorySanitizerPass(msan_options)); + } + }); + + llvm::ModulePassManager pm; + llvm::OptimizationLevel opt_level = GetOptimizationLevel(options_); if (opt_level == llvm::OptimizationLevel::O0) { pm.addPass(pb.buildO0DefaultPipeline(opt_level)); @@ -526,4 +583,54 @@ IrCompiler::build_target_machine() const { return target_machine_builder_(); } +void IrCompiler::InjectMsanEmulatedTls(llvm::Module& module) const { + llvm::LLVMContext& ctx = module.getContext(); + const llvm::DataLayout& dl = module.getDataLayout(); + llvm::Type* void_ptr_ty = llvm::PointerType::get(ctx, 0); + + auto inject_selector = [&](llvm::StringRef name, MsanTlsSelector selector) { + new llvm::GlobalVariable( + module, void_ptr_ty, /*isConstant=*/true, + llvm::GlobalValue::InternalLinkage, + llvm::Constant::getIntegerValue( + void_ptr_ty, llvm::APInt(dl.getPointerSizeInBits(), + static_cast(selector))), + name); + }; + + inject_selector("__emutls_v.__msan_param_tls", MsanTlsSelector::kParamTls); + inject_selector("__emutls_v.__msan_retval_tls", MsanTlsSelector::kRetvalTls); + inject_selector("__emutls_v.__msan_va_arg_tls", MsanTlsSelector::kVaArgTls); + inject_selector("__emutls_v.__msan_va_arg_overflow_size_tls", + MsanTlsSelector::kVaArgOverflowSizeTls); + inject_selector("__emutls_v.__msan_param_origin_tls", + MsanTlsSelector::kParamOriginTls); + inject_selector("__emutls_v.__msan_retval_origin_tls", + MsanTlsSelector::kRetvalOriginTls); + inject_selector("__emutls_v.__msan_va_arg_origin_tls", + MsanTlsSelector::kVaArgOriginTls); + inject_selector("__emutls_v.__msan_origin_tls", MsanTlsSelector::kOriginTls); + + llvm::FunctionType* emutls_get_addr_type = + llvm::FunctionType::get(void_ptr_ty, void_ptr_ty, /*isVarArg=*/false); + llvm::Function* emutls_get_addr_fn = llvm::cast( + module.getOrInsertFunction("__emutls_get_address", emutls_get_addr_type) + .getCallee()); + emutls_get_addr_fn->setLinkage(llvm::GlobalValue::InternalLinkage); + emutls_get_addr_fn->addFnAttr(kNoSanitizeMemoryAttr); + + llvm::FunctionCallee bridge_fn = module.getOrInsertFunction( + runtime::kMsanEmutlsGetAddressBridgeSymbolName, emutls_get_addr_type); + + llvm::BasicBlock* entry = + llvm::BasicBlock::Create(ctx, "entry", emutls_get_addr_fn); + llvm::IRBuilder<> builder(entry); + // LLVM's emutls calls __emutls_get_address with the *address* of the control + // variable. We need to dereference it to get the selector value. + auto* control_val = + builder.CreateLoad(void_ptr_ty, emutls_get_addr_fn->getArg(0)); + auto* call = builder.CreateCall(bridge_fn, {control_val}); + builder.CreateRet(call); +} + } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h index c170bc0755f9c7..ce67dc595800fe 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.h @@ -23,6 +23,10 @@ limitations under the License. #include #include +#include "absl/base/config.h" // IWYU pragma: keep +#ifdef ABSL_HAVE_MEMORY_SANITIZER +#include +#endif #include "absl/base/thread_annotations.h" #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" @@ -75,6 +79,19 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { bool disable_loop_unrolling = false; bool disable_platform_dependent_math = false; + // This should be the _only_ place where an #ifdef determines whether the + // generated code is msan-instrumented. This ensures an uninstrumented + // compiler can produce msan-instrumented AOT objects, and viceversa. +#ifdef ABSL_HAVE_MEMORY_SANITIZER + bool msan_enabled = true; + // Level of MSan origin tracking (0 = off, 1 = basic, 2 = full origins). + // In JIT compilation, this defaults to the host's origin tracking level. + int msan_track_origins = __msan_get_track_origins(); +#else + bool msan_enabled = false; + int msan_track_origins = 0; +#endif + bool dfsan_enabled = false; std::vector dfsan_abi_list_files; }; @@ -98,14 +115,16 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { static absl::StatusOr> InferTargetMachine(const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options); + const TargetMachineOptions& target_machine_options, + bool msan_enabled = false); // Returns a target machine builder that uses `InferTargetMachine` defined // above to infer the target machine for the given options. static TargetMachineBuilder InferTargetMachineBuilder( const llvm::TargetOptions& target_options, llvm::CodeGenOptLevel opt_level, - const TargetMachineOptions& target_machine_options); + const TargetMachineOptions& target_machine_options, + bool msan_enabled = false); // Compiles a `module` to an ObjectFile. llvm::Expected> operator()( @@ -140,6 +159,10 @@ class IrCompiler : public llvm::orc::IRCompileLayer::IRCompiler { // races when calling user provided compilation hooks. absl::Mutex mutex_; CompilationHooks hooks_ ABSL_GUARDED_BY(mutex_); + + // Injects MSAN emulated TLS symbols into the module. This is needed for + // supporting MSAN in JIT'ed and AOT'ed code. + void InjectMsanEmulatedTls(llvm::Module& module) const; }; } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc index 29a789f71d1dca..1c87c5c122ab8f 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc @@ -294,12 +294,13 @@ TEST(IrCompilerTest, EmitIntrinsicCall) { auto context = std::make_unique(); IrCompiler::CompilationHooks compilation_hooks; - std::unique_ptr ir_compiler = IrCompiler::Create( - llvm::TargetOptions(), - IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::Aggressive, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}, - compilation_hooks); + IrCompiler::Options options{/*opt_level=*/llvm::CodeGenOptLevel::Aggressive, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}; + options.msan_enabled = false; // Avoid msan interception of memcpy. + + std::unique_ptr ir_compiler = + IrCompiler::Create(llvm::TargetOptions(), options, compilation_hooks); TF_ASSERT_OK_AND_ASSIGN(auto ir_module, ParseModule(*context, kMemcpyCall, kModuleName)); @@ -353,6 +354,40 @@ INSTANTIATE_TEST_SUITE_P(IrCompilerParameterizedTestInstantiation, ::testing::Values("x86_64-grtev4-linux-gnu", "aarch64-unknown-linux-gnu")); +TEST(IrCompilerTest, MemorySanitizerTrackOrigins) { + auto context = std::make_unique(); + IrCompiler::CompilationHooks compilation_hooks; + + TargetMachineOptions target_machine_options(kTargetTripleForHost, + kTargetCpuForHost, ""); + + IrCompiler::Options options{ + /*opt_level=*/llvm::CodeGenOptLevel::None, + /*optimize_for_size=*/false, + target_machine_options, + }; + options.msan_enabled = true; + options.msan_track_origins = 2; + + std::unique_ptr ir_compiler = + IrCompiler::Create(llvm::TargetOptions(), options, compilation_hooks); + + ASSERT_OK_AND_ASSIGN(auto ir_module, + ParseModule(*context, kUnoptimizedIr, "test_module")); + + ASSERT_OK_AND_ASSIGN(auto target_machine, + ir_compiler->build_target_machine()); + + ir_module->setDataLayout(target_machine->createDataLayout()); + ir_module->setTargetTriple(target_machine->getTargetTriple()); + cantFail((*ir_compiler)(*ir_module)); + + auto ir = llvm_ir::DumpToString(ir_module.get()); + EXPECT_THAT(ir, HasSubstr("__msan_track_origins = weak_odr constant i32 2")); + EXPECT_THAT(ir, HasSubstr("__emutls_v.__msan_param_tls")); + EXPECT_THAT(ir, HasSubstr("@__emutls_get_address")); +} + } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc index ad3bf50ddf1acc..2c37fefcac7d26 100644 --- a/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/jit_compiler_test.cc @@ -87,6 +87,18 @@ static absl::StatusOr ParseModule( return llvm::orc::ThreadSafeModule(std::move(m), context); } +// Creates an IrCompiler for testing. We explicitly disable MSan instrumentation +// because unit tests in this file compile raw LLVM IR snippets without linking +// the XLA CPU runtime or BuiltinDefinitionGenerator. +static std::unique_ptr CreateTestIrCompiler() { + IrCompiler::Options options{/*opt_level=*/llvm::CodeGenOptLevel::None, + /*optimize_for_size=*/false, + TargetMachineOptions(GetDebugOptionsFromFlags())}; + options.msan_enabled = false; + return IrCompiler::Create(llvm::TargetOptions(), std::move(options), + IrCompiler::CompilationHooks()); +} + TEST(JitCompilerTest, Compile) { auto context = std::make_unique(); llvm::orc::ThreadSafeContext tsc(std::move(context)); @@ -102,12 +114,7 @@ TEST(JitCompilerTest, Compile) { thread_pool.Schedule(std::move(task)); }; - std::unique_ptr ir_compiler = IrCompiler::Create( - llvm::TargetOptions(), - IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}, - IrCompiler::CompilationHooks()); + std::unique_ptr ir_compiler = CreateTestIrCompiler(); TF_ASSERT_OK_AND_ASSIGN( auto compiler, @@ -201,12 +208,7 @@ TEST(JitCompilerTest, ExternalDefinitionGenerator) { return std::make_unique(); }; - std::unique_ptr ir_compiler = IrCompiler::Create( - llvm::TargetOptions(), - IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}, - IrCompiler::CompilationHooks()); + std::unique_ptr ir_compiler = CreateTestIrCompiler(); TF_ASSERT_OK_AND_ASSIGN( auto compiler, @@ -300,12 +302,7 @@ TEST(JitCompilerTest, CompileWithHighAlignment) { llvm::orc::ThreadSafeContext tsc(std::move(context)); JitCompiler::Options options; - std::unique_ptr ir_compiler = IrCompiler::Create( - llvm::TargetOptions(), - IrCompiler::Options{/*opt_level=*/llvm::CodeGenOptLevel::None, - /*optimize_for_size=*/false, - TargetMachineOptions(GetDebugOptionsFromFlags())}, - IrCompiler::CompilationHooks()); + std::unique_ptr ir_compiler = CreateTestIrCompiler(); TF_ASSERT_OK_AND_ASSIGN( auto compiler, diff --git a/third_party/xla/xla/backends/cpu/nanort/BUILD b/third_party/xla/xla/backends/cpu/nanort/BUILD index f67ac41f0addd5..fed38e3c85491d 100644 --- a/third_party/xla/xla/backends/cpu/nanort/BUILD +++ b/third_party/xla/xla/backends/cpu/nanort/BUILD @@ -94,6 +94,7 @@ xla_cc_test( "//xla/tsl/platform:test_benchmark", "//xla/tsl/platform:test_main", "@com_google_absl//absl/base", + "@com_google_absl//absl/base:config", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc index ec23293b9a8a8b..833d89c18bf12c 100644 --- a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc +++ b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include #include "absl/base/casts.h" +#include "absl/base/config.h" // IWYU pragma: keep #include "absl/container/inlined_vector.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" @@ -62,6 +63,10 @@ limitations under the License. #include "xla/xla_data.pb.h" #include "tsl/platform/casts.h" +#ifdef ABSL_HAVE_MEMORY_SANITIZER +#include +#endif + #define EIGEN_USE_THREADS #include "Eigen/ThreadPool" @@ -482,6 +487,54 @@ TEST_P(NanoRtClientTest, ProgramShapeKeepsLayout) { absl::Span({0, 1})); } +TEST_P(NanoRtClientTest, MsanTracksPoisonThroughKernel) { +#ifndef ABSL_HAVE_MEMORY_SANITIZER + GTEST_SKIP() << "This test requires an MSan build"; +#else + const char* kModuleStr = R"( + HloModule msan_shadow_test + + ENTRY e { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + ROOT sum = f32[4] add(p0, p1) + } + )"; + + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnUnverifiedModule(kModuleStr)); + XlaComputation computation(module->ToProto()); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr executable, + GetExecutable(computation, GetParam())); + + alignas(cpu::MinAlign()) float p0[4]; + p0[0] = 1.0f; + p0[1] = 2.0f; + // p0[2], p0[3] intentionally uninitialized (poisoned under MSan). + + alignas(cpu::MinAlign()) float p1[4] = {10.0f, 20.0f, 30.0f, 40.0f}; + alignas(cpu::MinAlign()) float result[4] = {}; + + Arguments arguments = {{p0, 4}, {p1, 4}}; + Results results = {{result, 4}}; + + auto event = executable->Execute(arguments, results, {}); + tsl::BlockUntilReady(event); + ASSERT_TRUE(event.IsConcrete()); + + EXPECT_EQ(__msan_test_shadow(&result[0], sizeof(float)), -1) + << "result[0] should be initialized (unpoisoned)"; + EXPECT_EQ(__msan_test_shadow(&result[1], sizeof(float)), -1) + << "result[1] should be initialized (unpoisoned)"; + + EXPECT_GE(__msan_test_shadow(&result[2], sizeof(float)), 0) + << "result[2] should be poisoned (p0[2] was uninitialized)"; + EXPECT_GE(__msan_test_shadow(&result[3], sizeof(float)), 0) + << "result[3] should be poisoned (p0[3] was uninitialized)"; +#endif +} + INSTANTIATE_TEST_SUITE_P(NanoRtClientTestSuite, NanoRtClientTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { diff --git a/third_party/xla/xla/backends/cpu/runtime/BUILD b/third_party/xla/xla/backends/cpu/runtime/BUILD index fc20ecd3f5fe5a..130bc7b2223b0e 100644 --- a/third_party/xla/xla/backends/cpu/runtime/BUILD +++ b/third_party/xla/xla/backends/cpu/runtime/BUILD @@ -47,6 +47,7 @@ filegroup( "dot_lib_f64.cc", "dot_lib_s32.cc", "dot_lib_s8.cc", + "msan_emulated_tls.cc", "rng_state_lib.cc", "sort_lib.cc", ], @@ -59,6 +60,7 @@ filegroup( "convolution_lib.h", "dot_lib.h", "kernel_c_api.h", + "msan_emulated_tls.h", "rng_state_lib.h", "sort_lib.h", "work_queue.h", @@ -166,6 +168,20 @@ cc_library( ], ) +cc_library( + name = "msan_emulated_tls", + srcs = ["msan_emulated_tls.cc"], + hdrs = ["msan_emulated_tls.h"], + deps = [ + "@com_google_absl//absl/base:config", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + ], + # Must always link because we might compile/run programs with msan enabled, + # regardless of whether the host compiler is built with msan or not. + alwayslink = True, +) + tf_proto_library( name = "thunk_proto", srcs = ["thunk.proto"], diff --git a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc new file mode 100644 index 00000000000000..c2079d29c03eff --- /dev/null +++ b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.cc @@ -0,0 +1,80 @@ +/* Copyright 2026 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. +==============================================================================*/ + +#include "xla/backends/cpu/runtime/msan_emulated_tls.h" + +#include // IWYU pragma: keep + +#include "absl/base/attributes.h" // IWYU pragma: keep +#include "absl/base/config.h" // IWYU pragma: keep +#include "absl/base/optimization.h" // IWYU pragma: keep +#include "absl/log/check.h" // IWYU pragma: keep + +#ifdef ABSL_HAVE_MEMORY_SANITIZER +extern "C" { +// Mark these initial-exec as compiler-rt does. +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_param_tls[]; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_param_origin_tls[]; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_retval_tls[]; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_retval_origin_tls; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint64_t __msan_va_arg_tls[]; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_va_arg_origin_tls[]; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uintptr_t + __msan_va_arg_overflow_size_tls; +extern ABSL_ATTRIBUTE_INITIAL_EXEC __thread uint32_t __msan_origin_tls; +} +#endif // ABSL_HAVE_MEMORY_SANITIZER + +static_assert( + static_cast(xla::cpu::MsanTlsSelector::kParamTls) == 1 && + static_cast(xla::cpu::MsanTlsSelector::kOriginTls) == 8, + "MsanTlsSelector must remain a contiguous 1..8 range"); + +extern "C" { + +ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void* __xla_cpu_runtime_emutls_get_address( + void* control) { +#ifdef ABSL_HAVE_MEMORY_SANITIZER + using xla::cpu::MsanTlsSelector; + // The control argument is already the selector value (not a pointer to it) + // because the internal __emutls_get_address wrapper in the LLVM module + // dereferences the selector before calling this bridge. + uintptr_t selector = reinterpret_cast(control); + DCHECK_GE(selector, static_cast(MsanTlsSelector::kParamTls)); + DCHECK_LE(selector, static_cast(MsanTlsSelector::kOriginTls)); + switch (static_cast(selector)) { + case MsanTlsSelector::kParamTls: + return __msan_param_tls; + case MsanTlsSelector::kRetvalTls: + return __msan_retval_tls; + case MsanTlsSelector::kVaArgTls: + return __msan_va_arg_tls; + case MsanTlsSelector::kVaArgOverflowSizeTls: + return &__msan_va_arg_overflow_size_tls; + case MsanTlsSelector::kParamOriginTls: + return __msan_param_origin_tls; + case MsanTlsSelector::kRetvalOriginTls: + return &__msan_retval_origin_tls; + case MsanTlsSelector::kVaArgOriginTls: + return __msan_va_arg_origin_tls; + case MsanTlsSelector::kOriginTls: + return &__msan_origin_tls; + } + ABSL_UNREACHABLE(); +#endif // ABSL_HAVE_MEMORY_SANITIZER + return nullptr; +} + +} // extern "C" diff --git a/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h new file mode 100644 index 00000000000000..5a592629ebaa9a --- /dev/null +++ b/third_party/xla/xla/backends/cpu/runtime/msan_emulated_tls.h @@ -0,0 +1,47 @@ +/* Copyright 2026 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. +==============================================================================*/ + +#ifndef XLA_BACKENDS_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ +#define XLA_BACKENDS_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ + +#include + +#include "absl/base/attributes.h" + +extern "C" { +// Returns the address of the host's MSAN TLS variables. +// See https://github.com/google/sanitizers/wiki/MemorySanitizerJIT. +ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void* __xla_cpu_runtime_emutls_get_address( + void* control); +} // extern "C" + +namespace xla::cpu { + +// Selectors for __emutls_get_address. +// All of these are needed for msan and msan-track-origins. +enum class MsanTlsSelector : uintptr_t { + kParamTls = 1, + kRetvalTls = 2, + kVaArgTls = 3, + kVaArgOverflowSizeTls = 4, + kParamOriginTls = 5, + kRetvalOriginTls = 6, + kVaArgOriginTls = 7, + kOriginTls = 8, +}; + +} // namespace xla::cpu + +#endif // XLA_BACKENDS_CPU_RUNTIME_MSAN_EMULATED_TLS_H_ diff --git a/third_party/xla/xla/service/compiler.h b/third_party/xla/xla/service/compiler.h index b7d37c0aac679f..82a06930ee57ea 100644 --- a/third_party/xla/xla/service/compiler.h +++ b/third_party/xla/xla/service/compiler.h @@ -438,6 +438,21 @@ class AotCompilationOptions { run_backend_only_ = run_backend_only; } + bool sanitize_memory() const { return sanitize_memory_; } + void set_sanitize_memory(bool sanitize_memory) { + sanitize_memory_ = sanitize_memory; + } + + int sanitize_memory_track_origins() const { + return sanitize_memory_track_origins_; + } + void set_sanitize_memory_track_origins(int track_origins) { + sanitize_memory_track_origins_ = track_origins; + if (track_origins > 0) { + sanitize_memory_ = true; + } + } + bool sanitize_dataflow() const { return sanitize_dataflow_; } void set_sanitize_dataflow(bool sanitize_dataflow) { sanitize_dataflow_ = sanitize_dataflow; @@ -485,6 +500,8 @@ class AotCompilationOptions { int64_t profile_version_ = 0; std::string cache_key_; bool run_backend_only_ = false; + bool sanitize_memory_ = false; + int sanitize_memory_track_origins_ = 0; bool sanitize_dataflow_ = false; std::vector sanitize_abilists_dataflow_; // Contains target-specific information required by AOT compilation. diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 99dc48e2494892..0d0410d6ce45ee 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -983,6 +983,7 @@ cc_library( "//xla/backends/cpu/collectives:cpu_cliques", "//xla/backends/cpu/collectives:cpu_collectives", "//xla/backends/cpu/collectives:in_process_collectives", + "//xla/backends/cpu/runtime:msan_emulated_tls", "//xla/backends/cpu/runtime:xfeed_manager", "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index 39a6808562aa7c..3fd1bcb5363d62 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -2272,6 +2272,9 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr hlo_module, IrCompiler::GetCodeGenOptLevel(hlo_module->config()); llvm::TargetOptions target_options = CompilerTargetOptions(hlo_module->config()); + if (options.sanitize_memory()) { + target_options.EmulatedTLS = true; + } auto target_machine_builder = [&]() { return absl::WrapUnique(target->createTargetMachine( triple, options.cpu_name(), options.features(), target_options, @@ -2348,6 +2351,8 @@ CpuCompiler::CompileAheadOfTimeThunks( options::DisableLoopUnrolling(module->config()), /*disable_platform_dependent_math=*/ options::DisablePlatformDependentMath(module->config()) || fast_compile, + /*msan_enabled=*/aot_options.sanitize_memory(), + /*msan_track_origins=*/aot_options.sanitize_memory_track_origins(), /*dfsan_enabled=*/aot_options.sanitize_dataflow(), /*dfsan_abilists_enabled=*/aot_options.sanitize_abilists_dataflow()}; diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h index c2b34ff178dc79..e96883b007eeca 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime.h +++ b/third_party/xla/xla/service/cpu/cpu_runtime.h @@ -38,8 +38,9 @@ namespace runtime { // Names of runtime functions. These get resolved from the generated code to the // right symbol at link time in one of two ways: -// 1. When using the JIT, the symbol resolver (xla::cpu::RuntimeSymbolGenerator) -// maps this symbol name to the actual symbol. +// 1. When using the JIT, the symbol resolver +// (xla::cpu::BuiltinDefinitionGenerator) maps this symbol name to the actual +// symbol. // 2. When using ahead-of-time compilation, the linker can resolve the name // because it is a symbol in the cpu_runtime library. inline constexpr absl::string_view kEigenMatMulF16SymbolName = @@ -138,6 +139,8 @@ inline constexpr absl::string_view kReduceScatterSymbolName = "__xla_cpu_runtime_ReduceScatter"; inline constexpr absl::string_view kHandleFfiCallSymbolName = "__xla_cpu_runtime_HandleFfiCall"; +inline constexpr absl::string_view kMsanEmutlsGetAddressBridgeSymbolName = + "__xla_cpu_runtime_emutls_get_address"; // All symbol names for XLA CPU runtime functions need to start with this // prefix. diff --git a/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc b/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc index aef57cadbb6b4f..5ab9ee1aa374bf 100644 --- a/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc +++ b/third_party/xla/xla/service/cpu/restricted/cpu_aot_compiler_test.cc @@ -18,6 +18,9 @@ limitations under the License. #include #include "absl/base/casts.h" +#ifdef ABSL_HAVE_MEMORY_SANITIZER +#include +#endif #include "absl/strings/match.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" @@ -93,6 +96,10 @@ ENTRY e { /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::BigPic); aot_options->set_executor(stream_exec); +#ifdef ABSL_HAVE_MEMORY_SANITIZER + aot_options->set_sanitize_memory(true); + aot_options->set_sanitize_memory_track_origins(__msan_get_track_origins()); +#endif auto test = [this, &compiler, aot_options = std::move(aot_options)]( absl::string_view test_name, absl::string_view hlo, int input, @@ -106,6 +113,9 @@ ENTRY e { TF_ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); +#ifdef ABSL_HAVE_MEMORY_SANITIZER + EXPECT_TRUE(absl::StrContains(serialized_aot_result, "__msan_")); +#endif TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, compiler->LoadAotCompilationResult(serialized_aot_result)); From 6bbfc2b8e48b252bc996514c205ee25db6996d91 Mon Sep 17 00:00:00 2001 From: Mason Chang Date: Fri, 28 Aug 2026 08:54:30 -0700 Subject: [PATCH 07/11] Delete GPU Record stack trace as we don't need it anymore. PiperOrigin-RevId: 972616295 --- third_party/xla/xla/service/gpu/BUILD | 21 -------- .../xla/xla/service/gpu/gpu_compiler.cc | 1 - .../xla/xla/service/gpu/gpu_compiler_test.cc | 37 -------------- third_party/xla/xla/service/gpu/metrics.cc | 48 ------------------ third_party/xla/xla/service/gpu/metrics.h | 9 ---- .../xla/xla/service/gpu/metrics_test.cc | 50 ------------------- 6 files changed, 166 deletions(-) delete mode 100644 third_party/xla/xla/service/gpu/metrics_test.cc diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index d4de5b0fd9462a..bf967c98dc36cb 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -3420,27 +3420,6 @@ cc_library( "//xla/tsl/lib/monitoring:counter", "//xla/tsl/lib/monitoring:gauge", "//xla/tsl/lib/monitoring:sampler", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - "@tsl//tsl/platform", - "@tsl//tsl/platform:stacktrace", - ], -) - -xla_cc_test( - name = "metrics_test", - srcs = ["metrics_test.cc"], - tags = [ - # Streamz recording doesn't work in OSS. - "no_oss", - ], - deps = [ - ":metrics", - "//xla/tests:xla_internal_test_main", - "//xla/tsl/lib/monitoring:collected_metrics", - "//xla/tsl/lib/monitoring:collection_registry", - "@com_google_googletest//:gtest", - "@tsl//tsl/platform:test", ], ) diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 15258447cc2958..cf971cf22f202e 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2957,7 +2957,6 @@ absl::StatusOr> GpuCompiler::RunBackend( module->name(), module->unique_id()); }}; - RecordGpuCompilerStacktrace(); if (module->config().has_static_device_assignment()) { const DeviceAssignment& da = module->config().static_device_assignment(); if (!da.IsIota() && !da.IsAll(0)) { diff --git a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc index 1fa627cfa97c74..6edf4434711d1b 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc @@ -230,43 +230,6 @@ ENTRY test_computation { "non-cyclical source-target pairs")); } -TEST_F(GpuCompilerTest, RecordsStreamzStackTrace) { - if (tsl::kIsOpenSource) { - GTEST_SKIP() << "Streamz is not supported in OSS."; - } - - const char* hlo_text = R"( -HloModule test - -ENTRY main { - p = f32[10]{0} parameter(0) - ROOT neg = f32[10]{0} negate(p) -} -)"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_text)); - - ASSERT_OK_AND_ASSIGN( - std::unique_ptr executable, - CreateExecutable(std::move(module), /*run_hlo_passes=*/false)); - - const std::string kGpuCompilerStacktraceMetricName = - "/xla/service/gpu/compiler_stacktrace_count"; - tsl::monitoring::CollectionRegistry::CollectMetricsOptions options; - std::unique_ptr metrics = - tsl::monitoring::CollectionRegistry::Default()->CollectMetrics(options); - - EXPECT_TRUE(metrics->point_set_map.find(kGpuCompilerStacktraceMetricName) != - metrics->point_set_map.end()); - - // Since Streamz is recorded every call, we expect at least one point. - // All other callers may increment the counter as well. - EXPECT_GT( - metrics->point_set_map[kGpuCompilerStacktraceMetricName]->points.size(), - 0); -} - TEST_F(GpuCompilerTest, GenerateDebugInfoForNonAutotuningCompilations) { const char* hlo_text = R"( HloModule test diff --git a/third_party/xla/xla/service/gpu/metrics.cc b/third_party/xla/xla/service/gpu/metrics.cc index 7c9e3ae6604da5..86aae1f6046f7f 100644 --- a/third_party/xla/xla/service/gpu/metrics.cc +++ b/third_party/xla/xla/service/gpu/metrics.cc @@ -16,18 +16,10 @@ limitations under the License. #include "xla/service/gpu/metrics.h" #include -#include -#include -#include "absl/strings/ascii.h" -#include "absl/strings/str_join.h" -#include "absl/strings/str_split.h" -#include "absl/strings/string_view.h" #include "xla/tsl/lib/monitoring/counter.h" #include "xla/tsl/lib/monitoring/gauge.h" #include "xla/tsl/lib/monitoring/sampler.h" -#include "tsl/platform/platform.h" -#include "tsl/platform/stacktrace.h" namespace xla { namespace { @@ -48,10 +40,6 @@ auto* xla_device_binary_size = tsl::monitoring::Gauge::New( "/xla/service/gpu/xla_device_binary_size", "The size of the XLA binary loaded onto the GPU device."); -auto* gpu_compiler_stacktrace_count = tsl::monitoring::Counter<1>::New( - "/xla/service/gpu/compiler_stacktrace_count", - "The number of times a compiler stacktrace was called.", "stacktrace"); - } // namespace void RecordHloPassesDuration(const uint64_t time_usecs) { @@ -101,40 +89,4 @@ void RecordXlaDeviceBinarySize(const int64_t size) { xla_device_binary_size->GetCell()->Set(size); } -void RecordGpuCompilerStacktrace() { - // Only record stack traces in google as streamz doesn't work in OSS. - if (tsl::kIsOpenSource) { - return; - } - - std::string tsl_stacktrace = tsl::CurrentStackTrace(); - - // tsl::CurrentStackTrace() adds a prefix and postfix lines, so remove them. - std::deque stack = absl::StrSplit(tsl_stacktrace, '\n'); - stack.pop_front(); - stack.pop_back(); - - const int kMaxStackDepth = 15; - if (stack.size() > kMaxStackDepth) { - stack.resize(kMaxStackDepth); - } - - // Stack traces with addresses would make too many unique streamz cells. - // We only care about the actual call stack. - // Format chars added by tsl::CurrentStackTrace(). - constexpr unsigned kFormatChars = 8; - constexpr unsigned kAddressFormat = kFormatChars + 2 * sizeof(void*); - for (int i = 0; i < stack.size(); ++i) { - stack[i] = std::string(absl::StripAsciiWhitespace( - absl::ClippedSubstr(stack[i], kAddressFormat))); - } - - std::string stacktrace = absl::StrJoin(stack, ";\n"); - gpu_compiler_stacktrace_count->GetCell(stacktrace)->IncrementBy(1); -} - -int GetGpuCompilerStacktraceCount(absl::string_view stacktrace) { - return gpu_compiler_stacktrace_count->GetCell(stacktrace)->value(); -} - } // namespace xla diff --git a/third_party/xla/xla/service/gpu/metrics.h b/third_party/xla/xla/service/gpu/metrics.h index 7995ac88b022db..88b308f080330d 100644 --- a/third_party/xla/xla/service/gpu/metrics.h +++ b/third_party/xla/xla/service/gpu/metrics.h @@ -18,8 +18,6 @@ limitations under the License. #include -#include "absl/strings/string_view.h" - namespace xla { // HLO passes (HLO -> HLO). @@ -50,13 +48,6 @@ int64_t GetCompiledProgramsCount(); // Records the size of the XLA device binary in bytes. void RecordXlaDeviceBinarySize(int64_t size); -// Records the stacktrace of the GPU compiler. -void RecordGpuCompilerStacktrace(); - -// Returns the number of times the GPU compiler was called with the given -// stacktrace. -int GetGpuCompilerStacktraceCount(absl::string_view stacktrace); - } // namespace xla #endif // XLA_SERVICE_GPU_METRICS_H_ diff --git a/third_party/xla/xla/service/gpu/metrics_test.cc b/third_party/xla/xla/service/gpu/metrics_test.cc deleted file mode 100644 index a6a1346b563894..00000000000000 --- a/third_party/xla/xla/service/gpu/metrics_test.cc +++ /dev/null @@ -1,50 +0,0 @@ -/* 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. -==============================================================================*/ - -#include "xla/service/gpu/metrics.h" - -#include -#include -#include - -#include -#include "xla/tsl/lib/monitoring/collected_metrics.h" -#include "xla/tsl/lib/monitoring/collection_registry.h" -#include "tsl/platform/test.h" - -namespace xla { -namespace gpu { -namespace { - -TEST(MetricsTest, RecordsGpuCompilerStacktrace) { - const std::string kGpuCompilerStacktraceMetricName = - "/xla/service/gpu/compiler_stacktrace_count"; - - RecordGpuCompilerStacktrace(); - - tsl::monitoring::CollectionRegistry::CollectMetricsOptions options; - std::unique_ptr metrics = - tsl::monitoring::CollectionRegistry::Default()->CollectMetrics(options); - - EXPECT_TRUE(metrics->point_set_map.find(kGpuCompilerStacktraceMetricName) != - metrics->point_set_map.end()); - EXPECT_EQ( - metrics->point_set_map[kGpuCompilerStacktraceMetricName]->points.size(), - 1); -} - -} // namespace -} // namespace gpu -} // namespace xla From 1b3c59ed7f8e274767f28e1cd4dedffb7c81a4d5 Mon Sep 17 00:00:00 2001 From: Seher Ellis Date: Fri, 28 Aug 2026 09:52:14 -0700 Subject: [PATCH 08/11] [XLA:SchedulingGroups] Do not drop the trivial groups that include ops for which "keep_trivial_sync_annotation" returns true. Delete the redundant code block which only checks groups with a single op. PiperOrigin-RevId: 972643528 --- .../legalize_scheduling_annotations.cc | 17 ++--------- .../service/legalize_scheduling_annotations.h | 2 +- .../legalize_scheduling_annotations_test.cc | 29 +++++++++++++++++++ 3 files changed, 33 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/service/legalize_scheduling_annotations.cc b/third_party/xla/xla/service/legalize_scheduling_annotations.cc index 97c0fc35553c0c..7dc8782c872c5c 100644 --- a/third_party/xla/xla/service/legalize_scheduling_annotations.cc +++ b/third_party/xla/xla/service/legalize_scheduling_annotations.cc @@ -491,26 +491,15 @@ bool LegalizeSchedulingAnnotations::RemoveTrivialGroups( std::vector instructions_across_comps; for (const auto& [comp, annotated_instructions] : comp_annotated_instructions) { - if (annotated_instructions.size() == 1 && - !config_.keep_trivial_sync_annotation(annotated_instructions[0])) { - // Remove annotations from synchronous operations (control flow, TC - // custom calls) since they won't do anything and will just get in the - // way of scheduling. - VLOG(2) << "Removing trivial group: " << group_id - << " from instruction: " << annotated_instructions[0]->name() - << " in computation: " << comp->name(); - changed |= RemoveSchedulingAnnotation(annotated_instructions[0]); - deleted_instructions.insert(annotated_instructions[0]); - continue; - } instructions_across_comps.insert(instructions_across_comps.end(), annotated_instructions.begin(), annotated_instructions.end()); } // Remove the groups without any async operations across all computations. - if (absl::c_none_of(instructions_across_comps, [](HloInstruction* instr) { + if (absl::c_none_of(instructions_across_comps, [&](HloInstruction* instr) { return IsSupportedAsyncOp(instr, /*supports_async_start=*/true, - /*check_sync_versions=*/true); + /*check_sync_versions=*/true) || + config_.keep_trivial_sync_annotation(instr); })) { for (HloInstruction* instr : instructions_across_comps) { VLOG(1) << "Removing group id: " << group_id diff --git a/third_party/xla/xla/service/legalize_scheduling_annotations.h b/third_party/xla/xla/service/legalize_scheduling_annotations.h index 1b835b1d4aa863..9d4b482e25b700 100644 --- a/third_party/xla/xla/service/legalize_scheduling_annotations.h +++ b/third_party/xla/xla/service/legalize_scheduling_annotations.h @@ -41,7 +41,7 @@ class LegalizeSchedulingAnnotations : public HloModulePass { public: struct Config { HloPredicate keep_sync_annotation = HloPredicateTrue; - HloPredicate keep_trivial_sync_annotation = HloPredicateTrue; + HloPredicate keep_trivial_sync_annotation = HloPredicateFalse; bool propagate_annotation = false; bool check_start_done_annotation_consistency = true; bool remove_loop_iteration_annotation_only = false; diff --git a/third_party/xla/xla/service/legalize_scheduling_annotations_test.cc b/third_party/xla/xla/service/legalize_scheduling_annotations_test.cc index 01226f95ef72cb..d98a4df945d268 100644 --- a/third_party/xla/xla/service/legalize_scheduling_annotations_test.cc +++ b/third_party/xla/xla/service/legalize_scheduling_annotations_test.cc @@ -1561,5 +1561,34 @@ ENTRY %entry (p0: f32[16,64]) -> (f32[64,64], f32[16,64]) { "1"); } } + +TEST_F(LegalizeSchedulingAnnotationsTest, KeepTrivialSyncAnnotationConfig) { + absl::string_view hlo_string = R"( +HloModule module, is_scheduled=true + +ENTRY %main.1 { + %p0 = f32[8,128]{1,0} parameter(0) + %add.1 = f32[8,128]{1,0} add(%p0, %p0), frontend_attributes={_scheduling_group_id="0"} + %add.2 = f32[8,128]{1,0} add(%p0, %p0), frontend_attributes={_scheduling_group_id="0"} + ROOT %tuple = (f32[8,128]{1,0}, f32[8,128]{1,0}) tuple(%add.1, %add.2) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); + LegalizeSchedulingAnnotations::Config config; + config.keep_trivial_sync_annotation = [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kAdd; + }; + + auto result = LegalizeSchedulingAnnotations(config).Run(hlo_module.get()); + EXPECT_IS_OK(result); + VLOG(1) << "module after: " << hlo_module->ToString(); + HloInstruction* add1 = FindInstruction(hlo_module.get(), "add.1"); + HloInstruction* add2 = FindInstruction(hlo_module.get(), "add.2"); + ASSERT_OK_AND_ASSIGN(auto annotation1, GetSchedulingAnnotation(add1)); + EXPECT_TRUE(annotation1.has_value()); + ASSERT_OK_AND_ASSIGN(auto annotation2, GetSchedulingAnnotation(add2)); + EXPECT_TRUE(annotation2.has_value()); +} } // namespace } // namespace xla From 2d48b97eb1fb15a93b8a8648b1022a827c0b911b Mon Sep 17 00:00:00 2001 From: Bhatu Date: Fri, 28 Aug 2026 10:42:52 -0700 Subject: [PATCH 09/11] Prevent float overflow by scaling kExp input bounds with downstream reduction sizes. Using isolated log bounds (such as [-85, 85] for F32) led to downstream +inf failures in reduction fusions and expressions like x * exp(x). This change: - Precomputes downstream addition-reduction element counts per kExp in reverse post-order. - Scales kExp bounds using base thresholds per data type (e.g. 4.0 for F32, 2.5 for BF16/F16), dynamically tightening when N > 1. - Adds a small positive floor (0.1) to avoid negative intervals and preserve positive test coverage around zero. PiperOrigin-RevId: 972671223 --- .../xla/xla/tests/constraint_propagator.cc | 132 +++++++++++++++++- .../xla/xla/tests/constraint_propagator.h | 17 +++ .../xla/tests/constraint_propagator_test.cc | 66 ++++++++- 3 files changed, 208 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 170a6787e6e331..3fd5b3b2426b8e 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -142,9 +142,9 @@ void SeedConstantInstruction( // the maximum finite representable value of that type. // // For any base b > 0 and exponent x: -// b^x = exp(x * ln(b)) <= max_val <=> x * ln(b) <= ln(max_val). -// Therefore, x <= ln(max_val) / ln(b). -double GetMaxLogForType(PrimitiveType type) { +// Returns the theoretical maximum log (ln(max_val)) for a given type, +// representing the limit of a single isolated exp(x) <= max_val. +double GetTheoreticalMaxLogForType(PrimitiveType type) { switch (type) { // 64-bit IEEE 754 Floating Point (F64): // max_val = 2^1024 * (1 - 2^-53) ≈ 1.7977e+308 @@ -211,6 +211,57 @@ double GetMaxLogForType(PrimitiveType type) { } } +// Returns the safe base max log for exp(x), providing headroom against +// post-multiplication (x * exp(x)) and operand differences (p0 - p1). +double GetBaseMaxLogForType(PrimitiveType type) { + switch (type) { + case F64: + // exp(20) ≈ 4.85e8; leaves ample headroom for large products and sums. + return 20.0; + case F32: + // exp(4) ≈ 54.6; prevents x * exp(x) and multi-operand overflow. + return 4.0; + case F16: + case BF16: + // exp(2.5) ≈ 12.2; safely accommodates accumulation while preserving + // positive range. + return 2.5; + case S64: + case U64: + return 10.0; + case S32: + case U32: + return 5.0; + case S16: + case U16: + return 3.0; + case S8: + case U8: + return 1.5; + default: + return 2.5; + } +} + +// Computes the safe max log bound for exp(x), dynamically scaling down +// when x contributes to a downstream addition reduction of N elements. +double GetMaxLogForType(PrimitiveType type, int64_t reduction_elements = 1) { + double base_max_log = GetBaseMaxLogForType(type); + if (reduction_elements <= 1) { + return base_max_log; + } + + // Ensure N * exp(x) <= max_val with a safety margin of 2.0. + double theoretical_max = GetTheoreticalMaxLogForType(type); + double reduction_budget = + theoretical_max - std::log(static_cast(reduction_elements)) - 2.0; + + // A small positive floor ensures bounds never turn negative, guaranteeing + // test coverage for positive values and zero across all types. + constexpr double kMinPositiveFloor = 0.1; + return std::max(kMinPositiveFloor, std::min(base_max_log, reduction_budget)); +} + // Seeds root constraints exclusively for 16-bit floating-point types (F16 and // BF16). Non-16-bit floating-point types (such as F32 or F64) are intentionally // not seeded. @@ -329,8 +380,77 @@ absl::Status ConstraintPropagator::Propagate( return absl::OkStatus(); } +void ConstraintPropagator::ComputeMaxAddReductionElementsPerExp( + const HloComputation* computation) { + // Tracks how many elements each instruction's output will be summed into + // across all downstream addition-reduction consumer paths. + absl::flat_hash_map + add_reduced_elements_downstream; + + auto instructions = computation->MakeInstructionPostOrder(); + for (auto it = instructions.rbegin(); it != instructions.rend(); ++it) { + const HloInstruction* instruction = *it; + + // How many elements this instruction's output contributes to downstream + // sums. + int64_t consumer_add_reduction_elements = 1; + if (auto it = add_reduced_elements_downstream.find(instruction); + it != add_reduced_elements_downstream.end()) { + consumer_add_reduction_elements = it->second; + } + + if (instruction->opcode() == HloOpcode::kExp) { + max_add_reduction_elements_per_exp_[instruction] = + consumer_add_reduction_elements; + } + + if (instruction->opcode() == HloOpcode::kReduce && + GetCanonicalReductionOpcode(*instruction->to_apply()) == + HloOpcode::kAdd) { + int64_t elements_in_add_reduction = 1; + for (int64_t dim : instruction->dimensions()) { + elements_in_add_reduction *= + instruction->operand(0)->shape().dimensions(dim); + } + + int64_t total_elements = + elements_in_add_reduction * consumer_add_reduction_elements; + + // Only data operands feed the sum (init values do not). + int64_t num_data_operands = instruction->operand_count() / 2; + for (int64_t i = 0; i < num_data_operands; ++i) { + const HloInstruction* operand = instruction->operand(i); + int64_t& downstream_elements = add_reduced_elements_downstream[operand]; + downstream_elements = std::max(downstream_elements, total_elements); + } + } else if (instruction->IsElementwise() || + instruction->opcode() == HloOpcode::kCopy || + instruction->opcode() == HloOpcode::kBitcast || + instruction->opcode() == HloOpcode::kReshape) { + // Elementwise and shape-preserving ops maintain 1-to-1 correspondence + // between input and output elements; forward the downstream count. + for (const HloInstruction* operand : instruction->operands()) { + int64_t& downstream_elements = add_reduced_elements_downstream[operand]; + downstream_elements = + std::max(downstream_elements, consumer_add_reduction_elements); + } + } + } +} + +int64_t ConstraintPropagator::GetMaxAddReductionElementsForExp( + const HloInstruction* exp_instruction) const { + if (auto it = max_add_reduction_elements_per_exp_.find(exp_instruction); + it != max_add_reduction_elements_per_exp_.end()) { + return it->second; + } + return 1; +} + absl::Status ConstraintPropagator::SeedConstraints( const HloComputation* computation) { + ComputeMaxAddReductionElementsPerExp(computation); + auto instructions = computation->MakeInstructionPostOrder(); // First pass: Seed all constants so they are available in states_ when @@ -365,7 +485,9 @@ absl::Status ConstraintPropagator::SeedConstraints( break; case HloOpcode::kExp: { // Safe domain [-max_log, max_log] prevents floating point overflow. - double max_log = GetMaxLogForType(inst->shape().element_type()); + int64_t reduction_elements = GetMaxAddReductionElementsForExp(inst); + double max_log = + GetMaxLogForType(inst->shape().element_type(), reduction_elements); states_[inst->operand(0)].AddConstraint( ConstraintInterval{-max_log, max_log, false}); break; @@ -570,7 +692,7 @@ absl::Status ConstraintPropagator::SeedConstraints( // // ML Patterns Handled: // 1. Guarded Division / Gradient & Activation Threshold Clipping: -// In deep learning models (e.g. GemFuse, diffusion models, transformers), +// In deep learning models (e.g. diffusion models, transformers), // gradients or activations are scaled down by a threshold tau > 0 when their // norm/magnitude exceeds tau: // scale(x) = where(x > tau, tau / x, 1.0) diff --git a/third_party/xla/xla/tests/constraint_propagator.h b/third_party/xla/xla/tests/constraint_propagator.h index 259bac7453c959..9ada9ec0953f01 100644 --- a/third_party/xla/xla/tests/constraint_propagator.h +++ b/third_party/xla/xla/tests/constraint_propagator.h @@ -55,6 +55,7 @@ IdentityElementType GetReductionIdentityElementType( // between operands are simplified into independent constraints to keep // computation scalable. // +class ConstraintPropagatorTest; // Limitations: // - Does not handle control flow. // - Uses weak heuristics for accumulation-intensive operations. @@ -67,11 +68,22 @@ class ConstraintPropagator { get_index_known_zeroes = nullptr); private: + friend class ConstraintPropagatorTest; + explicit ConstraintPropagator( std::function(const HloInstruction*, int64_t)> get_index_known_zeroes) : get_index_known_zeroes_(get_index_known_zeroes) {} + // Populates max_add_reduction_elements_per_exp_ by analyzing downstream + // addition-reduction chains across the computation. + void ComputeMaxAddReductionElementsPerExp(const HloComputation* computation); + + // Returns the maximum number of downstream addition-reduction elements for + // the given kExp instruction, or 1 if unreduced or not found. + int64_t GetMaxAddReductionElementsForExp( + const HloInstruction* exp_instruction) const; + // Propagates constraints in post-order throughout the given computation. absl::Status Propagate(const HloComputation* computation); @@ -142,6 +154,11 @@ class ConstraintPropagator { // Constraint states for each instruction in the module. absl::flat_hash_map states_; + + // Maps each kExp instruction to the maximum number of elements summed + // downstream across any addition reduction path. + absl::flat_hash_map + max_add_reduction_elements_per_exp_; }; } // namespace xla diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 9f22ba81062f3c..8a87fbe07d9240 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -16,16 +16,33 @@ limitations under the License. #include "xla/tests/constraint_propagator.h" #include +#include #include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/testlib/test.h" #include "xla/tests/constraint_state.h" #include "xla/tests/hlo_test_base.h" namespace xla { -namespace { +class ConstraintPropagatorTest : public HloTestBase { + protected: + ConstraintPropagator CreatePropagator() { + return ConstraintPropagator(nullptr); + } + + void ComputeMaxAddReductionElementsPerExp(ConstraintPropagator& propagator, + const HloComputation* comp) { + propagator.ComputeMaxAddReductionElementsPerExp(comp); + } + + int64_t GetMaxAddReductionElementsForExp( + const ConstraintPropagator& propagator, const HloInstruction* exp) { + return propagator.GetMaxAddReductionElementsForExp(exp); + } +}; -class ConstraintPropagatorTest : public HloTestBase {}; +namespace { TEST_F(ConstraintPropagatorTest, EmptyInterval) { ConstraintInterval a{0.0, 10.0, false}; @@ -1175,5 +1192,50 @@ ENTRY main { EXPECT_LE(p0_int.max, 150.0); } +TEST_F(ConstraintPropagatorTest, + MaxAddReductionElementsPerExpTracksDownstreamReductions) { + const char* hlo = R"( +HloModule TestReductionExpModule + +%add_reducer (a: f32[], b: f32[]) -> f32[] { + %a = f32[] parameter(0) + %b = f32[] parameter(1) + ROOT %sum = f32[] add(%a, %b) +} + +ENTRY %main { + %p0 = f32[2,256] parameter(0) + %exp_reduced = f32[2,256] exponential(%p0) + %mul = f32[2,256] multiply(%exp_reduced, %p0) + %c_zero = f32[] constant(0.0) + %reduce = f32[2] reduce(%mul, %c_zero), dimensions={1}, to_apply=%add_reducer + + %p1 = f32[2,256] parameter(1) + %exp_unreduced = f32[2,256] exponential(%p1) + + ROOT %tuple = (f32[2], f32[2,256]) tuple(%reduce, %exp_unreduced) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + HloComputation* entry = module->entry_computation(); + + ConstraintPropagator propagator = CreatePropagator(); + ComputeMaxAddReductionElementsPerExp(propagator, entry); + + const HloInstruction* exp_reduced = nullptr; + const HloInstruction* exp_unreduced = nullptr; + for (const HloInstruction* inst : entry->instructions()) { + if (inst->name() == "exp_reduced") { + exp_reduced = inst; + } else if (inst->name() == "exp_unreduced") { + exp_unreduced = inst; + } + } + + ASSERT_NE(exp_reduced, nullptr); + ASSERT_NE(exp_unreduced, nullptr); + EXPECT_EQ(GetMaxAddReductionElementsForExp(propagator, exp_reduced), 256); + EXPECT_EQ(GetMaxAddReductionElementsForExp(propagator, exp_unreduced), 1); +} } // namespace } // namespace xla From 2f66a4e9280cc7662493b7ce2b4490ab5b7d5d6b Mon Sep 17 00:00:00 2001 From: Shyamli Agrawal Date: Fri, 28 Aug 2026 10:43:06 -0700 Subject: [PATCH 10/11] Adapt xla aot test with new autotune cache format. Keep both legacy and new autotune cache formats in tests. PiperOrigin-RevId: 972671341 --- third_party/xla/xla/service/BUILD | 42 ++++++- .../xla/service/xla_aot_compile_gpu_test.cc | 23 +++- .../xla_aot_compile_test_autotune_cache.txtpb | 109 ++++++++++++++++++ third_party/xla/xla/service/xla_compile.bzl | 26 +++-- 4 files changed, 188 insertions(+), 12 deletions(-) create mode 100644 third_party/xla/xla/service/xla_aot_compile_test_autotune_cache.txtpb diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 443fdfb96c9a24..f1a7f7134f2de0 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -6473,42 +6473,79 @@ xla_aot_compile_cpu( xla_aot_compile_gpu( name = "xla_aot_compile_test_gpu_executable", + autotune_results = "xla_aot_compile_test_autotune_cache.txtpb", + gpu_targets = [ + "h100", + "b200", + ], + module = "xla_aot_compile_test.mlir", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=true", +) + +xla_aot_compile_gpu( + name = "xla_aot_compile_test_gpu_executable_legacy_cache", autotune_results = "xla_aot_compile_test_autotune_results.txtpb", gpu_targets = [ "h100", "b200", ], module = "xla_aot_compile_test.mlir", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=false", ) xla_aot_compile_gpu( name = "xla_aot_compile_test_gpu_executable_hlo", + autotune_results = "xla_aot_compile_test_autotune_cache.txtpb", + gpu_targets = [ + "h100", + "b200", + ], + module = "xla_aot_compile_test.hlo", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=true", +) + +xla_aot_compile_gpu( + name = "xla_aot_compile_test_gpu_executable_hlo_legacy_cache", autotune_results = "xla_aot_compile_test_autotune_results.txtpb", gpu_targets = [ "h100", "b200", ], module = "xla_aot_compile_test.hlo", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=false", ) xla_aot_compile_gpu( name = "xla_aot_compile_test_gpu_executable_constant", - autotune_results = "xla_aot_compile_test_autotune_results.txtpb", + autotune_results = "xla_aot_compile_test_autotune_cache.txtpb", gpu_targets = [ "h100", "b200", ], module = "xla_aot_compile_test_constant.mlir", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=true", ) xla_aot_compile_gpu( name = "xla_aot_compile_test_gpu_executable_convolution", + autotune_results = "xla_aot_compile_test_autotune_cache.txtpb", + gpu_targets = [ + "h100", + "b200", + ], + module = "xla_aot_compile_test_convolution.mlir", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=true", +) + +xla_aot_compile_gpu( + name = "xla_aot_compile_test_gpu_executable_convolution_legacy_cache", autotune_results = "xla_aot_compile_test_autotune_results.txtpb", gpu_targets = [ "h100", "b200", ], module = "xla_aot_compile_test_convolution.mlir", + xla_flags = "--xla_gpu_use_new_autotune_cache_format=false", ) xla_aot_compile_gpu_runtime_autotuning( @@ -6596,7 +6633,10 @@ xla_test( ":xla_aot_compile_test_gpu_executable", ":xla_aot_compile_test_gpu_executable_constant", ":xla_aot_compile_test_gpu_executable_convolution", + ":xla_aot_compile_test_gpu_executable_convolution_legacy_cache", ":xla_aot_compile_test_gpu_executable_hlo", + ":xla_aot_compile_test_gpu_executable_hlo_legacy_cache", + ":xla_aot_compile_test_gpu_executable_legacy_cache", ], tags = [ "cuda-only", diff --git a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc index 5cd597ced5812b..cac89d82602558 100644 --- a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc +++ b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc @@ -77,7 +77,9 @@ TEST_P(XlaAotCompileTest, LoadGpuExecutable) { INSTANTIATE_TEST_SUITE_P( TestingAotFormats, XlaAotCompileTest, ::testing::Values("xla_aot_compile_test_gpu_executable", - "xla_aot_compile_test_gpu_executable_hlo")); + "xla_aot_compile_test_gpu_executable_hlo", + "xla_aot_compile_test_gpu_executable_legacy_cache", + "xla_aot_compile_test_gpu_executable_hlo_legacy_cache")); TEST_F(XlaCompileTest, LoadGpuExecutableWithConstant) { Literal input = LiteralUtil::CreateR1({3.0f, 3.0f, 3.0f}); @@ -105,6 +107,25 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) { {&input1, &input2}, expected); } +TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolutionLegacyCache) { + Literal input1 = LiteralUtil::CreateR4( + {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}}, + {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}}, + {{21.0, 22.0}, {23.0, 24.0}, {25.0, 26.0}, {27.0, 28.0}}, + {{31.0, 32.0}, {33.0, 34.0}, {35.0, 36.0}, {37.0, 38.0}}}}); + Literal input2 = + LiteralUtil::CreateR4({{{{1.0}, {2.0}}, {{3.0}, {4.0}}}, + {{{5.0}, {6.0}}, {{7.0}, {8.0}}}, + {{{9.0}, {10.0}}, {{11.0}, {12.0}}}}); + Literal expected = LiteralUtil::CreateR4({{ + {{1310.0}, {1466.0}, {1622.0}}, + {{2090.0}, {2246.0}, {2402.0}}, + }}); + LoadAndRunExecutable( + "xla_aot_compile_test_gpu_executable_convolution_legacy_cache", + {&input1, &input2}, expected); +} + } // namespace } // namespace xla_compile } // namespace xla diff --git a/third_party/xla/xla/service/xla_aot_compile_test_autotune_cache.txtpb b/third_party/xla/xla/service/xla_aot_compile_test_autotune_cache.txtpb new file mode 100644 index 00000000000000..4c2562e7922e30 --- /dev/null +++ b/third_party/xla/xla/service/xla_aot_compile_test_autotune_cache.txtpb @@ -0,0 +1,109 @@ +# Copyright 2026 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. + +# proto-file: third_party/tensorflow/compiler/xla/autotune_cache.proto +# proto-message: xla.autotuner.AutotuneCache + +entries { + key { + target { + device: "699b92c0fb30cb16" + hlo_fingerprint: "74ccf3af7259becbcaa5dc0eaf995e2c" + } + } + value { + optimal_config { + backend: CUBLASLT + backend_config { + gemm { + algorithm: 13 + } + } + } + optimal_backend_version: "12.8.0" + } +} +entries { + key { + target { + device: "699b92c0fb30cb16" + hlo_fingerprint: "4cdf2fba5e5a5a4e010a212dafaa0c9a" + } + } + value { + optimal_config { + backend: CUDNN + backend_config { + algorithm { + algo_id: 28 + tuning_knobs { + key: 2 + value: 4 + } + tuning_knobs { + key: 3 + value: 0 + } + } + } + } + optimal_backend_version: "9.10.0" + } +} +entries { + key { + target { + device: "45a23d12be5bc693" + hlo_fingerprint: "74ccf3af7259becbcaa5dc0eaf995e2c" + } + } + value { + optimal_config { + backend: CUBLASLT + backend_config { + gemm { + algorithm: 13 + } + } + } + optimal_backend_version: "12.8.0" + } +} +entries { + key { + target { + device: "45a23d12be5bc693" + hlo_fingerprint: "4cdf2fba5e5a5a4e010a212dafaa0c9a" + } + } + value { + optimal_config { + backend: CUDNN + backend_config { + algorithm { + algo_id: 28 + tuning_knobs { + key: 2 + value: 4 + } + tuning_knobs { + key: 3 + value: 0 + } + } + } + } + optimal_backend_version: "9.10.0" + } +} diff --git a/third_party/xla/xla/service/xla_compile.bzl b/third_party/xla/xla/service/xla_compile.bzl index ce12171caeea9c..11096f37ed1ee3 100644 --- a/third_party/xla/xla/service/xla_compile.bzl +++ b/third_party/xla/xla/service/xla_compile.bzl @@ -97,32 +97,38 @@ def xla_aot_compile_gpu( name, module, gpu_targets, - autotune_results): + autotune_results, + xla_flags = ""): """Runs xla_compile to compile an MHLO, StableHLO or HLO module into an AotCompilationResult for GPU Args: name: The name of the build rule. module: The MHLO or StableHLO file to compile. gpu_targets: The list of gpu targets. - autotune_results: AOT AutotuneResults + autotune_results: AOT AutotuneResults or AutotuneCache file. + xla_flags: Additional XLA_FLAGS to set during compilation. """ res = [] for target in gpu_targets: # Run xla_compile to generate the file containing an AotCompilationResult. compiled_binary = name + "_" + target + cmd = ( + "$(location " + xla_compile_tool + ")" + + " --module_file=$(location " + module + ")" + + " --output_file=$(location " + compiled_binary + ")" + + " --platform=gpu" + + " --gpu_target_config=$(location " + gpu_target_config_map[target] + ")" + ) + flags = "--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")" + if xla_flags: + flags = flags + " " + xla_flags + cmd = "XLA_FLAGS=\"" + flags + "\" " + cmd native.genrule( name = "gen_" + name + "_" + target, srcs = [module, gpu_target_config_map[target], autotune_results], outs = [name + "_" + target], - cmd = ( - "$(location " + xla_compile_tool + ")" + - " --module_file=$(location " + module + ")" + - " --output_file=$(location " + compiled_binary + ")" + - " --platform=gpu" + - " --gpu_target_config=$(location " + gpu_target_config_map[target] + ")" + - " --autotune_results=$(location " + autotune_results + ")" - ), + cmd = cmd, tools = [xla_compile_tool], # copybara:comment_begin(oss-only) target_compatible_with = select({ From ec78c9a0ee15069823af6fd18a0db3a872b052fc Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Fri, 28 Aug 2026 11:04:28 -0700 Subject: [PATCH 11/11] Generalize PjRtStreamExecutorClient::Load to be a general function on CommonPjRtClient using InferDispatchInfo. PiperOrigin-RevId: 972683293 --- third_party/xla/xla/pjrt/BUILD | 4 +- .../xla/xla/pjrt/common_pjrt_client.cc | 169 ++++++++++- third_party/xla/xla/pjrt/common_pjrt_client.h | 53 ++-- third_party/xla/xla/pjrt/cpu/cpu_client.cc | 29 +- third_party/xla/xla/pjrt/cpu/cpu_client.h | 6 +- third_party/xla/xla/pjrt/dynamic_shapes.cc | 26 ++ third_party/xla/xla/pjrt/dynamic_shapes.h | 17 ++ .../pjrt/gpu/se_gpu_topology_description.cc | 20 ++ .../pjrt/gpu/se_gpu_topology_description.h | 3 + .../xla/xla/pjrt/infer_dispatch_info.cc | 43 +-- .../xla/xla/pjrt/infer_dispatch_info.h | 62 +++- third_party/xla/xla/pjrt/pjrt_compiler.h | 7 + third_party/xla/xla/pjrt/pjrt_executable.h | 5 + .../plugin/xla_cpu/cpu_topology_description.h | 3 +- third_party/xla/xla/pjrt/raw_pjrt_client.h | 11 + third_party/xla/xla/pjrt/se/BUILD | 1 + .../pjrt/se/pjrt_stream_executor_client.cc | 269 ++---------------- .../xla/pjrt/se/pjrt_stream_executor_client.h | 10 +- .../se/pjrt_stream_executor_client_test.cc | 4 +- .../xla/pjrt/se/stream_executor_executable.h | 4 + 20 files changed, 381 insertions(+), 365 deletions(-) diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index 5a3ba3547d0332..fc3285257f1494 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -206,6 +206,7 @@ cc_library( ":dynamic_shapes", ":host_callback", ":host_memory_spaces", + ":infer_dispatch_info", ":pjrt_client", ":pjrt_compiler", ":pjrt_executable", @@ -272,8 +273,9 @@ cc_library( ":friends", ]), deps = [ - ":common_pjrt_client", + ":dynamic_shapes", ":pjrt_client", + ":pjrt_compiler", ":pjrt_executable", ":pjrt_layout", ":utils", diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index f3da0d85167f07..eca3595ad77e00 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -336,6 +336,154 @@ CommonPjRtClient::LoadSerializedExecutable( return Load(std::move(executable), load_options); } +absl::StatusOr> CommonPjRtClient::Load( + std::shared_ptr executable, + const LoadOptions& load_options) { + return LoadInternal(std::move(executable), load_options, /*dump=*/false); +} + +absl::StatusOr> +CommonPjRtClient::LoadInternal(std::shared_ptr executable, + const LoadOptions& load_options, bool dump) { + tsl::profiler::TraceMe traceme("CommonPjRtClient::Load"); + VLOG(1) << "CommonPjRtClient::Load"; + ABSL_ASSIGN_OR_RETURN(const PjRtTopologyDescription* topology, + GetTopologyDescription()); + ABSL_ASSIGN_OR_RETURN(auto hlo_module, executable->GetHloModule()); + + ABSL_ASSIGN_OR_RETURN(CompileOptions compile_options, + executable->GetCompileOptions()); + ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides()); + if (IsEarlyExitCompilation(compile_options)) { + return InvalidArgument( + "Executable compiled with xla_early_exit_with_layouts cannot be " + "loaded."); + } + if (!IsGpuId(platform_id())) { + absl::StatusOr> runtime_abi_version = + RuntimeAbiVersion(); + if (!absl::IsUnimplemented(runtime_abi_version.status())) { + ABSL_RETURN_IF_ERROR(runtime_abi_version.status()); + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr executable_abi_version, + executable->GetAbiVersion()); + ABSL_RETURN_IF_ERROR( + (*runtime_abi_version)->IsCompatibleWith(*executable_abi_version)); + } + } + std::vector + addressable_device_logical_ids; + std::vector addressable_devices; + int num_replicas; + int num_partitions; + std::shared_ptr device_assignment; + ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions( + compile_options.compile_portable_executable, + &compile_options.executable_build_options, + [this, topology, &compile_options](int num_replicas, int num_partitions) { + return topology->GetDefaultDeviceAssignment( + process_index(), num_replicas, + /*num_replicas_per_slice=*/std::nullopt, num_partitions, + compile_options.multi_slice_config); + }, + &num_replicas, &num_partitions, &device_assignment)); + + // Find devices that are addressable by this client/task. + if (device_assignment != nullptr) { + int num_replicas = device_assignment->replica_count(); + int num_partitions = device_assignment->computation_count(); + addressable_device_logical_ids.reserve(num_replicas * num_partitions); + addressable_devices.reserve(num_replicas * num_partitions); + for (int replica = 0; replica < num_replicas; ++replica) { + for (int partition = 0; partition < num_partitions; ++partition) { + int64_t device_id = (*device_assignment)(replica, partition); + GlobalDeviceId global_device_id(device_id); + + ABSL_ASSIGN_OR_RETURN(PjRtDevice * device, LookupDevice(global_device_id)); + if (device->process_index() != process_index()) { + VLOG(3) << "Non-local device: " << device_id; + continue; + } + PjRtLoadedExecutable::LogicalDeviceIds logical_device_ids; + logical_device_ids.replica = replica; + logical_device_ids.partition = partition; + addressable_device_logical_ids.push_back(std::move(logical_device_ids)); + addressable_devices.push_back(device); + } + } + } + + const auto& ex_options = compile_options.executable_build_options; + const bool xla_dump_hlo_unoptimized_snapshots = + ex_options.has_debug_options() && + ex_options.debug_options().xla_dump_hlo_unoptimized_snapshots(); + if (dump) { + VLOG(1) << "Dumping deserialized executable"; + // Override the debug_options() embedded in the module with those + // explicitly passed in when deserializing. This allows options such as + // --xla_dump_to to be changed. Does not quite match the naming convention + // of the dump during compilation, which includes a backend-specific + // prefix. + DumpHloModuleIfEnabled( + *hlo_module, kAfterOptimizationsDumpName, + ex_options.has_debug_options() ? &ex_options.debug_options() : nullptr); + } + xla::Shape result_shape = hlo_module->result_shape(); + absl::Span result_shapes = + result_shape.IsTuple() ? absl::MakeSpan(result_shape.tuple_shapes()) + : absl::MakeSpan(&result_shape, 1); + for (auto& leaf_shape : result_shapes) { + if (leaf_shape.IsTuple()) { + return absl::InternalError( + absl::StrCat("Nested tuples are not supported with " + "PjRtStreamExecutorClient. got: ", + result_shape.ToString())); + } + } + + for (int result_index : compile_options.individually_defined_output_indices) { + if (result_index < 0 || + static_cast(result_index) >= result_shapes.size()) { + return InvalidArgument( + "Individually defined output index %d is out of range for %d " + "outputs", + result_index, result_shapes.size()); + } + } + + auto parameter_shapes = + GetParameterShapes(hlo_module->compute_computation_layout()); + using InputHloSnapshotBits = + CommonPjRtLoadedExecutable::DispatchInfo::InputHloSnapshotBits; + std::unique_ptr input_hlo_snapshot_bits; + if (xla_dump_hlo_unoptimized_snapshots) { + if (std::optional unoptimized_hlo_module_proto = + executable->GetUnoptimizedHloModule()) { + input_hlo_snapshot_bits = + std::make_unique(InputHloSnapshotBits{ + std::move(*unoptimized_hlo_module_proto), + compile_options.executable_build_options.debug_options()}); + } + } + + ABSL_ASSIGN_OR_RETURN( + auto dispatch_info, + InferDispatchInfo( + topology, std::move(parameter_shapes), std::move(result_shape), + hlo_module->input_output_alias_config(), std::move(device_assignment), + std::move(addressable_device_logical_ids), + std::move(addressable_devices), nullptr, + compile_options.parameter_is_tupled_arguments, + std::move(input_hlo_snapshot_bits))); + + auto load_state = raw_client()->MakeLoadState(); + ABSL_RETURN_IF_ERROR(load_state->Preload(executable.get())); + auto loaded_executable = std::make_unique( + this, raw_client()->ToAsyncExecutable(std::move(executable)), + std::move(dispatch_info), std::move(load_state)); + return std::unique_ptr(std::move(loaded_executable)); +} + absl::StatusOr CommonPjRtClient::LinearizeHostBufferInto( const void* data, PrimitiveType type, absl::Span dims, std::optional> byte_strides, @@ -1016,17 +1164,8 @@ CommonPjRtClient::BufferFromHostBuffer( absl::StatusOr CommonPjRtClient::GetOnDeviceBytesCount( int memory_space_kind, const xla::Shape& shape) const { - auto kind = GetDynamicShapeKind(memory_space_kind); - // PjRtShapeAndMetadataTransferRequirements::Get->ShapeUtil::ArraySize - // requires a layout. - if (!shape.IsToken() && !shape.has_layout()) { - return absl::FailedPreconditionError( - "Buffer's on-device shape has no layout. Cannot determine on-device " - "bytes count."); - } - auto requirements = - PjRtShapeAndMetadataTransferRequirements::Get(shape, kind); - return static_cast(requirements.size); + return PjRtGetOnDeviceBytesCount(shape, + GetDynamicShapeKind(memory_space_kind)); } absl::StatusOr> @@ -1645,8 +1784,12 @@ CommonPjRtClient::AllocateOutputBuffersWithInputReuse( } } if (memory_space == nullptr) { - return absl::InternalError( - absl::StrCat("No memory space found (kind_id: ", kind_id, ")")); + std::string silly; + for (PjRtMemorySpace* ms : device->memory_spaces()) { + absl::StrAppend(&silly, ms->kind_id(), ":", ms->kind(), ","); + } + return absl::InternalError(absl::StrCat( + "No memory space found (kind_id: ", kind_id, ")", silly)); } ABSL_ASSIGN_OR_RETURN(int64_t on_device_bytes, GetOnDeviceBytesCount(memory_space, leaf_shape)); diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h index f0705400f7dd50..1ec82406fe904c 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.h +++ b/third_party/xla/xla/pjrt/common_pjrt_client.h @@ -46,6 +46,7 @@ limitations under the License. #include "xla/pjrt/async_work_runner.h" #include "xla/pjrt/device_event.h" #include "xla/pjrt/dynamic_shapes.h" +#include "xla/pjrt/infer_dispatch_info.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/raw_buffer.h" #include "xla/pjrt/raw_pjrt_client.h" @@ -134,8 +135,8 @@ class CommonPjRtClient : public PjRtClient { // Gets the memory_space_kind for a particular XLA layout. virtual absl::StatusOr GetMemorySpaceKindForShape( const xla::Shape& shape) const { - return absl::UnimplementedError( - "GetMemorySpaceKindForShape is not supported."); + ABSL_ASSIGN_OR_RETURN(auto* topology, GetTopologyDescription()); + return topology->GetMemorySpaceKindForShape(shape); } // Allocates a raw buffer of a particular size after an optional @@ -527,7 +528,14 @@ class CommonPjRtClient : public PjRtClient { std::optional options, const LoadOptions& load_options) override; + absl::StatusOr> Load( + std::shared_ptr executable, + const LoadOptions& load_options) override; + protected: + absl::StatusOr> LoadInternal( + std::shared_ptr executable, + const LoadOptions& load_options, bool dump); // Returns the required alignment for device memory addresses when slicing. virtual absl::StatusOr GetDeviceAddressAlignment() const { return absl::UnimplementedError( @@ -553,39 +561,7 @@ class CommonPjRtClient : public PjRtClient { class CommonPjRtLoadedExecutable : public PjRtLoadedExecutable { public: - struct DispatchInfo { - std::vector parameter_device_shapes; - std::shared_ptr output_device_shape; - std::vector parameter_memory_space_kind_ids; - std::vector output_memory_space_kind_ids; - std::vector addressable_devices; - std::vector addressable_device_logical_ids; - std::shared_ptr device_assignment; - std::vector parameters_that_may_be_donated; - std::vector input_buffer_sizes_in_bytes; - // Executable shape information that is computable from the PjRtExecutable*. - struct Extras { - std::string name; - int num_partitions; - int num_replicas; - absl::StatusOr>> - parameter_layouts; - absl::StatusOr>> - output_layouts; - std::optional> parameter_shardings; - std::optional> output_shardings; - std::vector parameter_memory_kinds; - std::vector output_memory_kinds; - absl::StatusOr fingerprint; - HloInputOutputAliasConfig input_output_alias_config; - }; - struct InputHloSnapshotBits { - xla::HloModuleProto hlo_module; - xla::DebugOptions debug_options; - }; - std::unique_ptr input_hlo_snapshot_bits; - std::unique_ptr extras; - }; + using DispatchInfo = PjRtLoadedExecutableDispatchInfo; CommonPjRtLoadedExecutable( CommonPjRtClient* client, tsl::AsyncValueRef executable, @@ -807,8 +783,11 @@ class CommonPjRtLoadedExecutable : public PjRtLoadedExecutable { // side-effect of the execution. Derived classes may use custom logic. absl::Span ParametersThatMayBeDonated() const; - virtual const HloInputOutputAliasConfig& input_output_alias_config() - const = 0; + virtual const HloInputOutputAliasConfig& input_output_alias_config() const { + auto hlo_module = GetExecutable()->GetHloModule(); + CHECK_OK(hlo_module.status()); + return (*hlo_module)->input_output_alias_config(); + } // Checks that the input buffers passed in by the user have the correct size // on device for the compiled program. diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index 9fa31f3a93ae37..fe1a95aa466368 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -697,6 +697,12 @@ PjRtCpuClient::LoadInternal( std::move(load_state)); } +tsl::AsyncValueRef PjRtCpuRawClient::ToAsyncExecutable( + std::shared_ptr executable) const { + return tsl::MakeAvailableAsyncValueRef( + std::static_pointer_cast(executable)); +} + static absl::StatusOr> JitCompile( std::unique_ptr hlo_module, const ExecutableBuildOptions& build_options, @@ -1139,29 +1145,6 @@ PjRtCpuRawClient::CreateRawBufferChannel(PjRtMemorySpace* memory_space, return std::make_pair(std::move(raw_buffer), std::move(buffer_promise_cb)); } -absl::StatusOr PjRtCpuClient::GetMemorySpaceKindForShape( - const Shape& shape) const { - return topology().GetMemorySpaceKindForShape(shape); -} - -static std::vector GetParameterShapes(const ComputationLayout& layout) { - // For now, TPU programs compiled with multiple arguments cannot use tuples - // for any of their arguments, so we can assume that a tuple can only arise - // when there is a single argument. - std::vector shapes; - if (layout.parameter_count() == 1 && layout.parameter_shape(0).IsTuple()) { - shapes.reserve(layout.parameter_shape(0).tuple_shapes().size()); - absl::c_copy(layout.parameter_shape(0).tuple_shapes(), - std::back_inserter(shapes)); - } else { - shapes.reserve(layout.parameter_count()); - for (const ShapeLayout& sl : layout.parameter_layouts()) { - shapes.push_back(sl.shape()); - } - } - return shapes; -} - PjRtCpuExecutable::PjRtCpuExecutable( int num_replicas, int num_partitions, CompileOptions compile_options, std::unique_ptr cpu_executable, diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 01ca97aafd81d3..694c0fedfda696 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -179,6 +179,9 @@ class PjRtCpuRawClient : public PjRtRawClient { MaybeOwningMlirModule module, const CpuTopologyDescription& topology, int process_index, CompileOptions&& options); + tsl::AsyncValueRef ToAsyncExecutable( + std::shared_ptr executable) const override; + private: friend class PjRtCpuClient; friend class CpuExecutableLoadState; @@ -274,9 +277,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { &CommonPjRtClientImpl::topology()); } - absl::StatusOr GetMemorySpaceKindForShape( - const Shape& shape) const override; - bool BufferFromHostBufferSupportsZeroCopy( const void* data, PrimitiveType type, absl::Span dims, std::optional> byte_strides, const Shape& shape, diff --git a/third_party/xla/xla/pjrt/dynamic_shapes.cc b/third_party/xla/xla/pjrt/dynamic_shapes.cc index 9ae2b20a444fd8..fc39a352a67ef0 100644 --- a/third_party/xla/xla/pjrt/dynamic_shapes.cc +++ b/third_party/xla/xla/pjrt/dynamic_shapes.cc @@ -31,6 +31,32 @@ limitations under the License. namespace xla { +PjRtDynamicShapeKind GetPjRtDynamicShapeKind(const xla::Shape& shape) { + if (shape.is_static()) { + return PjRtDynamicShapeKind::kNotSupported; + } + if (shape.has_layout() && + shape.layout().dynamic_shape_metadata_prefix_bytes() > 0) { + return PjRtDynamicShapeKind::kPrefix; + } + return PjRtDynamicShapeKind::kSuffix; +} + +// Compute on-device size for a fully-specified shape. +absl::StatusOr PjRtGetOnDeviceBytesCount(const xla::Shape& shape, + PjRtDynamicShapeKind kind) { + // PjRtShapeAndMetadataTransferRequirements::Get->ShapeUtil::ArraySize + // requires a layout. + if (!shape.IsToken() && !shape.has_layout()) { + return absl::FailedPreconditionError( + "Buffer's on-device shape has no layout. Cannot determine on-device " + "bytes count."); + } + auto requirements = + PjRtShapeAndMetadataTransferRequirements::Get(shape, kind); + return static_cast(requirements.size); +} + PjRtShapeAndMetadataTransferRequirements PjRtShapeAndMetadataTransferRequirements::Get(const xla::Shape& shape, PjRtDynamicShapeKind kind) { diff --git a/third_party/xla/xla/pjrt/dynamic_shapes.h b/third_party/xla/xla/pjrt/dynamic_shapes.h index 386c3096f5f1e1..2c36a774c94521 100644 --- a/third_party/xla/xla/pjrt/dynamic_shapes.h +++ b/third_party/xla/xla/pjrt/dynamic_shapes.h @@ -32,11 +32,28 @@ enum class PjRtDynamicShapeKind { kSuffix, // Appended after the payload. }; +// Infer PjRtDynamicShapeKind from the shape. +PjRtDynamicShapeKind GetPjRtDynamicShapeKind(const xla::Shape& shape); + +// Compute on-device size for a fully-specified shape. +absl::StatusOr PjRtGetOnDeviceBytesCount(const xla::Shape& shape, + PjRtDynamicShapeKind kind); + +inline absl::StatusOr PjRtGetOnDeviceBytesCount( + const xla::Shape& shape) { + return PjRtGetOnDeviceBytesCount(shape, GetPjRtDynamicShapeKind(shape)); +} + // Offsets and bounds for extracting dynamic shape metadata. struct PjRtShapeAndMetadataTransferRequirements { static PjRtShapeAndMetadataTransferRequirements Get( const xla::Shape& shape, PjRtDynamicShapeKind kind); + static inline PjRtShapeAndMetadataTransferRequirements Get( + const xla::Shape& shape) { + return Get(shape, GetPjRtDynamicShapeKind(shape)); + } + size_t size = 0; size_t metadata_alignment = 0; size_t metadata_offset = 0; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc index 387b22917dd284..2b76674f95fbcc 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc @@ -390,4 +390,24 @@ StreamExecutorGpuTopologyDescription::FromProto( attributes, std::move(target_config)); } +absl::StatusOr +StreamExecutorGpuTopologyDescription::GetMemorySpaceKindForShape( + const xla::Shape& shape) const { + int kind = GetMemorySpaceKindIds()[0]; + if (shape.has_layout()) { + switch (shape.layout().memory_space()) { + case Layout::kHostMemorySpace: + return GetMemorySpaceKindIds()[1]; + break; + case Layout::kGenericFastMemorySpace: + case Layout::kDefaultMemorySpace: + break; + default: + return InvalidArgument("Unexpected memory space %d in output layout", + shape.layout().memory_space()); + } + } + return kind; +} + } // namespace xla diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.h b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.h index 4cf50f2c558c02..0bc572c3808cdf 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.h +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.h @@ -143,6 +143,9 @@ class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription { std::optional num_replicas_per_slice, int num_partitions, const MultiSliceConfig* multi_slice_config) const override; + absl::StatusOr GetMemorySpaceKindForShape( + const xla::Shape& shape) const override; + private: std::unique_ptr CreateDeviceDescription( int device_id) const; diff --git a/third_party/xla/xla/pjrt/infer_dispatch_info.cc b/third_party/xla/xla/pjrt/infer_dispatch_info.cc index 587edf9abd04f2..fde0780842797e 100644 --- a/third_party/xla/xla/pjrt/infer_dispatch_info.cc +++ b/third_party/xla/xla/pjrt/infer_dispatch_info.cc @@ -44,7 +44,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/mlir/utils/type_util.h" -#include "xla/pjrt/common_pjrt_client.h" +#include "xla/pjrt/dynamic_shapes.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/pjrt/pjrt_layout.h" @@ -77,16 +77,19 @@ std::vector GetParameterShapes(const ComputationLayout& layout) { return shapes; } -absl::StatusOr InferDispatchInfo( - CommonPjRtClient* client, std::vector parameter_device_shapes, - Shape output_device_shape, const HloInputOutputAliasConfig& alias_config, +absl::StatusOr InferDispatchInfo( + const PjRtTopologyDescription* topology, + std::vector parameter_device_shapes, Shape output_device_shape, + const HloInputOutputAliasConfig& alias_config, std::shared_ptr device_assignment, - std::vector + std::vector addressable_device_logical_ids, std::vector addressable_devices, - std::unique_ptr extras, - bool tuple_inputs) { - CommonPjRtLoadedExecutable::DispatchInfo result{ + std::unique_ptr extras, + bool tuple_inputs, + std::unique_ptr + input_hlo_snapshot_bits) { + PjRtLoadedExecutableDispatchInfo result{ .parameter_device_shapes = std::move(parameter_device_shapes), .output_device_shape = std::make_shared(std::move(output_device_shape)), @@ -94,10 +97,11 @@ absl::StatusOr InferDispatchInfo( .addressable_device_logical_ids = std::move(addressable_device_logical_ids), .device_assignment = std::move(device_assignment), + .input_hlo_snapshot_bits = std::move(input_hlo_snapshot_bits), .extras = std::move(extras), }; for (const auto& shape : result.parameter_device_shapes) { - ABSL_ASSIGN_OR_RETURN(int kind, client->GetMemorySpaceKindForShape(shape)); + ABSL_ASSIGN_OR_RETURN(int kind, topology->GetMemorySpaceKindForShape(shape)); result.parameter_memory_space_kind_ids.push_back(kind); } { @@ -107,7 +111,7 @@ absl::StatusOr InferDispatchInfo( : absl::MakeSpan(&*result.output_device_shape, 1); result.output_memory_space_kind_ids.reserve(shapes.size()); for (const auto& shape : shapes) { - ABSL_ASSIGN_OR_RETURN(int kind, client->GetMemorySpaceKindForShape(shape)); + ABSL_ASSIGN_OR_RETURN(int kind, topology->GetMemorySpaceKindForShape(shape)); result.output_memory_space_kind_ids.push_back(kind); } } @@ -121,9 +125,7 @@ absl::StatusOr InferDispatchInfo( result.parameter_device_shapes.size()); for (const Shape& shape : result.parameter_device_shapes) { DCHECK(!shape.IsTuple()); - ABSL_ASSIGN_OR_RETURN(int kind, client->GetMemorySpaceKindForShape(shape)); - ABSL_ASSIGN_OR_RETURN(int64_t size_in_bytes, - client->GetOnDeviceBytesCount(kind, shape)); + ABSL_ASSIGN_OR_RETURN(int64_t size_in_bytes, PjRtGetOnDeviceBytesCount(shape)); result.input_buffer_sizes_in_bytes.push_back(size_in_bytes); } return result; @@ -161,19 +163,18 @@ absl::StatusOr> GetShardShape( } // namespace -absl::StatusOr InferDispatchInfo( - CommonPjRtClient* client, mlir::ModuleOp mlir_module, +absl::StatusOr InferDispatchInfo( + const PjRtTopologyDescription* topology, mlir::ModuleOp mlir_module, const CompileOptions& options, std::shared_ptr device_assignment, - std::vector + std::vector addressable_device_logical_ids, std::vector addressable_devices, bool tuple_inputs) { if (!device_assignment) { return absl::UnimplementedError( "Async compilation requires a device_assignment"); } - auto extras = - std::make_unique(); + auto extras = std::make_unique(); extras->name = std::string(mlir_module.getSymName().value_or("?unknown program name?")); extras->num_partitions = device_assignment->replica_count(); @@ -215,8 +216,8 @@ absl::StatusOr InferDispatchInfo( xla::ShapeUtil::MakeShape(primitive_type, shard_shape); // TODO(parkers): Fix the nullptr layout. ABSL_ASSIGN_OR_RETURN(auto xla_shape, - client->MakeDefaultShapeForMemorySpace( - memory_space, xla_shard_shape, nullptr)); + topology->MakeCanonicalShapeForMemorySpace( + memory_space->kind_id(), xla_shard_shape, nullptr)); auto layout = std::make_shared(xla_shape.layout()); return std::make_tuple(xla_shape, layout); }; @@ -275,7 +276,7 @@ absl::StatusOr InferDispatchInfo( const auto& input_output_alias_config = extras->input_output_alias_config; ABSL_ASSIGN_OR_RETURN( auto result, - InferDispatchInfo(client, std::move(parameter_device_shapes), + InferDispatchInfo(topology, std::move(parameter_device_shapes), std::move(output_device_shape), input_output_alias_config, std::move(device_assignment), std::move(addressable_device_logical_ids), diff --git a/third_party/xla/xla/pjrt/infer_dispatch_info.h b/third_party/xla/xla/pjrt/infer_dispatch_info.h index 654861b9425363..0fd7ac62a63792 100644 --- a/third_party/xla/xla/pjrt/infer_dispatch_info.h +++ b/third_party/xla/xla/pjrt/infer_dispatch_info.h @@ -22,37 +22,75 @@ limitations under the License. #include "absl/status/statusor.h" #include "mlir/IR/BuiltinOps.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" -#include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" #include "xla/service/computation_layout.h" #include "xla/shape.h" namespace xla { +struct PjRtLoadedExecutableDispatchInfo { + std::vector parameter_device_shapes; + std::shared_ptr output_device_shape; + std::vector parameter_memory_space_kind_ids; + std::vector output_memory_space_kind_ids; + std::vector addressable_devices; + std::vector + addressable_device_logical_ids; + std::shared_ptr device_assignment; + std::vector parameters_that_may_be_donated; + std::vector input_buffer_sizes_in_bytes; + // Executable shape information that is computable from the PjRtExecutable*. + struct Extras { + std::string name; + int num_partitions; + int num_replicas; + absl::StatusOr>> + parameter_layouts; + absl::StatusOr>> + output_layouts; + std::optional> parameter_shardings; + std::optional> output_shardings; + std::vector parameter_memory_kinds; + std::vector output_memory_kinds; + absl::StatusOr fingerprint; + HloInputOutputAliasConfig input_output_alias_config; + }; + struct InputHloSnapshotBits { + xla::HloModuleProto hlo_module; + xla::DebugOptions debug_options; + }; + std::unique_ptr input_hlo_snapshot_bits; + std::unique_ptr extras; +}; + // Helper for extracting parameter shapes from GetParameterShapes. std::vector GetParameterShapes(const ComputationLayout& layout); -// Constructs CommonPjRtLoadedExecutable::DispatchInfo from both device lists +// Constructs PjRtLoadedExecutableDispatchInfo from both device lists // and metadata extracted from the final HloModule. -absl::StatusOr InferDispatchInfo( - CommonPjRtClient* client, std::vector parameter_device_shapes, - Shape output_device_shape, const HloInputOutputAliasConfig& alias_config, +absl::StatusOr InferDispatchInfo( + const PjRtTopologyDescription* topology, + std::vector parameter_device_shapes, Shape output_device_shape, + const HloInputOutputAliasConfig& alias_config, std::shared_ptr device_assignment, - std::vector + std::vector addressable_device_logical_ids, std::vector addressable_devices, - std::unique_ptr extras, - bool tuple_inputs); + std::unique_ptr extras, + bool tuple_inputs, + std::unique_ptr + input_hlo_snapshot_bits = nullptr); -// Constructs CommonPjRtLoadedExecutable::DispatchInfo from both device lists +// Constructs PjRtLoadedExecutableDispatchInfo from both device lists // and metadata extracted from the input mlir::ModuleOp. This may fail if all // information is not available yet. -absl::StatusOr InferDispatchInfo( - CommonPjRtClient* client, mlir::ModuleOp mlir_module, +absl::StatusOr InferDispatchInfo( + const PjRtTopologyDescription* topology, mlir::ModuleOp mlir_module, const CompileOptions& options, std::shared_ptr device_assignment, - std::vector + std::vector addressable_device_logical_ids, std::vector addressable_devices, bool tuple_inputs); diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h index 71457b80a8428c..69fa29c189d6f0 100644 --- a/third_party/xla/xla/pjrt/pjrt_compiler.h +++ b/third_party/xla/xla/pjrt/pjrt_compiler.h @@ -454,6 +454,13 @@ class PjRtTopologyDescription { "GetDefaultDeviceAssignment is not supported."); } + // Gets the memory_space_kind for a particular XLA layout. + virtual absl::StatusOr GetMemorySpaceKindForShape( + const xla::Shape& shape) const { + return absl::UnimplementedError( + "GetMemorySpaceKindForShape is not supported."); + } + // A list of all memory spaces kind_ids supported by this topology. virtual absl::Span GetMemorySpaceKindIds() const; diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h index 4a7c456c8457a1..5e3dd7288b45c2 100644 --- a/third_party/xla/xla/pjrt/pjrt_executable.h +++ b/third_party/xla/xla/pjrt/pjrt_executable.h @@ -376,6 +376,11 @@ class PjRtExecutable { virtual absl::StatusOr>> GetHloModules() const; + // Unoptimized hlo module. + virtual std::optional GetUnoptimizedHloModule() const { + return std::nullopt; + } + // Returns an output Shape per program, the size should be equal to // `GetHloModules()`. virtual absl::StatusOr> GetOutputShapes() const; diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.h b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.h index c1318a034d76aa..7e18923c167d96 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.h +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.h @@ -109,7 +109,8 @@ class CpuTopologyDescription : public PjRtTopologyDescription { return attributes_; } - absl::StatusOr GetMemorySpaceKindForShape(const Shape& shape) const; + absl::StatusOr GetMemorySpaceKindForShape( + const Shape& shape) const override; absl::StatusOr KindIdToKind(int kind) const; diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.h b/third_party/xla/xla/pjrt/raw_pjrt_client.h index d2b1663573a40a..ca4cf11cef9dd8 100644 --- a/third_party/xla/xla/pjrt/raw_pjrt_client.h +++ b/third_party/xla/xla/pjrt/raw_pjrt_client.h @@ -92,6 +92,10 @@ class PjRtExecutableLoadState virtual void Delete() = 0; virtual bool IsDeleted() const = 0; + virtual absl::Status Preload(PjRtExecutable* executable) { + return absl::OkStatus(); + } + virtual absl::StatusOr> LoadRawExecutable(tsl::AsyncValueRef executable, const ExecuteOptions& options, size_t host_callback_idx, @@ -194,6 +198,13 @@ class PjRtRawClient { return absl::UnimplementedError("RuntimeAbiVersion is not supported."); } + virtual tsl::AsyncValueRef ToAsyncExecutable( + std::shared_ptr executable) const = 0; + + virtual tsl::RCReference MakeLoadState() { + LOG(FATAL) << "Implement MakeLoadState()"; + } + virtual void ScheduleRemoteSend(PjRtMemorySpace* memory_space, PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD index 154374dbfde4e2..61eda90fc9c6cc 100644 --- a/third_party/xla/xla/pjrt/se/BUILD +++ b/third_party/xla/xla/pjrt/se/BUILD @@ -231,6 +231,7 @@ cc_library( "//xla/pjrt:dynamic_shapes", "//xla/pjrt:host_memory_allocator", "//xla/pjrt:host_memory_spaces", + "//xla/pjrt:infer_dispatch_info", "//xla/pjrt:layout_mode", "//xla/pjrt:maybe_owning_mlir_module", "//xla/pjrt:metrics", diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index f592ea5450426e..870d2482bdb027 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -121,6 +121,7 @@ limitations under the License. #include "xla/pjrt/dynamic_shapes.h" #include "xla/pjrt/host_memory_allocator.h" #include "xla/pjrt/host_memory_spaces.h" +#include "xla/pjrt/infer_dispatch_info.h" #include "xla/pjrt/layout_mode.h" #include "xla/pjrt/maybe_owning_mlir_module.h" #include "xla/pjrt/metrics.h" @@ -2153,7 +2154,7 @@ PjRtStreamExecutorClient::LoadSerializedExecutable( absl::string_view serialized, std::optional options, const LoadOptions& load_options) { ABSL_ASSIGN_OR_RETURN(auto executable, DeserializeExecutable(serialized, options)); - return LoadInternal(std::move(executable), /*dump=*/true); + return LoadInternal(std::move(executable), load_options, /*dump=*/true); } absl::StatusOr> @@ -2161,263 +2162,33 @@ PjRtStreamExecutorClient::LoadSerializedExecutable( const absl::Cord& serialized, std::optional options, const LoadOptions& load_options) { ABSL_ASSIGN_OR_RETURN(auto executable, DeserializeExecutable(serialized, options)); - return LoadInternal(std::move(executable), /*dump=*/true); + return LoadInternal(std::move(executable), load_options, /*dump=*/true); } -absl::StatusOr> -PjRtStreamExecutorClient::LoadInternal( - std::shared_ptr executable, bool dump) { - std::optional unoptimized_hlo_module_proto; - std::optional fingerprint = std::nullopt; - std::shared_ptr local_executable_ptr = nullptr; - CompileOptions compile_options; - { - auto se_executable = - std::static_pointer_cast(executable); - compile_options = se_executable->compile_options(); - - tsl::profiler::TraceMe traceme("PjRtStreamExecutorClient::Load"); - VLOG(1) << "PjRtStreamExecutorClient::Load"; - - ABSL_ASSIGN_OR_RETURN(local_executable_ptr, - se_executable->GetOrLoadExecutable(client())); - absl::StatusOr maybe_fingerprint = - se_executable->FingerprintExecutable(); - if (maybe_fingerprint.ok() && !maybe_fingerprint->empty()) { - fingerprint = *std::move(maybe_fingerprint); - } - - unoptimized_hlo_module_proto = - se_executable->unoptimized_hlo_module_proto(); - } - - ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides()); - std::vector - addressable_device_logical_ids; - std::vector addressable_devices; - ABSL_ASSIGN_OR_RETURN(auto device_assignment, - raw_client()->UpdateCompileOptions( - process_index(), topology(), &compile_options, - /*lookup_addressable_devices=*/true)); - - // Find devices that are addressable by this client/task. - if (device_assignment != nullptr) { - int num_replicas = device_assignment->replica_count(); - int num_partitions = device_assignment->computation_count(); - addressable_device_logical_ids.reserve(num_replicas * num_partitions); - addressable_devices.reserve(num_replicas * num_partitions); - for (int replica = 0; replica < num_replicas; ++replica) { - for (int partition = 0; partition < num_partitions; ++partition) { - int64_t device_id = (*device_assignment)(replica, partition); - GlobalDeviceId global_device_id(device_id); - - ABSL_ASSIGN_OR_RETURN(PjRtDevice * device, LookupDevice(global_device_id)); - if (device->process_index() != process_index()) { - VLOG(3) << "Non-local device: " << device_id; - continue; - } - PjRtLoadedExecutable::LogicalDeviceIds logica_device_ids; - logica_device_ids.replica = replica; - logica_device_ids.partition = partition; - addressable_device_logical_ids.push_back(std::move(logica_device_ids)); - addressable_devices.push_back(device); - } - } - } - - if (IsEarlyExitCompilation(compile_options)) { - return InvalidArgument( - "Executable compiled with xla_early_exit_with_layouts cannot be " - "loaded."); - } +absl::Status PjRtStreamExecutorExecutableLoadState::Preload( + PjRtExecutable* executable) { + return absl::down_cast(executable) + ->GetOrLoadExecutable(raw_client_->client()) + .status(); +} - const auto& ex_options = compile_options.executable_build_options; - const bool xla_dump_hlo_unoptimized_snapshots = - ex_options.has_debug_options() && - ex_options.debug_options().xla_dump_hlo_unoptimized_snapshots(); - if (dump) { - VLOG(1) << "Dumping deserialized executable"; - // Override the debug_options() embedded in the module with those - // explicitly passed in when deserializing. This allows options such as - // --xla_dump_to to be changed. Does not quite match the naming convention - // of the dump during compilation, which includes a backend-specific - // prefix. - if (local_executable_ptr->executable()->has_module()) { - DumpHloModuleIfEnabled(local_executable_ptr->executable()->module(), - kAfterOptimizationsDumpName, - ex_options.has_debug_options() - ? &ex_options.debug_options() - : nullptr); - } - } - ABSL_ASSIGN_OR_RETURN(PjRtMemorySpace* const default_memory_space, - this->addressable_devices()[0]->default_memory_space()); - xla::Shape result_shape = - local_executable_ptr->executable()->module().result_shape(); - std::vector output_memory_space_kind_ids; - { - absl::Span shapes = - result_shape.IsTuple() ? absl::MakeSpan(result_shape.tuple_shapes()) - : absl::MakeSpan(&result_shape, 1); - output_memory_space_kind_ids.reserve(shapes.size()); - for (const auto& shape : shapes) { - int kind = default_memory_space->kind_id(); - if (shape.has_layout()) { - switch (shape.layout().memory_space()) { - case Layout::kHostMemorySpace: - kind = PinnedHostMemorySpace::kKindId; - break; - case Layout::kGenericFastMemorySpace: - case Layout::kDefaultMemorySpace: - break; - default: - return InvalidArgument( - "Unexpected memory space %d in output layout", - shape.layout().memory_space()); - } - } - output_memory_space_kind_ids.push_back(kind); - } - } - if (result_shape.IsTuple()) { - for (auto& leaf_shape : result_shape.tuple_shapes()) { - if (leaf_shape.IsTuple()) { - return absl::InternalError( - absl::StrCat("Nested tuples are not supported with " - "PjRtStreamExecutorClient. got: ", - result_shape.ToString())); - } - } - } +tsl::RCReference +PjRtStreamExecutorRawClient::MakeLoadState() { + return tsl::MakeRef(this); +} - ComputationLayout computation_layout = - local_executable_ptr->executable()->compute_computation_layout(); - std::vector parameter_shapes; - parameter_shapes.reserve(computation_layout.parameter_count()); - for (int i = 0; i < computation_layout.parameter_count(); ++i) { - parameter_shapes.push_back(computation_layout.parameter_shape(i)); - } - std::vector parameter_memory_space_kind_ids; - { - absl::Span flat_parameter_shapes; - if (parameter_shapes.size() == 1 && parameter_shapes[0].IsTuple()) { - flat_parameter_shapes = parameter_shapes[0].tuple_shapes(); - } else { - flat_parameter_shapes = parameter_shapes; - } - parameter_memory_space_kind_ids.reserve(flat_parameter_shapes.size()); - for (const auto& shape : flat_parameter_shapes) { - int kind = default_memory_space->kind_id(); - if (shape.has_layout()) { - switch (shape.layout().memory_space()) { - case Layout::kHostMemorySpace: - kind = PinnedHostMemorySpace::kKindId; - break; - case Layout::kGenericFastMemorySpace: - case Layout::kDefaultMemorySpace: - break; - default: - return InvalidArgument( - "Unexpected memory space %d in output layout", - shape.layout().memory_space()); - } - } - parameter_memory_space_kind_ids.push_back(kind); - } - } - for (int result_index : compile_options.individually_defined_output_indices) { - if (result_index < 0 || static_cast(result_index) >= - output_memory_space_kind_ids.size()) { - return InvalidArgument( - "Individually defined output index %d is out of range for %d " - "outputs", - result_index, output_memory_space_kind_ids.size()); - } - } - auto load_state = - tsl::MakeRef(raw_client()); - ABSL_ASSIGN_OR_RETURN( - auto parameters_that_may_be_donated, - ComputeParametersThatMayBeDonated( - *tensorflow::down_cast(executable.get()) - ->hlo_module(), - compile_options.parameter_is_tupled_arguments)); - std::vector input_buffer_sizes_in_bytes; - { - if (device_assignment == nullptr) { - VLOG(3) << "PjRtStreamExecutorLoadedExecutable portable single-core"; - CHECK(addressable_devices.empty()); - } else { - VLOG(3) << "PjRtStreamExecutorLoadedExecutable device_assignment:\n" - << device_assignment->ToString(); - - if ((device_assignment->replica_count() > 1 || - device_assignment->computation_count() > 1) && - IsAllZeros(*device_assignment)) { - // This code path should only be triggered when we intentionally compile - // an HLO without having enough devices to actually run it. See the - // "--compile_only=true" option in - // tensorflow/compiler/xla/tools/multihost_hlo_runner/hlo_runner_main.cc. - // That will help us debug the XLA compiler locally. - LOG(INFO) - << "A workaround is in effect to allow compiling multi-device " - "HLOs on machines with fewer devices. Don't run this " - "executable."; - } else { - CHECK_LE(addressable_devices.size(), addressable_device_count()) - << "Inconsistent local device count."; - } - } - if (parameter_shapes.size() == 1 && parameter_shapes[0].IsTuple()) { - std::vector flat_parameter_shapes; - flat_parameter_shapes.reserve(parameter_shapes[0].tuple_shapes().size()); - for (const Shape& shape : parameter_shapes[0].tuple_shapes()) { - flat_parameter_shapes.push_back(shape); - } - std::swap(flat_parameter_shapes, parameter_shapes); - } - TransferManager* transfer_manager = client()->backend().transfer_manager(); - input_buffer_sizes_in_bytes.reserve(parameter_shapes.size()); - for (const Shape& shape : parameter_shapes) { - DCHECK(!shape.IsTuple()); - input_buffer_sizes_in_bytes.push_back( - transfer_manager->GetByteSizeRequirement(shape)); - } - } - using InputHloSnapshotBits = - CommonPjRtLoadedExecutable::DispatchInfo::InputHloSnapshotBits; - std::unique_ptr input_hlo_snapshot_bits; - if (xla_dump_hlo_unoptimized_snapshots && - unoptimized_hlo_module_proto.has_value()) { - input_hlo_snapshot_bits = - std::make_unique(InputHloSnapshotBits{ - std::move(*unoptimized_hlo_module_proto), - compile_options.executable_build_options.debug_options()}); - } - auto loaded_executable = std::make_unique( - this, - tsl::MakeAvailableAsyncValueRef( - std::static_pointer_cast(executable)), - CommonPjRtLoadedExecutable::DispatchInfo{ - std::move(parameter_shapes), - std::make_shared(result_shape), - std::move(parameter_memory_space_kind_ids), - std::move(output_memory_space_kind_ids), - std::move(addressable_devices), - std::move(addressable_device_logical_ids), - std::move(device_assignment), - std::move(parameters_that_may_be_donated), - std::move(input_buffer_sizes_in_bytes), - std::move(input_hlo_snapshot_bits), - }, - std::move(load_state)); - return std::unique_ptr(std::move(loaded_executable)); +tsl::AsyncValueRef +PjRtStreamExecutorRawClient::ToAsyncExecutable( + std::shared_ptr executable) const { + return tsl::MakeAvailableAsyncValueRef( + std::static_pointer_cast(executable)); } absl::StatusOr> PjRtStreamExecutorClient::Load(std::shared_ptr executable, const LoadOptions& load_options) { - auto loaded_executable = LoadInternal(std::move(executable), /*dump=*/false); + auto loaded_executable = + LoadInternal(std::move(executable), load_options, /*dump=*/false); for (const PjRtDevice* device : addressable_devices()) { LocalDeviceState* local_device_state = tensorflow::down_cast(device) diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h index ad11a2e68c32f3..196b80c0577445 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h @@ -293,6 +293,9 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { (executor_ == nullptr || !executor_->IsHostMemoryPinned(data, size)); } + tsl::AsyncValueRef ToAsyncExecutable( + std::shared_ptr executable) const override; + void ThenRecordEvent(BufferSequencingEventRef event, LocalDeviceState* local_device, EventPool::Handle device_event, se::Stream* stream); @@ -393,6 +396,8 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { virtual void RecordMemoryStats(LocalDeviceState* local_device_state) {} + tsl::RCReference MakeLoadState() override; + private: se::DeviceAddressAllocator* allocator_ = nullptr; std::unique_ptr owned_allocator_; @@ -428,6 +433,8 @@ class PjRtStreamExecutorExecutableLoadState : public PjRtExecutableLoadState { void Delete() override { is_deleted_.store(true); } bool IsDeleted() const override { return is_deleted_.load(); } + absl::Status Preload(PjRtExecutable* executable) override; + absl::StatusOr> LoadRawExecutable( tsl::AsyncValueRef executable, const ExecuteOptions& options, size_t host_callback_idx, @@ -540,9 +547,6 @@ class PjRtStreamExecutorClient : public CommonPjRtClientImpl { protected: friend class PjRtStreamExecutorRawBuffer; friend class PjRtStreamExecutorRawLoadedExecutable; - - absl::StatusOr> LoadInternal( - std::shared_ptr executable, bool dump); }; struct PjRtStreamExecutorExecutionOutput { diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index acbd860403a8db..e34af9ccc21186 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -212,7 +212,7 @@ absl::StatusOr> GetClient() { /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); std::vector> memory_spaces; memory_spaces.emplace_back(std::make_unique( - 0, devices.back().get(), "cpu", 0)); + 0, devices.back().get(), "device", tsl::Fingerprint32("device"))); devices.back()->AttachMemorySpace(memory_spaces.back().get(), /*is_default=*/true); auto topology = CreateCpuTopologyDescription(devices.size()); @@ -264,7 +264,7 @@ absl::StatusOr> GetClientWithDevices( /*process_index=*/0, /*process_index_in_partition=*/0, /*partition_index=*/0, "cpu")); memory_spaces.emplace_back(std::make_unique( - i, devices.back().get(), "cpu", 0)); + i, devices.back().get(), "device", tsl::Fingerprint32("device"))); devices.back()->AttachMemorySpace(memory_spaces.back().get(), /*is_default=*/true); } diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.h b/third_party/xla/xla/pjrt/se/stream_executor_executable.h index cb0175a729b5d2..6fac5f281fd9d4 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.h +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.h @@ -110,6 +110,10 @@ class StreamExecutorExecutable : public PjRtExecutable { return unoptimized_hlo_module_proto_; } + std::optional GetUnoptimizedHloModule() const override { + return unoptimized_hlo_module_proto_; + } + absl::StatusOr> GetAbiVersion() const override;