From 84966691e81f4fa74497894dece221c34a442e0a Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Mon, 12 Oct 2026 05:03:24 -0500 Subject: [PATCH] Linkable runtime: compile a kernel again in chunks when it would not fit XLA allocates the values an executable keeps as one temporary, and a raised kernel's batched parallel loops keep a value per iteration of many of them: past some problem size the executable cannot run (MFEM ceed BP3 order 4 at side 240: 35.79 GiB; BP4: 47.29 GiB). After compiling, the runtime now reads the executable's temporary bytes; when they exceed what the allocator has left (90%; else half the device; or REACTANT_EXEC_TEMP_BUDGET), it compiles the module again from the start asking the pattern list for parallel_while_to_batched_scatter_chunked(N) -- the loops batched N chunks at a time -- and without greedy_while_loop_batch_fission, which lifts ops out of a loop over all its iterations at once. N grows with the overshoot, again while the temporary still shrinks. Modules that fit are compiled exactly as before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD --- deps/ReactantExtra/API.cpp | 135 +++++++++++++++++++++++++++++++++++-- deps/ReactantExtra/BUILD | 3 + 2 files changed, 134 insertions(+), 4 deletions(-) diff --git a/deps/ReactantExtra/API.cpp b/deps/ReactantExtra/API.cpp index fdb7fc49ad..c0378b1e7a 100644 --- a/deps/ReactantExtra/API.cpp +++ b/deps/ReactantExtra/API.cpp @@ -76,6 +76,10 @@ void reactantReleaseAddressRange(void *base, size_t nbytes); #include "stablehlo/dialect/StablehloOps.h" #include "stablehlo/transforms/Passes.h" #include "stablehlo/transforms/optimization/Passes.h" +#if defined(REACTANT_CUDA) || defined(REACTANT_ROCM) +#include "xla/service/gpu/gpu_executable.h" +#endif +#include "xla/stream_executor/integrations/tf_allocator_adapter.h" #include "absl/log/globals.h" #include "absl/log/initialize.h" @@ -4095,7 +4099,8 @@ static void rawDeviceToHost(LinkableRuntime *__restrict__ lrt, void *dst, // The exec path's optimization of a kernel's module, before XLA: the // transform list (or REACTANT_EXEC_OPT=hlo-opt, the plain pass). -static void addExecOptimizationPasses(mlir::PassManager &pm) { +static void addExecOptimizationPasses(mlir::PassManager &pm, + int64_t chunks = 1) { // REACTANT_EXEC_OPT=hlo-opt runs the plain enzyme-hlo-opt pass instead. // REACTANT_EXEC_UNROLL sets the trip count up to which while loops are // unrolled. Unrolling the raised kernels' short loops (threshold 16) @@ -4143,6 +4148,19 @@ static void addExecOptimizationPasses(mlir::PassManager &pm) { {"convert_mul_convert;", "associative_binary_op_reordering<1>;"}) for (size_t at; (at = patterns.find(drop)) != std::string::npos;) patterns.erase(at, strlen(drop)); + // Asked for chunks (see reactantXLAExec): a tagged parallel loop is + // batched a chunk at a time, and fission, which lifts ops out of a loop + // over all its iterations at once, is left out. + if (chunks > 1) { + const std::string whole = "parallel_while_to_batched_scatter;"; + if (size_t at = patterns.find(whole); at != std::string::npos) + patterns.replace(at, whole.size(), + "parallel_while_to_batched_scatter_chunked(" + + std::to_string(chunks) + ");"); + const std::string fission = "greedy_while_loop_batch_fission;"; + if (size_t at = patterns.find(fission); at != std::string::npos) + patterns.erase(at, fission.size()); + } // The main list may introduce enzymexla ops (rotate, wrap, extend) // that XLA does not take; the Julia compiler lowers them before export // with this second list. @@ -4250,9 +4268,16 @@ static std::vector> returnLoopBounds(func::FuncOp funcOp) { // The executable `exec` specialized to the bound values `values`: the // bounds (the module's last results) become constants, and are returned // no more. +static int64_t tempBytes(xla::LocalExecutable *exec); +static int64_t tempBudget(LinkableRuntime *lrt); +// Asking for more chunks than this, the attempt with the smallest temporary +// is taken. +static constexpr int64_t kMaxChunkRequest = 1024; + static std::pair specializeRaw(LinkableRuntime *__restrict__ lrt, - LinkableRuntime::CachedExec &exec, ArrayRef values) { + LinkableRuntime::CachedExec &exec, ArrayRef values, + int64_t chunks = 1) { MLIRContext context(lrt->registry); RegisterDialects(wrap(&context)); mlir::ParserConfig config(&context, false); @@ -4281,10 +4306,25 @@ specializeRaw(LinkableRuntime *__restrict__ lrt, // Optimized again with the bounds known: the loops they bound unroll // (REACTANT_EXEC_UNROLL), and what that exposes simplifies. mlir::PassManager pm(&context); - addExecOptimizationPasses(pm); + addExecOptimizationPasses(pm, chunks); if (!mlir::succeeded(pm.run(*module))) ReactantThrowError("specialization: the passes failed"); auto [lexec, shapes, tupleResult] = compileRaw(lrt, *module); + // Too large to run (see reactantXLAExec): again with the parallel loops in + // more chunks, up to kMaxChunkRequest, keeping the attempt that takes the + // least. + if (int64_t temp = tempBytes(lexec), budget = tempBudget(lrt); + temp > budget) { + int64_t more = chunks * std::max(2, (temp + budget - 1) / budget); + if (more <= kMaxChunkRequest) { + auto other = specializeRaw(lrt, exec, values, more); + if (tempBytes(other.first) < temp) { + delete lexec; + return other; + } + delete other.first; + } + } return {lexec, tupleResult}; } @@ -4813,6 +4853,54 @@ REACTANT_ABI void reactantXLAFree(LinkableRuntime **__restrict__ lrtP, lrt->allocationBase.erase(id); } +// The bytes of the temporary allocation an executable makes when it runs +// (every value XLA keeps that is neither a parameter nor a result). +static int64_t tempBytes(xla::LocalExecutable *exec) { +#if defined(REACTANT_CUDA) || defined(REACTANT_ROCM) + auto *gpu = dynamic_cast(exec->executable()); + if (!gpu) + return 0; + int64_t bytes = 0; + for (const xla::BufferAllocation *a : gpu->GetAllocations()) + if (a->IsPreallocatedTempBuffer()) + bytes += a->size(); + return bytes; +#else + return 0; +#endif +} + +// Executables whose temporary allocation would not fit (in what the +// allocator has left, else half the device's memory; +// REACTANT_EXEC_TEMP_BUDGET bytes to override) are compiled again with the +// kernel's parallel loops batched in chunks: per module and arguments, the +// chunks asked for and the temporary bytes it last took. +struct ChunkRequest { + int64_t chunks = 1; + // the fewest temporary bytes an attempt took, and its chunks + int64_t bestTemp = INT64_MAX, bestChunks = 1; + bool settled = false; +}; +static std::mutex chunkRequestsLock; +static std::map>>, + ChunkRequest> + chunkRequests; + +static int64_t tempBudget(LinkableRuntime *lrt) { + if (const char *env = getenv("REACTANT_EXEC_TEMP_BUDGET")) + return atoll(env); + // What the allocator has left, the arguments already in it. + if (auto *multi = dynamic_cast(lrt->allocator)) + if (auto a = multi->GetAllocator(lrt->ordinal); a.ok()) + if (auto stats = (*a)->GetStats(); stats && stats->bytes_limit) + return (*stats->bytes_limit - stats->bytes_in_use) / 10 * 9; + int64_t free = 0, total = 0; + if (!lrt->local->backend().default_stream_executor()->DeviceMemoryUsage( + &free, &total)) + return INT64_MAX; + return total / 2; +} + REACTANT_ABI void reactantXLAExec(LinkableRuntime **__restrict__ lrtP, const char *modstr, int64_t argcnt, void **args, int64_t constcnt, @@ -4969,7 +5057,14 @@ REACTANT_ABI void reactantXLAExec(LinkableRuntime **__restrict__ lrtP, stablehlo::createStablehloCanonicalizeDynamismPass()); // The exec-time optimizer runs the pattern list the Julia compiler runs // on a traced program (the transform-dialect list with its defaults); - addExecOptimizationPasses(pm); + int64_t chunks = 1; + { + std::lock_guard guard(chunkRequestsLock); + if (auto found = chunkRequests.find({modstr, sizeKey}); + found != chunkRequests.end()) + chunks = found->second.chunks; + } + addExecOptimizationPasses(pm, chunks); if (getenv("REACTANT_EXEC_DUMP")) { llvm::errs() << "EXEC_DUMP before\n"; for (int64_t i = 0; i < argcnt; i++) @@ -5131,6 +5226,38 @@ REACTANT_ABI void reactantXLAExec(LinkableRuntime **__restrict__ lrtP, module->print(os); } std::tie(lexec, paramShapes, tupleResult) = compileRaw(lrt, *module); + // Too large to run: compiled again from the start with the parallel + // loops in chunks (more chunks may also batch loops a whole batch could + // not, so a larger request can take more), up to kMaxChunkRequest; + // past it, the attempt that took the least. + if (int64_t temp = tempBytes(lexec), budget = tempBudget(lrt); + temp > budget) { + bool again = false; + { + std::lock_guard guard(chunkRequestsLock); + auto &request = chunkRequests[{modstr, sizeKey}]; + if (temp < request.bestTemp) { + request.bestTemp = temp; + request.bestChunks = chunks; + } + if (!request.settled) { + int64_t more = + chunks * std::max(2, (temp + budget - 1) / budget); + if (more <= kMaxChunkRequest) { + request.chunks = more; + } else { + request.chunks = request.bestChunks; + request.settled = true; + } + again = request.chunks != chunks; + } + } + if (again) { + delete lexec; + free(written); + return reactantXLAExec(lrtP, modstr, argcnt, args, constcnt, consts); + } + } } else { exec = ClientCompileWithProto(lrt->client, wrap(module.get()), nullptr, 0); diff --git a/deps/ReactantExtra/BUILD b/deps/ReactantExtra/BUILD index 05b1651a8c..75f2ddfc62 100644 --- a/deps/ReactantExtra/BUILD +++ b/deps/ReactantExtra/BUILD @@ -1427,6 +1427,7 @@ cc_library( "@xla//xla/tsl/protobuf:histogram_proto_cc", "@xla//xla/tsl/protobuf:histogram_proto_cc_impl", "@xla//xla:xla_data_proto_cc_impl", + "@xla//xla/stream_executor/integrations:tf_allocator_adapter", "@xla//xla/tsl/platform:env", "@xla//xla/tsl/platform:errors", "@xla//xla/service:hlo_proto_cc_impl", @@ -1443,6 +1444,7 @@ cc_library( "@xla//xla/backends/profiler/gpu:device_tracer", "@xla//xla/pjrt/c:pjrt_c_api_gpu_internal", "@xla//xla/service/gpu:gpu_transfer_manager", + "@xla//xla/service/gpu:gpu_executable", "@xla//xla/service/gpu:nvptx_compiler", "@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl", "@xla//xla/service/gpu/model:hlo_op_profiles", @@ -1456,6 +1458,7 @@ cc_library( "@xla//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client", "@xla//xla/stream_executor:rocm_platform", "@xla//xla/service/gpu:amdgpu_compiler", + "@xla//xla/service/gpu:gpu_executable", "@xla//xla/backends/profiler/gpu:device_tracer", ]) + select({ # gloo tcp transport only builds on linux