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: 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/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 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 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/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 = ["//..."], 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, 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, 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; 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/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/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>> +} 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> -} 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 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/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. 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 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"