JIT: isolate JIT runtime symbols via explicit symbol mapping

[LeeLee26]

This PR refactors how the Codon JIT backend loads and resolves symbols from the runtime library (libcodonrt). It addresses symbol leakage by replacing global dynamic library loading with local isolation and explicit registration.

This PR is implemented based on #802, which can further mitigate symbol pollution during JIT compilation.

from __future__ import annotations

import ctypes
import os
import sys
from time import sleep
import traceback
import time


SYMBOLS = [
    "GC_malloc",
    "GC_init",
    "GC_get_version",
    "seq_alloc",
    "_ZTIN4llvm30DXILResourceBindingWrapperPassE", # libcodonc.so
    "openblas_read_env", # openblas symbol in libcodonrt.so
    "_Unwind_Resume", # symbol in libgcc_s.so.1
]


def section(title: str) -> None:
    print()
    print(f"== {title} ==")


def dlsym_default(symbol: str) -> int:
    libdl = ctypes.CDLL("libdl.so.2")
    libdl.dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
    libdl.dlsym.restype = ctypes.c_void_p
    return int(libdl.dlsym(ctypes.c_void_p(0), symbol.encode()) or 0)


def dump_visibility(stage: str) -> None:
    section(stage)
    print(f"dlopenflags: {sys.getdlopenflags()}")
    print(f"RTLD_GLOBAL bit: {bool(sys.getdlopenflags() & os.RTLD_GLOBAL)}")
    for symbol in SYMBOLS:
        addr = dlsym_default(symbol)
        print(f"{symbol}: {'VISIBLE' if addr else 'hidden'}" + (f" @ 0x{addr:x}" if addr else ""))


def main() -> int:
    dump_visibility("before import")

    section("import codon")
    try:
        import codon  # noqa: F401
        from codon.codon_jit import codon_library

        print("import codon: OK")
        print(f"codon_library(): {codon_library()!r}")
    except Exception:
        print("import codon: FAILED")
        traceback.print_exc()
        return 1

    dump_visibility("after import")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Apply this PR, the result for the case are as follows

== before import ==
dlopenflags: 2
RTLD_GLOBAL bit: False
GC_malloc: hidden
GC_init: hidden
GC_get_version: hidden
seq_alloc: hidden
_ZTIN4llvm30DXILResourceBindingWrapperPassE: hidden
openblas_read_env: hidden
_Unwind_Resume: hidden

== import codon ==
import codon: OK
codon_library(): '/xxx/.../lib/codon/libcodonc.so'

== after import ==
dlopenflags: 2
RTLD_GLOBAL bit: False
GC_malloc: hidden
GC_init: hidden
GC_get_version: hidden
seq_alloc: hidden
_ZTIN4llvm30DXILResourceBindingWrapperPassE: hidden
openblas_read_env: hidden
_Unwind_Resume: VISIBLE @ 0x7f2fa3479260

[arshajii]

Thanks again for the PR! Just to clarify, is #802 obsolete now with the introduction of this PR?

[BI71317]

I really appreciate for your follow-up PR. @LeeLee26

As I said #802, I am probably not familiar enough with this part of the JIT/runtime implementation to give a deep code review, so I tried to validate the behavior from the user side.

MRE Result in 802

== after import ==
dlopenflags: 2
RTLD_GLOBAL bit: False
GC_malloc: hidden
GC_init: hidden
GC_get_version: hidden
seq_alloc: hidden

== first jit call ==
jit result: 42

== after first jit ==
dlopenflags: 2
RTLD_GLOBAL bit: False
GC_malloc: hidden
GC_init: hidden
GC_get_version: hidden
seq_alloc: hidden

As you said, yes, no longer symbols in Runtime Library are visible through dlsym.

I also ran additional JIT probe that repeatedly allocates and resizes native codon types (list, dict, set ..),

to exercise whether they work well without runtime symbols visible.

MRE

from __future__ import annotations

import os
import sys
from typing import Tuple

import codon


@codon.jit
def mix(x: int) -> int:
    return ((x * 1103515245) + 12345) % 2147483647


def py_gc_wave(rounds: int, width: int) -> Tuple[int, int, int]:
    checksum = 0
    max_set_size = 0
    last_join_len = 0

    for step in range(rounds):
        # Many integer elements in a fresh list.
        values = [((step * width + i) * 17) % 1000 for i in range(width)]

        # Another list built from the first one.
        mixed = [mix(v) % 1000 for v in values]

        # Dict with string keys.
        table = {str(i): value for i, value in enumerate(mixed)}

        # Set derived from the list.
        bucket_ids = {value % 29 for value in mixed}

        # Tuples that include ints, strings and bools.
        triples = [(value, str(value), value % 7 == 0) for value in mixed]

        # String creation / joining.
        joined = "|".join(table.keys())

        checksum += sum(mixed) + len(table) + len(triples) + len(joined)
        if len(bucket_ids) > max_set_size:
            max_set_size = len(bucket_ids)
        last_join_len = len(joined)

    return (checksum, max_set_size, last_join_len)


@codon.jit
def gc_wave(rounds: int, width: int) -> Tuple[int, int, int]:
    checksum = 0
    max_set_size = 0
    last_join_len = 0

    for step in range(rounds):
        # Fresh list allocation.
        values = [((step * width + i) * 17) % 1000 for i in range(width)]

        # Another fresh list allocation.
        mixed = [mix(v) % 1000 for v in values]

        # Dict allocation with string keys.
        table = {str(i): value for i, value in enumerate(mixed)}

        # Set allocation.
        bucket_ids = {value % 29 for value in mixed}

        # List of tuples, with more string creation.
        triples = [(value, str(value), value % 7 == 0) for value in mixed]

        # Joined string from dict keys.
        joined = "|".join(table.keys())

        checksum += sum(mixed) + len(table) + len(triples) + len(joined)
        if len(bucket_ids) > max_set_size:
            max_set_size = len(bucket_ids)
        last_join_len = len(joined)

    return (checksum, max_set_size, last_join_len)


def main() -> int:
    print("== environment ==")
    print(f"python: {sys.executable}")
    print(f"python_version: {sys.version.split()[0]}")
    print(f"codon_module: {getattr(codon, '__file__', None)}")
    print(f"CODON_PATH: {os.environ.get('CODON_PATH')!r}")
    print()

    rounds = 40
    width = 64

    expected = py_gc_wave(rounds, width)
    got = gc_wave(rounds, width)

    print("== single run ==")
    print(f"python baseline: {expected!r}")
    print(f"codon jit      : {got!r}")
    print(f"match          : {got == expected}")
    print()

    print("== repeated calls ==")
    for i in range(3):
        result = gc_wave(rounds + i, width + i)
        print(f"run {i}: {result!r}")
    print()

    if got != expected:
        print("probe: FAIL")
        return 1

    print("probe: SUCCESS")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Result

...
== single run ==
python baseline: (1289052, 29, 181)
codon jit      : (1289052, 29, 181)
match          : True

== repeated calls ==
run 0: (1289052, 29, 181)
run 1: (1334274, 29, 184)
run 2: (1392988, 29, 187)

probe: SUCCESS

Seems works well too.

So, from the limited runtime checks I performed, this appears to fix the symbol leakage I was observing while still allowing JIT-compiled Codon functions that allocate runtime-managed objects to run correctly.

But, I have not deeply reviewed the implementation or validated all possible JIT/runtime cases, so please treat this as a limited behavioral confirmation rather than a full code review. @arshajii @LeeLee26 @inumanag

[LeeLee26]

Hi @arshajii,

is #802 obsolete now with the introduction of this PR?

I recommend merging PR #802 first, followed by this PR. Since #812 introduces more aggressive changes compared to #802 , I’m not yet certain about potential hidden risks. You can refer to the 19184f3 for details.

[LeeLee26]

Hi @BI71317,

Thank you so much for this thorough user‑side validation and extensive testing! It’s really helpful to confirm that PR #812 properly resolves the symbol leakage issue while keeping runtime‑managed objects working correctly under repeated JIT execution.

I really appreciate your careful probe covering list, dict, set allocations and GC‑related workloads. Your behavioral verification gives us great confidence in this change.

And I totally understand that this is a high‑level functional check rather than a full code review. Thanks again for your time and effort!

[inumanag]

Hi @LeeLee26

Thank you for your work here!

Can you please explain why exactly is RTLD_GLOBAL an issue? Aside from having a few extra symbols, I see no other downsides.

Before that, we used to have all sorts of JIT and symbol issues on different platforms; I am not sure if this PR fixes any of that. This PR also hardcodes lots of stuff (and looks like is not cross-platform as it relies on ".so" extension, fixed versions of shared libs, and so on), and also introduces a significant friction in JIT updates. Furthermore, future-proofing and being independent of LLVM ORC updates is also a major concern. Basically, I am afraid that merging this PR will introduce way more issues that it will solve.

[LeeLee26]

Hi @inumanag,

Thank you so much for your detailed feedback and valid concerns! I’d like to start by explaining why the RTLD_GLOBAL issue was a key problem we set out to resolve.

I first tackled this problem in PR #802 . Using RTLD_GLOBAL directly within Python loads numerous unneeded dependencies and symbols into the global scope, which may lead to symbol conflicts in production. For instance, symbols such as _ZTIN4llvm30DXILResourceBindingWrapperPassE from libcodonc.so are exposed globally despite never being used during LLVM JIT execution. Our main goal is to limit global symbol exposure and mitigate the risk of unexpected symbol collisions.

This PR builds upon the work and discussions from PR #802, introducing a full symbol isolation approach. I completely agree with your observations: this solution depends heavily on hardcoded settings, does not work well across platforms, and will add maintenance burden when LLVM ORC is updated down the line.

Taking all these factors into account, I share your concerns. I would suggest merging PR #802 first and keeping this PR as a backup option. Should we run into JIT symbol issues on Unix platforms later on, we can come back to review and assess this implementation again.

Please also know I’m happy to assist with investigation and root cause analysis if any Unix-specific JIT symbol problems appear in the future.