WIP: add nkigen-lite as a standalone IR-based kernel generation backend - #59
Draft
ymwangg wants to merge 107 commits into
Draft
WIP: add nkigen-lite as a standalone IR-based kernel generation backend#59ymwangg wants to merge 107 commits into
ymwangg wants to merge 107 commits into
Conversation
Migrates tensor_ir, nki_ir, and the direct lowering passes from nano-tensorizer/ir_lab into the nkipy workspace as a new package. The pipeline (canonicalize → decompose → layout_solver → direct_lower) produces legal NKI IR directly without intermediate passes.
Add nkigen-lite as a fully functional backend (backend="nkigen-lite") alongside hlo and nkigen. The pipeline traces Python kernels through nkigen_lite's tensor_ir Builder, lowers via the pass pipeline (canonicalize → decompose → layout_solver → direct_lower), and compiles to NEFF via the NKI kernel_builder API. nkipy integration: - backend/nkigen_lite.py: TraceContext, Tensor, IR adapter - ops/_nkigen_lite_impls.py: op implementations delegating to Builder - ops/_register_nkigen_lite.py: lazy op registration - trace.py: _specialize_nkigen_lite() dispatch - compile.py: _compile_nkigen_lite() via kernel_builder - knob.py, nki_op.py: backend-aware dispatch nkigen-lite enhancements: - Builder: add abs, sign, floor, ceil, power, floor_divide, mod ops - Interpreter: numpy dispatch for new ops, dtype-aware tensor_copy - Decompose pass: floor_divide/mod use divide-then-verify-and-correct strategy (matching neuronx-cc BIR), power→exp(b*log(a)), ceil→neg(floor(neg(x))), fixed-point iteration with max-iter guard - Direct lowering: abs/sign/sin via NisaActivationOp, floor via i32 truncation + sign correction, cast via tensor_copy, 1D reshape fix - docs/floor_divide_precision.md: documents the precision strategy Test results: 134/135 HLO-parity tests pass on trn2 hardware (99.3%).
Add "nkigen-lite" to the trace_mode fixture so all parametrized tests run with both backends. Add a pytest hook that marks NotImplementedError as xfail for nkigen-lite — ops not yet implemented show as expected failures and automatically start passing when added. Current results: - HLO: 741 passed, 4 xfailed, 42 skipped - nkigen-lite: ~340 passed, 161 xfailed (unimplemented ops), ~93 failed (partial implementations needing further work)
Add ReduceKeepdimsFalsePattern to decompose keepdims=False reductions into keepdims=True + reshape, which the layout solver and lowering require. Handle scalar (rank-0) tensors throughout the lowering pipeline by promoting them to (1,) at the NKI boundary since the hardware doesn't support rank-0 tensors. Also fix negative axis normalization in squeeze() and expand_dims().
- matmul: add 1D→2D promotion following NumPy semantics - squeeze: validate non-1 dims, normalize negative axis - reshape: handle int newshape argument - zeros/full: handle int shape argument - concatenate: handle single-tensor case, validate empty/axis bounds - split: validate axis bounds and unequal division - where: handle numpy array condition argument - _ensure_value: handle numpy array operands (uniform-fill)
- expand_dims: validate duplicate axes and out-of-bounds axis - Skip test_reduce_unsupported_op and test_topk_non_last_axis for non-HLO backends since they test HLO-specific internal behavior
NeuronCore hardware only supports Add and Max for cross_lane_reduce_arith. Implement MIN as -max(-x) transparently in the NKI IR builder so all existing P-dimension reduce codepaths work with min reductions.
Replace HLO-specific DeviceKernel.compile_and_load path with the shared on_device_test utility which handles input/output naming differences between backends automatically.
- broadcast_to: handle scalar (rank-0) source by loading the single element and broadcasting via tensor_scalar_arith with ones - emit_to_kb: auto-cast f16/bf16 operands to f32 around tensor_scalar_arith since the hardware scalar engine requires f32
Add NisaBitvecOp enum and tensor_tensor_bitvec builder method to NKI IR. Wire through the full pipeline: tensor IR opcodes, elementwise lowering, emit_to_kb mapping, and interpreter support. Replace the old arithmetic approximations (which only worked for booleans) with hardware bitwise instructions that work correctly on integer types.
Add NKI IR primitives for mixed-dtype operations: - tensor_tensor_compare: comparison ops (IsGT, IsGE, etc.) that accept float inputs and produce uint8 predicate output - tensor_scalar_bitvec: scalar bitvec ops (XOR for logical NOT, etc.) - Comparison and logical op variants in NisaArithOp enum Rewrite _emit_floor to use the NKI compiler's compare+select pattern: trunc→compare→conditional select in integer domain, avoiding float precision issues in the correction step.
…lice - emit_slice: add strides parameter; delegate to _emit_strided_slice for non-unit strides (element-by-element DMA for F-stride, row-by-row for P-stride) - dynamic_update_slice: handle numpy array value argument (uniform fill)
Add DType.FP8_E4M3_IEEE for the IEEE-standard float8_e4m3 format (distinct from the NaN-free float8_e4m3fn variant already supported). Wire through core, emit_to_kb, and compile dtype mappings.
Each pytest-xdist worker now claims a specific Neuron core via NEURON_RT_VISIBLE_CORES, enabling parallel test execution across all 64 available cores (~8.5x speedup).
- Comparison ops (equal, not_equal, greater, less, etc.) now produce same dtype as input (1.0/0.0 float) matching NKI convention, instead of DType.BOOL - where op lowered using NKI pattern: cond*x + (1-cond)*y with all float arithmetic — no mixed-dtype operations needed - Map DType.BOOL → uint8 in kernel builder and execution layer - Update tensor IR builder to remove BOOL requirement from where - Reduces xfail count from 162 → 125 (37 tests now passing)
Matches NKI compiler's approach: cos(x) = sin(x + π/2). The hardware sin activation instruction handles the computation.
Implement np.dot semantics as a composed op: - 1D/2D cases delegate directly to matmul - N-D × 1D delegates to matmul (batched matrix-vector) - N-D × M-D decomposes to reshape + matmul + reshape to achieve the outer-product batch semantics of np.dot
- arctan: wire native NISA ARCTAN activation through the Builder and direct-lower tables - invert/bitwise_not: composed_impl as XOR with all-ones (-1), matching the NKI compiler's implementation - logical_and/or/xor: composed_impl via 0/1 truthiness; also unblocks rint/round which depend on logical_and - constant: backend impl mirroring HLO (passthrough + uniform fill); non-uniform array constants raise NotImplementedError Fix a pre-existing bug in _emit_broadcast_scalar that fed a (1,1) tile to tensor_scalar_arith whose scalar operand partition dim must match the destination; replicate to (p_size, 1) via broadcast_partition. Also make test_ml_dtypes_constant_encoding's float8 xfails backend-aware so float8_e5m2 on nkigen-lite no longer reports XPASS.
The slice-based gather produced wrong output shapes: it ignored axis=None (no flatten), concatenated slices along the original axis instead of replacing it with indices.shape, and mishandled scalar and multi-dimensional index arrays. Rewrite to match numpy: out.shape == a.shape[:axis] + indices.shape + a.shape[axis+1:] - axis=None flattens the input first - negative indices are normalized modulo the axis dimension - each flat index becomes a width-1 slice; slices are concatenated then reshaped so the gathered axis is replaced by indices.shape (dropped entirely for a scalar index) Fixes 13 failing test_take_scalar / test_take_numpy_indices cases.
…) for nkigen-lite Wires the four distributed collectives through the full nkigen-lite stack: - tensor_ir Builder: collective ops with correct output-shape inference (all_gather grows the gather dim, reduce_scatter/all_to_all shrink/grow by world size) - nkipy lite impls + registration, mapping numpy reduce ufuncs to the collective reduce-op names - direct_lower: stage collectives through internal HBM scratch buffers (the compiler forbids collectives from reading/writing kernel IO tensors directly — "Collective instruction cannot read IO tensors") - nki_ir Builder: collective() side-effect node (HBM->HBM) - emit_to_kb: lower to nisa.all_reduce/all_gather/reduce_scatter/all_to_all via ExplicitReplicaGroupAttr + dma_compute_reduce_op The KB collective API only operates on the last (free) axis of 2D HBM tensors (cc_dim=0 raises std::bad_cast), so all_gather/reduce_scatter along other axes are staged via transpose-collective-transpose. Fixes the all_reduce/all_gather/reduce_scatter/all_to_all xfails (multiply-reduce variants stay xfailed for both backends — unsupported by the compiler).
The earlier transpose workaround for all_gather/reduce_scatter was based on a misdiagnosis: cc_dim=0 appeared to raise std::bad_cast, so collectives were staged through a transpose to operate on the last axis. Multi-core numerical verification showed that path silently dropped the remote rank's data (all_gather duplicated the local source; reduce_scatter ignored the per-rank scatter offset). Root cause: the KB nisa collective APIs forward cc_dim to the native builder un-converted, so a bare int 0 fails the int->enum cast. The NKI collectives contract also requires collective_dim=0 for HBM tensors. Fix: - emit_to_kb: convert the int dim to CollectiveDimension (DIM_0/DIM_1) before calling nisa.all_gather/reduce_scatter/all_to_all - drop the transpose workaround; gather/scatter along the requested dim directly Verified on 2 NeuronCores with distinct per-rank data: all_reduce, all_gather(dim0), reduce_scatter(dim0), and all_to_all all produce the correct cross-rank results.
Two bugs combined to make a // b and a % b off by one on the rare inputs where the true quotient is an exact integer: 1. The composed floor_divide impl (floor(divide(x, y))) ran at trace time, so the graph never contained a `floor_divide` opcode and the decompose pass's divide-then-verify-and-correct FloorDividePattern never fired. NeuronCore has no native divide -- it uses reciprocal multiply, which undershoots exact integers (2.0 -> 1.9999999), so plain floor gave N-1. Fix: register nkigen-lite-specific floor_divide/remainder impls that emit the native floor_divide/mod opcodes, so the correcting pattern runs. 2. Within FloorDividePattern, the up-correction used max(0, sign(|rem| - |b|)), which is 0 at the |rem| == |b| boundary (sign(0) == 0) -- exactly the exact-integer undershoot case. Replace with an inclusive greater_equal(|rem|, |b|); a genuine remainder is always strictly < |b|, so equality can only mean undershoot. Fixes the 3 failing floordiv/mod broadcasting cases. Exact-integer-boundary inputs remain inherently ambiguous under reciprocal division (numpy and the device can disagree by an ULP), but all test cases now pass.
Contributor
Author
Status Update: Replacing the HLO backend with nkigen-liteReported 2026-06-23 from full test suite run ( Where we areOver Jun 2–5, nkigen-lite was stood up as a third backend and brought to broad Current suite status: To finish retiring the
1. Hard failures — numerical bugs (2 cases) — fix firstThese already run but produce wrong results:
2. Unimplemented ops (58 cases)Each op unblocks its cluster of tests:
3. Capability gaps — broader subsystems (42 cases)
4. Out of scope — compiler limitation (2 cases — leave xfailed)
These cannot be cleared by nkigen-lite work — hardware/compiler limitation. Remaining work, in priority order
Clearing §1–§3 retires ~100 of the 102 hlo-only-passing cases; the 2 |
NRT does not carry FP8 dtypes through compiled neff metadata: e4m3/e5m2
surface as "int8" and e4m3fn surfaces as "unknown". The data round-trips
correctly on device, so the only blocker was spike's dtype validation,
which special-cased only e4m3/e5m2 and only accepted "int8".
Extend _check_dtype_compatibility to cover e4m3fn and accept either
placeholder ("int8"/"unknown"), gated behind the FP8 dtype set so
non-FP8 tensors are still validated strictly. Drop the matching
nkigen-lite xfail in test_ml_dtypes_constant_encoding.
The high-level tensor_ir.Builder had no index-ramp op, while the low-level nki_ir already exposed nisa.iota. Bridge that gap across all four layers: - tensor_ir/ir.py: Builder.iota(shape, dim, dtype) — out[..., i, ...] == i along dim, broadcast over other axes (np.arange-on-axis semantics). - core.py: numpy eval in eval_common_op (shared by both interpreters). - direct_lower.py: _emit_iota_op tiles under the canonical row-major layout and maps each axis to nisa.iota's pattern/channel_multiplier/offset (free -> step 1; partition -> channel_multiplier 1 + p_off; batch -> constant offset). Kept out of ELEMENTWISE_OPCODES since it is position-dependent. Unblocks tril/triu/diag/trace, which build index masks via iota. Adds TestIota HW coverage (per-axis, multi-tile partition, rank-3, and iota feeding an elementwise op).
Implement gap-8 triangular/diagonal ops on top of the new iota primitive, mirroring the HLO impls: build row/col index masks via iota + compare, then where(mask, x, 0). - tril/triu: keep row >= col-k (lower) or row <= col-k (upper); masks built over the last two axes, broadcast across batch dims. - diag 1D->2D: extend v to length N with a zero-pad on the side away from the diagonal, broadcast across columns, keep col == row+k. Avoids the HLO take-based gather (dynamic indexing is unsupported on nkigen-lite). - diag 2D->1D: mask the k-th diagonal and reduce-sum the off-axis to collapse to the diagonal vector; slice to diag_len. - trace: mask the diagonal (offset) and reduce-sum both axes. Flips 8 xfails to pass in tests/unit/test_tensor_api.py. Verified against numpy on non-square diag/tril/triu and trace-with-offset edge cases beyond the existing suite.
Implement gap-7 (pad) and gap-9 (flip/roll/tile/diff) as pure slice/concat data movement — no new primitives needed. - flip: reverse an axis by concatenating width-1 slices in descending order. - tile: concatenate r copies of the running result along each axis. - roll: cyclic shift via split at (n-shift) + swapped concat; supports axis=None (flatten), int axis, and tuple shift/axis. - diff: iterated x[1:] - x[:-1] along axis (n times). - pad: constant mode concatenates full-valued slabs; edge mode replicates the first/last slab. Handles scalar, per-axis, and asymmetric pad_width. Flips 14 xfails to pass in tests/unit/test_tensor_api.py. Verified against numpy on 3-D flip, multi-axis flip/roll, flattened roll, 3-D tile, and 3-D asymmetric pad beyond the existing suite.
Implement argmax/argmin via index masking on top of iota, mirroring the HLO algorithm: reduce to the extreme value along the axis, mark positions equal to it with their index (an iota ramp) and all others with a large sentinel, then min-reduce the indices — returning the first index that attains the extreme, matching numpy's tie-break. The whole computation runs in float32 (cast input up front, cast result to int32 at the end): min/max reductions initialize with +/-inf, which cannot be memset into an integer tile, so an integer input or index ramp would fail to compile. Handles axis=None (flatten), negative axis, and keepdims. Flips 8 xfails to pass; verified against numpy on int inputs and tie-breaking beyond the suite.
The lite builder only emits uniform fills, so non-uniform array constants previously raised NotImplementedError. Materialize them as a flat sequence of run-length fills, concatenate, and reshape — cheap for structured/small arrays, capped at 4096 runs to keep tracing bounded. Route the binary-operand path (_ensure_value) through the same logic. cumsum gets a dedicated nkigen-lite impl rather than relying on the composed fallback (which builds constant(np.triu(np.ones((N,N))))): for axis=None the flattened triangular matrix is (4096,4096), far too large as a literal. Build U[i,j] = (i<=j) via iota + compare instead, then cumsum = x_2d @ U. Handles axis=None/negative/middle axis and dtype. Flips 8 xfails to pass (4 cumsum, list/tuple constant, integer where-cond). Verified against numpy on 1-D/3-D/negative-axis cumsum and small structured constants beyond the suite.
nkigen-lite has no convolution primitive, so decompose N-D conv into im2col + a single matmul: gather each kernel position's strided window as a (N, Ci, out_pts) block, concat along the channel axis into (N, Ci*prod(K), out_pts), flatten the (transposed) weight to (Co, Ci*prod(K)), and do one batched matmul -> (N, Co, out_pts). A single fused matmul compiles ~35% faster than accumulating prod(K) separate matmuls (95.6s -> 61.5s on ic=16/oc=32/k=3); the official conv2d suite drops from 165s to 121s. Spatial padding is built from concat of zero slabs; strided/dilated windows use strided slice. groups != 1 is unsupported. Verified against PyTorch on conv2d (stride/padding/dilation/bias/1x1) and conv3d including a non-cubic (2,3,3) kernel. conv tests require torch as the oracle, now installed via the examples dependency group.
HardwareProfile/TRN2 was undocumented dead scaffolding: no cost model reads its fields, and lower_to_nki's target param is accepted but never used. Document both facts plainly instead of implying a live multi-target cost model. Also derive the tile-capacity fields (partition_max, sbuf/psum_per_partition_bytes, etc.) from nki_ir.ir's constants instead of duplicating them as separate literals, since those are what the graph verifier actually enforces.
…entwise segments Phase 0 of SBUF_FUSION_PLAN. Elementwise results that never escape their segment (not read by another segment, not a graph output) no longer get an HBM buffer or store. Segments tile to their widest dtype instead of the conservative F32 default (one shared dtype per segment so per-value slice offsets stay aligned with the rep loop). qwen3 MoE layer (L=8): 23017 -> 22833 dma_copy.
…ant concat inputs Two hotspot fixes from the qwen3 MoE profile: - gather_rows of wide rows (MoE expert weights, W=786432) emitted ceil(W/w_tile) single-lane indirect DMAs per row (~212 ops/expert). New packed path views the (N, W) table as (N*128, W/128), expands the dynamic index to the row's 128 sub-rows (idx*128 + lane via iota), and fetches the whole row as one partition-packed indirect DMA (~9 ops/row). Cost model keeps tall gathers (embedding-style) on the generic path. - concat inputs that are compile-time constants (the MoE router-weight assembly concats ~130 tiny constants per token) no longer materialize HBM buffers: emit_concat splat-fills their output window via memset + store, with a per-concat tile cache. Elementwise segments prune constant ops whose only consumer is such a concat; escape analysis in lower_graph skips their HBM allocation. qwen3 MoE layer (L=8): 79932 -> 45128 nki ops (17.2x -> 9.7x expansion), 22833 -> 11577 dma_copy, indirect DMAs 6720 -> 128.
Replace the global layout solver (solve_graph, five propagation phases) with per-segment layout decisions made at emission time. Every segment boundary round-trips through HBM, which is layout-agnostic (row-major), so a consumer segment can load a value in any layout regardless of how the producer stored it — there is no global layout problem until cross-segment SBUF residency exists (explicitly deferred). The solver's output was barely consumed: elementwise emission already used canonical row-major layouts, matmul emission never read it, and only _segment_ops (layout-flip group breaks that could only pessimize) and emit_reduce (axis classification) did. Its propagation phases also carried known direction bugs (transpose perm applied forward and backward with the same sign; phase-5 adapting through the producer op instead of the consumer). - passes/layout.py keeps the load-bearing parts: the Layout (I/P/F) dataclass, get_matmul_layouts (tensor-engine hard constraint), and default_layout (scored contiguous I|P|F split). - _segment_ops groups on collapsed-(P,F) shape compatibility only. - emit_reduce classifies reduce axes via default_layout of its input. - lower_graph(graph) / lower_to_nki(graph) drop the layouts dict. - Delete test_layout_solver.py golden baselines (locked in behavior nothing consumed); record the design in SBUF_FUSION_PLAN.md, including the Phase 2 anchor-driven predicate and fusion-boundary layout notes. Verified: full nkigen-lite suite (731 passed, 1 xfailed) and identical lowered op counts across 24 pattern graphs vs. the previous pipeline.
One view per concat instead of one per constant input (the router-weight assembly has ~130). Also records Phase 0.5 results in SBUF_FUSION_PLAN: qwen3 MoE layer 80332 -> 42805 nki ops (17.3x -> 9.2x).
…0/0.5 TTFT 4596 -> 3280 ms (1.40x), decode p50 2500 -> 2333 ms (1.07x) vs the pre-branch commit on the same box with fresh builds. Also documents that the checked-in benchmark_report.json (07-01) predates the basic-lowering-cleanup merge, which itself regressed decode ~1.5x.
# Conflicts: # nkigen-lite/SBUF_FUSION_PLAN.md # nkigen-lite/src/nkigen_lite/tensor_ir/passes/layout_solver.py
Document NEURON_RT_VISIBLE_CORES core pinning, how the test suite isolates cores per xdist worker, benchmarking a branch via git worktree (PYTHONPATH + core pinning for concurrent side-by-side runs), and ignored runtime env vars (NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS on RT 2.0).
A broadcast_to feeding only elementwise binary ops is pure waste: the vector/scalar engine broadcasts a size-1 free dim natively and a partition-1 operand is fanned by the load, so the consumer can read the un-broadcast source directly. New tensor_ir peephole (fold_broadcast.py) rewires collapse-safe broadcasts to their source and drops the broadcast. Runs before decompose, so div(a, broadcast_to(b)) folds to div(a,b) -> mul(a, reciprocal(b)) -- also removing the extra HBM round-trip that decompose.py flagged as the source of a residual ~1/65536 floor-divide precision error. Middle broadcasts (GQA head expansion (1,8,1,64)->(1,8,8,64)) are not collapse-safe and stay materialized. qwen3 MoE layer (L=8): broadcast_to 105 calls/1438 ops -> 7 calls/700 ops; total 42805 -> 41633 nki ops (-2.7%). All nkigen-lite tests pass (746 + fold suite).
A static-start, stride-1 slice preserves rank, so an elementwise consumer can read the sliced data by adding the slice's per-dim starts into its own tile-load offsets — no HBM buffer, no load-compute-store. lower_graph records these in slice_views and skips their allocation/emission; the generic elementwise path composes the offset (slice_srcs in _emit_ew_tile), while any other consumer (reshape/matmul/reduce/transpose/concat, the rank>=3 collapse path, or a graph output) calls _resolve to materialize the buffer + copy once (KB refuses .view() on a sliced tile, so only the generic path can compose). Non-unit strides keep the copy path. qwen3 MoE layer (L=8): 41633 -> 40993 nki ops (-1.5%), dma_copy 11369 -> 11113. The hot gate/up (384,)->(192,) split (x128) and top-k index slices fold into their consumers. All 695 tensor_ir tests pass plus 11 new slice-view tests (7 interpreter, 4 hardware).
transpose(transpose(x, p1), p2) is a single transpose(x, compose) with compose[i] = p1[p2[i]]. The chain materializes the intermediate through HBM (load -> remap -> store -> reload); the composed form skips it, and DCE drops the inner transpose's materialization when it has no other consumer. New tensor_ir peephole passes/fold_transpose.py, run before decompose beside fold_broadcast. The qwen3 attention path chains (0,2,1,3) then (0,1,3,2) on a (1,8,4096,128) tensor to feed QK^T; one of the two ~16 MB attention transposes folds away. qwen3 MoE layer (L=8): transpose 2864 -> 2096 nki ops; total 40993 -> 39968 (-2.5%); dma_copy 11113 -> 10601.
…re helper
Split direct_lower.py (~1000 lines lighter) into focused per-op modules —
direct_lower_{fusion,gather,iota,topk}.py — and consolidate the elementwise
load/store path.
- Segmentation + elementwise emission move to direct_lower_fusion.py; gather,
iota, and topk emitters get their own modules. direct_lower.py now dispatches
to them.
- New load_input_tile / store_output_tile / canonical_layout in
direct_lower_utils.py are the single elementwise tile load/store path,
replacing the two near-identical inlined branches in _emit_ew_tile. The
slice-view offset becomes an optional per-dim `offsets` descriptor composed
into the DimSlices — the seam a future transpose/reshape view composes
through.
- Dedup the hand-rolled batch-index unravel in broadcast/matmul/transpose
emitters onto the shared unravel() helper.
Behavior-neutral: qwen3 MoE layer stays at 39968 nki ops / 10601 dma_copy.
All 435 non-HW tensor_ir tests pass.
- qwen3/profile_layer.py: static per-op profiler for the fused MoE layer, attributing emitted nki ops + DMAs back to the high-level tensor_ir op. - qwen3/dump_layer_ir.py + qwen3_embedding/dump_layer_ir.py: dump the full tensor_ir/nki_ir for one fused transformer layer (whole unrolled MoE expert loop), complementing the per-building-block dump_nki_ir.py. Generated dumps and benchmark JSON stay untracked (matches qwen3_embedding).
Strip the segment-grouping and dead-store machinery from the direct lowering, trading the elementwise-fusion optimizations for a much simpler allocation and dispatch path. These optimizations are to be reintroduced uniformly (for all op classes, not just elementwise) once the planned SBUF-first residency model lands. - Remove _segment_ops / _collapsed_pf: lower one op per segment. Elementwise grouping (consecutive ops sharing one on-chip tile loop) is gone; it returns as the Phase-2 fusion-compatibility predicate. - Remove the segments / seg_of data structures: the emission loop iterates graph.ops directly. - Remove escapes / ew_results (dead-store elimination): every op result now gets an HBM buffer. This also disables the concat constant-splat path (a constant read only by concat was previously never materialized). - Fold the concat and view-slice special-cases into the per-op emitter table via closures, so the dispatch loop reads as "resolve views, then dispatch", with elementwise the one principled exception (it composes views without materializing). Deliberate op-count regression on the qwen3 MoE layer (L=8): 39968 -> 52686 nki ops, memset 2578 -> 4910 (dead stores + concat constants now materialized). All 435 non-HW and 330 HW (on-device) tensor_ir tests pass.
With dead-store elimination removed, every op result now gets an HBM buffer, so a constant read by a concat is already materialized by the time concat runs — the splat-fill path (memset straight into the output window) is redundant. Remove const_values and the splats plumbing; concat now loads all inputs uniformly from hbm_map. - direct_lower.py: drop the const_values dict and the concat closure; concat becomes a plain _OP_EMITTERS entry. - direct_lower_memory.py: emit_concat loses its splats parameter and splat branch; _emit_splat_window deleted. All 435 non-HW and 330 HW tensor_ir tests pass.
Remove dead code and the rank>=3 collapse fast path from the (renamed) elementwise emitter. - Drop the dead segment-pruning loop and the `out_name in hbm_map` store guard in _emit_elementwise_segment: with dead-store elimination gone every result has a buffer, so the prune kept every op and the guard was always true. - Remove _try_emit_collapsed_ew and _collapse_ew_shape (the rank>=3 leading-dims-onto-partition collapse), plus the shape_override plumbing in _emit_ew_tile only that path used. Each rank>=3 elementwise op now unrolls its leading axes via the generic per-tile path. This is a deliberate, sequence-length-scaling op-count regression (qwen3 MoE L=8: 53577 -> 56117 nki ops) to be reintroduced as the Phase-2 fusion-compatibility predicate. - fold_broadcast._collapses_cleanly keeps its own copy of the collapse-safety check (drop the stale cross-reference comment). - test_perf1_mixed_collapse_elementwise_no_blowup: keep the correctness check, drop the op-count guard that asserted the removed collapse; update the TestCollapsedElementwise docstring. All 435 non-HW and 330 HW tensor_ir tests pass.
Rewrite the elementwise emitter as a single linearize-and-tile flow: collapse
every operand's HBM buffer to 2D (P=prod(leading), F=last) via zero-copy views,
then tile P at 128 and F at the largest power-of-two <=512. Reuses the existing
emit_binary_op/emit_unary_op for on-chip (1,.)/(.,1) broadcast; slice-as-view
maps to a 2D (row_off, col_off) window. Soundness rests on the fold_broadcast
contract (operand collapses to P in {1,rep_P}, F in {1,rep_F}), asserted per
operand.
Deletes the old Layout/compute_tile_sizes/on_chip_shape/map_indices/hbm_slices/
canonical_layout/load_input_tile/store_output_tile machinery from
direct_lower_utils.py (net -88 lines). Renames the entry point
_emit_elementwise_segment -> _emit_elementwise_group ("segment" collided with
the layout-segment concept elsewhere).
All 765 tests pass on trn2.
Remove both data-movement peephole passes from the lowering pipeline (canonicalize → decompose → direct_lower) ahead of the SBUF-first rework. broadcast_to ops and chained transposes are now materialized through HBM again; the elementwise 2D-collapse soundness argument rests solely on that materialization (the trivially collapse-clean full-tensor case), so the emitter contract still holds. - Delete passes/fold_broadcast.py, passes/fold_transpose.py and their tests. - Update direct_lower_elementwise.py docstrings/comments that referenced the fold guarantee. - Note the revert in SBUF_FUSION_PLAN.md (historical entries left intact). - Mark test_view_assignment_semantics[nkigen-lite] xfail: nkigen-lite does not support assignment through a view. Verified: 688 tensor_ir tests + 299 tensor_ir HW tests (trn2) green. Full nkigen-lite backend suite passes modulo pre-existing order-dependent test_broadcasting flakiness (present on baseline with folds too).
Restore layout_solver.py (470-line solve_graph five-phase propagation) and test_layout_solver.py verbatim from 609e524^, which had deleted them in the segment-first refactor. Restored as-is: the current lowering pipeline does not call solve_graph, so this is currently unused alongside its 41 passing golden tests. The known propagation direction bugs noted in 609e524 (transpose perm sign; phase-5 adapting through the producer) are also restored.
…view Simplify the basic elementwise flow to one op per load→compute→store loop and remove the slice-as-view optimization. - _emit_elementwise_group -> _emit_elementwise_op: takes a single op, not an ops list. Drops the group_results dead-code, the ops loops, and the one-entry store_plan. LoadPlan is now a NamedTuple. - Remove slice_views / _resolve / _is_view_slice from lower_graph and the _slice_2d_offset offset mapping. Every static slice now materializes to its own HBM buffer via emit_slice (the existing copy path); elementwise consumers read that buffer normally. This drops the lazy-materialization machinery whose only payoff was ~1.5% ops on qwen3 MoE, and moots the 2D-window contiguity bug the offset mapping carried. - Fix stale module/test docstrings that described grouping and the removed offset-view path; mark slice-as-view reverted in SBUF_FUSION_PLAN.md. Full suite: 791 passed, 1 xfailed.
…ring Behavior-neutral cleanup of the basic lowering flow, plus a fix to the qwen3 profiler that instruments it. Dead code: - Delete build_out_slices (utils) — defined and imported but used nowhere. - Delete the _collapse_to_2d wrapper (memory) — a pass-through to collapse_view with an unused shape arg; inline collapse_view at 6 sites. - Drop emit_reduce's strategy param and unreachable "matmul" branch: the tensor-engine P-reduce is reached only via lower_p_reduce_matmul's force="p_matmul", never through emit_reduce. Thread strategy out of _lower_via_emit / lower_reduce too. - Rename unused src_order unpack to _src_order (transpose, 2 sites). Duplication: - Merge matmul _a_batch_slices/_b_batch_slices into one _batch_slices. - Use the shared iter_pf_tiles for the scatter backdrop copy (gather). - Factor topk's repeated "trim to first k cols" into a _keep_k helper. Stale comments: "per segment" / "segment's working set" -> per-op wording (segmentation was removed with single-op lowering). profile_layer.py: the attribution instrumentation patched module-global emitter names, several of which never existed (_emit_elementwise_segment, _emit_reduce_op, ...) — it crashed before lowering and, for the table- dispatched ops, never intercepted anything. Wrap the _OP_EMITTERS dict entries in place (keyed by opcode) and the elementwise direct call via the module global, so the attribution table populates correctly. Verified: full suite 791 passed / 1 xfailed; qwen3 MoE layer lowers to a byte-identical nki op count vs. before; qwen3 (30B TP=4) and qwen3-embedding both run correctly on trn2 via the nkigen-lite backend.
Promotes the tiling concept that every direct-lowering loop-nest reinvented (per-dim tile sizes, ceildiv iteration, clamped slice/extent) into a single data object, separating tiling from alloc/codegen. A TileSchedule knows only shapes/sizes/indices — never the Builder — so it stays cleanly separable. - New direct_lower_schedule.py: TileSchedule + TileIndex, with a .pf() policy constructor (the single place the data-movement tile budget is chosen) and .pf_tiles() for the flat-offset 2D loops. - Migrated all 7 iter_pf_tiles callers and all five reduce loop-nests (collapsed-f fast path, f-reduce, gpsimd/matmul p-reduce, mixed). The load-bearing accum-loop vs. no-loop branch in the doubly-nested emitters is preserved exactly. - Deleted the now-dead iter_pf_tiles / build_slices / clamped_extent. - Added test_direct_lower_schedule.py pinning TileSchedule against inlined copies of the old free functions (65 equivalence cases). Full tensor_ir suite: 794 passed, 1 xfailed.
Separates allocation from codegen in direct lowering. Every emitter allocated SBUF tiles and HBM scratch inline, mingling "what memory to allocate" with "what op to emit". Scratch collects both behind one named surface (sbuf/hbm/ load), wrapping the Builder and exposing only allocation. - New direct_lower_alloc.py: Scratch, with .load() fusing the pervasive dma_copy(alloc(shape, dtype, SBUF), src, slices) idiom (57 sites). - One Scratch threaded per graph from lower_graph through the emitter dispatch; standalone lower_* wrappers construct a local one (public emit_* default scratch=None and self-construct). - All HBM scratch routed through Scratch.hbm — the single audit choke-point for scratch HBM (the qwen3 OOM lever): collectives, topk scan/candidate/column buffers, reshape diff-f, broadcast_partition. - Removed the MemorySpace/ceildiv imports that went dead; no new ruff F401s. - Added test_direct_lower_alloc.py (space correctness + load==alloc+dma). Full suite: 860 passed, 1 xfailed.
Completes the tiling/alloc/codegen separation for the regular tile-loop emitters. Elementwise had its own hand-rolled 2D loop with a bespoke power-of-two free-tile policy (_free_tile); reduce already moved to TileSchedule in step 1. - Add TileSchedule.free_pow2(P, F, free_max): partition at 128, free = largest power of two <= min(F, 512). The single home of the elementwise tile policy. - Fold the elementwise emitter's loop onto free_pow2().pf_tiles(); delete _free_tile and the now-dead ceildiv/PARTITION_MAX imports. The emitter was already plan-once + per-tile body, so this finishes its body-split shape. - Collapse the 3 remaining nb.dma_copy(nb.alloc(...)) loads in broadcast's _emit_collapsed_broadcast to scratch.load (stride-0 fan-out + 3D middle broadcast) and its ones-multiply dst to scratch.sbuf. - Add test_free_pow2_matches_old_elementwise_loop equivalence test. Data-dependent / batch-unrolled loops (gather, iota, memory flat-range, generic broadcast) intentionally keep bespoke iteration — the flat tile_sizes model can't express their per-dim slice construction without changing IR. Full suite: 869 passed, 1 xfailed.
The class allocates SBUF working tiles and does loads, not only HBM scratch,
so "Scratch" undersold it. Rename the class Scratch -> Allocator and the
per-graph instance scratch -> alloc; reads as alloc.sbuf(...)/alloc.load(...)/
alloc.hbm(...).
Pure rename: concept prose ("HBM scratch round-trip") and pre-existing
scratch_* locals (scratch_hbm, val_scratch, the emit_to_kb scratch list) are
untouched. tensor_ir suite green (807 passed, 1 xfailed); no new ruff F401s.
Completes the allocation-seam invariant. Previously ~37 SBUF compute-output dst tiles (dst = nb.alloc(...); nb.tensor_reduce_arith(dst, ...)) and the per-op HBM result-buffer pre-pass bypassed Allocator, so "one surface for all allocation" wasn't actually true. - reduce/elementwise/broadcast: SBUF dst tiles -> alloc.sbuf. - direct_lower_utils.py arith helpers (emit_binary_op, emit_unary_op, _emit_floor, _materialize_broadcast): build a local Allocator(nb) — same pattern as broadcast_partition — and route their dst tiles through it. - lower_graph HBM result pre-pass -> alloc.hbm. - Removed 2 MemorySpace imports that went dead as a result. Remaining nb.alloc, by design: PSUM accumulators (matmul/transpose/reduce, x3 — systolic-array output, not poolable scratch) and the Allocator method bodies. nb.constant (alloc+memset) stays on the builder. Full suite: 869 passed, 1 xfailed. ruff F401: 5 (below prior baseline).
_emit_elementwise_op already shares the (nb, op, hbm_map, alloc) -> None contract with every _OP_EMITTERS entry, so the special-case branch was historical. Register the elementwise opcodes in the table and collapse the dispatch loop to one uniform lookup. Full tensor_ir suite: 807 passed, 1 xfailed.
Deletes the separate per-op-result HBM pre-pass in lower_graph, which had to know _is_view_reshape — coupling "what to allocate" to "how reshape lowers". Emitters now dispatch through _with_result_buffers, which allocates each result's HBM just-in-time before the emitter runs. Reshape is registered raw (via the _SELF_ALLOCATING set) so its zero-copy view case installs a view with no dead buffer (an expert-weight reshape result is ~100 MB); its real-copy case allocates its own buffer. lower_graph now has no per-opcode special cases. Generated IR is equivalent: identical opcode multiset, only alloc ordering differs (interleaved before each producer vs. all up front). Verified a view-reshape adds zero HBM allocs. Full monorepo suite: 1446 passed, 55 skipped, 15 xfailed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
nkigen-lite, a standalone IR-based kernel generation backend that lowers numpy-style tensor programs to NKI (Neuron Kernel Interface) code for NeuronCore targets.Architecture
The system is structured as a three-layer IR stack with a multi-phase lowering pipeline:
Core (
core.py)Shared SSA-based IR infrastructure used by both IRs:
Value,Op,Graph— SSA primitives with use-lists and mutation helpersDTypeenum covering f32/f16/bf16/tf32/fp8/int typesTensor IR (
tensor_ir/)High-level, hardware-agnostic IR operating on whole tensors:
NKI IR (
nki_ir/)Low-level IR that makes hardware concerns explicit:
Lowering Pipeline (
tensor_ir/passes/)The full pipeline:
tensor_ir → canonicalize → decompose → layout_solver → direct_lower → nki_irCanonicalize — recomposes primitive-op chains into high-level ops (e.g.,
div(1, sqrt(x))→rsqrt(x),mul(x, div(1, add(1, exp(neg(x)))))→silu(x))Decompose — lowers ops without direct NISA equivalents into supported primitives (e.g.,
div(a,b)→mul(a, reciprocal(b)),reduce(mean)→reduce(sum) * 1/N)Layout Solver — assigns each tensor dimension to one of three roles:
Propagates constraints across the graph to find a globally consistent assignment.
Direct Lower — converts tensor IR ops to tiled NKI IR:
Hardware Target (
passes/hardware.py)Parameterized hardware profiles (TRN2 defaults) defining partition limits, SBUF/PSUM sizes, and matmul constraints.
Status
🚧 Work in progress — not ready for review.
Test plan
uv run pytest nkigen-lite/tests/ -n auto)