From 41416365f0250e3076aad8511a730f8a3e3c9990 Mon Sep 17 00:00:00 2001 From: James Le Houx Date: Thu, 28 May 2026 15:50:18 +0000 Subject: [PATCH] session: disable AMReX arena preallocation to fit T4 with EB MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AMReX's default GPU arena preallocates a large chunk of VRAM upfront (~all available memory minus headroom). With EB enabled, the metadata allocations on top of this push the T4 (15 GB) over the limit, OOM-ing at AMReX init. Pass amrex.the_arena_init_size=0 to use lazy allocation instead — each kernel allocates what it needs, when it needs it. Small runtime overhead but unblocks Colab T4 users. --- python/bindings/module.cpp | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/python/bindings/module.cpp b/python/bindings/module.cpp index 1fa3cd8..8b757a1 100644 --- a/python/bindings/module.cpp +++ b/python/bindings/module.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -49,9 +50,20 @@ void init_config(py::module_& m); // ========================================================================= static void init_amrex() { if (!amrex::Initialized()) { - // Initialise with an empty argument list (no ParmParse input) - int argc = 0; - char** argv = nullptr; + // Initialise with ParmParse defaults that play nicely with EB on + // smaller GPUs (T4 on Colab has ~15 GB and AMReX's default arena + // tries to grab most of it upfront; with EB metadata the + // preallocation can OOM). Use lazy allocation instead. + std::vector argv_storage = { + "openimpala", + "amrex.the_arena_init_size=0", + }; + std::vector argv_ptrs; + for (auto& s : argv_storage) { + argv_ptrs.push_back(s.data()); + } + int argc = static_cast(argv_ptrs.size()); + char** argv = argv_ptrs.data(); amrex::Initialize(argc, argv); } }