nki_ctx, Scopes & the sema Handle

All addresses, line numbers, and interned identifiers on this page apply to neuronx_cc 2.24.5133.0+58f8de22, cp310 wheel (nki_ctx.cpython-310-x86_64-linux-gnu.so). cp311/cp312 ship byte-identical Cython logic at shifted addresses.

Abstract

Every NKI language op (nl.add, nl.load, nl.program_id, nl.affine_range, …) is a thin Python surface that, before it can emit a single IR node, must answer two questions: which IR-builder am I emitting into? and which scope region do my instructions land in? The whole answer is one tiny Cython module — neuronxcc/nki/compiler/backends/neuron/nki_ctx.py — compiled to a 197-line .so with exactly three functions and no instance state of its own. It is the ambient-accessor layer: it reads a process-global cell, validates that a kernel trace is actually in progress, and hands back the active codegen context.

The cell is TraceContext.global_ctx. The object it holds is the active codegen context — the NeuronCodegen instance from KernelBuilder.so that owns .builder (the pelican/BIR IRBuilder C++ binding) and the .cur_scope chain. nki_ctx() returns that object or raises; get_cur_scope() returns nki_ctx().cur_scope; nki_ctx_or_none() returns the raw cell unchecked. The scope chain itself is a stack of @contextmanager region emitters on NeuronCodegen — kernel_scope, loop_scope, if_scope, while_scope, allocation_scope, stage_scope — each opening a Penguin ScopeRegion of a tagged kind. The guard that fires when no trace is active, err_nki_api_outside_of_nki_kernel, does not live here at all; it is a function in the sibling sema module, imported by name.

This page documents the fetch chain bottom-up: the three accessors and their exact decompiled control flow; the is-None guard and the error path; the two distinct scope vocabularies (TraceContext's abstract-interp scope classes vs NeuronCodegen's IR region kinds) and how cur_scope bridges them; and the ScopeRegion-kind set every traced op selects into.

For reimplementation, the contract is:

• The singleton cell: TraceContext.global_ctx, set once per kernel trace, cleared on teardown; the only mutable state the accessors read.
• The three accessors and their exact guard semantics: nki_ctx() (None-check + raise), nki_ctx_or_none() (raw read), get_cur_scope() (truthy-check + .cur_scope).
• The outside-kernel guard: sema.err_nki_api_outside_of_nki_kernel(), fired on global_ctx is None.
• The ScopeRegion-kind set an op's cur_scope can be — KernelScope / FunctionScope / LoopScope / IfScope / WhileScope / AllocationScope (+ ScopeRegion generic) — and the @contextmanager emitters that push them.

The Singleton Cell — TraceContext.global_ctx

Purpose

The entire ambient-context mechanism is one class attribute: TraceContext.global_ctx. It is None between kernel traces and holds the active codegen context during a trace. TraceContext.new_ctx(opts) (TraceContext.so, classmethod) constructs the context and installs it; trace teardown clears it back to None. The accessors in nki_ctx.py never own a context — they only read this cell.

CORRECTION (W11) — report D-W11 describes the object in global_ctx as a TraceContext instance with a .sema attribute, and writes the guard as if ctx is None or ctx.sema is None. The binary disagrees on both points. (1) TraceContext.so interns cur_scope, global_ctx, new_ctx, nki_func — but not sema and not builder as instance attributes. The object actually fetched by nki_ctx() and called with .affine_range(...)/.if_scope(...)/.cur_scope is the NeuronCodegen codegen context (KernelBuilder.so), which does intern builder, cur_scope, and affine_range. (2) nki_ctx() performs a single is None test on global_ctx, not a .sema check — see §Algorithm below. The .sema referenced in the disassembly is the imported sema module, off which the error helper is fetched — not a context attribute.

Considerations

global_ctx is a single mutable cell, so NKI tracing is single-active-trace per interpreter — there is no per-thread context stack at this layer despite the module docstring's "(thread local) global context" wording (the thread-locality is enforced upstream by TraceContext's main_interpreter_id guard, not by nki_ctx.py). Nested traced functions reuse the same global_ctx; what changes per nested call is cur_scope (a new FunctionScope pushed on the chain), not the cell.

nki_ctx() — the guarded accessor

Purpose

The function nearly every NKI intrinsic calls first. Returns the active codegen context, or raises if no kernel trace is in progress. It is a METH_NOARGS wrapper — it takes no arguments.

CORRECTION (W11) — D-W11 gives the signature as nki_ctx(allow_none=False) with an allow_none keyword selecting raise-vs-return-None. There is no allow_none parameter: the wrapper at 0x7700 is METH_NOARGS (PyObject *__pyx_self, PyObject *unused), and the string/name tables contain no allow_none identifier. The "return None instead of raising" behavior is provided by the separate function nki_ctx_or_none(), not by a keyword on nki_ctx().

Entry Point

nl.<op>(...)                          ── any W01/W02 NKI intrinsic
  └─ nki_ctx()                        ── __pyx_pw_..._nki_ctx_1nki_ctx @0x7700
       ├─ from .TraceContext import TraceContext     (ImportModuleLevelObject + ImportFrom)
       ├─ GetAttr(TraceContext, "global_ctx")        (tp_getattro, py:17)
       └─ if None → sema.err_nki_api_outside_of_nki_kernel()   (py:19)

Algorithm

// __pyx_pw_..._nki_ctx_1nki_ctx @0x7700   src nki_ctx.py:14
PyObject *nki_ctx(void):
    TraceContext = ImportFrom(import ".TraceContext", "TraceContext")   // py:15
    ctx = getattr(TraceContext, "global_ctx")        // py:17  (tp_getattro, n_s global_ctx)
    if ctx != Py_None:
        return ctx                                   // common path: trace in progress
    // ctx is None → not inside a traced kernel:
    err = GetBuiltinName("sema").err_nki_api_outside_of_nki_kernel   // py:19
    return err()                                     // raises; never returns a value

The decompiled body makes the None-branch unambiguous: after getattr(ctx, "global_ctx") it tests if ( Attr != &Py_NoneStruct ) goto LABEL_18 (return ctx). Only when Attr == Py_None does it resolve sema (via _Pyx_GetBuiltinName(__pyx_n_s_sema)), fetch its err_nki_api_outside_of_nki_kerne[l] attribute through the tp_call slot at +144, and call it with the empty tuple. The error helper unconditionally raises, so nki_ctx() has exactly two outcomes: return the live context, or raise "calling NKI API outside of NKI kernels is not supported.".

NOTE — sema is resolved as a builtin/module global name, not as ctx.sema. nki_ctx.py does from neuronxcc.nki.compiler.backends.neuron import sema (and from .TraceContext import TraceContext); the import-from for sema appears in __pyx_pymod_exec_nki_ctx. The error helper is sema.err_nki_api_outside_of_nki_kernel, confirmed present in sema.cpython-310-*.so with the message string. This is why a stale/None global_ctx after trace teardown still produces a clean, attributable error rather than an AttributeError.

Function Map

nki_ctx_or_none() — the soft accessor

Purpose

Returns TraceContext.global_ctx raw — None if no trace is active, the context otherwise. No guard, no error. This is the accessor for code paths that legitimately run both inside and outside a trace (e.g. helpers that branch on "am I tracing right now?").

Algorithm

// __pyx_pw_..._nki_ctx_3nki_ctx_or_none @0x7440   src nki_ctx.py:23
PyObject *nki_ctx_or_none(void):
    TraceContext = ImportFrom(import ".TraceContext", "TraceContext")   // py:24
    return getattr(TraceContext, "global_ctx")       // py:26  — may be None, no check

The decompiled body is a strict prefix of nki_ctx() minus the None-branch: import TraceContext, getattr(..., "global_ctx"), return it. There is no Py_None comparison and no sema resolution — the function has a single return site for the attribute. This is the true allow_none=True equivalent that D-W11 ascribed to a keyword on nki_ctx().

get_cur_scope() — the scope accessor

Purpose

The function the op emitters call to learn which region to emit into. Returns nki_ctx().cur_scope — the active scope on the codegen context's scope chain. Because it routes through nki_ctx() (not nki_ctx_or_none()), calling it outside a kernel raises the same outside-kernel error.

Algorithm

// __pyx_pw_..._nki_ctx_5get_cur_scope @0x6f60   src nki_ctx.py:29
PyObject *get_cur_scope(void):
    f = GetBuiltinName("nki_ctx")                    // py:30  module-global nki_ctx
    ctx = f()                                        // calls nki_ctx() → raises if outside
    if PyObject_IsTrue(ctx):                          // py:31  truthy guard before getattr
        return getattr(ctx, "cur_scope")             // py:31  (n_s cur_scope)
    return ctx                                        // ctx falsy → return it as-is

The body resolves the module global nki_ctx (not a direct internal call), invokes it, then runs PyObject_IsTrue on the result. For a real context this is true and it does getattr(ctx, "cur_scope"); the False/None/True fast-paths short-circuit the truthiness check. In practice nki_ctx() either raises or returns a truthy context object, so the falsy return is a defensive tail.

QUIRK — get_cur_scope() calls nki_ctx by name through the module dict (_Pyx_GetBuiltinName(__pyx_n_s_nki_ctx)), not as a private C call. A reimplementation that inlines the context fetch would diverge if user code monkey-patched the module-level nki_ctx; the binary preserves the indirection, so the public nki_ctx is the single choke-point even for the in-module caller.

The Scope Chain — what cur_scope points at

cur_scope is the active node of a scope stack, but two different objects wear the name "scope" in this subsystem, and the page that conflates them will mis-model the IR. They live in different binaries and serve different layers.

TraceContext scope classes (abstract-interpretation side)

TraceContext.so defines the Python-level scope objects that the abstract interpreter walks the kernel body under. Confirmed interned class names: KernelScope, FunctionScope, StmtScope. These track the Python nesting (which kernel, which traced function, which statement) and carry parent_scope, name_and_loc, stack_frame. cur_scope on the context is set to one of these as the interpreter enters/leaves functions and statements.

NeuronCodegen scope region kinds (IR-builder side)

KernelBuilder.so's NeuronCodegen is the codegen context — the object nki_ctx() actually returns. Its scopes are @contextmanager emitters that open a Penguin ScopeRegion of a tagged kind. Confirmed kind strings in .rodata:

The axis kinds AffineAxis / DynamicAxis are not scope kinds — they tag the LoopAxis node a LoopScope is opened around. The task brief's "loop / if / while / kernel / allocation / stage" kinds correspond to loop_scope/if_scope/while_scope/kernel_scope/allocation_scope/stage_scope here; stage_scope opens a generic ScopeRegion (the pipeline-stage container), with the per-stage marker carried by pipeline_stage(stage_id, execution_order).

NOTE — there is no StmtScope region kind in NeuronCodegen. StmtScope exists only as a TraceContext class (statement-level abstract-interp bookkeeping); it does not produce an IR ScopeRegion. A reimplementer mapping cur_scope → IR region must skip StmtScope nodes — only Kernel/Function/Loop/If/While/Allocation scopes back a Penguin region.

Algorithm — new_scope, the generic region opener

// NeuronCodegen.new_scope  (KernelBuilder.so) — @contextmanager generator
contextmanager new_scope(self, kind, **payload):
    block = self.create_scope_region_block(self.curstmt, parent, kind, attrs)  // builder.makeScopeRegion
    prev = self.cur_scope
    self.cur_scope = block            // push: ops emitted now land in `block`
    try:
        yield block
    finally:
        self.cur_scope = prev         // pop on __exit__

Each named scope (if_scope, loop_scope, allocation_scope, …) is create_*_block (the node factory) + new_scope(kind=…) (the push/pop). cur_scope is the mutated cell: it is what get_cur_scope() reads, so an op emitted inside a with self.if_scope(pred): sees the IfScope block as its cur_scope and the op is inserted in_cur_scope. (Region-construction internals — create_scope_region_block, create_affine_axis_block, create_while_block — are documented in 6.5.5 scope emitters; predicate lowering for if_scope/while_scope in 6.4.4 sema legality engine.)

End-to-End — how one traced op reaches its builder and scope

nl.add(a, b)                                   (W01/W02 intrinsic surface)
 ├─ ctx   = nki_ctx()                           @0x7700  → active NeuronCodegen (or raise)
 ├─ scope = get_cur_scope()  (= nki_ctx().cur_scope)  @0x6f60  → active ScopeRegion
 ├─ builder = ctx.builder                        (pelican/BIR IRBuilder binding)
 └─ builder.<emit>(... , scope=scope)            → appends a Penguin IR node into `scope`

The single mutable cell is TraceContext.global_ctx, installed by TraceContext.new_ctx(opts) at trace start and cleared on finalize_kernel. Between traces it is None, so a stray nl.* call resolves nki_ctx() → global_ctx is None → sema.err_nki_api_outside_of_nki_kernel() → "calling NKI API outside of NKI kernels is not supported.". The cur_scope the op selects into is whatever region is on top of the NeuronCodegen scope stack at emit time — the kernel region by default, narrowed by every active with loop_scope()/if_scope()/allocation_scope() the trace has entered.

Adversarial Self-Verification

The five strongest claims, re-challenged against the binary:

1. Exactly three public functions, no allow_none, no nested generator. nki_ctx.so __pyx_pw_* set = {nki_ctx@0x7700, nki_ctx_or_none@0x7440, get_cur_scope@0x6f60}; zero __pyx_gb_* (no generator); name table has no allow_none. HOLDS.
2. The guard tests global_ctx is None, not ctx.sema. Decompiled nki_ctx does getattr(ctx,"global_ctx") then if (Attr != &Py_NoneStruct) goto return; only the None-branch resolves sema. TraceContext.so interns no sema/builder instance attr. HOLDS — issued as CORRECTION (W11).
3. The error helper lives in sema, not in the context. sema is _Pyx_GetBuiltinName-resolved; err_nki_api_outside_of_nki_kernel is confirmed in sema.cpython-310-*.so with message "calling NKI API outside of NKI kernels is not supported.". HOLDS.
4. nki_ctx() returns the NeuronCodegen codegen context. The iterators trampoline does nki_ctx().affine_range(...); NeuronCodegen (KernelBuilder.so) interns affine_range, cur_scope, builder. So the object held in global_ctx is the codegen context. HOLDS (STRONG — the object is the codegen context; whether global_ctx is literally typed NeuronCodegen vs a thin trace-context wrapper that is the codegen is not byte-pinned, but the .affine_range/.cur_scope/.builder surface is confirmed on NeuronCodegen).
5. ScopeRegion-kind set = Kernel/Function/Loop/If/While/Allocation (+ generic ScopeRegion); StmtScope is not a region. KernelBuilder.so .rodata confirms KernelScope/FunctionScope/LoopScope/IfScope/WhileScope/AllocationScope/ScopeRegion; StmtScope appears only in TraceContext.so. HOLDS.

Nothing in the five required claims failed; the two W11 misreadings (the allow_none keyword and the ctx.sema guard) are corrected in place above.

Cross-References

• TraceContext — 6.1.1, the abstract-interpreter that owns global_ctx, new_ctx, and the scope classes (KernelScope/FunctionScope/StmtScope)
• NeuronCodegen Control / Scope / Predicate Emitters — 6.5.5, NeuronCodegen @contextmanager region emitters and create_*_block node factories
• sema legality engine — 6.4.4, the sema module that hosts err_nki_api_outside_of_nki_kernel and predicate lowering for if_scope/while_scope