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