[Build] Target torch 2.10 + add the Triton codegen route (draft) - #305
Draft
YWHyuk wants to merge 35 commits into
Draft
[Build] Target torch 2.10 + add the Triton codegen route (draft)#305YWHyuk wants to merge 35 commits into
YWHyuk wants to merge 35 commits into
Conversation
Wire an opt-in second codegen route for the npu device: Inductor's own Triton
backend produces the kernel, and triton-npu (tnpu) lowers it to a RISC-V ELF,
in place of the hand-written MLIR emission in PyTorchSimFrontend/mlir.
Off by default. TORCHSIM_TRITON_CODEGEN=1 selects it at device registration;
the MLIR route is unchanged and still passes tests/ops/elementwise/test_add.py.
Flow, one module per seam:
scheduling.py TritonNPUScheduling.define_kernel intercepts the generated
Triton source; TritonNPUKernel emits a plain call instead
of triton's .run(grid=..., stream=...)
kernel_spec.py Inductor kernel -> tnpu KernelSpec: collects arg roles /
dtypes / numels while V.graph is live, strips the
triton_heuristics decorator and torch imports, pins the
block sizes
codecache.py one directory per source hash, mirroring extension_codecache
tnpu_bridge.py runs tnpu out of process
_triton_compat.py torch 2.8 Inductor against triton 3.6, no GPU
Measured today on x + y (1024 elements): Inductor codegen, spec generation and
tnpu stages 1-5 all succeed and link 05-triton_npu_fused_add_0.elf. The lowering
is structurally right -- three togsim.transfer ops, and Inductor's xmask arrives
as a masked DMA (masked_axes / masked_fill). The run then stops in
TritonNPULauncher.__call__, which is not wired: the functional launch needs
tensor marshalling into runtime/*.raw and the timing launch needs the build_tog
adapters. Both are tracked, with the rest of the gaps, in
PyTorchSimFrontend/triton_backend/README.md.
Three decisions are recorded there rather than in commit history: block sizes
are fixed at codegen time (Inductor defers the grid to a runtime autotuner that
this route has no equivalent of), tnpu runs in its own process (LLVM 23 vs 20
bindings cannot share an interpreter), and the frontend is shimmed onto triton
3.6 rather than the reverse (3.6 pins LLVM 23, which the IR seam requires).
…ton 3.6 triton-npu pins triton 3.6 because 3.6 pins LLVM 23, and both sides of its textual IR seam must be the same LLVM. torch 2.10 is the first release whose Inductor targets 3.6, so moving the image to it removes the version skew the Triton route had to shim around on 2.8. Dockerfile.base: torch 2.8.0 -> 2.10.0, torchvision 0.23.0 -> 0.25.0 (resolver-confirmed pair) add triton==3.6.0 -- Inductor imports triton while GENERATING a kernel, and it is not a dependency of the CPU torch wheels Editing Dockerfile.base is enough to re-tag the base image: the tag is the sha256 of thirdparty/github-releases.json plus that file, so the pin moves 1bb94039a884 -> dde0ff29a141 and CI builds the new base. The GitHub release pins themselves (gem5 / riscv-llvm / spike) are tied to ubuntu 22.04 and python 3.11, not to torch, and are unchanged. The rest of the pip set resolves clean against 2.10 (numpy 1.26.4, transformers 4.43.4, diffusers 0.34.0). _triton_compat shrinks to what is actually still needed. The triton_key injection is now guarded by _torch_handles_triton() and does not run on 2.10, which reaches the symbol through torch._inductor.runtime.triton_compat. What remains is not a version shim: triton_hash_with_backend() asks the triton runtime for the current target and raises "0 active drivers" on a box with no GPU, so it is short-circuited to a deterministic cache key. TritonNPUKernel.call_kernel takes **kwargs: 2.10 added a deallocate_ws parameter, and none of these apply to a route with no triton launcher. Verified against torch 2.10.0+cpu in an isolated venv: device C++ extension builds with 0 errors; eager add on npu:0 gives max_err 0 the MLIR production route still passes tests/ops/elementwise/test_add.py the Triton route still compiles through to 05-triton_npu_fused_add_0.elf Device-side API review found no conflicts: DeviceGuardImplInterface has the same 8 pure virtuals as on 2.8, c10::Allocator changed only additively, PrivateUse1HooksInterface only a comment typo, and every registration macro we use still exists. On the python side rename_privateuse1_backend, generate_methods_for_privateuse1_backend, register_interface_for_device and the DeviceInterface method set are unchanged; register_backend_for_device gained only optional parameters.
Stands up CI for the Triton codegen route, sourcing the harness from
PSAL-POSTECH/triton-npu.
Kept off the main CI path. The route is WIP, and its toolchain (LLVM 23, the
CONFIG_DESC-capable spike, the triton runtime) is ~1.8 GiB compressed that no
other job needs, so it is a separate image and a separate workflow rather than
part of torchsim_base.
thirdparty/triton-npu.json pins, deliberately NOT in github-releases.json:
that file plus Dockerfile.base is hashed into the
torchsim_base tag, so keeping these apart means
the main image is not rebuilt when tnpu moves
scripts/ci/tnpu_base_pin.sh sha256(manifest + Dockerfile.tnpu), same scheme
scripts/ci/tnpu_asset_env.sh resolves the three release asset ids
Dockerfile.tnpu torchsim_base -> torchsim_tnpu_base
triton_npu.yml preflight, image build, then the jobs
Jobs: tnpu-baselines runs the harness's own add/mul/relu/gemm/bmm through Spike
and gates, since nothing downstream means anything if the toolchain regressed.
triton-route runs test_triton_codegen.py and reports without gating, because its
launch is deliberately unimplemented. mlir-route-regression runs test_add.py --
this layer adds a second LLVM and a second triton to the image, and that job
checks the production path did not notice.
Dockerfile.tnpu mirrors setup/restore.sh --prebuilt rather than calling it, and
says why in place: that script clones a triton_shared fork that no longer
exists, and setup/package.sh collects the triton tree with
`find python -name '*.so'`, which misses python/triton/backends/{amd,nvidia,
triton_shared} -- symlinks the build creates, so present in neither the tarball
nor git, and without them `import triton` dies in entry-point discovery. Both
are worth fixing upstream; both are worked around here.
TWO PREREQUISITES, BOTH OUTSIDE THIS REPO, neither satisfiable from here:
1. secrets.TNPU_TOKEN -- PSAL-POSTECH/triton-npu is private and the default
Actions token is scoped to this repository, so it can neither clone the
repo nor read its releases. The existing gem5 / riscv-llvm / spike pins
need no secret because those repos are public.
2. A release tagged toolchain-llvm23 on PSAL-POSTECH/triton-npu carrying
llvm23-install.tar.gz, spike-install.tar.gz, triton-runtime.tar.gz. That
repo currently has NO releases -- the assets exist only on the fork it came
from and have to be mirrored across.
The preflight job reports which of the two is missing instead of letting a
docker build fail deep. Asset resolution was verified end to end against the
repo that does carry the release: it resolves all three ids, and against
PSAL-POSTECH it fails with the message above.
The two things Dockerfile.tnpu worked around are fixed upstream now (PSAL-POSTECH/triton-npu e686799): restore.sh no longer clones a triton_shared fork that does not exist, and it recreates the triton backend symlinks the release tarballs cannot carry. So the image just clones the harness and runs its own setup/restore.sh --prebuilt. Every pin lives in that repo's setup/versions.env, which also removes the duplicated triton pin and the asset-id resolution here (scripts/ci/tnpu_asset_env.sh deleted -- restore.sh resolves them). thirdparty/triton-npu.json pins `ref` to a commit rather than a branch, so an upstream change moves this image's tag.
The toolchain-llvm23 release now exists on PSAL-POSTECH/triton-npu, so secrets.TNPU_TOKEN is the only thing left to set up.
It prints nothing. The method walks the IR and attaches TOG nodes; `bfs` and `display` do the printing. The name came from the C++ pass this file ports, where one method does both -- the docstring now records that so the correspondence is still findable.
A Triton kernel describes a single program instance: the tile loop is not in the kernel, it is the launch grid outside it. The trace producer already wants that same split -- togsim_kernel_tile per work-item, enumerated by togsim_kernel (docs/design/togsim_cpp_trace.md sec 9.3) -- so the two models agree and only the enumeration was missing. This teaches the pipeline to accept a kernel in that shape, rather than requiring the loop nest PyTorchSim's codegen emits. build_tog _find_kernel falls back to the module's only func.func; the name comes from the Triton kernel, not the fixed "kernel". _build roots at top-level loops carrying a ROLE attribute, and treats the whole body as one work-item when there are none. Keying on the attribute matters on its own: bank_vectorize also leaves a bare top-level affine.for for the tile's vector work, and rooting at that one made every DMA a sibling the traversal never reached. A DMA's tensor identity comes from the producer's dram_arg attribute when the operand is a view of the argument rather than the argument. Inferring it would mean chasing memref view ops back, and there are nine of them with no ViewLikeOpInterface in the python bindings to ask generically. DMA nodes are recorded in a list so _collect_dma_nodes can seed from them; it only descended from loop nodes, and a DMA outside any loop was dropped before reaching the skeleton. lower_to_emitc WorkItem + _materialize_grid_loop supply the grid, on the trace artifact only: the body is wrapped in a loop per axis, tagged outer_loop, with each program-id argument replaced by its induction variable. It must run before _rewrite_signature, which erases the arguments and first asserts none are still used. Everything after is unchanged -- _parallel_loop_chain finds the tagged loop, the outliner threads the induction variable through iv[], and the loop left behind becomes the dispatch enumeration. Two things about building that nest are easy to get wrong and only show at rank >= 2: a nested loop is created before the enclosing yield (an InsertionPoint on a block appends, and an scf.for body is already terminated), and every bound is created before the first loop, so that a bound made after an outer loop does not end up below it while an inner loop uses it. A parallel loop may be scf.for as well as affine.for; the role is carried by the attribute, not the dialect. _strip_aux keeps the kernel the caller resolved instead of matching on a name. Every change is a fallback: the existing conditions are tested first, so PyTorchSim's own codegen takes exactly the path it did before.
Closes the timing loop: one torch.compile now produces a cycle count. The
launcher emits the trace producer from tnpu's post-vcix IR and hands it to
TOGSim, reusing PyTorchSim's existing trace pipeline unchanged.
timing.py
emit_trace 04-custom.mlir -> build_skeleton -> trace.so + trace_cycles.tsv
run_togsim hand the kernel directory to TOGSimulator.run_standalone
work_item_for derives the program-id argument positions from the signature
layout (pointers, user scalars, then triton-shared's own
gridX,Y,Z / pidX,Y,Z) and the grid from the pinned block sizes
codecache persists meta.json beside the artifacts so the timing step can run
standalone, and TritonNPULauncher.__call__ simulates instead of raising.
kernel_spec._grid becomes grid_of: the timing path needs the same extents to
enumerate work-items, so it is computed in one place.
The test drives a 2-D grid: that it verifies as MLIR, nests one loop per axis,
and dispatches both indices. Checking the module and not only the C++ it becomes
is the point -- the emitc lowering hoists constants to a flat scope, so it hides
a bound that does not dominate its use. Inductor cannot reach this path here (it
uses y/z only when x would overflow), so it gets a test rather than waiting for
a kernel to exercise it.
Measured on Inductor's `x + y` (1024 elements, XBLOCK 128, grid 8): TOGSim
totals 573 cycles, and channel-0 DRAM traffic of 16 reads x 32 B x 16 channels
is 8192 B -- exactly the 8 work-items x 2 loads x 512 B the kernel should move,
so every dispatch in the enumeration really ran. The cycle table is a
placeholder until gem5 sampling lands, and says so on every emit.
Output tensors are NOT written; marshalling them through Spike is the remaining
functional half. The launcher logs that on every call rather than letting an
undefined value pass for a computed one, and the test asserts the timing
artifacts exist instead of comparing values.
The cycle table held a placeholder, so TOGSim modelled the DMA but charged nothing for compute. It is now a measurement. measure_tile_cycles chains three pieces that already existed: build_tog's sample mode annotates the post-vcix IR (inline-asm markers around each compute node, every loop rewritten to one trip, so what runs is one tile), `python -m tnpu.cycle` lowers that to a RISC-V binary in tnpu's own process -- the Gemmini/VCIX lowering and its LLVM live there -- and CycleSimulator runs it under gem5, reading one numCycles per marker pair. build_cycle_table then turns the list into the tsv, keyed by tile_id. Sampling runs before build_skeleton because both consume the same post-vcix IR and build_skeleton rewrites it in place. Failure is not fatal: any step that does not produce a measurement falls back to the placeholder table and says so, since a kernel that simulates with the wrong compute cost is more useful than one that will not simulate -- as long as it announces which it is. Measured on `x + y` (1024 elements, grid 8): the tile samples at 19 cycles and TOGSim's total moves 573 -> 650. The MLIR route reports 251 on the same computation; the remaining gap is double buffering, which tnpu does not emit yet, so nothing overlaps.
The grid machinery handles N axes and is tested to, but grid_of only ever
computed one from xnumel, so work_item_for could only ever build a 1-D
WorkItem -- general plumbing behind a caller that never used it.
grid_of now walks every parallel axis, and fixed_config_for pins a block per
axis. Parallel vs reduction is Inductor's own test (a prefix starting with "r"
is looped inside the kernel, not gridded), and the block name is
f"{prefix.upper()}BLOCK", so neither is guessed.
Two orderings meet here and they are not the same: grid_of returns axes
OUTERMOST first (z, y, x -- x is Inductor's contiguous axis), while
triton-shared always appends the program ids as pidX, pidY, pidZ. work_item_for
therefore builds the argument list per axis instead of as a range; zipping the
two blindly would pair the outermost loop with the wrong id.
Block sizes: the outermost axis gets the lane count, because that is the tile
dimension bank_vectorize spreads over the lanes. The rest get 1, which leaves
the tile exactly the verified [lanes] shape and lets the grid cover everything
else. That is correct but pathological -- an inner block of 1 makes each
work-item move a strided column -- so a multi-axis kernel logs a warning saying
it is not a tiling worth measuring. Choosing real tile sizes is the block-size
policy gap in the README.
Verified: axes/grid/parallel_args come out ['x']/(8,)/[pidX],
['y','x']/(2,1024)/[pidY,pidX] and ['z','y','x']/[pidZ,pidY,pidX]. End-to-end on
a real multi-axis kernel is still untested -- Inductor reaches for y/z only when
x would overflow, which the shapes this route handles do not.
Only the NUMBER of grid axes has to be compiled in -- how many loops to nest and how many iv[] slots to fill. The trip counts are just values, and the producer ABI already carries them: togsim_kernel(ctx, shape_args, n). Baking them in was what forced a recompile per shape. A WorkItem extent of None now means "read it from shape_args". The bound cannot be wired when the loop is built -- shape_args does not exist until _rewrite_signature adds it -- so the loop takes a placeholder and _bind_runtime_bounds replaces it once the signature is there. The loops stay in the entry function (the outliner moves only their bodies), so the read is in scope where the bound is used. timing.write_shape computes the grid per launch and writes trace_shape.txt. The launch already knows the real extents: Inductor appends the numels after the tensor arguments, so the trailing values are them. Only the PARALLEL numels ride along -- a reduction axis is looped inside the kernel and never passed, so counting it would misalign the mapping. TOGSim reads that sidecar the same way it reads trace_cycles.tsv, from the kernel directory. main.cc passed nullptr for shape_args; absent file still means nullptr, so a producer with its bounds baked in is unaffected. Measured. One trace.so (md5 identical across all three), torch.compile( dynamic=True), a single kernel directory reused: n=1024 grid 8 650 cycles n=2048 grid 16 1316 cycles n=4096 grid 32 2586 cycles The functional path still compiles per shape -- tnpu's spec bakes the tensor extents into the flat memref view -- so this opens the timing half only.
A persistent reduction declares R0_BLOCK in the kernel BODY (R0_BLOCK: tl.constexpr = 64), not in its signature -- Inductor has already chosen it, and there is nothing left to pin. Passing it anyway put a name in the spec that the kernel does not take, and tnpu's stage 1 died on "'R0_BLOCK' is not in list" -- a malformed spec masquerading as a reduction problem. fixed_config entries absent from the signature are now skipped, so the "block size is unset" guard fires only when the block really is ours to choose. What that uncovers is the actual boundary: a reduction now reaches tnpu and fails THERE. triton-shared hands over a linalg.reduce plus a linalg.transpose, and tnpu has no lane-aware reduction to lower them with -- the scratchpad is lane-banked, so the reduced axis has to stay inside a lane. Measured on x.sum(dim=1): 02-ttshared carries linalg.reduce, linalg.transpose and tts.scatter, and stage 3 stops. The test pins that boundary rather than leaving it undocumented. It passes while a reduction is REFUSED; a reduction that compiles would mean either the lane path landed (delete the check) or the kernel is simulating compute the hardware would not do.
YWHyuk
force-pushed
the
feature/triton-codegen
branch
from
July 27, 2026 14:43
2208d64 to
8da4bbc
Compare
triton-shared keeps the user scalars in the lowered signature ahead of its own six grid/pid arguments: func.func @k(%arg0..2: memref<*xf32>, %arg3: i32 xnumel, %arg4,5,6: gridX,Y,Z, %arg7,8,9: pidX,Y,Z) tnpu's wrapper reads them from spec.extra["scalar_args"], and the generated spec had no extra at all, so the call passed six i32 where the kernel takes seven. Every argument after the tensors landed one slot early: pidX got pid_y, which the grid loop never varies, so program 0 ran eight times and only the first block of the output was written -- the rest stayed zero. Measured on x + y, 1024 elements: 896 of 1024 wrong, and the wrong part was silently zero rather than garbage. The values are compile-time constants, like the grid, so they come from the same numels the grid is computed from. On x + y this now gives 0.0 max error over the whole tensor under Spike. This only ever affected the functional path: the timing path reads the argument positions from the lowered MLIR signature, which was always right.
The launcher simulated the kernel but left the caller's output tensors untouched, so a compiled graph returned undefined values and the test could only assert that the timing artifacts existed. functional.py marshals the launch: every argument is written to runtime/<arg>.raw -- outputs as zeros, since the wrapper loads and dumps by argv position -- tnpu's stage 6 runs the ELF on Spike, and the output files are read back into the caller's tensors. Spike runs before TOGSim so the values survive a timing failure; the two halves are independent. Gated on pytorchsim_functional_mode, as the MLIR route is. With it off the launcher says so rather than letting an undefined value pass for a computed one. The binary is compiled for ONE shape: the spec bakes the grid, the scalar values and the memref extents in. A launch whose shapes differ is rejected with ShapeMismatch naming the compiled extent, instead of running against the wrong bounds. The timing path has no such limit -- it takes the grid at run time -- so the error points at pytorchsim_functional_mode: False for cycle-only studies. Measured: x + y and the fused (x + y) * 2 - x both match torch to 0.0 max error over 1024 elements. The test now checks values instead of reporting the gap.
Picks up tnpu.spike, the entry point the functional launch calls. Without it the launcher raises on every kernel once functional mode is on, which is the default.
triton_backend/README.md is the working reference for the modules: what each seam does and how to run it. What it does not answer is the question anyone reviewing the route asks first -- how it differs from the path already in production, and whether the numbers hold. The report covers that: where the two paths diverge (the kernel is a loop nest vs one tile) and where they rejoin (the same trace.so + trace_cycles.tsv, so TOGSim cannot tell them apart, which is why no hardware model changed), an item-by-item contrast, and the measurements behind each claim. Two findings worth having written down. The Triton route is AHEAD on dynamic shape: develop's trace path calls trace_to_tilegraph(..., nullptr, 0) and passes no shape arguments at all, so it recompiles per shape, and PR #269 is still open to fix that -- the Triton route takes the grid at run time and one trace serves every shape. It is BEHIND on DMA overlap: the generated IR carries three togsim.transfer and zero togsim.wait, which is the whole of the 251 vs 650 cycle gap, and the MLIR route already has the barrier machinery to copy. The two documents cross-reference rather than overlap.
The report read as though one person had built the whole route. Two thirds of what it describes is not ours: the lowering pass -- Triton IR through linalg/memref and the tts-level backend passes down to a RISC-V ELF -- is 이정민's, and this work is the PORT of that pass onto the existing PyTorchSim stack. The document now says so up front, in the pipeline diagram, in the comparison diagram and in the contrast table, so a reader cannot mistake the boundary wherever they enter. Naming follows the same split. The lowering layer is called the "PyTorchSim lowering pass" in prose, since that is what it is within this project. Concrete identifiers keep reading tnpu -- tnpu/passes/, tnpu.spike, the CI job names -- because the document has to stay navigable back to the code. triton-npu#1 is described for what it is: not a change to the pass, but the three entry points the port needed to call it. tnpu.cycle to time one tile under gem5, dram_arg so the TOG builder is told which argument a DMA belongs to instead of inferring it, and tnpu.spike so stage 6 can run on the caller's own tensors.
The report ended with four sections that were either backstory (design calls already reverted, PR bookkeeping) or a wish list. What a reader actually needs next is narrower: the MLIR route is held up by the op suite under tests/ops, and the question is how far the Triton route gets on the same ground. So those four are gone and one section replaces them, built from running the representative cases rather than predicting them. Two results were not what the guesses would have been. a @ b never enters the route at all -- Inductor emits an extern aten.mm, so the run passes with 0.0 error while touching no simulator. A probe that only checks values would have recorded that as coverage. The table therefore carries a "routed" column, taken from whether a work directory appeared. x.t() + 1 returns exactly x + 1: the transpose is dropped, 4030 of 4096 elements wrong, and nothing raises. This one is ours, not the lowering pass's. Inductor folds the transpose into the output buffer's stride and emits an identity index, so the kernel is right; functional.py is what assumes logical order equals storage order -- contiguous() on the way in, view_as on the way out. True for contiguous tensors, which is exactly the set that passed so far. Reordered so the contrast with the MLIR route comes first: what the route is differs from is easier to hold than how it is built, and the pipeline reads better once you know which half of it is new.
Six rounds of edits had left it patched rather than written: the boundary stated three times in three voices, a hook table stranded where a deleted section used to reference it, a summary that still read as though everything worked while the section below reported a silent wrong answer. Same material, rewritten straight through. The bug found while wiring the launch now sits inside the measurements it belongs to rather than standing as its own section, the two remaining top-level sections are the boundary and where the route actually reaches, and the coverage caveat is in the summary instead of only in the section that measures it. Nine sections became five.
…name Two lines still called it "tnpu stage 6" in running text, which reads as a separate tool rather than the lowering stage of this route. What stays is the set of names that exist in the code: tnpu_bridge.py, tnpu/passes/, tnpu.spike, tnpu.cycle, strip_for_tnpu. Renaming those in prose would make the document point at files nobody can open. The terminology note now lists them, so a reader knows the two vocabularies map onto one thing.
Naming it after PyTorchSim read as though the frontend owned it, which is the opposite of the boundary the document is drawing. It lowers Triton IR for this NPU and belongs to whoever owns that lowering, so name it for the target. Prose only. tnpu_bridge.py, tnpu/passes/, tnpu.spike, tnpu.cycle and strip_for_tnpu are real names and stay quoted as they are.
Section 3 named WorkItem and _materialize_grid_loop but never said what they reconcile, so the one piece of wiring the port actually turns on read as an implementation detail. Triton keeps the grid on the KernelSpec: grid_of computes the per-axis ceil-div from Inductor's numels and the pinned block sizes, and the kernel body only learns which slice it is, through pidX/Y/Z. build_tog expects the opposite -- it roots the TOG at a top-level affine.for carrying a role attribute, so it reads work-items out of a loop. A Triton kernel has no such loop, and without a root the TOG comes out empty. So the grid is stood back up as that loop: one nest per axis, the body moved inside, every use of a pid argument replaced by the induction variable, and outer_loop set on the result -- which is the marker build_tog looks for. Both sides then denote the same thing and the rest of the pipeline is shared unchanged.
The remaining work was framed as getting the op suite to pass, which is the milestone rather than the point. What the route has to reach is the models the MLIR path already runs. Section 5 now states that goal and splits the way there in two. Ops first, because a model is a composition of them and one blocked op stops the model at its first kernel. Models second, ordered by difficulty -- MLP, then MobileNet/ResNet, then ViT/Transformer, then Llama. The second stage gets its own list, because what it surfaces is not in the op tests at all: compile time once there are hundreds of kernels each spawning a lowering pass subprocess, buffer reuse between kernels, shapes that do not divide evenly, mixed f16/bf16, and the backward kernels on the training path. Marked as expectation, not measurement -- unlike the op table above it, none of it has been run yet.
The remaining-work section had drifted into implementation notes -- the two functional.py calls that mishandle stride, the pass that refuses a reduction, the helper the spec writer names when it stops. A reader reaching for what is left does not have that context and does not need it to understand the shape of the work. Each item is now one sentence: what is broken and roughly where. The op probe table keeps its verdicts but drops the error strings. Anyone who needs the detail has the code and the sections above. Section 5 goes from 85 lines to 55.
240 commits: the checkout layout drops the hexagon-mlir nesting, triton_shared repoints at the PSAL fork, and the lane-axis and reduction passes are rewritten. tnpu.spike and tnpu.cycle -- the two entry points this route calls -- are byte-identical to the pinned version, so functional launch and timing are unaffected. The pin also asks spike for zvfp8, which the toolchain release does not carry; an unknown extension stops spike at startup, failing every kernel rather than the fp8 one. Both jobs pin TNPU_SPIKE and TNPU_SPIKE_ISA=rv64gcv_zfh to the released build, which tnpu's own doctor() suggests. Costs kernels/coverage/ops/ops_fp8_roundtrip.py, which CI does not run. Verified on the new pin with the released spike: doctor clean, add/mul/relu err 0, gemm/bmm 1.52588e-05.
The triton checkout moves to /workspace/triton-src. A directory named `triton` shadowed the package for anything run from /workspace, so doctor died on AttributeError inside the image build -- Dockerfile.tnpu runs it from there.
Dockerfile.tnpu runs doctor as its last build step, and that run had no TNPU_SPIKE -- so the image build failed on the fp8 spike being absent even though every job that used the image passed it. Putting both variables in the image covers the build and the jobs at once, and removes the duplicate.
TORCHSIM_TRITON_CODEGEN is read once at device registration, so every test under tests/ is already a test of this route -- nothing in the tests needed to change. What was missing was a runner and somewhere for the failures to land. triton_route_sweep.py runs the suite with the variable set and splits the outcome three ways. triton_route_passing.txt is the gate: the tests that pass today, so coverage can only grow. The rest is a report, bucketed by the layer that owns the failure and by which tnpu stage the kernel reached. A test that passes WITHOUT emitting a kernel -- CPU-only, eager fallback, or an op Inductor sends to an extern call -- is kept out of the allowlist, since counting it would overstate coverage; matmul, bmm and topk all land there. Every failure leaves the Inductor kernel that was rejected, the stage IR it produced, stage.log and the error, so a report needs no rerun. That required one reorder in codecache: write_spec_file raises for exactly the kernels worth keeping (triton_helpers, SpecIncomplete), and it ran before kernel.py was written, so the interesting sources were the ones being thrown away. First measurement, 69 tests: 11 pass through the route, 5 pass without using it. Of the failures, 17 are missing test deps present in the CI image; the rest are 8 tnpu_stage, 7 triton_helpers, 7 spec_incomplete, 6 wrapper_gap, 3 device_op, 2 togsim. Jobs move to the PSAL Slurm runner farm: runs-on must carry the slurm label, image builds and the sweep on the big bucket.
Three diagnostics were losing the cause on the way out. tnpu writes its stage table to stdout and the diagnostic itself to stage.log, so TnpuError reported "exit 1" and nothing else. It now reads stage.log and carries the failing line, which is how six of the eight tnpu_stage failures turned out to be one bug. That bug: strip_for_tnpu drops `from torch...`, and Inductor takes libdevice and tl_math from torch._inductor.runtime.triton_helpers -- but those are re-exports of triton's own symbols. Dropping the import took valid triton names with it and the kernels died as NameError inside stage 1. tl_math is rebound from triton.language; test_pointwise now gets through fourteen ops and as far as the trace producer instead of failing on the first. libdevice cannot be rebound -- its members are @core.extern with no triton_shared implementation, so a call returns None -- and is now named the way triton_helpers is. The sweep also gets -j: tests are independent subprocesses with their own dump dir, Inductor cache and TOGSim FIFO, so they parallelise with no coordination. 69 tests go from ~50 minutes to 5 at -j 10. Coverage is unchanged at 11 of 69, but the failures are better attributed: tnpu_stage 8 -> 2, spec_incomplete 7 -> 13.
11 of 69 tests pass through the route, with every failure attributed to an owning layer and the pipeline stage it reached. Ranked next steps come from the measured unblock counts, not estimates.
Adds the softmax kernel end to end as a worked example, the Python -> ttir -> diagnostic chain for both tnpu pass rejections, and test_cat's togsim.transfer ops. Also records that the two Spike failures cannot be reproduced by hand today: tnpu.spike reports only exit 255, and write_inputs rewrites the .raw files per launch, so a by-hand replay uses stale inputs.
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.
Targets torch 2.10 and adds an opt-in second codegen route: Inductor's own
Triton backend produces the kernel, and triton-npu lowers it to a RISC-V ELF.
Draft. The route now runs end to end on elementwise kernels — values and
cycles both — but the gap list below is still long, and the environment change
needs a base image rebuild and a full test-suite run against torch 2.10.
Why torch 2.10
triton-npu pins triton 3.6 because 3.6 pins LLVM 23, and both sides of its
textual IR seam must be the same LLVM. torch 2.10 is the first release whose
Inductor targets 3.6, so the version skew the Triton route had to shim around on
2.8 disappears.
Editing
Dockerfile.basemoves the base-image tag automatically(
sha256(thirdparty/github-releases.json + Dockerfile.base)), so CI buildstorchsim_base:thirdparty-dde0ff29a141on this PR. The GitHub release pins(gem5 / riscv-llvm / spike) are tied to ubuntu 22.04 and python 3.11, not to
torch, and are unchanged. The rest of the pip set resolves clean (numpy 1.26.4,
transformers 4.43.4, diffusers 0.34.0).
Device-side review
No conflicts found, verified against torch 2.10.0+cpu in an isolated venv:
addonnpu:0gives max_err 0DeviceGuardImplInterfacehas the same 8 pure virtuals as on 2.8c10::Allocatorchanged only additively;PrivateUse1HooksInterfaceonly acomment typo; every registration macro we use still exists
rename_privateuse1_backend,generate_methods_for_privateuse1_backend,register_interface_for_deviceand the
DeviceInterfacemethod set are unchanged;register_backend_for_devicegained only optional parameterstests/ops/elementwise/test_add.pyOnly one signature moved:
TritonKernel.call_kernelgaineddeallocate_ws, soour override takes
**kwargs.The Triton route
Off by default;
TORCHSIM_TRITON_CODEGEN=1selects it at device registration.One module per seam, described in
PyTorchSimFrontend/triton_backend/README.md.Measured on
x + y(1024 elements): Inductor codegen, spec generation and tnpustages 1-5 all succeed and link
05-triton_npu_fused_add_0.elf. The lowering isstructurally right — three
togsim.transferops, and Inductor'sxmaskarrivesas a masked DMA.
The launch has both halves:
trace.so+trace_cycles.tsvthatTOGSim consumes. gem5 samples the tile at 19 cycles; TOGSim reports 650,
with DRAM traffic of exactly the 8 work-items × 2 loads × 512 B the kernel
should move. The MLIR route on the same
x + yreports 251 — the same order,higher here because tnpu emits synchronous DMA, so nothing overlaps.
.raw, Spike runs theELF, and the outputs are read back into the caller's tensors.
x + yand thefused
(x + y) * 2 - xboth match torch to 0.0 max error.A Triton kernel body is one program instance, so the grid that enumerates
instances is supplied by
lower_to_emitc.WorkItemrather than read out of thekernel. The timing path takes it at run time, so one trace serves every shape.
The functional binary does not: tnpu bakes the grid, the scalar values and the
memref extents in, so a dynamic-shape launch is rejected with
ShapeMismatchnaming the compiled extent instead of running against the wrong bounds.
Still open, in order: shape-specialised functional launch, double buffering,
triton_helpers, reductions (blocked in tnpu — no lane-aware reduction path),and matmul timing.
tests/system/test_triton_codegen.pypins the reductionboundary so it fails loudly rather than producing wrong numbers.
Depends on PSAL-POSTECH/triton-npu#1 for the timing and functional entry points;
thirdparty/triton-npu.jsonpins the commit.CI for the route
.github/workflows/triton_npu.yml, separate from the main CI: the toolchainlayer is ~1.8 GiB no other job needs. It will fail at
preflightuntilsecrets.TNPU_TOKENis set —PSAL-POSTECH/triton-npuis private and thedefault Actions token is scoped to this repository. That is expected on this PR.
What to expect from CI here
docker-image.ymlbuilds the new base image and runs the full suite threetimes (3 TOGSim configs) against torch 2.10. Only
test_addwas verifiedlocally, so other jobs may surface 2.10 regressions — that signal is the
reason this PR is open.
triton_npu.ymlbuilds the toolchain image and runstests/system/test_triton_codegen.py.🤖 Generated with Claude Code