[ORCJIT] Add keep_module_alive to ExecutionSession.load_module

Mirror tvm_ffi.load_module's keep_module_alive option on the ORCJIT session
so callers can pin a JIT module in the runtime's process-global registry.
Without pinning, JIT-allocated heap Objects (String, Array, Map, Tuple, ...)
that escape past their owning Module deref a deleter pointer into unmapped
JIT memory and crash the process at teardown.

- session.py: forward keep_module_alive through to ModuleGlobalsAdd; Notes
  document the transitive session-pin and slab-granularity RSS trade-off.
- examples/quick-start: drop std::string-returning concat (it hit the
  escape hazard without opting into the pin); example is now int-only.
- tests/sources/cc/test_containers.cc: new source covering String / Array
  / Map / Tuple returns, one function per container.
- tests/test_basic.py: load() gains keep_module_alive; new parametrized
  tests verify each container return survives del mod when pinned.

Author: cyx-6

addons/tvm_ffi_orcjit/examples/quick-start/add.cc

@@ -40,7 +40,3 @@ int fib_impl(int n) {
   return fib_impl(n - 1) + fib_impl(n - 2);
 }
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fibonacci, fib_impl);
-
-// String concatenation example
-std::string concat_impl(std::string a, std::string b) { return a + b; }
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(concat, concat_impl);

addons/tvm_ffi_orcjit/examples/quick-start/run.py

@@ -98,13 +98,6 @@ def _run_tests(obj_file: str, lang: str) -> None:
     print(f"fibonacci(10) = {result}")
     assert result == 55, f"Expected 55, got {result}"
 
-    if lang == "cpp":
-        # String concatenation only available in C++ variant (uses std::string)
-        print("\n=== Testing concat function ===")
-        result = mod.concat("Hello, ", "World!")
-        print(f"concat('Hello, ', 'World!') = '{result}'")
-        assert result == "Hello, World!", f"Expected 'Hello, World!', got '{result}'"
-
     print("\n" + "=" * 50)
     print("All tests passed successfully!")
     print("=" * 50)

addons/tvm_ffi_orcjit/python/tvm_ffi_orcjit/session.py

@@ -119,6 +119,7 @@ def load_module(
         self,
         objects: str | Path | bytes | bytearray | Sequence[str | Path | bytes | bytearray],
         name: str = "",
+        keep_module_alive: bool = False,
     ) -> Module:
         """Load one or more object files into a fresh module.
 
@@ -132,19 +133,39 @@ def load_module(
             Object files to load, given as paths and/or in-memory object-file
             images. A single item is accepted as shorthand for a one-element list.
         name : str
-            Optional name for the underlying library. If empty, a unique name is
-            generated.
+            Optional name for the underlying library. Auto-generated if empty.
+        keep_module_alive : bool
+            If True, pin the module in the runtime's process-global registry
+            (see Notes). Defaults to False.
 
         Returns
         -------
         Module
             A :class:`tvm_ffi.Module` whose imports and library context are fully
-            wired. Dropping every reference to it frees the JIT memory.
+            wired.
+
+        Notes
+        -----
+        ``keep_module_alive`` mirrors :func:`tvm_ffi.load_module`'s option of
+        the same name. When True, the module is inserted into the runtime's
+        process-global module registry, so its JITDylib — and every function
+        pointer, deleter, and static allocation it owns — stays mapped until
+        the interpreter unloads ``libtvm_ffi``. Use this when Objects produced
+        by the module may outlive the local ``mod`` reference (e.g., a
+        JIT-allocated ``String`` or ``Array`` returned to Python and held past
+        ``del mod``). When False (default), the caller owns the module's
+        lifetime and its JIT memory is reclaimed on drop; callers must ensure
+        no JIT-produced Object outlives the returned module.
+
+        Pinning is transitive: the module holds a strong reference to its
+        :class:`ExecutionSession`, so pinning one module also keeps that
+        session (and its slab-pool memory manager) alive for the process
+        lifetime. Reclaim is slab-granular — a slab shared between a pinned
+        and an unpinned module stays mapped until the pinned module is gone.
 
         Examples
         --------
         >>> session = default_session()
-        >>> mod = session.load_module("dylib0.o")
         >>> mod = session.load_module(["a.o", "b.o", Path("c.o").read_bytes()])
         >>> mod.my_function(1, 2)
 
@@ -162,7 +183,10 @@ def load_module(
                     "load_module objects must be a path (str or Path) or object-file "
                     f"bytes, but got {type(obj).__name__}"
                 )
-        return _ffi_api.SessionLoadModule(self, normalized, name)  # type: ignore
+        mod = _ffi_api.SessionLoadModule(self, normalized, name)  # type: ignore
+        if keep_module_alive:
+            tvm_ffi._ffi_api.ModuleGlobalsAdd(mod)  # type: ignore
+        return mod
 
     def clear_free_slabs(self) -> int:
         """Release drained slabs (no live JIT allocations) back to the OS.

addons/tvm_ffi_orcjit/tests/CMakeLists.txt

@@ -84,6 +84,7 @@ if (NOT WIN32)
   add_test_object(sources/cc/test_link_order_base.cc)
   add_test_object(sources/cc/test_link_order_caller.cc)
   add_test_object(sources/cc/test_error.cc)
+  add_test_object(sources/cc/test_containers.cc)
 endif ()
 
 # Pure C object files — built on all platforms (no C++ runtime deps)

addons/tvm_ffi_orcjit/tests/sources/cc/test_containers.cc

@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+// One JIT-allocated FFI container per test. Every returned Object is freshly
+// constructed inside this translation unit, so its header.deleter points into
+// JIT-mapped memory — the escape hazard that keep_module_alive=True pins the
+// JITDylib against.
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/string.h>
+
+#include <string>
+
+using tvm::ffi::Array;
+using tvm::ffi::Map;
+using tvm::ffi::String;
+using tvm::ffi::Tuple;
+
+// Character replace in a String.
+String test_string_replace_impl(String s, String from, String to) {
+  std::string out(s);
+  char c_from = std::string(from)[0];
+  char c_to = std::string(to)[0];
+  for (char& c : out) {
+    if (c == c_from) c = c_to;
+  }
+  return String(out);
+}
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_string_replace, test_string_replace_impl);
+
+// Concatenate two Array<int64_t>s.
+Array<int64_t> test_array_concat_impl(Array<int64_t> a, Array<int64_t> b) {
+  Array<int64_t> out;
+  for (int64_t x : a) out.push_back(x);
+  for (int64_t x : b) out.push_back(x);
+  return out;
+}
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_array_concat, test_array_concat_impl);
+
+// Hand-built Map<String, int64_t>.
+Map<String, int64_t> test_map_build_impl() {
+  Map<String, int64_t> out;
+  out.Set(String("a"), 1);
+  out.Set(String("b"), 2);
+  return out;
+}
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_map_build, test_map_build_impl);
+
+// Reverse a four-element Tuple.
+Tuple<String, int64_t, String, int64_t> test_tuple_reverse_impl(
+    Tuple<int64_t, String, int64_t, String> t) {
+  return Tuple<String, int64_t, String, int64_t>(t.get<3>(), t.get<2>(), t.get<1>(), t.get<0>());
+}
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_tuple_reverse, test_tuple_reverse_impl);

addons/tvm_ffi_orcjit/tests/test_basic.py

@@ -63,11 +63,18 @@ def obj(name: str) -> str:
     return str(path)
 
 
-def load(*obj_names: str, session: ExecutionSession | None = None, name: str = "") -> Module:
+def load(
+    *obj_names: str,
+    session: ExecutionSession | None = None,
+    name: str = "",
+    keep_module_alive: bool = False,
+) -> Module:
     """Load one or more object files into a fresh module on ``session``."""
     if session is None:
         session = ExecutionSession()
-    return session.load_module([obj(o) for o in obj_names], name)
+    return session.load_module(
+        [obj(o) for o in obj_names], name, keep_module_alive=keep_module_alive
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -111,6 +118,16 @@ def ctor_dtor_obj(self) -> str:
         """Return path prefix for test_ctor_dtor object."""
         return f"{self.subdir}/test_ctor_dtor"
 
+    def containers_obj(self) -> str:
+        """Return path prefix for test_containers object.
+
+        C++-only source; C variants return an empty prefix so ``obj()`` sees
+        a missing path and skips the parametrized case for that variant.
+        """
+        if not self.subdir.startswith("cc"):
+            return ""
+        return f"{self.subdir}/test_containers"
+
     def fn(self, base_name: str) -> str:
         """Return the function name for a given base name."""
         return base_name
@@ -558,6 +575,60 @@ def test_drop_all_modules_then_session(v: Variant) -> None:
     # surface as a crash under pytest.
 
 
+# ---------------------------------------------------------------------------
+# FFI-container returns under keep_module_alive.
+#
+# JIT-emitted code allocates each returned Object through inline templates
+# (make_object / make_inplace_array_object), so the object's header.deleter
+# points into JIT-mapped memory. If the owning Module were torn down before
+# the escaped Object, the deleter would dangle. keep_module_alive=True pins
+# the module in the runtime's process-global registry, keeping the JITDylib
+# mapped for the process lifetime. Each test loads with the flag, drops the
+# local mod reference, forces GC, and then reads the returned container —
+# proving the deleter is still reachable after the local mod is gone.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_containers_string(v: Variant) -> None:
+    """String return survives del mod when pinned."""
+    mod = load(v.containers_obj(), keep_module_alive=True)
+    result = mod.test_string_replace("hello world", "l", "L")
+    del mod
+    gc.collect()
+    assert str(result) == "heLLo worLd"
+
+
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_containers_array(v: Variant) -> None:
+    """Array<int64_t> return survives del mod when pinned."""
+    mod = load(v.containers_obj(), keep_module_alive=True)
+    result = mod.test_array_concat([1, 2, 3], [4, 5, 6])
+    del mod
+    gc.collect()
+    assert list(result) == [1, 2, 3, 4, 5, 6]
+
+
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_containers_map(v: Variant) -> None:
+    """Map<String, int64_t> return survives del mod when pinned."""
+    mod = load(v.containers_obj(), keep_module_alive=True)
+    result = mod.test_map_build()
+    del mod
+    gc.collect()
+    assert dict(result) == {"a": 1, "b": 2}
+
+
+@pytest.mark.parametrize("v", _all_variants, ids=_variant_id)
+def test_containers_tuple(v: Variant) -> None:
+    """Tuple<...> return survives del mod when pinned."""
+    mod = load(v.containers_obj(), keep_module_alive=True)
+    result = mod.test_tuple_reverse((1, "a", 2, "b"))
+    del mod
+    gc.collect()
+    assert tuple(result) == ("b", 2, "a", 1)
+
+
 # ---------------------------------------------------------------------------
 # Slab-pool growth (Stage B).
 #