From 2a7bc1737f8b9582c26a8742276ad2b684eb95e5 Mon Sep 17 00:00:00 2001 From: Rishipal Singh Bhatia Date: Wed, 2 Sep 2026 12:12:25 -0700 Subject: [PATCH 01/13] [XLA:TPU] Rename AllocationValue defining_position and remove defining_instruction. PiperOrigin-RevId: 975257608 --- .../memory_space_assignment/algorithm.cc | 176 +++++++++--------- .../memory_space_assignment/algorithm.h | 2 +- .../allocation_value.cc | 4 +- .../allocation_value.h | 15 +- .../memory_space_assignment.h | 6 +- .../memory_space_assignment_test.cc | 23 ++- 6 files changed, 106 insertions(+), 120 deletions(-) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index a30bbe7990b0ae..84a9d65f6cd55d 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -634,9 +634,9 @@ bool MsaAlgorithm::MatchesPrefetchContext( const PrefetchContext& context, absl::string_view producer_name, ShapeIndex producer_shape_index, absl::string_view consumer_name) const { return context.request->use->hlo_use.instruction->name() == consumer_name && - context.request->allocation_value->defining_position() - .instruction->name() == producer_name && - context.request->allocation_value->defining_position().index == + context.request->allocation_value->position().instruction->name() == + producer_name && + context.request->allocation_value->position().index == producer_shape_index; } @@ -747,7 +747,7 @@ AllocationValue* FindAllocationValueForAsyncOperationStateUse( HloPosition source_position = GetNonTrivialSourcePosition( HloPosition{use.instruction->mutable_operand(0), use.operand_index}); for (AllocationValue& allocation_value : candidate_allocation_values) { - if (allocation_value.defining_position() == source_position) { + if (allocation_value.position() == source_position) { return &allocation_value; } } @@ -779,7 +779,7 @@ AllocationValue* FindLatestAllocationValueForUse( it != candidate_allocation_values.rend(); ++it) { AllocationValue* allocation_value = &(*it); const HloInstruction* defining_instruction = - allocation_value->defining_instruction(); + allocation_value->position().instruction; // Skip definitions from different computations. if (allocation_value->computation() != use_computation) { @@ -906,7 +906,7 @@ void MsaAlgorithm::CreateAllocationValues( for (int i = beginning_idx; i < allocation_values.size(); ++i) { AllocationValue& allocation_value = allocation_values.at(i); if (IsAsyncOperationStateDefinition( - allocation_value.defining_instruction())) { + allocation_value.position().instruction)) { for (const AllocationValue::Use& use : allocation_value.uses()) { HloInstruction* use_instruction = use.hlo_use.instruction; CHECK(use_instruction->opcode() == HloOpcode::kAsyncUpdate || @@ -921,15 +921,15 @@ void MsaAlgorithm::CreateAllocationValues( << "Unexpected use_instruction opcode: " << HloOpcodeString(use_instruction->opcode()) << " (" << use_instruction->ToString() << ") for async def: " - << allocation_value.defining_instruction()->ToString(); + << allocation_value.position().instruction->ToString(); } } bool is_async_operation_state = IsAsyncOperationStateDefinition( - allocation_value.defining_instruction()) || + allocation_value.position().instruction) || (has_async_pipelined_while_loops_ && IsBufferAliasedToAsyncPipelinedWhileLoop(allocation_value.value()) && - (IsAsyncPipelinedWhilePosition(allocation_value.defining_position()) || + (IsAsyncPipelinedWhilePosition(allocation_value.position()) || absl::c_any_of(allocation_value.uses(), [](const AllocationValue::Use& use) { return IsAsyncOperationStateUse(use.hlo_use); @@ -954,7 +954,7 @@ void MsaAlgorithm::CreateAllocationValues( "done/update use)."; allocation_value.set_requires_contiguous_allocation(true); } else if (options_.position_requires_contiguous_allocation_fn( - allocation_value.defining_position())) { + allocation_value.position())) { VLOG(3) << "Mark " << allocation_value.ToShortString() << " to require contiguous allocation because of options."; allocation_value.set_requires_contiguous_allocation(true); @@ -970,7 +970,7 @@ void MsaAlgorithm::FindAliases( std::vector> values_by_defining_inst; for (AllocationValue& value : *allocation_values) { - values_by_defining_inst[value.defining_instruction()].push_back(&value); + values_by_defining_inst[value.position().instruction].push_back(&value); } auto maybe_add_alias_with_instruction = [&](const HloInstruction* instruction, AllocationValue::Use* use) { @@ -980,16 +980,15 @@ void MsaAlgorithm::FindAliases( // When aliasing while loop boundaries, ensure that only matching tuple // shape indexes are linked together as aliases. if (use->hlo_use.instruction->opcode() == HloOpcode::kWhile && - !aliased_value->defining_position().index.empty() && - aliased_value->defining_position().index != - use->hlo_use.operand_index) { + !aliased_value->position().index.empty() && + aliased_value->position().index != use->hlo_use.operand_index) { continue; } - if (absl::c_find(use->aliases, aliased_value->defining_position()) == + if (absl::c_find(use->aliases, aliased_value->position()) == use->aliases.end()) { VLOG(3) << "Adding aliasing for use " << use->hlo_use.ToString() << " to " << aliased_value->ToShortString(); - use->aliases.push_back(aliased_value->defining_position()); + use->aliases.push_back(aliased_value->position()); } } } @@ -2268,7 +2267,7 @@ void MsaAlgorithm::CreateAllocationValuesForJointProcessedValues( proposal.allocation_values.end(), [this](AllocationValue& allocation_value) { return !IsInstructionPendingReplacements( - allocation_value.defining_instruction()); + allocation_value.position().instruction); }); NicePrintAllocationValues(proposal.allocation_values, /*log_level=*/3); @@ -5196,13 +5195,13 @@ std::vector MsaAlgorithm::GetInefficientAllocationSites( // The logic below is used mostly for testing, allowing a test case to inject // some custom logic for this method. if (options_.get_inefficient_allocation_sites_fn) { - std::vector defining_positions; - defining_positions.reserve(allocation_values.size()); + std::vector positions; + positions.reserve(allocation_values.size()); for (const AllocationValue& value : allocation_values) { - defining_positions.push_back(value.defining_position()); + positions.push_back(value.position()); } return options_.get_inefficient_allocation_sites_fn( - absl::MakeSpan(defining_positions)); + absl::MakeSpan(positions)); } if (!options_.cost_analysis || @@ -5338,7 +5337,7 @@ void MsaAlgorithm::CreateAllocationValuesFromColocatedIntervals( // when we try to allocate the AllocationValue, we would think they overlap. auto create_instruction_vector = [](const AllocationValue& allocation_value) { std::vector instruction_vector; - instruction_vector.push_back(allocation_value.defining_instruction()); + instruction_vector.push_back(allocation_value.position().instruction); for (const AllocationValue::Use& use : allocation_value.uses()) { instruction_vector.push_back(use.hlo_use.instruction); } @@ -5348,8 +5347,8 @@ void MsaAlgorithm::CreateAllocationValuesFromColocatedIntervals( for (int j = i + 1; j < new_allocation_values.size(); ++j) { const AllocationValue& allocation_value_1 = new_allocation_values[i]; const AllocationValue& allocation_value_2 = new_allocation_values[j]; - if (allocation_value_1.defining_position().index == - allocation_value_2.defining_position().index && + if (allocation_value_1.position().index == + allocation_value_2.position().index && create_instruction_vector(allocation_value_1) == create_instruction_vector(allocation_value_2)) { VLOG(3) << "Allocation values " << allocation_value_1.ToShortString() @@ -5431,7 +5430,7 @@ void MsaAlgorithm::MaybeSplitAllocationValues( bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( AllocationValue& allocation_value) const { if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, allocation_value.defining_position(), + options_.msa_tensor_overrides, allocation_value.position(), allocation_value.size())) { return true; } @@ -5443,9 +5442,9 @@ bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( void MsaAlgorithm::AssignDefaultMemIfNotAllowedInAlternateMem( AllocationValue& allocation_value, int64_t definition_time) { if (!options_.is_position_allowed_in_alternate_mem_fn( - allocation_value.defining_position()) || + allocation_value.position()) || MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - options_.msa_tensor_overrides, allocation_value.defining_position(), + options_.msa_tensor_overrides, allocation_value.position(), allocation_value.size())) { std::optional existing_req = RequiredMemoryAssignmentAt(allocation_value.value(), definition_time); @@ -5453,21 +5452,20 @@ void MsaAlgorithm::AssignDefaultMemIfNotAllowedInAlternateMem( // memory requirement, preserve that requirement and log a warning instead // of forcing an incompatible assignment in default memory. if (RequiresNoCopyAlternateMemAllocation(allocation_value) || - IsPositionColoredInAlternateMemory( - allocation_value.defining_position()) || - IsPositionColoredInAlternateMemoryAtTime( - allocation_value.defining_position(), definition_time) || + IsPositionColoredInAlternateMemory(allocation_value.position()) || + IsPositionColoredInAlternateMemoryAtTime(allocation_value.position(), + definition_time) || (existing_req.has_value() && existing_req->memory_space == MemorySpace::kAlternate)) { LOG(WARNING) << "The value " << allocation_value.value()->ToShortString() << " is pre-colored for alternate memory but the position " - << allocation_value.defining_position().ToString() + << allocation_value.position().ToString() << " is not allowed in the alternate memory. Respecting the " "color " "but this may break things later in compilation."; } else { AddRequiredAssignment(allocation_value.value(), - allocation_value.defining_instruction(), + allocation_value.position().instruction, static_cast(MemorySpace::kDefault), static_cast(definition_time), RequiredMemoryAssignment::Source:: @@ -5499,7 +5497,7 @@ MsaAlgorithm::GenerateAllocationSegmentContexts( value_indices_by_sync_inst.at(primary_use.hlo_use.instruction)) { AllocationValue& sync_destination = allocation_values.at(sync_destination_idx); - if (sync_destination.defining_instruction() == + if (sync_destination.position().instruction == primary_use.hlo_use.instruction) { VLOG(3) << "Adding secondary uses related to allocation value " << sync_destination.ToShortString() @@ -5617,7 +5615,7 @@ bool MsaAlgorithm::GetUpdatedRequireNoCopyAlternateMemForAsyncPipelinedWhile( } // Check 3: we're only looking to override async DUS and async dynamic slice. - const HloPosition& def_pos = allocation_value_to_update.defining_position(); + const HloPosition& def_pos = allocation_value_to_update.position(); const HloInstruction* def_instr = def_pos.instruction; if (def_instr->IsAsynchronous() && def_instr->async_wrapped_opcode() != HloOpcode::kDynamicUpdateSlice && @@ -5650,7 +5648,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( value_indices_by_sync_inst; for (size_t idx = 0; idx < allocation_values.size(); ++idx) { const HloInstruction* inst = - allocation_values.at(idx).defining_instruction(); + allocation_values.at(idx).position().instruction; if (IsInstructionPendingReplacements(inst)) { value_indices_by_sync_inst[inst].push_back(idx); } @@ -5713,7 +5711,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( << allocation_value.ToShortString(); if (IsInstructionPendingReplacements( - allocation_value.defining_instruction())) { + allocation_value.position().instruction)) { VLOG(3) << "Skip allocating allocation value " << allocation_value.ToShortString(); continue; @@ -5724,13 +5722,13 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( if (RequiresNoCopyAlternateMemAllocation(allocation_value) && allocation_value.size() > available_heap_size()) { if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, - allocation_value.defining_position(), allocation_value.size())) { + options_.msa_tensor_overrides, allocation_value.position(), + allocation_value.size())) { return absl::ResourceExhaustedError(absl::StrCat( "Cannot allocate pinned tensor in alternate memory: tensor size (", allocation_value.size(), " bytes) exceeds available heap size (", available_heap_size(), " bytes) for defining instruction '", - allocation_value.defining_instruction()->name(), "'")); + allocation_value.position().instruction->name(), "'")); } VLOG(3) << "Skip " << allocation_value.value()->ToShortString() << " because the buffer is larger than the heap size."; @@ -5758,7 +5756,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( &allocation_value_to_update)) { definition_time_for_allocation_value[&allocation_value_to_update] = hlo_live_range_.instruction_schedule().at( - allocation_value_to_update.defining_instruction()); + allocation_value_to_update.position().instruction); AssignDefaultMemIfNotAllowedInAlternateMem( allocation_value_to_update, definition_time_for_allocation_value.at( &allocation_value_to_update)); @@ -5774,8 +5772,8 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( AliasedOffset* preferred_offset = nullptr; if (has_async_pipelined_while_loops_) { const HloBuffer& hlo_buffer = alias_analysis_.GetUniqueBufferAt( - allocation_value_to_update.defining_position().instruction, - allocation_value_to_update.defining_position().index); + allocation_value_to_update.position().instruction, + allocation_value_to_update.position().index); auto buf_it = pipelined_while_buffer_id_to_aliased_offset_.find( hlo_buffer.id()); if (buf_it != pipelined_while_buffer_id_to_aliased_offset_.end()) { @@ -5806,14 +5804,14 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( preferred_offset = comp_it->second; } else if (has_async_pipelined_while_loops_ && IsAsyncPipelinedWhilePosition( - allocation_value_to_update.defining_position())) { + allocation_value_to_update.position())) { auto comp_it = pipelined_while_preferred_offset_for_computation_.find( allocation_value_to_update.computation()); if (comp_it != pipelined_while_preferred_offset_for_computation_.end()) { auto index_it = comp_it->second.find( - allocation_value_to_update.defining_position().index); + allocation_value_to_update.position().index); if (index_it != comp_it->second.end()) { preferred_offset = index_it->second; } @@ -5870,7 +5868,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( (request.require_no_copy_alternate_mem_allocation && MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( options_.msa_tensor_overrides, - allocation_value_to_update.defining_position(), + allocation_value_to_update.position(), allocation_value_to_update.size()))) { if (allocate_segment_result != AllocationResult::kSuccess) { std::string reason = ResultToString(allocate_segment_result); @@ -5897,7 +5895,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( ->shape() .ToString(), ", defining instruction '", - allocation_value_to_update.defining_instruction()->name(), + allocation_value_to_update.position().instruction->name(), "': ", reason, " (requested schedule time: ", request.preferred_prefetch_time.has_value() ? absl::StrCat(*request.preferred_prefetch_time) @@ -6021,7 +6019,7 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( "memory, which could not be satisfied. This typically happens " "because more pinned buffers are live than the alternate memory " "capacity.", - allocation_value.defining_position().ToString()); + allocation_value.position().ToString()); LOG(ERROR) << failed_precondition; return failed_precondition; } @@ -6134,7 +6132,7 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( // the value is used in a position or the earliest use time of the updated // allocation value. We find the minimum of these two times. int64_t min_time = - GetCorrectedUseTime(allocation_value.defining_instruction()); + GetCorrectedUseTime(allocation_value.position().instruction); int64_t earliest_position_time = std::numeric_limits::max(); for (auto& position : allocation_value.value()->positions()) { auto position_time = GetCorrectedUseTime(position.instruction); @@ -6224,9 +6222,9 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( auto is_required_in_alt_mem = [&]() { return require_no_copy_alternate_mem_allocation || IsPositionColoredInAlternateMemory( - allocation_value_to_update.defining_position()) || + allocation_value_to_update.position()) || IsPositionColoredInAlternateMemoryAtTime( - allocation_value_to_update.defining_position(), use_time) || + allocation_value_to_update.position(), use_time) || IsUseColoredInAlternateMemory(hlo_use) || (existing_req.has_value() && existing_req->memory_space == MemorySpace::kAlternate); @@ -6380,7 +6378,7 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( if (MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( options_.msa_tensor_overrides, hlo_use, allocation_value.size()) || MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - options_.msa_tensor_overrides, allocation_value.defining_position(), + options_.msa_tensor_overrides, allocation_value.position(), allocation_value.size())) { allow_prefetch = false; allow_no_copy_alternate_mem_allocation = false; @@ -6391,7 +6389,7 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory); } if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, allocation_value.defining_position(), + options_.msa_tensor_overrides, allocation_value.position(), allocation_value.size())) { require_no_copy_alternate_mem_allocation = true; allow_prefetch = false; @@ -6405,7 +6403,7 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( options_.msa_tensor_overrides, options_.preferred_prefetch_overrides, allocation_value.size(), hlo_use, instruction_schedule, live_range_start_time, - latest_prefetch_time, allocation_value.defining_position()); + latest_prefetch_time, allocation_value.position()); CHECK_OK(prefetch_override_info.status()); if (prefetch_override_info.value().has_value()) { VLOG(1) << "Overriding prefetch for " << hlo_use.instruction->name() @@ -6576,8 +6574,7 @@ void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse( allocation_sequence->rbegin(), allocation_sequence->rend(), [&](const auto& allocation) { return allocation->memory_space() == MemorySpace::kDefault && - allocation->defining_position() == - allocation_value.defining_position(); + allocation->defining_position() == allocation_value.position(); }); if (prev_allocation_in_default_mem_it != allocation_sequence->rend()) { VLOG(3) << "Found a prev allocation in default mem for while use: " @@ -6589,20 +6586,19 @@ void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse( auto body_allocation_value_it = absl::c_find_if(allocation_values, [&](const AllocationValue& value) { return value.computation() == hlo_use.instruction->while_body() && - value.defining_instruction()->opcode() == + value.position().instruction->opcode() == HloOpcode::kParameter && - value.defining_position().index == hlo_use.operand_index; + value.position().index == hlo_use.operand_index; }); CHECK_NE(body_allocation_value_it, allocation_values.end()); VLOG(3) << "Body allocation value: " << body_allocation_value_it->ToShortString(); int64_t body_parameter_time = instruction_schedule.at( - body_allocation_value_it->defining_instruction()); + body_allocation_value_it->position().instruction); body_allocation_value_it->mutable_allocation_sequence()->push_back( std::make_unique( **prev_allocation_in_default_mem_it, hlo_use.instruction, - body_allocation_value_it->defining_position(), - body_parameter_time)); + body_allocation_value_it->position(), body_parameter_time)); VLOG(3) << "Created: " << body_allocation_value_it->allocation_sequence() ->back() @@ -6610,8 +6606,8 @@ void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse( auto after_while_allocation_value_it = absl::c_find_if(allocation_values, [&](const AllocationValue& value) { - return value.defining_instruction() == hlo_use.instruction && - value.defining_position().index == hlo_use.operand_index; + return value.position().instruction == hlo_use.instruction && + value.position().index == hlo_use.operand_index; }); CHECK_NE(after_while_allocation_value_it, allocation_values.end()); VLOG(3) << "After while allocation value: " @@ -6735,10 +6731,10 @@ void MsaAlgorithm::CreateMirroredAllocations( // set of processed allocation values throughout conditionals, to avoid // re-processing. for (AllocationValue& allocation_val : allocation_values) { - // 1. Find allocation values whose defining positions lie inside the live + // 1. Find allocation values whose positions lie inside the live // range of the conditional. int64_t position_time = hlo_live_range_.instruction_schedule().at( - allocation_val.defining_instruction()); + allocation_val.position().instruction); int64_t last_use_time = position_time; for (const AllocationValue::Use& use : allocation_val.uses()) { last_use_time = std::max( @@ -6785,7 +6781,7 @@ void MsaAlgorithm::CreateMirroredAllocations( "result in a copy from alternate memory to alternate memory." << allocation_val.ToString(); allocation_val.mutable_allocation_sequence()->push_back( - std::make_unique(allocation_val.defining_position(), + std::make_unique(allocation_val.position(), MemorySpace::kDefault, kDummyChunk, position_time, last_use_time)); for (const AllocationValue::Use& use : allocation_val.uses()) { @@ -6803,7 +6799,7 @@ void MsaAlgorithm::CreateMirroredAllocations( // 4. We create a mirrored allocation in the alternate memory for the // allocation value. allocation_val.mutable_allocation_sequence()->push_back( - std::make_unique(allocation_val.defining_position(), + std::make_unique(allocation_val.position(), *last_allocation, position_time, last_use_time)); for (const AllocationValue::Use& use : allocation_val.uses()) { @@ -6840,7 +6836,7 @@ bool MsaAlgorithm::IsEvictionRequiredForPreviousUseAtConditional( for (AllocationValue& allocation_val : allocation_values) { int64_t position_time = hlo_live_range_.instruction_schedule().at( - allocation_val.defining_instruction()); + allocation_val.position().instruction); int64_t last_use_time = position_time; for (const AllocationValue::Use& use : allocation_val.uses()) { last_use_time = std::max( @@ -6890,8 +6886,7 @@ bool MsaAlgorithm::IsEvictionRequiredForPreviousUseAtConditional( // Require eviction if the set of allocation values jointly processed // within the conditional live range span across different buffers. const HloBuffer* destination_buffer = &alias_analysis_.GetUniqueBufferAt( - allocation_val.defining_position().instruction, - allocation_val.defining_position().index); + allocation_val.position().instruction, allocation_val.position().index); if (buffer_in_alt_mem != destination_buffer) { CHECK(!IsAsyncConversionCandidate( allocation_val.value()->defining_instruction())); @@ -8635,7 +8630,7 @@ void MsaAlgorithm::FreeAlternateMemoryColoringReservedAllocations( return; } const HloPosition& defining_position = - request.allocation_value_to_update->defining_position(); + request.allocation_value_to_update->position(); const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( defining_position.instruction, defining_position.index); auto reserved_allocations_it = @@ -8668,7 +8663,7 @@ void MsaAlgorithm::UpdateRequestWithAlternateMemoryColoringRequirements( return; } const HloPosition& defining_position = - request.allocation_value_to_update->defining_position(); + request.allocation_value_to_update->position(); int64_t inclusive_start_time = request.inclusive_start_time; int64_t use_time = request.end_time; @@ -8699,7 +8694,7 @@ void MsaAlgorithm::UpdateRequestWithDefaultMemoryColoringRequirements( return; } const HloPosition& defining_position = - request.allocation_value_to_update->defining_position(); + request.allocation_value_to_update->position(); int64_t inclusive_start_time = request.inclusive_start_time; int64_t use_time = request.end_time; @@ -8739,8 +8734,7 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { return AllocationResult::kSuccess; } - const HloPosition& defining_position = - request.allocation_value->defining_position(); + const HloPosition& defining_position = request.allocation_value->position(); VLOG(2) << "Finding allocation for " << request.allocation_value->ToShortString() << " [" << request.inclusive_start_time << ", " << request.end_time @@ -8754,12 +8748,11 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { } CHECK_LE(request.inclusive_start_time, request.end_time); if (VLOG_IS_ON(3) && options_.cost_analysis) { - const HloPosition& defining_position = - request.allocation_value->defining_position(); + const HloPosition& defining_position = request.allocation_value->position(); const HloUse& use = request.use->hlo_use; VLOG(3) << "Definition benefit = " << options_.cost_analysis->GetAlternateMemoryBenefit( - request.allocation_value->defining_position()) + request.allocation_value->position()) << " use benefit = " << options_.cost_analysis->GetAlternateMemoryBenefit( request.use->hlo_use); @@ -8933,7 +8926,7 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { // a fallback for this purpose. CHECK(allocation_result == AllocationResult::kSuccess); } else if (!IsAsyncConversionCandidate( - request.allocation_value_to_update->defining_position() + request.allocation_value_to_update->position() .instruction)) { // If the start of an allocation is in alternate memory and the allocation // sequence is not empty: @@ -9308,8 +9301,8 @@ bool MsaAlgorithm::ViolatesMaximumOutstandingAsyncCopies( AllocationResult MsaAlgorithm::ForceAlternateMemoryAllocationForMinTime( const AllocationRequest& request) { - CHECK_EQ(request.allocation_value->defining_position(), - request.allocation_value_to_update->defining_position()); + CHECK_EQ(request.allocation_value->position(), + request.allocation_value_to_update->position()); MsaBufferInterval alternate_mem_interval = MsaBufferInterval{ /*buffer=*/request.allocation_value->value(), @@ -9328,8 +9321,7 @@ AllocationResult MsaAlgorithm::ForceAlternateMemoryAllocationForMinTime( AddToPendingChunks(alternate_mem_interval, chunk_candidate); - const HloPosition& defining_position = - request.allocation_value->defining_position(); + const HloPosition& defining_position = request.allocation_value->position(); request.allocation_value->mutable_allocation_sequence()->push_back( std::make_unique( defining_position, MemorySpace::kAlternate, chunk_candidate, @@ -9379,8 +9371,7 @@ AllocationResult MsaAlgorithm::AllocateInAlternateMemoryNoCopy( return AllocationResult::kFailPrevAllocationNotInAlternateMem; } - const HloPosition& defining_position = - request.allocation_value->defining_position(); + const HloPosition& defining_position = request.allocation_value->position(); // If prefer_no_copy_alternate_mem_allocation is true, bypass the live range // duration checks. if (!request.require_no_copy_alternate_mem_allocation && @@ -9520,11 +9511,11 @@ AllocationResult MsaAlgorithm::Evict(const AllocationRequest& request, int64_t eviction_end_time = prev_allocation->end_time(); CHECK(eviction_exclusive_start_time <= eviction_end_time); - int64_t preferred_eviction_end_time = std::max( - options_.prefetch_interval_picker->PreferredEvictionEndTime( - request.allocation_value_to_update->defining_position().shape(), - eviction_exclusive_start_time, request.end_time), - eviction_end_time); + int64_t preferred_eviction_end_time = + std::max(options_.prefetch_interval_picker->PreferredEvictionEndTime( + request.allocation_value_to_update->position().shape(), + eviction_exclusive_start_time, request.end_time), + eviction_end_time); // Evictions must complete by the time of this use. preferred_eviction_end_time = std::min(preferred_eviction_end_time, request.latest_prefetch_time); @@ -9595,7 +9586,7 @@ AllocationResult MsaAlgorithm::Evict(const AllocationRequest& request, options_.cost_analysis ? options_.cost_analysis->GetAsyncCopyElapsed( options_.cost_analysis->GetShapeSizeBytes( - request.allocation_value->defining_position().shape())) + request.allocation_value->position().shape())) : 0.1; bool eviction_interval_too_short = @@ -10733,8 +10724,7 @@ std::vector MsaAlgorithm::FindBestChunkCandidates( // Then find the latest use that can be allocated contiguously without // copies. - const Shape& shape = - request.allocation_value_to_update->defining_position().shape(); + const Shape& shape = request.allocation_value_to_update->position().shape(); const int64_t shape_size = GetShapeSizeBytes(options_.cost_analysis, shape); for (; (use_time_it + 1) != use_times.end() && diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.h b/third_party/xla/xla/service/memory_space_assignment/algorithm.h index 799e302f1ae212..71b7f2f69eac19 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.h +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.h @@ -1169,7 +1169,7 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap { AllocationValue& allocation_value) const; // Adds a required assignment in default memory, at the given time, if - // allocation_value's defining position is not allowed in alternate memory. + // allocation_value's position is not allowed in alternate memory. void AssignDefaultMemIfNotAllowedInAlternateMem( AllocationValue& allocation_value, int64_t time); diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc index d1df3be48dccde..010c5737311e18 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc @@ -27,7 +27,7 @@ std::string AllocationValue::ToString() const { absl::StrAppend( &out, (requires_contiguous_allocation_ ? " (contiguous alloc)" : "")); absl::StrAppend(&out, "\n position:\n"); - absl::StrAppend(&out, " ", defining_position_.ToString(), "\n"); + absl::StrAppend(&out, " ", position_.ToString(), "\n"); absl::StrAppend(&out, " uses:\n"); for (const Use& use : uses_) { absl::StrAppend(&out, " ", use.hlo_use.ToString(), "\n"); @@ -37,7 +37,7 @@ std::string AllocationValue::ToString() const { std::string AllocationValue::ToShortString() const { return absl::StrCat("computation = ", computation()->name(), - ", position = ", defining_position_.ToString(), + ", position = ", position_.ToString(), ", value = ", value_->ToShortString(), (requires_contiguous_allocation_ ? " (cont alloc)" : "")); } diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h index 502335244f6b37..f3e8f7d3dca438 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h @@ -35,9 +35,9 @@ namespace memory_space_assignment { // (trivial positions are considered Tuple, GetTupleElement, and Bitcast). An // HloValue may include positions and uses that alias with each other across // multiple computations. We use this class to break these HloValues such that -// every AllocationValue has one defining position (that may alias with other +// every AllocationValue has one position (that may alias with other // AllocationValues). The uses field of the AllocationValue contains only the -// direct uses of the AllocationValue's defining position. +// direct uses of the AllocationValue's position. // // For example, consider the following HLO snippet: // @@ -129,22 +129,19 @@ class AllocationValue { AllocationValue(const HloValue* value, const HloPosition& position, int64_t size) : value_(value), - defining_position_(position), + position_(position), size_(size), requires_contiguous_allocation_(false), split_shape_(std::nullopt) {} - const HloPosition& defining_position() const { return defining_position_; } - const HloInstruction* defining_instruction() const { - return defining_position().instruction; - } + const HloPosition& position() const { return position_; } int64_t size() const { return size_; } void set_size(int64_t size) { size_ = size; } const std::vector& uses() const { return uses_; } std::vector& uses() { return uses_; } const HloValue* value() const { return value_; } const HloComputation* computation() const { - return defining_instruction()->parent(); + return position().instruction->parent(); } AllocationSequence* mutable_allocation_sequence() { return &allocation_sequence_; @@ -174,7 +171,7 @@ class AllocationValue { private: const HloValue* value_; - HloPosition defining_position_; + HloPosition position_; int64_t size_; // If true, there must be a contiguous allocation for this buffer without // any copies. diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h index f0d5c78703e8e7..2de772dca1499b 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h @@ -95,21 +95,21 @@ Classes * We do not create AllocationValues for trivial HloPositions, e.g., ones defined by Tuple, GetTupleElement, and Bitcast instructions. * The HloPosition used to define the AllocationValue is referred to as the - AllocationValue's defining position. + AllocationValue's position. * Typically, this is also the defining position of the HloValue. However, it may not be. For example, we would create an AllocationValue with an HloPosition of a read-only while loop parameter, but the HloValue corresponding to that HloPosition would have a different defining position. * The uses of an AllocationValue are limited to the direct uses of the - AllocationValue's defining position. + AllocationValue's position. * An AllocationValue is associated with an AllocationSequence, describing what to do with the underlying tensor, in memory, over the lifetime of the AllocationValue. - (Use) Segment: Each AllocationValue and its uses are separated into periods of time called use segments. The first use segment is from the (inclusive) - time of the AllocationValue's defining position to its first use + time of the AllocationValue's position to its first use (inclusive). The second use segment is from the first use (inclusive) to the second use (inclusive), etc. diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc index 7500e842ea22eb..0610693b4ad4db 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc @@ -1550,8 +1550,8 @@ ENTRY entry { options_result_modifier.allocation_result_modifier_testing_fn = [](const AllocationRequest& request, AllocationResult& result, int64_t retry_number) { - if (request.allocation_value_to_update->defining_instruction() - ->name() == "p0" && + if (request.allocation_value_to_update->position() + .instruction->name() == "p0" && request.use->hlo_use.instruction->name() == "add0") { result = AllocationResult::kFailRequiresUncommit; } @@ -1573,7 +1573,7 @@ ENTRY entry { options_request_modifier.max_retries = 1; options_request_modifier .allocation_request_modifier_testing_fn = [](AllocationRequest& request) { - if (request.allocation_value_to_update->defining_instruction()->name() == + if (request.allocation_value_to_update->position().instruction->name() == "p0" && request.use->hlo_use.instruction->name() == "add0") { // Schedule the copy-done before negate4 (scheduled at 6). @@ -1647,7 +1647,7 @@ ENTRY entry { options.allocation_result_modifier_testing_fn = [](const AllocationRequest& request, AllocationResult& result, int64_t retry_number) { - if (request.allocation_value->defining_instruction()->name() == + if (request.allocation_value->position().instruction->name() == "p0_copy" && request.use->hlo_use.instruction->name() == "concat") { result = AllocationResult::kFailRequiresUncommit; @@ -8382,11 +8382,11 @@ TEST_F(MemorySpaceAssignmentTest, // default memory, and creates a new one which is wrong. bool marked_inefficient = false; options.get_inefficient_allocation_sites_fn = - [&](absl::Span defining_positions) + [&](absl::Span positions) -> std::vector> { - if (absl::c_find(defining_positions, + if (absl::c_find(positions, HloPosition{FindInstruction(module.get(), "while1"), - {1}}) != defining_positions.end() && + {1}}) != positions.end() && !marked_inefficient) { LOG(INFO) << "Marking the use inefficient."; marked_inefficient = true; @@ -8423,11 +8423,10 @@ TEST_F(MemorySpaceAssignmentTest, InefficientAllocationRetryWithoutProgress) { Options options = DefaultMemorySpaceOptions(); // The hook flags the add instruction's use of p0 as inefficient. options.get_inefficient_allocation_sites_fn = - [&](absl::Span defining_positions) + [&](absl::Span positions) -> std::vector> { - if (absl::c_find(defining_positions, - HloPosition{FindInstruction(module.get(), "p0"), {}}) != - defining_positions.end()) { + if (absl::c_find(positions, HloPosition{FindInstruction(module.get(), "p0"), + {}}) != positions.end()) { return {HloUse{FindInstruction(module.get(), "add"), 1}}; } return {}; @@ -17851,7 +17850,7 @@ ENTRY entry { memory_space_options.allocation_result_modifier_testing_fn = [](const AllocationRequest& request, AllocationResult& result, int64_t retry_number) { - if (request.allocation_value->defining_instruction()->name() == + if (request.allocation_value->position().instruction->name() == "negate0" && retry_number <= 0 && result == AllocationResult::kSuccess) { result = AllocationResult::kFailOutOfMemory; From bdfeb5324de5f33aaefaaec73792699866273692 Mon Sep 17 00:00:00 2001 From: Bill Varcho Date: Wed, 2 Sep 2026 13:34:42 -0700 Subject: [PATCH 02/13] Reverts e22f2b27decb17d5396811f71a7e1f45cdbc36c9 PiperOrigin-RevId: 975296535 --- .../stablehlo_round_trip/shard_map_export.cc | 52 +++++++------------ ...stablehlo_round_trip_shard_map_export.mlir | 40 -------------- 2 files changed, 20 insertions(+), 72 deletions(-) diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc index 0366d458356da7..6a072cab73678b 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc @@ -184,25 +184,16 @@ void setOpManualAxes(Operation* op, ManualAxesAttr manualAxes, op->setAttr(kManualAxes, manualAxes); } -void setManualAxesForOpsInBody( - ManualComputationOp op, const mlir::SymbolTable& symbolTable, - ManualComputationToParentManualAxes& parentManualCompAxes); - -void setFuncManualAxesRecursively( - FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, - const mlir::SymbolTable& symbolTable, - ManualComputationToParentManualAxes& parentManualCompAxes); - -mlir::WalkResult setManualAxes( - Operation* op, ManualAxesAttr manualAxes, Attribute meshOrRef, - const mlir::SymbolTable& symbolTable, - ManualComputationToParentManualAxes& parentManualCompAxes) { - if (auto manualCompOp = mlir::dyn_cast(op)) { - // Record parent manual axes for this manualCompOp and process its body. - SmallVector& parentAxes = parentManualCompAxes[manualCompOp]; - parentAxes.assign(manualAxes.getValue().begin(), - manualAxes.getValue().end()); - setManualAxesForOpsInBody(manualCompOp, symbolTable, parentManualCompAxes); +void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, + Attribute meshOrRef, + const mlir::SymbolTable& symbolTable); + +mlir::WalkResult setManualAxes(Operation* op, ManualAxesAttr manualAxes, + Attribute meshOrRef, + const mlir::SymbolTable& symbolTable) { + if (mlir::isa(op)) { + // Skip `ManualComputationOp`s and their nested operations, they will + // be handled separately. return mlir::WalkResult::skip(); } if (!mlir::isa(op)) { @@ -211,16 +202,14 @@ mlir::WalkResult setManualAxes( if (CallOp callOp = mlir::dyn_cast(op)) { FuncOp funcOp = symbolTable.lookup(callOp.getCallee()); CHECK(funcOp) << "Failed to lookup function: " << callOp.getCallee().str(); - setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable, - parentManualCompAxes); + setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable); } return mlir::WalkResult::advance(); } -void setFuncManualAxesRecursively( - FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, - const mlir::SymbolTable& symbolTable, - ManualComputationToParentManualAxes& parentManualCompAxes) { +void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, + Attribute meshOrRef, + const mlir::SymbolTable& symbolTable) { llvm::SmallVector funcArgAttrs; funcArgAttrs.reserve(funcOp.getNumArguments()); for (int argNum = 0; argNum < funcOp.getNumArguments(); argNum++) { @@ -261,15 +250,15 @@ void setFuncManualAxesRecursively( // Walk in preorder of blocks in order to stop walks on manual computations. funcOp->walk([&](Operation* op) { - return setManualAxes(op, manualAxes, meshOrRef, symbolTable, - parentManualCompAxes); + return setManualAxes(op, manualAxes, meshOrRef, symbolTable); }); } // Sets the manual axes of all operations in `op`'s body. void setManualAxesForOpsInBody( - ManualComputationOp op, const mlir::SymbolTable& symbolTable, - ManualComputationToParentManualAxes& parentManualCompAxes) { + ManualComputationOp op, + const ManualComputationToParentManualAxes& parentManualCompAxes, + const mlir::SymbolTable& symbolTable) { TensorShardingAttr sharding = getFirstSharding(op); if (!sharding) { // If there are no in/out shardings, op.getManualAxes() must be empty. We do @@ -291,8 +280,7 @@ void setManualAxesForOpsInBody( // Set the manual axes of all operations in the body. op.getBody().front().walk( [&](Operation* opInBody) { - return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable, - parentManualCompAxes); + return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable); }); } @@ -500,7 +488,7 @@ class ShardMapExportPass parentAxes.insert(parentAxes.end(), parentOp.getManualAxes().begin(), parentOp.getManualAxes().end()); } - setManualAxesForOpsInBody(op, symbolTable, parentManualCompAxes); + setManualAxesForOpsInBody(op, parentManualCompAxes, symbolTable); }); // Need to do a separate post order walk to inline the diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir index 5ca34525922de5..8a7d5aa3cbd6b4 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir @@ -575,43 +575,3 @@ func.func private @bar(%arg0: tensor<8xi32>) -> tensor<8xi32> { // CHECK-NEXT: %1 = call @bar(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8xi32>) -> tensor<8xi32> // CHECK-NEXT: return %arg0 : tensor<4xi32> // CHECK-NEXT: } - -// ----- -sdy.mesh @mesh = <["a"=2, "b"=4]> - -// Tests that when a function containing a `ManualComputationOp` (callee) is -// called from within another `ManualComputationOp` (caller), the callee's -// `ManualComputationOp` correctly inherits parent manual axes from the caller. -// -// 1. Caller passes parent manual axes {"a"} to callee's inputs/outputs: -// CHECK-LABEL: func private @called_func_with_manual_comp( -// CHECK-SAME: %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) -// CHECK-SAME: -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) { -// CHECK-NEXT: %[[COPY:.*]] = mhlo.copy %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x8xf32> -// -// 2. Full-to-shard inside callee inherits caller axis "a" + callee axis "b" -> manual axes {"a", "b"}: -// CHECK-NEXT: %[[FULL_TO_SHARD:.*]] = stablehlo.custom_call @SPMDFullToShardShape(%[[COPY]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x2xf32> -// -// 3. Inner body call runs with both axes manual: -// CHECK-NEXT: %[[CALL:.*]] = call @xla.sdy.inlinable_manual_computation_body(%[[FULL_TO_SHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x2xf32> -// CHECK-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x2xf32> -// -// 4. Shard-to-full returns to caller scope, with manual axis {"a"} preserved: -// CHECK-NEXT: %[[SHARD_TO_FULL:.*]] = stablehlo.custom_call @SPMDShardToFullShape(%[[COPY_1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x8xf32> -// CHECK-NEXT: return %[[SHARD_TO_FULL]] : tensor<8x8xf32> -func.func private @called_func_with_manual_comp(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{}, {"b"}]>] out_shardings=[<@mesh, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x2xf32>) { - %1 = stablehlo.add %arg1, %arg1 : tensor<8x2xf32> - sdy.return %1 : tensor<8x2xf32> - } : (tensor<8x8xf32>) -> tensor<8x8xf32> - return %0 : tensor<8x8xf32> -} - -// CHECK-LABEL: func @manual_comp_calls_func_with_manual_comp( -func.func @manual_comp_calls_func_with_manual_comp(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> { - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<8x8xf32>) { - %1 = func.call @called_func_with_manual_comp(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32> - sdy.return %1 : tensor<8x8xf32> - } : (tensor<16x8xf32>) -> tensor<16x8xf32> - return %0 : tensor<16x8xf32> -} From 45703171bfaeb52409ed948763a391f024532431 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Wed, 2 Sep 2026 14:22:15 -0700 Subject: [PATCH 03/13] Move IsCpu to the topology and make PjRtCpuClient a trivial subclass of CommonPjRtClientImpl. PiperOrigin-RevId: 975326716 --- third_party/xla/xla/pjrt/BUILD | 1 + .../xla/xla/pjrt/common_pjrt_client.cc | 75 +++++++++++++++++++ third_party/xla/xla/pjrt/common_pjrt_client.h | 27 +++++-- third_party/xla/xla/pjrt/cpu/cpu_client.cc | 31 ++++---- third_party/xla/xla/pjrt/cpu/cpu_client.h | 23 ------ third_party/xla/xla/pjrt/gpu/BUILD | 1 + .../pjrt/gpu/se_gpu_topology_description.cc | 8 +- .../pjrt/gpu/se_gpu_topology_description.h | 2 + third_party/xla/xla/pjrt/pjrt_compiler.h | 4 + .../plugin/xla_cpu/cpu_topology_description.h | 4 + .../xla/pjrt/se/pjrt_stream_executor_client.h | 8 ++ 11 files changed, 139 insertions(+), 45 deletions(-) diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index fc3285257f1494..4becb8e978a00b 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -221,6 +221,7 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", + "//xla/backends/cpu:alignment", "//xla/error:error_codes", "//xla/hlo/ir:hlo", "//xla/pjrt/c:pjrt_c_api_device_event_hdrs", diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index 6a63332acaf787..13df6c962b4810 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -53,6 +53,7 @@ limitations under the License. #include "llvm/Support/MathExtras.h" #include "riegeli/bytes/cord_reader.h" #include "riegeli/bytes/string_reader.h" +#include "xla/backends/cpu/alignment.h" #include "xla/error/error_codes.h" #include "xla/executable_run_options.h" #include "xla/future.h" @@ -98,6 +99,63 @@ limitations under the License. namespace xla { +PjRtDynamicShapeKind CommonPjRtClient::GetDynamicShapeKind( + int memory_space_kind_id) const { + if (!IsTpuId(platform_id()) || + UnpinnedHostMemorySpace::kKindId == memory_space_kind_id) { + return PjRtDynamicShapeKind::kSuffix; + } + return PjRtDynamicShapeKind::kPrefix; +} + +absl::StatusOr> +CommonPjRtClient::GetHloCostAnalysis() const { + return std::make_unique( + [](const xla::Shape& shape) -> int64_t { + if (shape.IsTuple()) { + return ShapeUtil::TupleElementCount(shape) * sizeof(uint32_t); + } + auto result = PjRtGetOnDeviceBytesCount(shape); + CHECK_OK(result.status()) << shape.ToString(true); + return *result; + }); +} + +bool CommonPjRtClient::IsOnCpu(PjRtMemorySpace* memory_space) { + auto topology = GetTopologyDescription(); + CHECK_OK(topology.status()); + return (*topology)->IsMemorySpaceOnCpu(memory_space->kind_id()); +} + +bool CommonPjRtClient::BufferFromHostBufferSupportsZeroCopy( + const void* data, PrimitiveType type, absl::Span dims, + std::optional> byte_strides, const Shape& shape, + PjRtMemorySpace* memory_space, const Layout* device_layout) const { + if (!IsCpuId(platform_id()) && + memory_space->kind_id() != UnpinnedHostMemorySpace::kKindId) { + return false; + } + if (byte_strides && !HasMajorToMinorLayout(type, dims, *byte_strides)) { + return false; + } + // Packed arrays are unpacked on host and packed on device. + if (primitive_util::IsSubByteNonPredType(type)) { + return false; + } + if (shape.has_layout() && + !LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) { + return false; + } + + // If the input buffer has a default layout and is sufficiently aligned, we + // can simply point to the input array's data without any further copies. At + // the time of writing we require a 16-byte alignment because XLA may generate + // code which requires it. + if ((absl::bit_cast(data) & (cpu::MinAlign() - 1)) != 0) { + return false; + } + return true; +} void CommonPjRtClient::TrackFuture(PjRtMemorySpace* memory_space, absl::string_view debug_info, const Future<>& future) {} @@ -4045,6 +4103,23 @@ CommonPjRtClientImpl::CommonPjRtClientImpl( kv_store_(std::move(kv_store)), raw_client_(std::move(raw_client)) { CHECK(topology_) << " topology is required."; + auto set_bool_attr_from_plugin_attrs = [&](absl::string_view key, bool& out) { + if (!plugin_attributes_) { + return; + } + auto it = plugin_attributes_->attributes.find(key); + if (it != plugin_attributes_->attributes.end()) { + if (const bool* b = std::get_if(&it->second)) { + out = *b; + } + } + }; + set_bool_attr_from_plugin_attrs("allow_fallback_for_donation", + allow_fallback_for_donation_); + set_bool_attr_from_plugin_attrs("supports_two_phase_launch", + supports_two_phase_launch_); + set_bool_attr_from_plugin_attrs("supports_predetermined_error", + supports_predetermined_error_); } void CommonPjRtClientImpl::AttachDevices( diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h index 1ec82406fe904c..4f7e479c5de92a 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.h +++ b/third_party/xla/xla/pjrt/common_pjrt_client.h @@ -109,9 +109,7 @@ class CommonPjRtClient : public PjRtClient { virtual void CallOomHandlers() const {} virtual PjRtDynamicShapeKind GetDynamicShapeKind( - int memory_space_kind_id) const { - return PjRtDynamicShapeKind::kNotSupported; - } + int memory_space_kind_id) const; virtual void LaunchOnDevice(PjRtDevice* device, absl::AnyInvocable execute_fn) const { @@ -124,6 +122,9 @@ class CommonPjRtClient : public PjRtClient { return false; } + absl::StatusOr> GetHloCostAnalysis() + const override; + // Computes the memory requirements for storing shape on memory_space. absl::StatusOr GetOnDeviceBytesCount(int memory_space_kind, const xla::Shape& shape) const; @@ -304,7 +305,7 @@ class CommonPjRtClient : public PjRtClient { const xla::Shape& shape, PjRtMemorySpace* src_memory_space, PjRtMemorySpace* dst_memory_space); - virtual bool IsOnCpu(PjRtMemorySpace* memory_space) { return false; } + virtual bool IsOnCpu(PjRtMemorySpace* memory_space); virtual bool use_stream_based_compaction() const { return false; } absl::StatusOr> BufferFromHostBuffer( @@ -371,9 +372,7 @@ class CommonPjRtClient : public PjRtClient { virtual bool BufferFromHostBufferSupportsZeroCopy( const void* data, PrimitiveType type, absl::Span dims, std::optional> byte_strides, const Shape& shape, - PjRtMemorySpace* memory_space, const Layout* device_layout) const { - return false; - } + PjRtMemorySpace* memory_space, const Layout* device_layout) const; virtual absl::StatusOr LinearizeHostBufferInto( const void* data, PrimitiveType type, absl::Span dims, @@ -1091,6 +1090,16 @@ class CommonPjRtClientImpl : public CommonPjRtClient { std::shared_ptr kv_store, std::optional plugin_attributes = std::nullopt); + bool allow_fallback_for_donation() const override { + return allow_fallback_for_donation_; + } + bool supports_two_phase_launch() const override { + return supports_two_phase_launch_; + } + bool supports_predetermined_error() const override { + return supports_predetermined_error_; + } + private: const PjRtPlatformId platform_id_; const std::string platform_name_; @@ -1116,6 +1125,10 @@ class CommonPjRtClientImpl : public CommonPjRtClient { std::shared_ptr kv_store_; std::unique_ptr raw_client_; + + bool allow_fallback_for_donation_ = false; + bool supports_two_phase_launch_ = true; + bool supports_predetermined_error_ = true; }; } // namespace xla diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index fe1a95aa466368..0aae628700402a 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -423,6 +423,21 @@ PjRtCpuRawClient::PjRtCpuRawClient( tsl::Env::Default(), "XLAPjRtCpuClient", num_threads)) {} PjRtCpuRawClient::~PjRtCpuRawClient() {} + +PjRtPluginAttributes GetDefaultCpuPluginAttributes() { + PjRtPluginAttributes attrs; + attrs.pjrt_c_api_major_version = 0; + attrs.pjrt_c_api_minor_version = 0; + attrs.attributes["serialize_with_sdy"] = true; + attrs.attributes["allow_fallback_for_donation"] = true; + // This is needed because CPU currently doesn't have per-device dispatching + // threads for Execute() so two-phase launch can run into thread starvation. + attrs.attributes["supports_two_phase_launch"] = false; + // TODO(parkers): implement proper predetermined error support. + attrs.attributes["supports_predetermined_error"] = false; + return attrs; +} + PjRtCpuClient::PjRtCpuClient( int process_index, std::vector> devices, std::unique_ptr raw_client, @@ -430,7 +445,8 @@ PjRtCpuClient::PjRtCpuClient( : CommonPjRtClientImpl( xla::CpuPlatformId(), std::string(xla::CpuPlatformName()), std::string(xla::CpuPlatformVersion()), process_index, - std::move(topology), std::move(raw_client), /*kv_store=*/nullptr) { + std::move(topology), std::move(raw_client), /*kv_store=*/nullptr, + GetDefaultCpuPluginAttributes()) { std::vector> generic_devices; generic_devices.reserve(devices.size()); std::vector> memory_spaces; @@ -466,11 +482,6 @@ PjRtCpuClient::PjRtCpuClient( PjRtCpuClient::~PjRtCpuClient() { VLOG(1) << "PjRtCpuClient destroyed."; } -absl::StatusOr> -PjRtCpuClient::GetHloCostAnalysis() const { - return std::make_unique(cpu::CpuExecutable::ShapeSizeBytes); -} - // Find the root instruction of the entry computation. static const InstructionValueSet& GetRootValueSet( const BufferAssignment& assignment, const HloModule& module) { @@ -1039,14 +1050,6 @@ absl::StatusOr PjRtCpuRawClient::CreateDeviceEvent( return ToCpuEvent(std::move(dependency)); } -bool PjRtCpuClient::BufferFromHostBufferSupportsZeroCopy( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, const Shape& shape, - PjRtMemorySpace* memory_space, const Layout* device_layout) const { - return AbstractCpuBuffer::BufferFromHostBufferSupportsZeroCopy( - data, type, dims, byte_strides, shape); -} - absl::StatusOr PjRtCpuExecutable::GetCompiledMemoryStats() const { const auto& buffer_assignment = cpu_executable_->buffer_assignment(); diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 694c0fedfda696..6c72e89b619bea 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -251,38 +251,15 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { CommonPjRtClientImpl::raw_client()); } - bool allow_fallback_for_donation() const override { return true; } - // This is needed because CPU currently doesn't have per-device dispatching - // threads for Execute() so two-phase launch can run into thread starvation. - bool supports_two_phase_launch() const override { return false; } - // TODO(parkers): implement proper predetermined error support. - bool supports_predetermined_error() const override { return false; } - - PjRtDynamicShapeKind GetDynamicShapeKind( - int memory_space_kind_id) const override { - return PjRtDynamicShapeKind::kSuffix; - } - - absl::StatusOr> GetHloCostAnalysis() - const override; - absl::StatusOr> Load( std::shared_ptr executable, const LoadOptions& load_options) override; - bool IsOnCpu(PjRtMemorySpace* memory_space) override { return true; } - const xla::CpuTopologyDescription& topology() const { return *absl::down_cast( &CommonPjRtClientImpl::topology()); } - bool BufferFromHostBufferSupportsZeroCopy( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, const Shape& shape, - PjRtMemorySpace* memory_space, - const Layout* device_layout) const override; - private: friend class PjRtCpuLoadedExecutable; friend class CpuPjRtRawLoadedExecutable; diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD index 91833397fdf4a8..9d9b3d048f492f 100644 --- a/third_party/xla/xla/pjrt/gpu/BUILD +++ b/third_party/xla/xla/pjrt/gpu/BUILD @@ -996,6 +996,7 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", + "//xla/pjrt:host_memory_spaces", "//xla/pjrt:pjrt_common", "//xla/pjrt:pjrt_compiler", "//xla/pjrt:pjrt_device_description", 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 2b76674f95fbcc..020d669b9f9e3d 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 @@ -31,6 +31,7 @@ limitations under the License. #include "absl/types/span.h" #include "xla/layout.h" #include "xla/layout_util.h" +#include "xla/pjrt/host_memory_spaces.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_device_description.h" #include "xla/pjrt/pjrt_device_dimensions.h" @@ -295,10 +296,15 @@ absl::Span StreamExecutorGpuTopologyDescription::GetMemorySpaceKindIds() const { static const int kGpuMemorySpaceKindIds[] = { static_cast(tsl::Fingerprint32("device")), - static_cast(tsl::Fingerprint32("pinned_host"))}; + PinnedHostMemorySpace::kKindId}; return absl::MakeConstSpan(kGpuMemorySpaceKindIds); } +bool StreamExecutorGpuTopologyDescription::IsMemorySpaceOnCpu( + int memory_space_kind_id) const { + return memory_space_kind_id == PinnedHostMemorySpace::kKindId; +} + absl::StatusOr StreamExecutorGpuTopologyDescription::ChipBounds() const { return PjRtDeviceDimensions{gpu_topology_->num_partitions(), 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 0bc572c3808cdf..8e94cb488b9a10 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 @@ -146,6 +146,8 @@ class StreamExecutorGpuTopologyDescription : public PjRtTopologyDescription { absl::StatusOr GetMemorySpaceKindForShape( const xla::Shape& shape) const override; + bool IsMemorySpaceOnCpu(int memory_space_kind_id) const override; + private: std::unique_ptr CreateDeviceDescription( int device_id) const; diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h index 69fa29c189d6f0..239af2e09ceacd 100644 --- a/third_party/xla/xla/pjrt/pjrt_compiler.h +++ b/third_party/xla/xla/pjrt/pjrt_compiler.h @@ -467,6 +467,10 @@ class PjRtTopologyDescription { // GetMemorySpaceKindIds()[0] should be the default memory space id. int GetDefaultMemorySpaceKindId() const { return GetMemorySpaceKindIds()[0]; } + virtual bool IsMemorySpaceOnCpu(int memory_space_kind_id) const { + return false; + } + virtual absl::StatusOr ToProto() const { return absl::UnimplementedError("ToProto is unsupported."); } 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 7e18923c167d96..bd8beab21d5e35 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 @@ -134,6 +134,10 @@ class CpuTopologyDescription : public PjRtTopologyDescription { std::optional num_replicas_per_slice, int num_partitions, const MultiSliceConfig* multi_slice_config) const override; + bool IsMemorySpaceOnCpu(int memory_space_kind_id) const override { + return true; + } + private: const PjRtPlatformId platform_id_; const std::string platform_name_; 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 196b80c0577445..359820015aedd9 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 @@ -519,6 +519,14 @@ class PjRtStreamExecutorClient : public CommonPjRtClientImpl { return PjRtDynamicShapeKind::kSuffix; } + bool BufferFromHostBufferSupportsZeroCopy( + const void* data, PrimitiveType type, absl::Span dims, + std::optional> byte_strides, const Shape& shape, + PjRtMemorySpace* memory_space, + const Layout* device_layout) const override { + return false; + } + bool ShouldPerformZeroCopyLinearize( const void* data, const xla::Shape& device_shape, PrimitiveType type, absl::Span dims, From 413c89d02548a5df057b8ce83b061b909f727bd1 Mon Sep 17 00:00:00 2001 From: Aditya Sharma Date: Wed, 2 Sep 2026 14:38:34 -0700 Subject: [PATCH 04/13] Support GCS paths for saving Riegeli profiles PiperOrigin-RevId: 975336433 --- .../xla/xla/tsl/profiler/rpc/client/BUILD | 5 +- .../tsl/profiler/rpc/client/save_profile.cc | 44 ++++++++++--- .../profiler/rpc/client/save_profile_test.cc | 66 ++++++++++++++++++- 3 files changed, 101 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/BUILD b/third_party/xla/xla/tsl/profiler/rpc/client/BUILD index 1a8a2cf8700ec3..ad627cf9c3f8f5 100644 --- a/third_party/xla/xla/tsl/profiler/rpc/client/BUILD +++ b/third_party/xla/xla/tsl/profiler/rpc/client/BUILD @@ -75,12 +75,14 @@ cc_library( "//xla/tsl/platform:env", "//xla/tsl/platform:logging", "//xla/tsl/profiler/utils:file_system_utils", + "@com_google_absl//absl/cleanup", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", "@riegeli//riegeli/bytes:fd_writer", "@riegeli//riegeli/records:record_writer", + "@tsl//tsl/platform:path", "@tsl//tsl/profiler/protobuf:profiler_service_proto_cc", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", ], @@ -203,10 +205,11 @@ tsl_cc_test( "//xla/tsl/platform:test", "//xla/tsl/util/proto:parse_text_proto", "//xla/tsl/util/proto:proto_matchers", + "@com_google_absl//absl/cleanup", "@com_google_absl//absl/status:status_matchers", - "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", "@riegeli//riegeli/bytes:fd_reader", + "@riegeli//riegeli/bytes:string_reader", "@riegeli//riegeli/records:record_reader", "@tsl//tsl/platform:path", "@tsl//tsl/profiler/protobuf:profiler_service_proto_cc", diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/save_profile.cc b/third_party/xla/xla/tsl/profiler/rpc/client/save_profile.cc index fa1ff78cb6f75a..1d17ebadf008a1 100644 --- a/third_party/xla/xla/tsl/profiler/rpc/client/save_profile.cc +++ b/third_party/xla/xla/tsl/profiler/rpc/client/save_profile.cc @@ -15,7 +15,6 @@ limitations under the License. #include "xla/tsl/profiler/rpc/client/save_profile.h" -#include #include #include #include @@ -23,6 +22,7 @@ limitations under the License. #include #include +#include "absl/cleanup/cleanup.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/strings/match.h" @@ -40,6 +40,7 @@ limitations under the License. #include "xla/tsl/platform/file_system.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/profiler/utils/file_system_utils.h" +#include "tsl/platform/path.h" #include "tsl/profiler/protobuf/profiler_service.pb.h" #include "tsl/profiler/protobuf/xplane.pb.h" @@ -231,10 +232,30 @@ absl::Status SaveXSpaceChunks( ABSL_RETURN_IF_ERROR(record_options.FromString("brotli:6")); SetPadding(record_options); - std::string temp_path = - absl::StrCat(out_path, ".tmp.", Env::Default()->GetProcessId(), "_", - Env::Default()->NowMicros()); + absl::string_view scheme; + absl::string_view host_part; + absl::string_view path_part; + io::ParseURI(out_path, &scheme, &host_part, &path_part); + // An empty URI scheme indicates a local filesystem path. + const bool is_local_path = scheme.empty(); + + std::string temp_path; + if (is_local_path) { + temp_path = absl::StrCat(out_path, ".tmp.", Env::Default()->GetProcessId(), + "_", Env::Default()->NowMicros()); + } else if (!Env::Default()->LocalTempFilename(&temp_path)) { + return absl::InternalError( + absl::StrCat("Failed to create local temp filename for: ", out_path)); + } + + absl::Cleanup cleanup_temp_file = [&temp_path] { + Env::Default()->DeleteFile(temp_path).IgnoreError(); + }; + riegeli::RecordWriter writer(riegeli::FdWriter<>(temp_path), record_options); + if (!writer.ok()) { + return writer.status(); + } for (tensorflow::profiler::XSpace& xspace : xspaces) { std::string plane_names = GetPlaneNames(xspace); VLOG(1) << "SaveXSpaceChunks " @@ -249,15 +270,18 @@ absl::Status SaveXSpaceChunks( } xspaces.clear(); if (!writer.Close()) { - Env::Default()->DeleteFile(temp_path).IgnoreError(); return writer.status(); } - absl::Status status = - Env::Default()->RenameFile(temp_path, out_path, /*overwrite=*/true); - if (!status.ok()) { - Env::Default()->DeleteFile(temp_path).IgnoreError(); - return status; + if (is_local_path) { + ABSL_RETURN_IF_ERROR( + Env::Default()->RenameFile(temp_path, out_path, /*overwrite=*/true)); + } else { + absl::Status status = Env::Default()->CopyFile(temp_path, out_path); + if (!status.ok()) { + Env::Default()->DeleteFile(out_path).IgnoreError(); + return status; + } } return absl::OkStatus(); } diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/save_profile_test.cc b/third_party/xla/xla/tsl/profiler/rpc/client/save_profile_test.cc index c4dbef8491a052..4f716332ded139 100644 --- a/third_party/xla/xla/tsl/profiler/rpc/client/save_profile_test.cc +++ b/third_party/xla/xla/tsl/profiler/rpc/client/save_profile_test.cc @@ -17,12 +17,14 @@ limitations under the License. #include #include +#include #include #include +#include "absl/cleanup/cleanup.h" #include "absl/status/status_matchers.h" -#include "absl/strings/string_view.h" #include "riegeli/bytes/fd_reader.h" +#include "riegeli/bytes/string_reader.h" #include "riegeli/records/record_reader.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/status_matchers.h" // IWYU pragma: keep @@ -38,7 +40,6 @@ namespace profiler { namespace { using ::absl_testing::IsOk; -using ::tensorflow::profiler::XPlane; using ::tensorflow::profiler::XSpace; using ::testing::ElementsAre; using ::testing::Not; @@ -82,7 +83,66 @@ TEST(SaveProfileTest, SaveXSpaceChunksVectorSuccessAndPadding) { read_space.Clear(); } ASSERT_OK(reader.status()); - reader.Close(); + EXPECT_TRUE(reader.Close()); + + EXPECT_THAT(read_spaces, ElementsAre(Partially(EqualsProto(R"pb( + hostnames: "host1" + planes { name: "plane1" } + )pb")), + Partially(EqualsProto(R"pb( + hostnames: "host2" + planes { name: "plane2" } + )pb")))); +} + +TEST(SaveProfileTest, SaveXSpaceChunksRemoteFileSystemCopySuccess) { + std::string repo_root = "ram://test_remote_repo"; + absl::Cleanup cleanup = [&repo_root] { + int64_t undeleted_files = 0; + int64_t undeleted_dirs = 0; + Env::Default() + ->DeleteRecursively(repo_root, &undeleted_files, &undeleted_dirs) + .IgnoreError(); + }; + std::string run = "test_run_remote"; + std::string host = "test_host_remote"; + + XSpace space1 = ParseTextProtoOrDie(R"pb( + hostnames: "host1" + planes { name: "plane1" } + )pb"); + + XSpace space2 = ParseTextProtoOrDie(R"pb( + hostnames: "host2" + planes { name: "plane2" } + )pb"); + + std::vector spaces = {space1, space2}; + + ASSERT_OK(SaveXSpaceChunks(repo_root, run, host, spaces)); + EXPECT_TRUE(spaces.empty()); + + std::string file_path = + io::JoinPath(repo_root, run, "test_host_remote.xplane.riegeli"); + EXPECT_OK(Env::Default()->FileExists(file_path)); + + uint64_t file_size = 0; + ASSERT_OK(Env::Default()->GetFileSize(file_path, &file_size)); + EXPECT_GT(file_size, 0); + EXPECT_EQ(file_size % (64 * 1024), 0); + + std::string contents; + ASSERT_OK(tsl::ReadFileToString(Env::Default(), file_path, &contents)); + riegeli::RecordReader> reader{ + riegeli::StringReader<>(std::move(contents))}; + std::vector read_spaces; + XSpace read_space; + while (reader.ReadRecord(read_space)) { + read_spaces.push_back(std::move(read_space)); + read_space.Clear(); + } + ASSERT_OK(reader.status()); + EXPECT_TRUE(reader.Close()); EXPECT_THAT(read_spaces, ElementsAre(Partially(EqualsProto(R"pb( hostnames: "host1" From 358d659bbaa5208041c5100338e9e71f5c79edf2 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 2 Sep 2026 14:41:42 -0700 Subject: [PATCH 05/13] Refactor CUPTI Hardware Event System (HES) initialization. Add CuptiTracer::EnableHES() static method to support early initialization of CUPTI activity hardware trace before CUDA context creation. Support enabling HES in GPU tracer via the XPROF_ENABLE_HES environment variable. PiperOrigin-RevId: 975338177 --- .../xla/xla/backends/profiler/gpu/BUILD | 1 + .../xla/backends/profiler/gpu/cupti_tracer.cc | 96 ++++++++++++------- .../xla/backends/profiler/gpu/cupti_tracer.h | 4 + .../profiler/gpu/device_tracer_cuda.cc | 18 +++- .../gpu/profile_with_cuda_kernels_test.cc | 6 ++ 5 files changed, 86 insertions(+), 39 deletions(-) diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD index 7c7ea026cfe4a6..919e4b8faa26fd 100644 --- a/third_party/xla/xla/backends/profiler/gpu/BUILD +++ b/third_party/xla/xla/backends/profiler/gpu/BUILD @@ -963,6 +963,7 @@ xla_test( "h100", "b200", ], + shard_count = 3, tags = [ "cuda-only", "no_mac", diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc index 15f30480b471bd..a15577dc2b2e46 100644 --- a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc +++ b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.cc @@ -22,7 +22,6 @@ limitations under the License. #include #include #include -#include #include #include #include @@ -31,8 +30,10 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/const_init.h" #include "absl/base/no_destructor.h" #include "absl/base/optimization.h" +#include "absl/base/thread_annotations.h" #include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -59,7 +60,6 @@ limitations under the License. #include "xla/backends/profiler/gpu/cupti_utils.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/tsl/profiler/backends/cpu/annotation_stack.h" #include "xla/tsl/profiler/utils/per_thread.h" #include "xla/tsl/profiler/utils/xplane_builder.h" @@ -1120,13 +1120,6 @@ const char* GetCuptiErrorString(CuptiInterface* cupti_interface, return err_str; } -bool& IsCuptiHardwareEventSystemEnabled() { - // This flag can not flip to true once per process. Once enabled, it will stay - // enabled until the process is terminated. - static bool is_enabled = false; - return is_enabled; -} - } // namespace CuptiTracer::CuptiTracer(CuptiInterface* cupti_interface) @@ -1679,31 +1672,6 @@ absl::Status CuptiTracer::EnableActivityTracing() { << err; } } - if (option_->enable_activity_hardware_tracing) { - if (IsCuptiHardwareEventSystemEnabled()) { - LOG(INFO) << "CUPTI activity HW trace already enabled."; - } else { - auto err = cupti_interface_->ActivityEnableHWTrace(true); - if (err == CUPTI_ERROR_NOT_SUPPORTED) { - LOG(INFO) - << "CUPTI activity HW trace not enabled due to not supported on " - "this platform!"; - } else if (err != CUPTI_SUCCESS) { - LOG(WARNING) - << "Fail to enable CUPTI activity HW trace, CUPTI ERROR CODE:" - << err << " (" << GetCuptiErrorString(cupti_interface_, err) - << ")"; - } else { - LOG(INFO) << "CUPTI activity HW trace successfully enabled."; - IsCuptiHardwareEventSystemEnabled() = true; - } - } - } else { - if (IsCuptiHardwareEventSystemEnabled()) { - LOG(INFO) - << "CUPTI activity HW trace already enabled, continue with it."; - } - } if (using_v2_subscriber_api_) { RETURN_IF_CUPTI_ERROR(ActivityRegisterCallbacksV2( @@ -2073,6 +2041,66 @@ absl::Status CuptiTracer::ProcessActivityBuffer(CUcontext context, return ""; } +/*static*/ absl::Status CuptiTracer::EnableHES() { + static absl::Mutex mu(absl::kConstInit); + static bool is_hes_enabled ABSL_GUARDED_BY(mu) = false; + + absl::MutexLock lock(mu); + if (is_hes_enabled) { + LOG(INFO) << "CUPTI activity HW trace already enabled."; + return absl::OkStatus(); + } + + CUresult cu_err = cuInit(0); + if (cu_err != CUDA_SUCCESS) { + return absl::InternalError(absl::StrCat( + "cuInit(0) failed with error code: ", static_cast(cu_err))); + } + + CUcontext ctx = nullptr; + if (cuCtxGetCurrent(&ctx) == CUDA_SUCCESS && ctx != nullptr) { + return absl::FailedPreconditionError( + "Cannot enable HES: a CUDA context is already active on the current " + "thread."); + } + + int gpu_count = NumGpus(); + for (int i = 0; i < gpu_count; ++i) { + CUdevice dev; + if (cuDeviceGet(&dev, i) == CUDA_SUCCESS) { + unsigned int flags = 0; + int active = 0; + if (cuDevicePrimaryCtxGetState(dev, &flags, &active) == CUDA_SUCCESS && + active) { + return absl::FailedPreconditionError(absl::StrCat( + "Cannot enable HES: active primary CUDA context found on device ", + i)); + } + } + } + + CuptiInterface* cupti_interface = GetCuptiInterface(); + auto err = cupti_interface->ActivityEnableHWTrace(true); + if (err == CUPTI_ERROR_NOT_SUPPORTED) { + LOG(INFO) + << "CUPTI activity HW trace not enabled due to not supported on this " + "platform!"; + return absl::UnimplementedError( + "CUPTI activity HW trace not supported on this platform."); + } + if (err != CUPTI_SUCCESS) { + LOG(WARNING) << "Fail to enable CUPTI activity HW trace, CUPTI ERROR CODE: " + << err << " (" << GetCuptiErrorString(cupti_interface, err) + << ")"; + return absl::InternalError( + absl::StrCat("Fail to enable CUPTI activity HW trace: ", + GetCuptiErrorString(cupti_interface, err))); + } + LOG(INFO) << "CUPTI activity HW trace successfully enabled."; + is_hes_enabled = true; + return absl::OkStatus(); +} + std::vector CuptiTracer::GatherCallbackAnnotationsAndEvents(bool stop_recording) { // Note that it is OK to call PerThread's StartRecording() multiple times diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.h b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.h index 2f84f8a65b6837..2bf1653cac6118 100644 --- a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.h +++ b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer.h @@ -158,6 +158,10 @@ class CuptiTracer { // Returns the error (if any) when using libcupti. static std::string ErrorIfAny(); + // Enables activity hardware events tracing using HES (Hardware Event System). + // Once enabled, it stays enabled for the process lifetime. + static absl::Status EnableHES(); + // Returns true if the number of annotation strings is too large. The input // count is the per-thread count. bool TooManyAnnotationStrings(size_t count) const; diff --git a/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc b/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc index 93ae16cc4c9fe6..380447fc6f91da 100644 --- a/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc +++ b/third_party/xla/xla/backends/profiler/gpu/device_tracer_cuda.cc @@ -48,6 +48,18 @@ using tensorflow::ProfileOptions; using tensorflow::profiler::XSpace; using tsl::ReadBoolFromEnvVar; +static void MaybeEnableHES() { + bool enable_hes = false; + tsl::ReadBoolFromEnvVar("TF_GPU_CUPTI_ENABLE_ACTIVITY_HW_TRACING", false, + &enable_hes) + .IgnoreError(); + if (enable_hes) { + if (auto status = CuptiTracer::EnableHES(); !status.ok()) { + LOG(WARNING) << "Failed to enable HES: " << status.message(); + } + } +} + // GpuTracer for GPU. class GpuTracer : public tsl::profiler::ProfilerInterface { public: @@ -103,11 +115,6 @@ absl::Status GpuTracer::DoStart() { options_.activities_selected.push_back(CUPTI_ACTIVITY_KIND_OVERHEAD); options_.activities_selected.push_back(CUPTI_ACTIVITY_KIND_MEMSET); - // TODO: Change default to true once we have more confidence in HES. - ReadBoolFromEnvVar("TF_GPU_CUPTI_ENABLE_ACTIVITY_HW_TRACING", false, - &options_.enable_activity_hardware_tracing) - .IgnoreError(); - // CUDA/CUPTI 10 have issues (leaks and crashes) with cuptiFinalize. #if CUDA_VERSION >= 11000 options_.cupti_finalize = true; @@ -243,6 +250,7 @@ std::unique_ptr CreateGpuTracer( } auto register_gpu_tracer_factory = [] { + MaybeEnableHES(); RegisterProfilerFactory(&CreateGpuTracer); return 0; }(); diff --git a/third_party/xla/xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc b/third_party/xla/xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc index 88964e97f29aa2..59813aec95810c 100644 --- a/third_party/xla/xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc +++ b/third_party/xla/xla/backends/profiler/gpu/profile_with_cuda_kernels_test.cc @@ -110,6 +110,12 @@ void HandleRecords(PmSamples* samples) { void SimpleAddSubWithProfilerTest(bool enable_activity_hardware_tracing, bool enable_pm_sampling) { + if (enable_activity_hardware_tracing) { + if (auto status = CuptiTracer::EnableHES(); !status.ok()) { + LOG(WARNING) << "Failed to enable HES: " << status.message(); + } + } + uint32_t cupti_version = 0; cuptiGetVersion(&cupti_version); LOG(INFO) << "RUNTIME CUPTI version " << cupti_version From bb973de15184c82b42400dbd321bb5f41bd08e8a Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Wed, 2 Sep 2026 14:50:59 -0700 Subject: [PATCH 06/13] Test component-wise ULP distance for complex tanh in ArrayElementwiseOpTest - Add execution tests checking component-wise ULP distance with xla::UlpDistance instead of Euclidean relative error. - Enforce <= 2 ULPs for complex64 across all platforms. - Enforce <= 2 ULPs for complex128 on platforms with native float64 units (CPU, GPU), and <= 100 ULPs on TPU due to software-emulated double-double transcendentals. PiperOrigin-RevId: 975342933 --- .../xla/tests/array_elementwise_ops_test.cc | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/third_party/xla/xla/tests/array_elementwise_ops_test.cc b/third_party/xla/xla/tests/array_elementwise_ops_test.cc index 2a8ea692a1fb74..9f0a17faf3c19e 100644 --- a/third_party/xla/xla/tests/array_elementwise_ops_test.cc +++ b/third_party/xla/xla/tests/array_elementwise_ops_test.cc @@ -24,6 +24,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -111,6 +112,9 @@ class ArrayElementwiseOpTest : public ClientLibraryTestRunnerMixin< static constexpr double kEpsF64 = std::numeric_limits::epsilon(); ErrorSpec error_spec_{60 * kEpsF32, 60 * kEpsF32}; ErrorSpec strict_error_spec_{100 * kEpsF64, 100 * kEpsF64}; + + template + void TestComplexTanhUlps(int64_t max_ulps = 2); }; class ArrayElementwiseOpTestParamCount @@ -2949,6 +2953,91 @@ TEST_F(ArrayElementwiseOpTest, TanhF64sVector) { ComputeAndCompare(&builder, {}, strict_error_spec_); } +template +std::vector> GetComplexTanhTestInputs() { + return { + // Small inputs where (exp(a))^2 - (exp(-a))^2 suffered precision loss + // due to exp2 cancellation on older TPUs: + {T(0.0017180424), T(0.0017180424)}, + {T(-0.0017180424), T(0.0017180424)}, + {T(0.0017180424), T(-0.0017180424)}, + {T(-0.0017180424), T(-0.0017180424)}, + + // Other small magnitudes: + {T(1e-5), T(1e-5)}, + {T(-1e-5), T(1e-5)}, + {T(1e-4), T(1e-4)}, + {T(-1e-4), T(-1e-4)}, + {T(1e-3), T(1e-3)}, + {T(-1e-3), T(-1e-3)}, + {T(1e-2), T(1e-2)}, + {T(0.05), T(0.05)}, + + // Points near or on axes: + {T(0.0), T(0.0)}, + {T(1e-3), T(0.0)}, + {T(0.0), T(1e-3)}, + {T(-1e-3), T(0.0)}, + {T(0.0), T(-1e-3)}, + + // Moderate inputs: + {T(0.5), T(0.5)}, + {T(-0.5), T(0.5)}, + {T(1.0), T(1.0)}, + {T(2.0), T(-1.5)}, + + // Large inputs (overflow handling, Re(z) > 15 region where tanh(z) -> +/- + // 1): + {T(15.0), T(0.5)}, + {T(-15.0), T(0.5)}, + {T(20.0), T(1.0)}, + {T(-20.0), T(1.0)}, + }; +} + +template +void ArrayElementwiseOpTest::TestComplexTanhUlps(int64_t max_ulps) { + using RealT = typename ComplexT::value_type; + std::vector xs = GetComplexTanhTestInputs(); + XlaBuilder builder(TestName()); + auto a = ConstantR1(&builder, xs); + Tanh(a); + ASSERT_OK_AND_ASSIGN(Literal actual, this->ExecuteAndTransfer(&builder, {})); + for (int64_t i = 0; i < xs.size(); ++i) { + ComplexT act = actual.Get({i}); + std::complex zd(static_cast(xs[i].real()), + static_cast(xs[i].imag())); + std::complex ref = std::tanh(zd); + ComplexT exp(static_cast(ref.real()), + static_cast(ref.imag())); + + auto real_ulps = UlpDistance(act.real(), exp.real()); + auto imag_ulps = UlpDistance(act.imag(), exp.imag()); + ASSERT_TRUE(real_ulps.has_value()) + << "NaN/Inf mismatch on real part for input " << xs[i]; + ASSERT_TRUE(imag_ulps.has_value()) + << "NaN/Inf mismatch on imag part for input " << xs[i]; + EXPECT_LE(*real_ulps, max_ulps) + << "Real part ULP error exceeded for input " << xs[i] + << ": actual=" << act.real() << ", expected=" << exp.real() + << ", ulp_distance=" << *real_ulps; + EXPECT_LE(*imag_ulps, max_ulps) + << "Imag part ULP error exceeded for input " << xs[i] + << ": actual=" << act.imag() << ", expected=" << exp.imag() + << ", ulp_distance=" << *imag_ulps; + } +} + +TEST_F(ArrayElementwiseOpTest, TanhC64s) { TestComplexTanhUlps(); } + +TEST_F(ArrayElementwiseOpTest, TanhC128s) { + // Float64 transcendentals on TPU use software emulation, which has an error + // bound of up to 100 ULPs (matching strict_error_spec_). On platforms with + // native float64 units (CPU, GPU), enforce <= 2 ULPs. + const int64_t max_ulps = test::DeviceTypeIs(test::kTpu) ? 100 : 2; + TestComplexTanhUlps(max_ulps); +} + TEST_F(ArrayElementwiseOpTest, ExpF32sVector) { // The input tensor is large enough to exercise the vectorized exp // implementation on XLA CPU. From 5d5e899dd384ad27b12db5469c799ad41cfe0e25 Mon Sep 17 00:00:00 2001 From: Alexandros Theodoridis Date: Wed, 2 Sep 2026 14:57:27 -0700 Subject: [PATCH 07/13] PR #48140: [ROCm] Dynamically select rocm repo runfiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48140 ๐Ÿ“ Summary of Changes Symlink all the runfiles from lit test deps to execution path ๐ŸŽฏ Justification This change is required to support bzlmod builds in lit tests, we now switch rocm ci builds to bzlmod ๐Ÿš€ Kind of Contribution Please remove what does not apply: โ™ป๏ธ Cleanup ๐Ÿ“Š Benchmark (for Performance Improvements) Not relevant ๐Ÿงช Unit Tests: CI ๐Ÿงช Execution Tests: CI Copybara import of the project: -- 2ee21fed04c255ec7aeb2ec09215aaa0c2a746e3 by Alexandros Theodoridis : Dynamically select rocm runfiles dirs in workspace and bzlmod -- 03e79fa529e898ec2eb2227f88e2958d73eca106 by Alexandros Theodoridis : Clean up -- 2c8c71e46e10d5349b2f44f638c49a4801234c71 by Alexandros Theodoridis : Clean-up -- 408262b41467d7cff1cf43e848428894440aa290 by Alexandros Theodoridis : Address review comments Merging this change closes #48140 PiperOrigin-RevId: 975346419 --- third_party/xla/tensorflow.bazelrc | 6 ++-- .../xla/third_party/gpus/rocm/BUILD.tpl | 6 ++-- third_party/xla/xla/sh_test_with_runfiles.py | 34 ++++++++++++------- 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index 0c18ed64c47782..12006c287621e5 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -292,20 +292,20 @@ common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_rocm=True common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_cuda=False common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_sycl=False common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_hermetic_cc=True -common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace +common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic common:rocm_clang_hermetic --strategy=CppLink=local common:rocm_clang_hermetic --@rules_ml_toolchain//common:static_libcxx=False common:rocm --config=rocm_clang_hermetic common:rocm_ci --config=rocm -common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace +common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic common:rocm_ci_hermetic --dynamic_mode=off common:rocm_ci_hermetic --config=rocm_clang_hermetic common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a" common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a" common:rocm_ci_hermetic --repo_env=SYSROOT_DIST=linux_glibc_2_31 -common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace +common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic # This config option is used for SYCL as GPU backend. # SYCL Configuration (non-hermetic) diff --git a/third_party/xla/third_party/gpus/rocm/BUILD.tpl b/third_party/xla/third_party/gpus/rocm/BUILD.tpl index d05070d8d74223..ae5c129bba196d 100644 --- a/third_party/xla/third_party/gpus/rocm/BUILD.tpl +++ b/third_party/xla/third_party/gpus/rocm/BUILD.tpl @@ -129,16 +129,16 @@ cc_library( name = "rocm_rpath", linkopts = select({ ":build_hermetic": [ - "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib", + "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib", ], ":link_only": [ ], ":multiple_rocm_paths": [ - "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib", + "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib", "-Wl,-rpath=%{rocm_lib_paths}", ], "//conditions:default": [ - "-Wl,-rpath,external/%{rocm_repo_name}/rocm/%{rocm_root}/lib", + "-Wl,-rpath,../%{rocm_repo_name}/rocm/%{rocm_root}/lib", "-Wl,-rpath,/opt/rocm/lib", ], }), diff --git a/third_party/xla/xla/sh_test_with_runfiles.py b/third_party/xla/xla/sh_test_with_runfiles.py index dca6c6caca952a..9362b3a5ecb11f 100644 --- a/third_party/xla/xla/sh_test_with_runfiles.py +++ b/third_party/xla/xla/sh_test_with_runfiles.py @@ -27,17 +27,21 @@ def execute(self, test, lit_config): created_symlinks = [] if runfiles_env: rf_path = pathlib.Path(runfiles_env) - runfiles_dir = rf_path / "xla" - if runfiles_dir.is_dir(): - dst = pathlib.Path(test.getExecPath()).parent - dst.mkdir(parents=True, exist_ok=True) - for item in runfiles_dir.iterdir(): - target = dst / item.name - try: - target.symlink_to(item, target_is_directory=item.is_dir()) - created_symlinks.append(target) - except FileExistsError: - pass + + test_exec_dir = pathlib.Path(test.getExecPath()).parent + test_exec_dir.mkdir(parents=True, exist_ok=True) + + # Symlink all directories from runfiles root to test_exec_dir.parent + # RUNPATH has "../+rocm_configure_ext+local_config_rocm/..." patterns + for item in rf_path.iterdir(): + if item.is_dir(): + test_exec_symlink = test_exec_dir.parent / item.name + if not test_exec_symlink.exists(): + try: + test_exec_symlink.symlink_to(item, target_is_directory=True) + created_symlinks.append(test_exec_symlink) + except FileExistsError: + pass # Dynamically resolve hermetic cuda_nvcc in runfiles if present. # Static relative paths (e.g. %S/../../..) break for deeply nested targets @@ -71,6 +75,12 @@ def execute(self, test, lit_config): ) result = super().execute(test, lit_config) + + # Clean up created symlinks for target in created_symlinks: - target.unlink() + try: + target.unlink() + except FileNotFoundError: + pass + return result From 5468766b7f4a2c9f077ada127e2faa652be07327 Mon Sep 17 00:00:00 2001 From: Bhatu Date: Wed, 2 Sep 2026 16:20:07 -0700 Subject: [PATCH 08/13] Teach test_utils to find constrained reduction identity uses through data-formatting ops and fusions. - Update FindConstrainedUses to recursively trace through transparent data-formatting ops (copy, reshape, transpose, slice, bitcast, get_tuple_element) and fusions using a single-pass DFS. - Add IdentityElementType support for kMinimum and kMaximum. - Add unit tests verifying min/max reduction identity element generation through copies and fusions. PiperOrigin-RevId: 975388177 --- third_party/xla/xla/tests/BUILD | 9 +- third_party/xla/xla/tests/test_utils.cc | 178 ++++++++----------- third_party/xla/xla/tests/test_utils_test.cc | 65 ++++++- 3 files changed, 149 insertions(+), 103 deletions(-) diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 55382cafa7ae96..5b4d1905423859 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -117,13 +117,16 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", + "//xla/service:hlo_module_config", + "//xla/service:hlo_value", "//xla/service:hlo_verifier", - "//xla/tsl/platform:statusor", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:protobuf", ], @@ -3245,11 +3248,13 @@ xla_test( # There is nothing backend specific in this test, so just pick an arbitrary backend. backends = ["cpu"], deps = [ - ":hlo_pjrt_test_base", + ":hlo_test_base", ":test_utils", ":xla_internal_test_main", "//xla:literal", + "//xla:literal_util", "//xla:shape_util", + "//xla:types", "//xla:xla_data_proto_cc", "//xla:xla_proto_cc", "//xla/hlo/builder:xla_builder", diff --git a/third_party/xla/xla/tests/test_utils.cc b/third_party/xla/xla/tests/test_utils.cc index c846b0df2734d3..be58d91c2de931 100644 --- a/third_party/xla/xla/tests/test_utils.cc +++ b/third_party/xla/xla/tests/test_utils.cc @@ -26,10 +26,13 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/hlo_casting_utils.h" @@ -37,12 +40,13 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/service/hlo_module_config.h" +#include "xla/service/hlo_value.h" #include "xla/service/hlo_verifier.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tests/constraint_propagator.h" #include "xla/tests/constraint_state.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -62,49 +66,76 @@ bool NeedsInitValue(const HloUse& use) { op_num >= instruction->operand_count() / 2)); } -// Generate random values that are constrained to the input_shape minus the -// output_shape so as not to produce wrapping slices, for instance. -Literal MakeRandomIndex(int64_t index_bound, std::minstd_rand0* engine) { - std::uniform_int_distribution generator(0, index_bound); - return LiteralUtil::CreateR0(generator(*engine)); -} - -// Returns true if `dest' is reachable from `src' through data-formatting and -// custom call instructions within the same computation. -bool ReachableViaDataFormatting(const HloInstruction* src, - const HloInstruction* dest, - bool treat_gte_as_data_formatting) { - if (src == dest) { - return true; - } - switch (dest->opcode()) { +bool IsDataFormattingOp(const HloInstruction* instruction, + bool treat_gte_as_data_formatting) { + switch (instruction->opcode()) { + case HloOpcode::kConvert: + case HloOpcode::kReducePrecision: + case HloOpcode::kCopy: case HloOpcode::kReshape: case HloOpcode::kTranspose: - case HloOpcode::kCopy: case HloOpcode::kSlice: - break; + case HloOpcode::kBitcast: + return true; case HloOpcode::kCustomCall: - if (dest->custom_call_target() == "AssumeGatherIndicesInBound") { - break; - } - return false; - // TODO(b/249417724): a workaround for tuple param. + return instruction->custom_call_target() == "AssumeGatherIndicesInBound"; case HloOpcode::kGetTupleElement: - if (treat_gte_as_data_formatting) { - break; - } else { - return false; - } + return treat_gte_as_data_formatting; default: return false; } - for (const auto* operand : dest->operands()) { - if (ReachableViaDataFormatting(src, operand, - treat_gte_as_data_formatting)) { - return true; +} + +void FindConstrainedUsesHelper( + const HloDataflowAnalysis& dataflow, const HloInstruction& instruction, + bool treat_gte_as_data_formatting, + absl::flat_hash_set& visited, + std::vector& constrained_uses) { + auto [it, inserted] = visited.insert(&instruction); + if (!inserted) { + return; + } + + for (const auto& pair : dataflow.GetInstructionValueSet(&instruction)) { + for (const HloValue* value : pair.second.values()) { + for (const HloUse& use : value->GetUses()) { + HloInstruction* const user = use.instruction; + const HloOpcode opcode = user->opcode(); + const int64_t op_num = use.operand_number; + + if ((opcode == HloOpcode::kDynamicSlice && op_num >= 1) || + (opcode == HloOpcode::kDynamicUpdateSlice && op_num >= 2)) { + constrained_uses.push_back(use); + } else if ((opcode == HloOpcode::kGather || + opcode == HloOpcode::kScatter) && + op_num == 1) { + constrained_uses.push_back(use); + } else if (opcode == HloOpcode::kFusion) { + const HloInstruction* const to_analyze = + user->fused_parameter(op_num); + FindConstrainedUsesHelper(dataflow, *to_analyze, + treat_gte_as_data_formatting, visited, + constrained_uses); + } else if (NeedsInitValue(use)) { + constrained_uses.push_back(use); + } else if (IsDataFormattingOp(user, treat_gte_as_data_formatting)) { + FindConstrainedUsesHelper(dataflow, *user, + treat_gte_as_data_formatting, visited, + constrained_uses); + } else if (opcode == HloOpcode::kSort && + (user->operand_count() >= 2 || + Cast(user)->is_stable()) && + op_num == 0) { + // Operand 0 of sort is the array of keys used for key/value + // (two-operand) kSort instructions. Since sort stability is not + // guaranteed, constrain keys of key-value sort not to have + // duplicates, since otherwise the value order may legitimately + // differ. + constrained_uses.push_back(use); + } + } } } - return false; } // Use dataflow analysis on each parameter to see if there are uses that would @@ -116,62 +147,9 @@ std::vector FindConstrainedUses(const HloDataflowAnalysis& dataflow, const HloInstruction& param, bool treat_gte_as_data_formatting) { std::vector constrained_uses; - for (const auto& pair : dataflow.GetInstructionValueSet(¶m)) { - const HloValue& value = dataflow.GetUniqueValueAt(¶m, pair.first); - for (const HloUse& use : value.GetUses()) { - HloInstruction* instruction = use.instruction; - const HloOpcode opcode = instruction->opcode(); - const int64_t op_num = use.operand_number; - if ((opcode == HloOpcode::kDynamicSlice && op_num >= 1) || - (opcode == HloOpcode::kDynamicUpdateSlice && op_num >= 2)) { - constrained_uses.push_back(use); - } else if ((opcode == HloOpcode::kGather || - opcode == HloOpcode::kScatter) && - op_num == 1) { - constrained_uses.push_back(use); - } else if (opcode == HloOpcode::kFusion) { - const HloInstruction* const to_analyze = - instruction->fused_parameter(op_num); - auto fused_uses = FindConstrainedUses(dataflow, *to_analyze, - treat_gte_as_data_formatting); - constrained_uses.insert(constrained_uses.end(), fused_uses.begin(), - fused_uses.end()); - } else if (NeedsInitValue(use)) { - constrained_uses.push_back(use); - } else if (opcode == HloOpcode::kConvert || - opcode == HloOpcode::kReducePrecision) { - auto converted_uses = FindConstrainedUses(dataflow, *instruction, - treat_gte_as_data_formatting); - constrained_uses.insert(constrained_uses.end(), converted_uses.begin(), - converted_uses.end()); - } else if (opcode == HloOpcode::kSort && - (instruction->operand_count() >= 2 || - Cast(instruction)->is_stable()) && - op_num == 0) { - // Operand 0 of sort is the array of keys used for key/value - // (two-operand) kSort instructions. Since sort stability is not - // guaranteed, constrain keys of key-value sort not to have - // duplicates, since otherwise the value order may legitimately - // differ. - constrained_uses.push_back(use); - } - } - } - - for (auto* instruction : param.parent()->instructions()) { - const HloOpcode opcode = instruction->opcode(); - if (opcode == HloOpcode::kGather || opcode == HloOpcode::kScatter) { - if (instruction->operand(1) == ¶m) { - // Above already covers this case. - continue; - } - if (ReachableViaDataFormatting(¶m, instruction->operand(1), - treat_gte_as_data_formatting)) { - constrained_uses.push_back( - HloUse{instruction, /*operand_number=*/1, ShapeIndex{}}); - } - } - } + absl::flat_hash_set visited; + FindConstrainedUsesHelper(dataflow, param, treat_gte_as_data_formatting, + visited, constrained_uses); return constrained_uses; } @@ -187,8 +165,8 @@ absl::StatusOr CreateLiteralForConstrainedUses( bool generate_aligned_ds_indices, GetIndexKnownZeroesFn get_index_known_zeroes = nullptr) { int64_t index_bound = INT64_MAX; - // Used for operations like DUS / DS which need to be aligned when they appear - // in a fusion. + // Used for operations like DUS / DS which need to be aligned when they + // appear in a fusion. std::optional index_alignment = std::nullopt; bool no_duplicates = false; bool needs_constant = false; @@ -305,7 +283,8 @@ absl::StatusOr CreateLiteralForConstrainedUses( needs_sorted_indices, no_duplicates, use_large_range, max_bits_of_precision, index_alignment, index_known_zeroes, /*float_generator=*/nullptr); - } else if (needs_constant) { + } + if (needs_constant) { switch (identity_type) { case IdentityElementType::kZero: return LiteralUtil::Zero(param_shape.element_type()); @@ -327,14 +306,13 @@ absl::StatusOr CreateLiteralForConstrainedUses( /*index_known_zeroes=*/std::nullopt, /*float_generator=*/nullptr); } - } else { - return MakeFakeLiteral(param_shape, engine, /*limit=*/std::nullopt, - /*is_sorted=*/needs_sorted_indices, no_duplicates, - use_large_range, max_bits_of_precision, - /*index_alignment=*/std::nullopt, - /*index_known_zeroes=*/std::nullopt, - /*float_generator=*/nullptr); } + return MakeFakeLiteral(param_shape, engine, /*limit=*/std::nullopt, + /*is_sorted=*/needs_sorted_indices, no_duplicates, + use_large_range, max_bits_of_precision, + /*index_alignment=*/std::nullopt, + /*index_known_zeroes=*/std::nullopt, + /*float_generator=*/nullptr); } // Given a module entry parameter, use the dataflow analysis to see if a diff --git a/third_party/xla/xla/tests/test_utils_test.cc b/third_party/xla/xla/tests/test_utils_test.cc index 6c331b7ae30b00..1308d2a4db91eb 100644 --- a/third_party/xla/xla/tests/test_utils_test.cc +++ b/third_party/xla/xla/tests/test_utils_test.cc @@ -29,13 +29,15 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" +#include "xla/literal_util.h" #include "xla/service/hlo_runner_interface.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/tests/hlo_pjrt_test_base.h" +#include "xla/tests/hlo_test_base.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" +#include "xla/types.h" #include "xla/xla.pb.h" #include "xla/xla_data.pb.h" @@ -620,5 +622,66 @@ ENTRY %module (param: f4e2m1fn[1024]) -> f4e2m1fn[1024] { EXPECT_EQ(values.size(), num_possible_values); } +// Tests that max reduction uses MinValue as the identity element through copy +// pass-through. +TEST_F(TestUtilsTest, ReduceMaxIdentityElement) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule ReduceMaxIdentityModule + +max_BF16 (lhs: bf16[], rhs: bf16[]) -> bf16[] { + lhs = bf16[] parameter(0) + rhs = bf16[] parameter(1) + ROOT maximum = bf16[] maximum(lhs, rhs) +} + +ENTRY entry { + param_0 = bf16[4,5,128,256] parameter(0) + param_1 = bf16[] parameter(1) + copy = bf16[] copy(param_1) + ROOT reduce-window = bf16[3,4,128,256] reduce-window(param_0, copy), + window={size=2x2x1x1 pad=0_0x0_0x0_0x0_0}, to_apply=max_BF16 +} +)")); + + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get())); + ASSERT_EQ(args.size(), 2); + EXPECT_EQ(args[1].Get({}), + LiteralUtil::MinValue(BF16).Get({})); +} + +// Tests that min reduction uses MaxValue as the identity element through copy +// pass-through and fusion. +TEST_F(TestUtilsTest, ReduceMinIdentityElement) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule ReduceMinIdentityModule + +min_F32 (lhs: f32[], rhs: f32[]) -> f32[] { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT minimum = f32[] minimum(lhs, rhs) +} + +fused_computation (param_0: f32[4,5,128,256], param_1: f32[]) -> f32[3,4,128,256] { + param_0 = f32[4,5,128,256] parameter(0) + param_1 = f32[] parameter(1) + ROOT reduce-window = f32[3,4,128,256] reduce-window(param_0, param_1), + window={size=2x2x1x1 pad=0_0x0_0x0_0x0_0}, to_apply=min_F32 +} + +ENTRY entry { + param_0 = f32[4,5,128,256] parameter(0) + param_1 = f32[] parameter(1) + copy = f32[] copy(param_1) + ROOT fusion = f32[3,4,128,256] fusion(param_0, copy), kind=kOutput, calls=fused_computation +} +)")); + + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get())); + ASSERT_EQ(args.size(), 2); + EXPECT_EQ(args[1].Get({}), LiteralUtil::MaxValue(F32).Get({})); +} + } // namespace } // namespace xla From 5b92fd9e3219a551c30ecd9b5c65ce0684926750 Mon Sep 17 00:00:00 2001 From: Zac Mustin Date: Wed, 2 Sep 2026 16:29:18 -0700 Subject: [PATCH 09/13] Remove unneeded `xla_gpu_internal_packages` users. PiperOrigin-RevId: 975392160 --- third_party/xla/xla/pjrt/gpu/package_groups.bzl | 5 ----- 1 file changed, 5 deletions(-) diff --git a/third_party/xla/xla/pjrt/gpu/package_groups.bzl b/third_party/xla/xla/pjrt/gpu/package_groups.bzl index 51c29e5d8c3073..8b6dba8ef320c7 100644 --- a/third_party/xla/xla/pjrt/gpu/package_groups.bzl +++ b/third_party/xla/xla/pjrt/gpu/package_groups.bzl @@ -26,11 +26,6 @@ def xla_gpu_internal_packages(name = "xla_gpu_internal_packages"): packages = ["//..."], ) - native.package_group( - name = "legacy_gpu_internal_users", - packages = ["//..."], - ) - native.package_group( name = "legacy_se_gpu_pjrt_compiler_users", packages = ["//..."], From f31cee3d95e4fcfe676ff9c34a883fb266549040 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 2 Sep 2026 16:34:53 -0700 Subject: [PATCH 10/13] Fix bazel-out stripping regex. Sometime while resolving export issues I made manual changes and didn't rerun my validation steps; I've triple-checked this time. PiperOrigin-RevId: 975394784 --- tensorflow/compiler/aot/tfcompile.bzl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index b6facb7eb12e4e..11d63ab758d253 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -378,7 +378,8 @@ def _tf_library( "-e \"s|{{TFCOMPILE_HEADER}}|$(location " + header_file + ")|g\" " + "-e \"s|{{TFCOMPILE_CPP_CLASS}}|" + cpp_class + "|g\" " + "-e \"s|{{TFCOMPILE_NAME}}|" + no_ns_name + "|g\" " + - "-e \"s!bazel-out/[^/]*/(bin|genfiles)/!!g\" " + "-e \"s!bazel-out/[^/]*/bin/!!g\" " + + "-e \"s!bazel-out/[^/]*/genfiles/!!g\" " ) if gen_test: From 7a7748ce83827156efc687e798442e8638776d86 Mon Sep 17 00:00:00 2001 From: Jian Cai Date: Wed, 2 Sep 2026 16:39:59 -0700 Subject: [PATCH 11/13] [XLA][HLO Value Tracking] Preserve OriginalValue in InfeedTokenPropagation Update InfeedTokenPropagation to preserve and adjust instruction OriginalValue when tuplifying instruction shapes and inserting tokens into tuples. PiperOrigin-RevId: 975397101 --- .../xla/xla/hlo/transforms/collectives/BUILD | 2 + .../collectives/infeed_token_propagation.cc | 115 +++++++++++---- .../infeed_token_propagation_test.cc | 135 ++++++++++++++++++ 3 files changed, 225 insertions(+), 27 deletions(-) diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD index b9219888f34422..b9fe2ada53b71b 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/BUILD +++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD @@ -572,6 +572,7 @@ cc_library( "//xla/hlo/transforms/simplifiers:tuple_simplifier", "//xla/service:call_graph", "//xla/service:tuple_util", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -591,6 +592,7 @@ xla_cc_test( srcs = ["infeed_token_propagation_test.cc"], deps = [ ":infeed_token_propagation", + "//xla:shape_util", "//xla/hlo/analysis:hlo_ordering", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", diff --git a/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation.cc b/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation.cc index eaf8f308630929..158c9a1b09263f 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation.cc @@ -16,9 +16,11 @@ limitations under the License. #include "xla/hlo/transforms/collectives/infeed_token_propagation.h" #include +#include #include #include +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -33,6 +35,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/hlo/transforms/simplifiers/tuple_simplifier.h" @@ -171,6 +175,10 @@ absl::StatusOr InsertTokenIntoTuple(HloInstruction* tuple, // Trying to use tuple->ReplaceAllUsesWith(original_tuple) cause a cycle. std::vector original_users = tuple->users(); HloInstruction* original_tuple = TupleUtil::Duplicate(tuple); + if (tuple->original_value() && + !tuple->original_value()->is_synthetic_call()) { + original_tuple->set_original_value(tuple->original_value()); + } for (HloInstruction* original_user : original_users) { for (int64_t idx : original_user->operand_indices(tuple)) { // We expect the shape to be same, but checking that it is the same is @@ -180,8 +188,14 @@ absl::StatusOr InsertTokenIntoTuple(HloInstruction* tuple, } } + int64_t old_tuple_size = tuple->shape().tuple_shapes().size(); // Append the token to the parameter tuple. *tuple->mutable_shape()->add_tuple_shapes() = ShapeUtil::MakeTokenShape(); + absl::flat_hash_map old_to_new_tuple_idx; + for (int64_t i = 0; i < old_tuple_size; ++i) { + old_to_new_tuple_idx[i] = i; + } + CopyOriginalValue(original_tuple, tuple, old_to_new_tuple_idx); if (add_token_operand) { tuple->AppendOperand( computation->AddInstruction(HloInstruction::CreateToken())); @@ -198,6 +212,25 @@ absl::StatusOr InsertTokenIntoTuple(HloInstruction* tuple, tuple, tuple->shape().tuple_shapes().size() - 1)); return input_token_gte; } + +void TuplifyInstructionShape(HloInstruction* instruction) { + if (instruction->shape().IsTuple()) { + return; + } + *instruction->mutable_shape() = + ShapeUtil::MakeTupleShape({instruction->shape()}); + if (instruction->has_sharding()) { + instruction->set_sharding( + HloSharding::Tuple(instruction->shape(), {instruction->sharding()})); + } + if (instruction->original_value() && + !instruction->original_value()->is_synthetic_call()) { + auto new_ov = std::make_shared(instruction->shape()); + new_ov->mutable_tree()->CopySubtreeFrom( + instruction->original_value()->tree(), {}, {0}); + instruction->set_original_value(new_ov); + } +} } // namespace absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { @@ -207,15 +240,13 @@ absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { // Tuplify the branch parameter if needed. HloInstruction* parameter = branch->parameter_instruction(0); if (!parameter->shape().IsTuple()) { - *parameter->mutable_shape() = - ShapeUtil::MakeTupleShape({parameter->shape()}); - if (parameter->has_sharding()) { - HloSharding sharding = - HloSharding::Tuple(parameter->shape(), {parameter->sharding()}); - parameter->set_sharding(sharding); - } + TuplifyInstructionShape(parameter); HloInstruction* original = branch->AddInstruction( HloInstruction::CreateGetTupleElement(parameter, 0)); + if (parameter->original_value()) { + original->set_original_value( + OriginalValue::CreateFromInstruction(original)); + } ABSL_RETURN_IF_ERROR(parameter->ReplaceAllUsesWithDifferentShape(original)); } @@ -226,6 +257,10 @@ absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { if (!branch_tuple->shape().IsTuple()) { branch_tuple = conditional->parent()->AddInstruction( HloInstruction::CreateTuple({branch_tuple})); + if (branch_tuple->operand(0)->original_value()) { + branch_tuple->set_original_value( + OriginalValue::CreateFromInstruction(branch_tuple)); + } ABSL_RETURN_IF_ERROR(conditional->ReplaceOperandWithDifferentShape( branch_operand_idx, branch_tuple)); } @@ -233,7 +268,11 @@ absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { // Explicitly disjoin computation parameters from branch inputs, so we can // insert tokens into the input tuple. if (branch_tuple->opcode() == HloOpcode::kParameter) { + HloInstruction* old_branch_tuple = branch_tuple; branch_tuple = TupleUtil::Duplicate(branch_tuple); + if (old_branch_tuple->original_value()) { + branch_tuple->set_original_value(old_branch_tuple->original_value()); + } // We expect the shape to be same, but checking that it is the same is // expensive. ABSL_RETURN_IF_ERROR(conditional->ReplaceOperandWithDifferentShape( @@ -243,7 +282,11 @@ absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { // Explicitly make the root of the branch a tuple. HloInstruction* root = branch->root_instruction(); if (root->opcode() != HloOpcode::kTuple) { + HloInstruction* old_root = root; root = TupleUtil::Duplicate(root); + if (old_root->original_value()) { + root->set_original_value(old_root->original_value()); + } branch->set_root_instruction(root); } } @@ -255,7 +298,11 @@ absl::Status CanonicalizeConditionalInstruction(HloInstruction* conditional) { // Explicitly disjoin the conditional from being a computation root, so that // we can insert tokens into, while preserving the original computation shape. if (conditional->IsRoot()) { + HloInstruction* old_conditional = conditional; HloInstruction* new_root = TupleUtil::Duplicate(conditional); + if (old_conditional->original_value()) { + new_root->set_original_value(old_conditional->original_value()); + } conditional->parent()->set_root_instruction(new_root); } @@ -270,15 +317,13 @@ absl::Status CanonicalizeWhileInstruction(HloInstruction* loop) { // Tuplify the body parameter if needed. HloInstruction* body_parameter = body->parameter_instruction(0); if (!body_parameter->shape().IsTuple()) { - *body_parameter->mutable_shape() = - ShapeUtil::MakeTupleShape({body_parameter->shape()}); - if (body_parameter->has_sharding()) { - HloSharding sharding = HloSharding::Tuple(body_parameter->shape(), - {body_parameter->sharding()}); - body_parameter->set_sharding(sharding); - } + TuplifyInstructionShape(body_parameter); HloInstruction* original = body->AddInstruction( HloInstruction::CreateGetTupleElement(body_parameter, 0)); + if (body_parameter->original_value()) { + original->set_original_value( + OriginalValue::CreateFromInstruction(original)); + } ABSL_RETURN_IF_ERROR(body_parameter->ReplaceAllUsesWithDifferentShape(original)); } @@ -286,34 +331,34 @@ absl::Status CanonicalizeWhileInstruction(HloInstruction* loop) { HloInstruction* root = body->root_instruction(); if (!root->shape().IsTuple()) { root = body->AddInstruction(HloInstruction::CreateTuple({root})); + if (root->operand(0)->original_value()) { + root->set_original_value(OriginalValue::CreateFromInstruction(root)); + } body->set_root_instruction(root, /*accept_different_shape=*/true); } // Tuplify the condition parameter if needed. HloInstruction* cond_parameter = cond->parameter_instruction(0); if (!cond_parameter->shape().IsTuple()) { - *cond_parameter->mutable_shape() = - ShapeUtil::MakeTupleShape({cond_parameter->shape()}); - if (cond_parameter->has_sharding()) { - HloSharding sharding = HloSharding::Tuple(cond_parameter->shape(), - {cond_parameter->sharding()}); - cond_parameter->set_sharding(sharding); - } + TuplifyInstructionShape(cond_parameter); HloInstruction* original = cond->AddInstruction( HloInstruction::CreateGetTupleElement(cond_parameter, 0)); + if (cond_parameter->original_value()) { + original->set_original_value( + OriginalValue::CreateFromInstruction(original)); + } ABSL_RETURN_IF_ERROR(cond_parameter->ReplaceAllUsesWithDifferentShape(original)); } // Tuplify the while instruction if needed. if (!loop->shape().IsTuple()) { - *loop->mutable_shape() = ShapeUtil::MakeTupleShape({loop->shape()}); - if (loop->has_sharding()) { - HloSharding sharding = - HloSharding::Tuple(loop->shape(), {loop->sharding()}); - loop->set_sharding(sharding); - } + TuplifyInstructionShape(loop); HloInstruction* original = loop->parent()->AddInstruction( HloInstruction::CreateGetTupleElement(loop, 0)); + if (loop->original_value()) { + original->set_original_value( + OriginalValue::CreateFromInstruction(original)); + } ABSL_RETURN_IF_ERROR(loop->ReplaceAllUsesWithDifferentShape(original)); } @@ -322,13 +367,21 @@ absl::Status CanonicalizeWhileInstruction(HloInstruction* loop) { if (!loop_tuple->shape().IsTuple()) { loop_tuple = loop->parent()->AddInstruction( HloInstruction::CreateTuple({loop_tuple})); + if (loop_tuple->operand(0)->original_value()) { + loop_tuple->set_original_value( + OriginalValue::CreateFromInstruction(loop_tuple)); + } ABSL_RETURN_IF_ERROR(loop->ReplaceOperandWithDifferentShape(0, loop_tuple)); } // Explicitly disjoin computation parameters from loop inputs, so we can // insert tokens into the input tuple. if (loop_tuple->opcode() == HloOpcode::kParameter) { + HloInstruction* old_loop_tuple = loop_tuple; loop_tuple = TupleUtil::Duplicate(loop_tuple); + if (old_loop_tuple->original_value()) { + loop_tuple->set_original_value(old_loop_tuple->original_value()); + } // We expect the shape to be same, but checking that it is the same is // expensive. ABSL_RETURN_IF_ERROR(loop->ReplaceOperandWithDifferentShape(0, loop_tuple)); @@ -336,14 +389,22 @@ absl::Status CanonicalizeWhileInstruction(HloInstruction* loop) { // Explicitly make the root of the body a tuple. if (root->opcode() != HloOpcode::kTuple) { + HloInstruction* old_root = root; root = TupleUtil::Duplicate(root); + if (old_root->original_value()) { + root->set_original_value(old_root->original_value()); + } body->set_root_instruction(root); } // Explicitly disjoin the loop from being a computation root, so that // we can insert tokens into, while preserving the original computation shape. if (loop->IsRoot()) { + HloInstruction* old_loop = loop; HloInstruction* new_root = TupleUtil::Duplicate(loop); + if (old_loop->original_value()) { + new_root->set_original_value(old_loop->original_value()); + } loop->parent()->set_root_instruction(new_root); } diff --git a/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation_test.cc b/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation_test.cc index a5d25636dc6d89..5dc458e9454207 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/infeed_token_propagation_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/shape_util.h" #include "xla/tsl/platform/statusor.h" namespace op = xla::testing::opcode_matchers; @@ -406,6 +407,140 @@ ENTRY main { EXPECT_TRUE(cond_param->shape().tuple_shapes()[0].IsToken()); } +TEST_F(InfeedTokenPropagationTest, WhileInfeedPreservesOriginalValue) { + constexpr absl::string_view kHlo = R"( +HloModule main + +comp { + arg.0 = s32[] parameter(0), origin={{"prev.1"}} + token.0 = after-all() + infeed.0 = (s32[], token[]) infeed(token.0) + ROOT res.0 = s32[] constant(0), origin={{"res.0"}} +} + +cond { + arg.0 = s32[] parameter(0), origin={{"cond_arg.0"}} + ROOT true.0 = pred[] constant(true) +} + +ENTRY main { + init.0 = s32[] constant(0), origin={{"init.0"}} + ROOT while.0 = s32[] while(init.0), condition=cond, body=comp, origin={{"while.0"}} +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHlo)); + InfeedTokenPropagation itp; + ASSERT_OK_AND_ASSIGN(bool changed, itp.Run(module.get())); + EXPECT_TRUE(changed); + + // The while loop should have been tuplified and had token appended. + HloInstruction* loop = FindInstruction(module.get(), "while.0"); + ASSERT_NE(loop, nullptr); + EXPECT_TRUE(loop->shape().IsTuple()); + EXPECT_EQ(loop->shape().tuple_shapes().size(), 2); + ASSERT_NE(loop->original_value(), nullptr); + EXPECT_TRUE( + loop->original_value()->tree().element(ShapeIndex({0})).has_value()); + EXPECT_EQ( + loop->original_value()->tree().element(ShapeIndex({0}))->instruction_name, + "while.0"); + EXPECT_FALSE( + loop->original_value()->tree().element(ShapeIndex({1})).has_value()); + + // The body parameter should have matching original_value. + HloComputation* body_comp = FindComputation(module.get(), "comp"); + HloInstruction* body_param = body_comp->parameter_instruction(0); + ASSERT_NE(body_param, nullptr); + EXPECT_TRUE(body_param->shape().IsTuple()); + EXPECT_EQ(body_param->shape().tuple_shapes().size(), 2); + ASSERT_NE(body_param->original_value(), nullptr); + EXPECT_TRUE(body_param->original_value() + ->tree() + .element(ShapeIndex({0})) + .has_value()); + EXPECT_EQ(body_param->original_value() + ->tree() + .element(ShapeIndex({0})) + ->instruction_name, + "prev.1"); + EXPECT_FALSE(body_param->original_value() + ->tree() + .element(ShapeIndex({1})) + .has_value()); + + // The condition parameter should have matching original_value. + HloComputation* cond_comp = FindComputation(module.get(), "cond"); + HloInstruction* cond_param = cond_comp->parameter_instruction(0); + ASSERT_NE(cond_param, nullptr); + EXPECT_TRUE(cond_param->shape().IsTuple()); + EXPECT_EQ(cond_param->shape().tuple_shapes().size(), 2); + ASSERT_NE(cond_param->original_value(), nullptr); + EXPECT_TRUE(cond_param->original_value() + ->tree() + .element(ShapeIndex({0})) + .has_value()); + EXPECT_EQ(cond_param->original_value() + ->tree() + .element(ShapeIndex({0})) + ->instruction_name, + "cond_arg.0"); + EXPECT_FALSE(cond_param->original_value() + ->tree() + .element(ShapeIndex({1})) + .has_value()); +} + +TEST_F(InfeedTokenPropagationTest, ConditionalInfeedPreservesOriginalValue) { + constexpr absl::string_view kHlo = R"( +HloModule main + +true_comp { + arg.0 = s32[] parameter(0), origin={{"true_arg.0"}} + token.0 = after-all() + infeed.0 = (s32[], token[]) infeed(token.0) + gte.0 = get-tuple-element(infeed.0), index=0 + ROOT res.0 = (s32[]) tuple(gte.0) +} + +false_comp { + arg.0 = s32[] parameter(0), origin={{"false_arg.0"}} + ROOT res.1 = (s32[]) tuple(arg.0) +} + +ENTRY main { + pred.0 = pred[] constant(true) + arg.1 = s32[] constant(0), origin={{"arg.1"}} + arg.2 = s32[] constant(1), origin={{"arg.2"}} + ROOT cond.0 = (s32[]) conditional(pred.0, arg.1, arg.2), true_computation=true_comp, false_computation=false_comp, origin={({"cond.0"})} +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHlo)); + InfeedTokenPropagation itp; + ASSERT_OK_AND_ASSIGN(bool changed, itp.Run(module.get())); + EXPECT_TRUE(changed); + + HloComputation* true_comp = FindComputation(module.get(), "true_comp"); + HloInstruction* true_param = true_comp->parameter_instruction(0); + ASSERT_NE(true_param, nullptr); + EXPECT_TRUE(true_param->shape().IsTuple()); + ASSERT_NE(true_param->original_value(), nullptr); + EXPECT_TRUE(true_param->original_value() + ->tree() + .element(ShapeIndex({0})) + .has_value()); + EXPECT_EQ(true_param->original_value() + ->tree() + .element(ShapeIndex({0})) + ->instruction_name, + "true_arg.0"); + EXPECT_FALSE(true_param->original_value() + ->tree() + .element(ShapeIndex({1})) + .has_value()); +} + TEST_F(InfeedTokenPropagationTest, WhileOutfeed) { constexpr absl::string_view kHlo = R"( HloModule main From 87dfcf526aa6fd34850f6604a894680944826c17 Mon Sep 17 00:00:00 2001 From: Bixia Zheng Date: Wed, 2 Sep 2026 16:44:51 -0700 Subject: [PATCH 12/13] Added createStablehloCanonicalizeFromHloImportPass() to the beginning of addStablehloImportPipeline() in stablehlo_import.cc. This canonicalizes tuple-returning operations into multi-result operations before HLO sharding import. Added FileCheck unit tests to stablehlo_import_pipeline.mlir verifying HLO sharding import on tuple-returning custom calls. PiperOrigin-RevId: 975399339 --- .../service/spmd/shardy/stablehlo_round_trip/BUILD | 1 + .../shardy/stablehlo_round_trip/stablehlo_import.cc | 3 +++ .../spmd/shardy/test/stablehlo_import_pipeline.mlir | 11 +++++++++++ 3 files changed, 15 insertions(+) diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/BUILD index d4e4fd40d4d276..59c37be8059f81 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/BUILD +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/BUILD @@ -128,6 +128,7 @@ cc_library( "//xla/hlo/ir:named_sharding", "//xla/hlo/ir:tile_assignment", "//xla/hlo/translate/mhlo_to_hlo:attribute_exporter", + "//xla/mlir_hlo:stablehlo_extension_passes", "//xla/service/spmd/shardy:constants", "//xla/service/spmd/shardy/sdy_round_trip:pipelines", "@com_google_absl//absl/algorithm:container", diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc index 1b9f62cacf1dc6..9354c3db106e82 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc @@ -63,6 +63,7 @@ limitations under the License. #include "xla/hlo/ir/named_sharding.h" #include "xla/hlo/ir/tile_assignment.h" #include "xla/hlo/translate/mhlo_to_hlo/attribute_exporter.h" +#include "xla/mlir_hlo/stablehlo_ext/transforms/passes.h" #include "xla/service/spmd/shardy/constants.h" #include "xla/service/spmd/shardy/sdy_round_trip/pipelines.h" #include "xla/shape.h" @@ -744,6 +745,8 @@ void addStablehloImportPipeline(mlir::OpPassManager& pm, ArrayRef allowPropagationToArgs, ArrayRef allowPropagationToResults, bool enableHloShardingV3) { + pm.addNestedPass( + mlir::stablehlo_ext::createStablehloCanonicalizeFromHloImportPass()); pm.addPass(createImportShardingsPass(allowPropagationToArgs, allowPropagationToResults)); addSdyRoundTripImportPipeline(pm, /*enableConstantImport=*/true, diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_import_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_import_pipeline.mlir index 25cfdea88d7b7e..9bae1067db8e19 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_import_pipeline.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_import_pipeline.mlir @@ -125,3 +125,14 @@ func.func @import_sharding_group_with_unused_result(%arg0: tensor<8x8xf32>) -> t %0 = stablehlo.custom_call @xla.sdy.ShardingGroup(%arg0) {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding_group_id = "21 : i64"}} : (tensor<8x8xf32>) -> tuple<> return %arg0 : tensor<8x8xf32> } + +// ----- + +// CHECK-LABEL: func @custom_call_tuple_result_sharding +func.func @custom_call_tuple_result_sharding(%arg0: tensor<8x8xf32>) -> tuple, tensor<8x8xf32>> { + // CHECK-NEXT: %[[CUSTOM_CALL:.*]]:2 = stablehlo.custom_call @foo(%arg0) + // CHECK-SAME: {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"_axis_0"}, {}]>, <@mesh, [{}, {"_axis_0"}]>]>} + // CHECK-SAME: : (tensor<8x8xf32>) -> (tensor<8x8xf32>, tensor<8x8xf32>) + %0 = stablehlo.custom_call @foo(%arg0) {mhlo.sharding = "{{devices=[2,1]<=[2]},{devices=[1,2]<=[1,2]T(1,0)}}"} : (tensor<8x8xf32>) -> tuple, tensor<8x8xf32>> + return %0 : tuple, tensor<8x8xf32>> +} From 6d3956e2d2bdfa6d5eef6e511fa6d3354450b828 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Wed, 2 Sep 2026 16:49:44 -0700 Subject: [PATCH 13/13] [IFRT Proxy] Migrate RemapPlan callers to directly construct input_devices_for_output_map. Migrate RemapPlan callers in IFRT proxy client to directly construct `input_devices_for_output_map` without deprecated `mappings`. In `Array::RemapArrays`, use `plan.output_specs()[i].layout` directly instead of deriving output layouts from `plan.mappings()`. In `array_test.cc`, construct `RemapPlan` with `input_devices_for_output_map` instead of `mappings`. PiperOrigin-RevId: 975401333 --- .../xla/xla/python/ifrt_proxy/client/BUILD | 1 + .../xla/xla/python/ifrt_proxy/client/array.cc | 14 +------------ .../python/ifrt_proxy/client/array_test.cc | 20 +++++++++++-------- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/third_party/xla/xla/python/ifrt_proxy/client/BUILD b/third_party/xla/xla/python/ifrt_proxy/client/BUILD index 948d73ccc43502..3c78bf8293ab30 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/BUILD +++ b/third_party/xla/xla/python/ifrt_proxy/client/BUILD @@ -328,6 +328,7 @@ ifrt_proxy_cc_test( "//xla/tsl/concurrency:ref_count", "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/python/ifrt_proxy/client/array.cc b/third_party/xla/xla/python/ifrt_proxy/client/array.cc index 104bc4f15ecf3f..eceae5358f4caa 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/array.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/array.cc @@ -541,18 +541,6 @@ absl::StatusOr> Array::RemapArrays( req->add_array_handles(handle.handle); } - std::vector> output_layouts( - plan.output_specs().size()); - for (const auto& mapping : plan.mappings()) { - if (output_layouts[mapping.out_array] == nullptr) { - const xla::ifrt::ArrayRef& rcref = arrays[mapping.in_array]; - Array* array = cast(rcref.get()); - ABSL_ASSIGN_OR_RETURN(std::shared_ptr layout, - array->pjrt_layout()); - output_layouts[mapping.out_array] = std::move(layout); - } - } - std::vector result; result.reserve(plan.output_specs().size()); for (int i = 0; i < plan.output_specs().size(); ++i) { @@ -561,7 +549,7 @@ absl::StatusOr> Array::RemapArrays( result.push_back(xla::ifrt::ArrayRef(tsl::MakeRef( client, rpc_helper, plan.output_specs()[i].dtype, plan.output_specs()[i].shape, plan.output_specs()[i].sharding, - ArrayHandle{h}, std::move(output_layouts[i])))); + ArrayHandle{h}, plan.output_specs()[i].layout))); } rpc_helper->RemapArrays(std::move(req)); return result; diff --git a/third_party/xla/xla/python/ifrt_proxy/client/array_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/array_test.cc index 7377ffae0c005e..efa6150f69d038 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/array_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/array_test.cc @@ -20,6 +20,7 @@ #include #include +#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" @@ -327,21 +328,24 @@ TEST_F(ArrayTest, RemapArraysSuccess) { std::vector> arrays; arrays.push_back(array_1); arrays.push_back(array_2); - std::vector mappings; - mappings.push_back({/*in_array=*/0, /*out_array=*/1}); - mappings.push_back({/*in_array=*/1, /*out_array=*/0}); + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map[0].push_back( + {/*in_array=*/1, sharding_->devices()}); + input_devices_for_output_map[1].push_back( + {/*in_array=*/0, sharding_->devices()}); std::vector input_specs; input_specs.push_back(xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), Shape({}), sharding_, kLayout1}); input_specs.push_back(xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), Shape({}), sharding_, kLayout2}); std::vector output_specs; - output_specs.push_back( - xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), Shape({}), sharding_}); - output_specs.push_back( - xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), Shape({}), sharding_}); + output_specs.push_back(xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), + Shape({}), sharding_, kLayout2}); + output_specs.push_back(xla::ifrt::ArraySpec{DType(DType::Kind::kBF16), + Shape({}), sharding_, kLayout1}); RemapPlan plan(std::move(input_specs), std::move(output_specs), - std::move(mappings)); + std::move(input_devices_for_output_map)); absl::StatusOr>> result = Array::RemapArrays(mock_client_.get(), rpc_helper_, plan,