NkiTypeSystem and Traced-Tile Operator Overloading

All symbols, addresses, and strings on this page apply to neuronx_cc 2.24.5133.0+58f8de22, cp310 wheel. The tracer lives in three Cython extension modules under neuronxcc/nki/compiler/backends/neuron/: NkiTypeSystem.cpython-310-x86_64-linux-gnu.so (the n-ary tracer), NkiTypeSystemCmpOp.…so (comparison synthesis), and NkiTypeSystemLogicalOp.…so (logical synthesis), all Cython 3.0.10 with debug_info — class and method names below are real __pyx_pw_… / __pyx_n_s_… symbols. cp311/cp312 are byte-twins. Addresses are pinned to the cp310 artifact's .so; per-method offsets shift between Cython outputs, so a figure such as 0x10b70 for the base module's generator trampoline belongs to a different build — in the cp310 wheel that trampoline is at 0xbe30. Treat offsets as module-relative landmarks and the symbols and strings as the hard anchors.

Abstract

When a NKI kernel writes lo <= x < hi, or i < N, or (a < b) & (c >= 0) over traced operands, Python never computes a bool. The comparison and logical operators on traced tiles, scalars, and tile-indices are synthesized at class-construction time by three decorator modules and dispatched at trace time by operand type category into one of two outcomes: an affine predicate / mask object (the deferred symbolic structure documented in §6.2.4 Mask / Predicate Algebra), or — when a real tensor is involved — an elementwise tensor op whose result dtype is np.int8. This is the entire reason a NKI mask= argument or affine_select predicate can be built from ordinary Python comparison syntax: the syntax is the same, but the operators are overloaded to produce structure, not truth.

The machinery has three layers. NkiTypeSystem (the "TypeTraceContext" tracer) holds the only hand-written method bodies: it lowers Python chained comparisons (a < b < c) and in / not in into n-ary folds over the per-operator methods, and it provides logical_and_nary / logical_or_nary reductions. NkiTypeSystemCmpOp synthesizes __eq__/__ne__/__lt__/__le__/__gt__/__ge__ onto traced-value classes via two class decorators, choosing a result type from a TypeCategory lattice. NkiTypeSystemLogicalOp synthesizes __and__/__or__/__invert__ (and the logical_and/logical_or/logical_not the n-ary folds consume) over an extended TypeCategory lattice with a promotion map that coerces heterogeneous logical operands (mask, predicate, scalar, bool) to a common logical type before combining.

The single docstring fragment recovered from the base module states the purpose: "This file define the TypeTraceContext class, it defines a tracer that implement NKI's core type system" (.rodata). This page reconstructs the synthesis factories, the two TypeCategory lattices, the promotion lattice, and the exact dispatch each operator performs.

What a reimplementer must reproduce

• The synthesis decorators — three class decorators that, applied at class-body evaluation, inject the six rich-comparison dunders and the three logical dunders (plus their logical_* aliases) onto every traced-value class.
• The two TypeCategory lattices — a 4-member one in CmpOp (OTHER/SCALAR/TILE_INDEX/TENSOR) and a 6-member one in LogicalOp (adds MASK/PREDICATE), each with an isinstance-chain classifier get_type_category.
• The comparison dispatch — given (cat_a, cat_b), pick EQTileMask/EQScalarPredicate (for ==/!=), TileMaskIntersection/ScalarPredicateIntersection (for <,<=,>,>=), or a tensor._binop(…, dtype=np.int8) elementwise op (when a tensor is present).
• The logical dispatch — highest_category, promote_to_logical_type via PROMOTION_MAP, then operator.and_/or_/invert over the promoted predicate/mask objects, with bool/number operands promoted to AlwaysTrue/AlwaysFalsePredicate and tensor operands routed to an elementwise np.int8 op (or rejected, for ~).
• The chained-comparison lowering — compare_nary rewriting a OP1 b OP2 c into (a OP1 b) & (b OP2 c) & … AND-folded by logical_and_nary.

The one claim to internalise. A comparison or logical operation on a traced NKI value is never a Python bool. i < N evaluates to an EQTileMask / TileMaskIntersection object; p & q to a combined predicate/mask object; a tensor compare to an int8-dtype tensor op. The "type system" is precisely the dispatch that guarantees this: every operator is replaced, at class-construction, with a handler that classifies its operands and constructs structure. This page is the construction side; §6.2.4 is the algebra and lowering of the objects it produces.

1. Three modules, one decorated class set

The traced-value classes — scalar, tensor, tile_index, and the predicate/mask classes — are built by the metaclass machinery (§6.3.2) and then decorated with three class decorators that rewrite their operator slots. The decorators are applied at module-exec time of whatever module defines the traced-value class; the base NkiTypeSystem class itself is decorated by all three (its module-exec looks up the three synthesize_* names and applies them as class decorators), which is why NkiTypeSystem also gains logical_and/logical_or that its own logical_and_nary/logical_or_nary consume.

                       class body of a traced-value class
                                     │
            ┌────────────────────────┼────────────────────────┐
            ▼                        ▼                         ▼
 @synthesize_logical_operations   @synthesize_equality_   @synthesize_ordered_
   (NkiTypeSystemLogicalOp)         comparisons              comparisons
                                   (NkiTypeSystemCmpOp)     (NkiTypeSystemCmpOp)
            │                        │                         │
   injects __and__ __or__          injects __eq__            injects __lt__ __le__
   __invert__ + logical_and        __ne__                    __gt__ __ge__
   logical_or logical_not

Each injected dunder is a thin make_binary_method / make_unary_method wrapper around a handler closure produced by one of three factories (make_equality_operation, make_comparison_operation, make_logical_operation / make_logical_not). The handler is where the type-category dispatch lives. Every factory and wrapper symbol appears verbatim in the modules' _strings.json:

NkiTypeSystemCmpOp:      make_equality_operation   make_comparison_operation
                         make_binary_method        make_commutative_binary_method
                         synthesize_equality_comparisons  synthesize_ordered_comparisons
NkiTypeSystemLogicalOp:  make_logical_operation     make_logical_not
                         make_binary_method         make_unary_method
                         synthesize_logical_operations

NOTE — the same factory shape repeats across all three layers. Every operator is setattr(cls, dunder, wrapper(handler)) where handler = factory(op, op_name, numpy_op[, pred_op]). The op is the stdlib operator.<x> (used as the actual reducing operation on already-promoted operands), op_name is the dunder stem used for error messages and the tensor _binop selector, numpy_op is the np.<x> ufunc passed to the tensor path, and (for ordered comparisons) pred_op is the affine-predicate constructor pred_lt/le/gt/ge. A reimplementer can build all nine operators from one parameterized factory family.

2. The TypeCategory lattices and get_type_category

The dispatch key is an IntEnum category assigned to each operand by an isinstance chain. There are two lattices — CmpOp has four members, LogicalOp has six. Both member sets are readable directly: every member name appears as a __pyx_n_s_<MEMBER> string in the respective module, and each module ships a TypeCategory.__str__ symbol (pretty enum printing).

2.1 CmpOp lattice (comparisons)

// NkiTypeSystemCmpOp — TypeCategory(IntEnum), members read from _strings.json
//   {OTHER, SCALAR, TILE_INDEX, TENSOR}      // exactly four — no MASK/PREDICATE here
enum TypeCategory { OTHER, SCALAR, TILE_INDEX, TENSOR };   // integer values INFERRED

// get_type_category  (__pyx_pw_…_18NkiTypeSystemCmpOp_1get_type_category)
TypeCategory get_type_category(PyObject *x) {
    if (isinstance(x, scalar))      return SCALAR;       // neuron `scalar`
    if (isinstance(x, tensor))      return TENSOR;       // neuron `tensor`
    if (isinstance(x, tile_index))  return TILE_INDEX;   // `tile_index`
    return OTHER;                                        // numbers / np scalars / anything else
}

2.2 LogicalOp lattice (logical combine)

The logical module extends the lattice with MASK and PREDICATE so that already-built mask/predicate objects classify distinctly from the index/scalar/tensor inputs. The isinstance chain order is recovered directly from the decompiled body of get_type_category (__pyx_pw_…_22NkiTypeSystemLogicalOp_1get_type_category, file …_1get_type_categor_0x15cb0_…c), which probes the module globals in exactly this sequence: scalar → tensor → tile_index → nki_mask → predicate, falling through to OTHER:

// NkiTypeSystemLogicalOp — TypeCategory(IntEnum), all six members read from the pool
//   `__pyx_k_MASK` and `__pyx_n_s_PREDICATE` are real constants in this module's pool
enum TypeCategory { OTHER, SCALAR, TILE_INDEX, MASK, PREDICATE, TENSOR };  // ints INFERRED

// get_type_category  — decompiled isinstance chain, in this order
TypeCategory get_type_category(PyObject *x) {
    if (isinstance(x, scalar))      return SCALAR;
    if (isinstance(x, tensor))      return TENSOR;
    if (isinstance(x, tile_index))  return TILE_INDEX;
    if (isinstance(x, nki_mask))    return MASK;        // a built tile-mask object
    if (isinstance(x, predicate))   return PREDICATE;   // a built scalar-predicate object
    return OTHER;                                        // bool / number / etc.
}

GOTCHA — the isinstance probe order is not the precedence order. get_type_category returns the first matching class, so its chain order is just classification, not dominance. The dominance order — which category "wins" when two operands differ — is the IntEnum integer ordering consumed by highest_category(a,b) = max(get_type_category(a), get_type_category(b)). With members declared OTHER, SCALAR, TILE_INDEX, MASK, PREDICATE, TENSOR, the auto-assigned IntEnum values (0..5) put TENSOR highest, so any operation touching a real tensor takes the tensor path. The exact integer values are INFERRED from declaration order — IntEnum member values are interned at import and not readable from a fixed struct — but the relative ordering (TENSOR dominant, OTHER least) is corroborated by the dispatch itself: the tensor branch is checked first in every handler, and promote_to_logical_type explicitly refuses TENSOR.

Alongside it sit a get_target_type_category function (…_3get_target_type_c_0xfb50_…) and a PROMOTION_TARGET_MAP: together they map an (catA, catB) pair to the target logical category (MASK vs PREDICATE) that the promotion should produce. The precise contents of PROMOTION_TARGET_MAP are INFERRED — the map is built at import, so only its name and the two target categories are readable from the binary.

TYPE_COMBINATIONS = itertools.product(TypeCategory, repeat=2) — the names product, repeat, and TYPE_COMBINATIONS are present in both modules — is the full Cartesian set of category pairs, used to drive the synthesis loop and/or validate dispatch coverage.

3. Comparison synthesis (NkiTypeSystemCmpOp)

Two factories build the comparison handlers; two class decorators install them.

3.1 make_equality_operation — == and !=

// make_equality_operation(op, op_name, numpy_op) -> handler(self, a, b)
//   op       = operator.eq / operator.ne
//   op_name  = 'eq' / 'ne'                 (also selects tensor _binop op)
//   numpy_op = np.equal / np.not_equal
// closure refs: SCALAR/TENSOR/TILE_INDEX, EQTileMask, EQScalarPredicate,
//               _binop, return_tensor_or_extracted_scalar, "Unexpected type "
PyObject *handler(self, a, b) {
    cat_a = get_type_category(a);
    cat_b = get_type_category(b);

    if (cat_a == TENSOR || cat_b == TENSOR) {            // a real tensor is present
        // elementwise tensor compare -> int8 boolean mask tensor
        result = <tensor>._binop(numpy_op_or_op_name, other, /*dtype=*/np.int8);
        return return_tensor_or_extracted_scalar(result);   // scalar-collapse if 0-D
    }
    else if (cat_a == TILE_INDEX || cat_b == TILE_INDEX)
        return EQTileMask(a, b, op);                     // tile-shaped equality mask
    else if (/* SCALAR involved */)
        return EQScalarPredicate(a, b, op);             // scalar equality predicate
    else
        raise TypeError("Unexpected type " ...);        // .rodata literal
}

3.2 make_comparison_operation — <, <=, >, >=

Identical category dispatch, but the predicate/mask results are the Intersection variants, the tile/scalar path fetches combine_tile_with as a method on the operand, and the affine relation is encoded by pred_op:

// make_comparison_operation(op, op_name, numpy_op, pred_op) -> handler(self, a, b)
//   pred_op = pred_lt / pred_le / pred_gt / pred_ge   (affine predicate constructor)
// closure refs: TileMaskIntersection, ScalarPredicateIntersection, combine_tile_with,
//               _binop, return_tensor_or_extracted_scalar, pred_lt/le/gt/ge
PyObject *handler(self, a, b) {
    cat_a = get_type_category(a);
    cat_b = get_type_category(b);

    if (cat_a == TENSOR || cat_b == TENSOR) {
        result = <tensor>._binop(numpy_op_or_op_name, other, /*dtype=*/np.int8);
        return return_tensor_or_extracted_scalar(result);
    }
    else if (cat_a == TILE_INDEX || cat_b == TILE_INDEX)
        return TileMaskIntersection(a.combine_tile_with(b), pred_op);
    else if (/* SCALAR involved */)
        return ScalarPredicateIntersection(a.combine_tile_with(b), pred_op);
    else
        raise TypeError("Unexpected type " ...);
}

QUIRK — ordered comparisons build an Intersection, not a bare leaf. i < N does not return a single EQTileMask; it returns a TileMaskIntersection wrapping combine_tile_with of the two operands and the relation pred_lt. This is what makes a chained compare lo <= x < hi fold cleanly — every link is already an intersection-shaped mask, and compare_nary just &-folds them (§5). Equality (==/!=) takes the simpler EQTileMask(a,b,op) leaf path. The split shows up as disjoint result-class strings in the two factories.

3.3 Wrappers and decorators

make_binary_method(op) wraps a handler into a real method(self, b) that re-runs the dispatch on the bound self. make_commutative_binary_method(op) additionally pulls in operator.itemgetter (the itemgetter string is in the pool) to order the operand pair by category, so a OP b and b OP a route identically. The two decorators install them:

// synthesize_equality_comparisons(cls)   (class decorator)
for ((name, op, numpy_op) in [('eq', operator.eq, np.equal),
                              ('ne', operator.ne, np.not_equal)]) {
    handler = make_equality_operation(op, name, numpy_op);
    setattr(cls, "__" + name + "__", make_binary_method(handler));   // __eq__, __ne__
}

// synthesize_ordered_comparisons(cls)    (class decorator)
for ((name, op, numpy_op, pred_op) in
        [('lt', operator.lt, np.less,          pred_lt),
         ('le', operator.le, np.less_equal,    pred_le),
         ('gt', operator.gt, np.greater,       pred_gt),
         ('ge', operator.ge, np.greater_equal, pred_ge)]) {
    handler = make_comparison_operation(op, name, numpy_op, pred_op);
    setattr(cls, "__" + name + "__", make_binary_method(handler));   // __lt__ __le__ __gt__ __ge__
}

Both op-name lists are readable in full: equal/not_equal/less/less_equal/greater/greater_equal (the np ufunc .lower() names) and eq/ne/lt/le/gt/ge (dunder stems) all appear in _strings.json, alongside lower (the .lower() call) and the pred_* constructors.

4. Logical synthesis and the promotion lattice (NkiTypeSystemLogicalOp)

The logical module is where heterogeneous operands are unified. p & q, p | q, ~p first compute highest_category(a,b), then promote both operands to a common logical type via PROMOTION_MAP, then apply the stdlib operator over the now-homogeneous predicate/mask objects.

4.1 Promotion helpers

// promote_other_to_predicate(a)
//   bool / number  ->  a constant predicate
PyObject *promote_other_to_predicate(PyObject *a) {
    if (!PyObject_IsTrue(a)) return AlwaysFalsePredicate();   // falsey -> ⊥
    return AlwaysTruePredicate();                              // truthy -> ⊤
}

// promote_to_logical_type(value, highest_category)
PyObject *promote_to_logical_type(PyObject *value, TypeCategory highest_category) {
    if (get_type_category(value) == TENSOR)
        raise(...);                                  // tensors are not promotable here
    promoter = PROMOTION_MAP[highest_category];      // PyObject_GetItem(PROMOTION_MAP, …)
    return promoter(value);                          // promote_to_mask / promote_to_predicate
    // failure path emits "Cannot promote from <…> to <…>"   (.rodata "Cannot promote from ")
}

The PROMOTION_MAP entries are the module functions promote_to_mask and promote_to_predicate (both present as names in the pool) plus promote_other_to_predicate for the OTHER slot. The exact key→value pairing of PROMOTION_MAP is INFERRED — the dict is constructed at import, so only the participating function names and the GetItem lookup are readable from the binary. The structurally-implied map:

4.2 make_logical_operation — & and |

// make_logical_operation(op, op_name, numpy_op) -> handler(self, a, b)
//   op = operator.and_ / operator.or_ ;  numpy_op = np.logical_and / np.logical_or
// body refs: highest_category, TENSOR branch, promote_to_logical_type, a_promoted/b_promoted
PyObject *handler(self, a, b) {
    cat = highest_category(a, b);
    if (cat == TENSOR)                               // tensor & tensor -> elementwise int8 mask
        return <tensor>._binop(numpy_op_or_op_name, other, /*dtype=*/np.int8);
    a_promoted = promote_to_logical_type(a, cat);    // both coerced to common logical type
    b_promoted = promote_to_logical_type(b, cat);
    return op(a_promoted, b_promoted);               // predicate/mask  &/|  predicate/mask
}

a_promoted / b_promoted are the handler's real local-variable names, preserved in the decompiled body.

4.3 make_logical_not — ~

// make_logical_not(op_name) -> handler(self, a)   (unary invert)
PyObject *handler(self, a) {
    cat = highest_category(a, a);
    if (cat == TENSOR)
        raise err_ambiguous_tensor_truth_value(...); // inverting a tensor is ambiguous
    a_promoted = promote_to_logical_type(a, cat);
    return operator.invert(a_promoted);              // De Morgan-aware invert on the promoted object
}

GOTCHA — a tensor has no logical truth value in tracing. & / | over two real tensors lower to an elementwise np.int8 op (a tensor result, not a predicate), but ~tensor and any bool-test of a tensor raise err_ambiguous_tensor_truth_value — a real string and a real sema error factory. This mirrors NumPy's "truth value of an array is ambiguous" and is the tracer's way of forcing the user to write an explicit comparison (which does produce a predicate) before combining. The asymmetry — binary logical over tensors is allowed but unary invert is not — is intentional: a & b has a defined elementwise meaning, ~a over an int8 tensor would silently mean bitwise-not, which is not a mask negation.

4.4 The decorator

// synthesize_logical_operations(cls)   (class decorator)
for ((name, op, numpy_op) in [('and', operator.and_, np.logical_and),
                              ('or',  operator.or_,  np.logical_or)]) {
    handler = make_logical_operation(op, name, numpy_op);
    method  = make_binary_method(handler);
    setattr(cls, "__" + name + "__", method);     // __and__, __or__
    setattr(cls, "logical_" + name, method);      // cls.logical_and, cls.logical_or   <-- consumed by §5
}
not_handler = make_logical_not('not');
setattr(cls, "__invert__",  make_unary_method(not_handler));   // ~
setattr(cls, "logical_not", make_unary_method(not_handler));

This is what gives NkiTypeSystem its logical_and / logical_or (the n-ary folds in §5 call self.logical_and), and gives every traced predicate/mask value its & / | / ~. nki_method decoration (sema integration) is applied at factory time; nki_method is imported by both CmpOp and LogicalOp.

5. The tracer: chained comparisons and n-ary folds (NkiTypeSystem)

The base NkiTypeSystem class (the "TypeTraceContext") is the only module with hand-written method bodies. It holds a single attribute (self.ctx, the trace context) and provides the lowerings for Python constructs that the synthesized binary operators cannot express alone: chained comparisons and membership.

class NkiTypeSystem {                       // decorated by all three synthesize_* (§1)
    void __init__(self, ctx) { self.ctx = ctx; }   // only attribute stored

    // ---- membership:  x in coll  /  x not in coll ----
    PyObject *in_(self, a, b) {
        return operator.contains(a, b);     // GetBuiltin 'operator', getattr 'contains'
    }                                        //   -> yields a predicate, not a bool
    PyObject *not_in(self, a, b) {
        return !self.in_(a, b);             // reconstructed: body refs in_, PyObject_Not
    }

    // ---- chained comparison:  a OP1 b OP2 c  ->  (a OP1 b) & (b OP2 c) & … ----
    PyObject *compare_nary(self, ops, operands) {
        assert(len(ops) == len(operands) - 1);            // two PyObject_Size, v==v+1
        return self.logical_and_nary(                     // AND-fold the per-link results
            (op(a, b)
             for (op, (a, b)) in zip(ops, zip(operands[:-1], operands[1:]))));
        //   operands[:-1] / operands[1:] = pairwise; the genexpr applies each
        //   comparison op to consecutive operand pairs (slices, zip, genexpr)
    }

    // ---- n-ary logical reductions (consume the synthesized logical_and/_or) ----
    PyObject *logical_and_nary(self, operands) {
        return functools.reduce(self.logical_and, operands);   // reduce + self.logical_and
    }
    PyObject *logical_or_nary(self, operands) {
        return functools.reduce(self.logical_or, operands);    // reduce + self.logical_or
    }
}

So lo <= x < hi, which Python would normally evaluate as (lo <= x) and (x < hi) (short-circuiting to a bool), is instead routed through compare_nary([le, lt], [lo, x, hi]), producing (lo <= x) & (x < hi) — two TileMaskIntersection/ScalarPredicateIntersection objects AND-folded by self.logical_and into a single combined mask/predicate. The genexpr trampoline is present in the base module (Pyx_Generator_Next at 0xbe30), and compare_nary calling self.logical_and_nary is the recovered tail of its body.

NOTE — why the tracer, not the operators, owns chained compares. Python evaluates a < b < c with built-in short-circuit and, which would force a bool out of a < b before the second comparison runs — defeating the entire predicate machinery. NKI's front-end therefore rewrites chained comparisons at trace time (the nisa validators in the W-strand call compare_nary when they encounter a chained relation on traced operands) so that each link is evaluated independently to a predicate/mask object and then &-folded. A reimplementer must intercept chained comparisons at the AST/trace level — they cannot be recovered from the __lt__/__le__ dunders alone.

6. End-to-end dispatch summary

Putting the three layers together, here is what each surface-syntax construct resolves to over traced operands:

The predicate/mask objects produced here flow into the algebra of §6.2.4 Mask / Predicate Algebra, whose EQTileMask / TileMaskIntersection / ScalarPredicateIntersection / AlwaysTrue* / AlwaysFalse* classes are exactly the result types named above — the two pages describe the construction side (this page) and the combination + lowering to AffinePredicate side (mask-predicate-algebra) of one object graph. The dtype of the tensor-compare result (np.int8) connects to the dtype model in §6.3.3 dtype; the metaclasses that build scalar/tensor/tile_index (and thus make get_type_category's isinstance chain meaningful) are §6.3.2 metaclasses.

7. Evidence summary

Limits of this reading

Per-method offsets are a weak anchor for the base module. NkiTypeSystem.…so is a plain Cython output with no dedicated _strings.json — only disasm and decompiled sidecars — and its function offsets move between Cython versions, so figures such as compare_nary @ 0xbf80 or a generator trampoline at 0x10b70 do not survive a rebuild; in the cp310 wheel the trampoline sits at 0xbe30. Cite the symbols and strings, not the absolute offsets. The CmpOp / LogicalOp offsets (0x15cb0, 0x16a10, 0x13750, 0xfb50) are a different matter: they match the cp310 decompiled filenames exactly and are dependable for those two modules.