perf: cache interpreter registry and statement trait lookups

[kaihsin]

Summary

Two independent caches for values that are invariant but recomputed constantly. Both were found by profiling a compile-heavy workload where ~19s of a 38s run was inspect.getmembers.

1. Cache the interpreter registry on the dialect group

Registry.interpreter builds its Signature -> BoundedDef table by walking every method table of every dialect with inspect.getmembers. Interpreter.__post_init__ requests that table for every interpreter instance, and Method.__call__ constructs a fresh Interpreter on every call:

def __call__(self, *args, **kwargs):
    interp = Interpreter(self.dialects)   # rebuilds the whole registry
    _, ret = interp.run(self, *args, **kwargs)
    return ret

So calling a kernel from Python costs O(size of the dialect group) per call, independent of what the kernel does. With the basic prelude (26 dialects) that dominates: in a profile of 20k calls, inspect.getmembers alone accounted for 429,903 calls and ~4.8s of self time.

The table depends only on (dialect group, interpreter keys), so it is now cached on the group. Dialect.register can add method tables after a group exists, so it bumps a module-level epoch that invalidates every cached table.

One subtlety worth flagging for review: the cached table is shared, not copied. Interpreter only reads it (in, .get, []). I first returned a defensive copy, and copying per call is itself expensive enough to erase the win entirely, so interpreter() now documents the mapping as read-only.

2. Memoize Statement trait lookups

has_trait / get_trait / get_present_trait each scan cls.traits with an ABC isinstance per entry. Rewrite passes hammer these — every purity check in DeadCodeElimination goes through is_pure — and a single Fold over a 56-method call graph performs ~305k lookups.

traits is an immutable ClassVar[frozenset] fixed at class creation, so the result depends only on (cls, trait). has_trait and get_present_trait now route through get_trait so all three share one cache and one definition of matching.

Benchmark (MWE)

import gc, time, statistics
from kirin.ir import Method
from kirin.passes import Fold
from kirin.prelude import basic, basic_no_opt

@basic
def kernel(x: int) -> int:
    y = x + 1
    z = y * 2
    return z - 1

def _leaf() -> Method:
    @basic_no_opt
    def leaf(x: int) -> int:
        y = x + 1
        z = y * 2
        return z - 1
    return leaf

def _node(a: Method, b: Method, c: Method) -> Method:
    @basic_no_opt
    def node(x: int) -> int:
        p = a(x)
        q = b(p)
        r = c(q)
        return p + q + r
    return node

def build(width=8, layers=7) -> Method:
    """`layers` layers of `width` kernels; each calls three below it."""
    level = [_leaf() for _ in range(width)]
    for _ in range(layers):
        level = [
            _node(level[i % len(level)],
                  level[(i + 1) % len(level)],
                  level[(i + 2) % len(level)])
            for i in range(width)
        ]
    return _node(level[0], level[1], level[2])

def bench(label, make, run, repeats=7):
    samples = []
    for _ in range(repeats):
        payload = make()
        gc.collect()
        t0 = time.perf_counter()
        run(payload)
        samples.append(time.perf_counter() - t0)
    print(f"  {label:24s} best={min(samples):7.3f}s median={statistics.median(samples):7.3f}s")

bench("20000 kernel calls", lambda: None,
      lambda _: [kernel(3) for _ in range(20000)])
bench("Fold (56 methods)", lambda: build(),
      lambda mt: Fold(mt.dialects)(mt))

	main	this PR
20,000 Method.__call__	5.155s	0.456s (11.3x)
Fold over 56-method call graph	1.237s	1.194s (~3%)

The Fold row is the trait cache measured on its own; it is a modest win and is a separate commit so it can be dropped independently. The registry cache is the headline.

As an incidental datapoint, kirin's own test suite goes from 7.40s to 2.69s.

Correctness

• traits is an immutable ClassVar[frozenset] and is never reassigned anywhere in src/; the cache is keyed per (class, trait) so sibling classes cannot share entries, and a cached None is distinguished from "not cached" by a sentinel.
• DialectGroup.data is a frozenset assigned once in __init__, so a per-group cache cannot go stale from group mutation. The one real staleness path, Dialect.register, is handled by the epoch counter and covered by a test.
• 11 new tests; full suite passes (715 passed, 13 xfailed, 1 xpassed), pyright clean, pre-commit clean.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

☂️ Code Coverage

current status: ✅

Overall Coverage

Statements	Covered	Coverage	Threshold	Status
11915	10753	90%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
src/kirin/ir/dialect.py	92%	🟢
src/kirin/ir/group.py	91%	🟢
src/kirin/ir/nodes/stmt.py	87%	🟢
src/kirin/registry.py	93%	🟢
TOTAL	91%	🟢

updated for commit: 6b92f82 by action🐍

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-09-02 15:08 UTC

[zhenrongliew]

LGTM

[github-actions]

Backport results for 2481110

Succeeded:

• release-0-22: Backport PR #719 to release-0-22 #722