From 4c077b471445d18d00b0668bf08455ca0e694395 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 17:37:52 +0900 Subject: [PATCH 01/35] [Frontend] Add the Triton codegen route (Inductor triton + tnpu passes) 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). --- PyTorchSimDevice/torch_openreg/__init__.py | 29 +- PyTorchSimFrontend/extension_config.py | 15 + PyTorchSimFrontend/triton_backend/README.md | 104 ++++++ PyTorchSimFrontend/triton_backend/__init__.py | 37 ++ .../triton_backend/_triton_compat.py | 132 +++++++ .../triton_backend/codecache.py | 80 +++++ .../triton_backend/kernel_spec.py | 337 ++++++++++++++++++ .../triton_backend/scheduling.py | 88 +++++ .../triton_backend/tnpu_bridge.py | 82 +++++ .../triton_backend/wrapper_codegen.py | 21 ++ tests/system/test_triton_codegen.py | 65 ++++ 11 files changed, 984 insertions(+), 6 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/README.md create mode 100644 PyTorchSimFrontend/triton_backend/__init__.py create mode 100644 PyTorchSimFrontend/triton_backend/_triton_compat.py create mode 100644 PyTorchSimFrontend/triton_backend/codecache.py create mode 100644 PyTorchSimFrontend/triton_backend/kernel_spec.py create mode 100644 PyTorchSimFrontend/triton_backend/scheduling.py create mode 100644 PyTorchSimFrontend/triton_backend/tnpu_bridge.py create mode 100644 PyTorchSimFrontend/triton_backend/wrapper_codegen.py create mode 100644 tests/system/test_triton_codegen.py diff --git a/PyTorchSimDevice/torch_openreg/__init__.py b/PyTorchSimDevice/torch_openreg/__init__.py index e8158391..2667ac70 100644 --- a/PyTorchSimDevice/torch_openreg/__init__.py +++ b/PyTorchSimDevice/torch_openreg/__init__.py @@ -18,13 +18,30 @@ sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')) import PyTorchSimFrontend.extension_config # noqa: F401 +from PyTorchSimFrontend import extension_config as _extension_config from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen -from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling -torch._inductor.codegen.common.register_backend_for_device( - "npu", - lambda scheduling: MLIRScheduling(scheduling), - ExtensionWrapperCodegen -) + +# Two mutually exclusive codegen routes for `npu`, chosen here because Inductor +# registers a backend per device, once. +# MLIR (default) hand-written MLIR emission, PyTorchSimFrontend/mlir +# Triton (opt-in) Inductor's own Triton codegen + the triton-npu passes, +# TORCHSIM_TRITON_CODEGEN=1. WIP; see +# PyTorchSimFrontend/triton_backend/README.md +if _extension_config.CONFIG_TRITON_CODEGEN: + from PyTorchSimFrontend.triton_backend import ( + TritonNPUScheduling, TritonNPUWrapperCodegen) + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: TritonNPUScheduling(scheduling), + TritonNPUWrapperCodegen + ) +else: + from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: MLIRScheduling(scheduling), + ExtensionWrapperCodegen + ) torch_openreg.openreg.init() sys.modules['torch.npu'] = torch_openreg.openreg diff --git a/PyTorchSimFrontend/extension_config.py b/PyTorchSimFrontend/extension_config.py index 2d706bbc..09e5168d 100644 --- a/PyTorchSimFrontend/extension_config.py +++ b/PyTorchSimFrontend/extension_config.py @@ -11,6 +11,21 @@ CONFIG_TORCHSIM_DUMP_MLIR_IR = int(os.environ.get("TORCHSIM_DUMP_MLIR_IR", default=False)) CONFIG_TORCHSIM_DUMP_LLVM_IR = int(os.environ.get("TORCHSIM_DUMP_LLVM_IR", default=False)) +# --- Triton codegen route (WIP, opt-in) -------------------------------------- +# Replaces the hand-written MLIR emission in PyTorchSimFrontend/mlir with +# Inductor's own Triton codegen, lowered to the NPU by the triton-npu (tnpu) +# pass pipeline. OFF by default: the MLIR route stays the production path until +# this one is complete. See PyTorchSimFrontend/triton_backend/README.md. +CONFIG_TRITON_CODEGEN = bool(int(os.environ.get("TORCHSIM_TRITON_CODEGEN", default=0))) +# The triton-npu checkout that owns stages 1-5 (ttir -> ttshared -> tnpu passes +# -> RISC-V ELF). It is a SEPARATE repository, deliberately not vendored. +CONFIG_TNPU_DIR = os.environ.get( + "TNPU_DIR", default=os.path.join(CONFIG_TORCHSIM_DIR, "triton-npu")) +# tnpu runs in its own process: its passes need LLVM 23's MLIR bindings while +# this process holds LLVM 20's, and `mlir` is a namespace package, so the two +# cannot coexist in one interpreter (tnpu/config.py:activate_bindings). +CONFIG_TNPU_PYTHON = os.environ.get("TNPU_PYTHON", default=sys.executable) + def get_dump_path(): """Resolve TORCHSIM_DUMP_PATH and re-point Inductor's cache dir at it. diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md new file mode 100644 index 00000000..ed286c0f --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -0,0 +1,104 @@ +# Triton codegen route (WIP) + +Replaces the hand-written MLIR emission in `PyTorchSimFrontend/mlir/` with +**Inductor's own Triton codegen**, lowered to this NPU by the **triton-npu** +(`tnpu`) pass pipeline. Opt-in and off by default; the MLIR route is untouched +and stays the production path. + +```bash +TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py +``` + +## Why + +The MLIR route does not just emit loops — it hand-implements the whole hardware +mapping (tiling, vectorization, DMA, scratchpad, lane distribution) as ~5,500 +lines of Python string emission, which entangles *what to compute* with *how to +map it*. See `docs/linalg-codegen-migration.md` for the long form. + +This route keeps Inductor for the first and triton-npu for the second: + +| | owns | +|---|---| +| Inductor (upstream) | fusion, index expressions, masking, reductions, the kernel source | +| triton-shared | Triton IR -> `linalg` / `tts` pointer descriptors | +| tnpu passes | `tts` -> `togsim.transfer` DMA, scratchpad, lane-banked vectors, systolic array | + +## Flow + +``` +torch.compile + └ TritonNPUScheduling.define_kernel scheduling.py + │ Inductor's triton kernel SOURCE TEXT + collected metadata + ▼ + triton_npu_compile(src, meta, name) codecache.py + │ a tnpu kernel file (KernelSpec) kernel_spec.py + ▼ + run.py --to binary (subprocess) tnpu_bridge.py + │ 01-ttir → 02-ttshared → 03-adapted → 04-lowered → 05-*.elf + ▼ + TritonNPULauncher.__call__ ← NOT WIRED YET +``` + +Artifacts land in one directory per source hash under the dump path +(`outputs/triton_/`), alongside the unmodified Inductor source +(`kernel.py`) so the rewrite is diffable. + +## What works today (measured) + +`x + y`, 1024 elements, on `npu:0`: + +- Inductor generates the Triton kernel and our `define_kernel` intercepts it +- `kernel_spec` pins `XBLOCK` = lane count, computes `grid = (8,)`, writes the spec +- tnpu runs stages 1–5 and links **`05-triton_npu_fused_add_0.elf`** (20 B/lane spad) +- the lowering is correct in shape: `tl.load/store` became three + `togsim.transfer` ops, and Inductor's `xmask` came through as a **masked DMA** + (`masked_axes = [0]`, `masked_fill`), which tnpu already supports +- the run stops in `TritonNPULauncher.__call__`, by design + +## Gap list, in order + +1. **Launch (functional).** Marshal the caller's tensors into + `runtime/*.raw`, run Spike on the ELF, read outputs back. tnpu's stage 6 does + this for its own kernels but generates inputs from the spec; here the tensors + come from the caller. +2. **Launch (timing).** Emit `trace.so` + `trace_cycles.tsv` and hand them to + TOGSim. Blocked on the `build_tog` adapters — the tnpu IR is structurally + invisible to it today (no top-level `affine.for`, `scf.for` instead of + `affine.for`, vcix as LLVM intrinsics rather than dialect ops, DMA addresses + as `arith` chains rather than `affine.apply`, grid outside the IR). +3. **`triton_helpers`.** Any kernel using `triton_helpers.*` (reductions, + clamps, `maximum`/`minimum`) cannot compile: the module lives in torch and + the tnpu venv has none. `strip_for_tnpu` raises and names the helper. Needs a + minimal vendored copy. +4. **Reductions.** Independently blocked in tnpu itself — no lane-aware + reduction path; see `triton-npu/kernels/reduce.py`. +5. **Block-size policy.** `fixed_config_for` pins `XBLOCK` to the lane count and + deliberately leaves reduction blocks unset. Real tile selection (the MLIR + route's autotuner / `codegen_mapping_strategy`) has no equivalent here yet. +6. **Dynamic shapes.** `collect_meta` resolves numels through `size_hint`; a + genuinely dynamic dim gives `None` and `_grid` raises. + +## Three design decisions + +**Block sizes are fixed at codegen time.** Inductor defers the grid to +`triton_heuristics` at runtime (`grid = cdiv(xnumel, XBLOCK)` after autotuning). +tnpu compiles one binary ahead of time and walks the grid as a sequential loop in +generated C, so there is nothing to autotune later and no runtime `grid=` +callable. Pinning the config is what makes the launch shape statically +describable — the premise of this route, not a shortcut. (`kernel_spec.fixed_config_for`) + +**tnpu runs in its own process.** Its passes need LLVM 23's MLIR bindings while +this process holds LLVM 20's, and `mlir` is a namespace package, so two LLVMs in +one interpreter silently merge. The seam between them is a file, and that is +measured to work: LLVM 23 prints IR that LLVM 20's bindings parse without +complaint. (`tnpu_bridge`) + +**torch 2.8's Inductor is shimmed onto triton 3.6.** Inductor targets the +triton ~3.3 API; tnpu pins 3.6 because 3.6 pins LLVM 23, and both sides of the +IR seam must be the same LLVM. So the frontend bends: `triton_key` is injected +back into `triton.compiler.compiler`, and `triton_hash_with_backend` is +short-circuited because it asks a GPU driver for the current target and there is +no GPU. The alternative — a second, torch-compatible triton in the driver +interpreter purely for codegen — stays open if the API drift widens beyond this +one symbol. (`_triton_compat`) diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py new file mode 100644 index 00000000..f5cf1b60 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -0,0 +1,37 @@ +"""The Triton codegen route: Inductor's Triton backend + the tnpu lowering passes. + +WHAT THIS REPLACES +------------------ +The production path emits MLIR by hand (`PyTorchSimFrontend/mlir/`, ~5,500 lines +of string emission that also decides tiling, vectorization, DMA and scratchpad +placement). This route keeps Inductor's OWN Triton codegen for "what to compute" +and hands the resulting Triton kernel to triton-npu for "how to map it onto the +NPU": + + Inductor -> TritonNPUScheduling.define_kernel (scheduling.py) + | triton kernel SOURCE TEXT + v + -> TritonNPUCodeCache.load (codecache.py) + | a tnpu KernelSpec file (kernel_spec.py) + v + -> triton-npu, in a subprocess (tnpu_bridge.py) + ttir -> ttshared -> tnpu passes -> RISC-V ELF + v + -> Spike (functional) / TOGSim (timing) + +The two routes are mutually exclusive and chosen at device-registration time by +`extension_config.CONFIG_TRITON_CODEGEN` (env `TORCHSIM_TRITON_CODEGEN=1`). +Default OFF -- nothing here is on the production path yet. + +STATUS: scaffolding. The seams are wired and each stage names precisely what it +still owes; see README.md for the gap list. Expect failures, not results. +""" + +from . import _triton_compat + +# Before anything imports Inductor's Triton codegen: it needs `triton` in THIS +# interpreter, and torch 2.8 expects a triton API older than the 3.6 tnpu pins. +_triton_compat.install() + +from .scheduling import TritonNPUScheduling # noqa: E402,F401 +from .wrapper_codegen import TritonNPUWrapperCodegen # noqa: E402,F401 diff --git a/PyTorchSimFrontend/triton_backend/_triton_compat.py b/PyTorchSimFrontend/triton_backend/_triton_compat.py new file mode 100644 index 00000000..c97fe850 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -0,0 +1,132 @@ +"""Make torch 2.8's Inductor codegen work against tnpu's triton 3.6, CPU-only. + +TWO SKEWS, BOTH STRUCTURAL +-------------------------- +1. VERSION. Inductor in torch 2.8 targets the triton ~3.3 API. tnpu pins triton + 3.6 and that pin is not negotiable: 3.6 is what pins LLVM 23, and both sides + of the textual IR seam must be the same LLVM (triton-npu/setup/versions.env). + So the frontend has to bend, not the backend. + +2. NO GPU. `triton_hash_with_backend()` asks the triton runtime driver for the + *current target*, which on a machine with no GPU has nothing to answer with. + We never launch through triton's runtime -- the kernel is compiled ahead of + time to a RISC-V ELF -- so the value is only a cache-key ingredient. + +Both are handled by replacing `torch.utils._triton.triton_hash_with_backend` +with a deterministic string. It is a monkeypatch, and it is the cheap half of a +real choice: the alternative is installing a torch-2.8-compatible triton (3.3.x) +in the driver interpreter purely for codegen, keeping 3.6 in the tnpu venv for +stage 1. That works because the two interpreters exchange only SOURCE TEXT and +never share objects -- but it means two tritons to keep straight, so it is worth +doing only if the API drift turns out to be wider than this one symbol. +""" + +import functools +import hashlib +import importlib +import os +import sys + +_installed = False + + +def triton_src_dir(): + """Where tnpu's triton checkout lives (its editable install points here). + + Read out of tnpu's own `setup/versions.env` rather than guessed, so the two + repos cannot drift: that file is the single place the checkout layout is + pinned (HEXAGON_MLIR_ROOT). + """ + from PyTorchSimFrontend import extension_config + override = os.environ.get("TNPU_TRITON_SRC") + if override: + return override + + root = "/workspace/hexagon-mlir" + versions = os.path.join(extension_config.CONFIG_TNPU_DIR, "setup", "versions.env") + try: + with open(versions) as f: + for line in f: + if line.startswith("HEXAGON_MLIR_ROOT="): + root = line.split("=", 1)[1].strip() + break + except OSError: + pass + return os.path.join(root, "triton", "python") + + +def ensure_triton_importable(): + """`import triton` in THIS interpreter, borrowing tnpu's checkout if needed. + + Inductor's Triton codegen imports triton at codegen time (for metadata and + hashing), so the driver needs it even though it never compiles with it. + """ + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + cand = triton_src_dir() + if os.path.isdir(os.path.join(cand, "triton")): + sys.path.insert(0, cand) + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + return False + + +def _stable_backend_hash(): + try: + import triton + version = triton.__version__ + except Exception: # noqa: BLE001 + version = "unknown" + key = f"pytorchsim-tnpu-{version}" + return hashlib.sha256(key.encode("utf-8")).hexdigest().upper() + + +def _needs_hash_patch(): + """True when triton_hash_with_backend cannot work here.""" + try: + mod = importlib.import_module("triton.compiler.compiler") + except Exception: # noqa: BLE001 + return True + return not hasattr(mod, "triton_key") + + +def install(): + """Idempotently apply the shims. Returns a short report for logging.""" + global _installed + notes = [] + if not ensure_triton_importable(): + raise ModuleNotFoundError( + f"the Triton codegen route needs `triton` importable in this " + f"interpreter (Inductor imports it during codegen). Not found, and " + f"no checkout at {triton_src_dir()}. Set TNPU_TRITON_SRC, or install " + f"triton into this environment.") + if _installed: + return notes + + if _needs_hash_patch(): + # `triton_key` is imported from triton.compiler.compiler by SEVERAL torch + # call sites (codecache.CacheBase.get_system, _triton.triton_hash_with_ + # backend, ...), each with its own local import. Supplying the symbol on + # the triton side satisfies all of them at once instead of chasing every + # call site; it is a cache-key ingredient, so any stable string will do. + mod = importlib.import_module("triton.compiler.compiler") + mod.triton_key = _stable_backend_hash + notes.append("injected triton.compiler.compiler.triton_key " + "(removed in triton 3.6; torch 2.8 still imports it)") + + # Separately: triton_hash_with_backend also asks the triton runtime driver + # for the current target, which needs a GPU. We compile ahead of time to a + # RISC-V ELF and never use triton's runtime, so short-circuit it. + import torch.utils._triton as _t + _t.triton_hash_with_backend = functools.cache(_stable_backend_hash) + notes.append("patched torch.utils._triton.triton_hash_with_backend " + "(no GPU target to query)") + + _installed = True + return notes diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py new file mode 100644 index 00000000..b4817cc6 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -0,0 +1,80 @@ +"""Compile cache for the Triton route -- the counterpart of extension_codecache. + +`triton_npu_compile` is what the generated wrapper calls, exactly where the MLIR +route calls `custom_async_compile.mlir(...)`. It compiles the Triton kernel via +tnpu and returns the callable the wrapper then invokes per launch. + + define_kernel -> triton_npu_compile(src, meta, kernel_name) -> launcher + call site -> launcher(arg0, arg1, ..., xnumel) + +Layout mirrors the MLIR route so the two are comparable: one directory per source +hash under the dump path, holding the generated tnpu kernel file and every tnpu +artifact (01-ttir.mlir ... 05-*.elf). +""" + +import os + +from filelock import FileLock +from torch._inductor.codecache import get_hash + +from PyTorchSimFrontend import extension_config +from . import kernel_spec, tnpu_bridge + +logger = extension_config.setup_logger() + +LOCK_TIMEOUT = 600 + + +def _write_path(src_code): + return os.path.join(extension_config.get_dump_path(), + "triton_" + get_hash(src_code.strip())[1:12]) + + +class TritonNPULauncher: + """What a compiled kernel name is bound to in the generated wrapper. + + Holds the compile result; each call is one launch of the whole grid. + """ + + def __init__(self, kernel_name, workdir, meta): + self.kernel_name = kernel_name + self.workdir = workdir + self.meta = meta + self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") + + def __call__(self, *args): + raise NotImplementedError( + f"{self.kernel_name}: compiled to {self.elf}, but the launch is not " + f"wired yet. Two pieces are missing and both are tracked in " + f"triton_backend/README.md:\n" + f" 1. functional -- marshal the caller's tensors into " + f"{self.workdir}/runtime/*.raw, run Spike on the ELF, read the " + f"outputs back into the caller's tensors;\n" + f" 2. timing -- emit trace.so + trace_cycles.tsv from the tnpu IR " + f"and hand them to TOGSim (needs the build_tog adapters).\n" + f"Compilation itself succeeded, so the codegen half of this route " + f"is exercised by getting this far.") + + +def triton_npu_compile(src_code, meta, kernel_name): + """Compile one Inductor-generated Triton kernel through tnpu. + + Called from the generated wrapper at module import time (same point as + `custom_async_compile.mlir`). Synchronous for now: the MLIR route's thread + pool buys nothing until the pipeline itself is proven. + """ + write_path = _write_path(src_code) + os.makedirs(write_path, exist_ok=True) + + lock = FileLock(os.path.join(write_path, ".compile.lock"), timeout=LOCK_TIMEOUT) + with lock: + spec_path = os.path.join(write_path, f"{kernel_name}_spec.py") + elf = os.path.join(write_path, f"05-{kernel_name}.elf") + if not os.path.isfile(elf): + kernel_spec.write_spec_file(src_code, meta, spec_path, + tnpu_bridge.tnpu_dir()) + with open(os.path.join(write_path, "kernel.py"), "w") as f: + f.write(src_code) # the unmodified Inductor source, for diffing + tnpu_bridge.run_pipeline(spec_path, write_path, to_stage="binary") + logger.info("[triton-npu] %s -> %s", kernel_name, write_path) + return TritonNPULauncher(kernel_name, write_path, meta) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py new file mode 100644 index 00000000..ef779913 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -0,0 +1,337 @@ +"""Inductor kernel -> tnpu KernelSpec. + +Two jobs, both of which exist because Inductor's Triton output is written for a +GPU launcher and tnpu's is written for a static, ahead-of-time pipeline: + +1. `collect_meta` -- pull everything tnpu needs out of the Inductor kernel while + we still have `V.graph`: argument names/roles/dtypes/sizes, the constexprs, + and the numels the grid is computed from. This runs at codegen time; by the + time the compile callable fires, `V.graph` is gone. + +2. `write_spec_file` -- turn the Triton source + that metadata into a kernel file + tnpu can load (`tnpu.spec.load_spec`). + +WHY THE SOURCE HAS TO BE REWRITTEN +---------------------------------- +Inductor emits, above the kernel: + + from torch._inductor.runtime import triton_heuristics + @triton_heuristics.pointwise(size_hints={'x': 1024}, ..., inductor_meta=...) + @triton.jit + def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): + +Neither line can survive into tnpu: + + * the tnpu triton venv has NO torch (deliberately -- tnpu/spec.py), so the + `torch._inductor.runtime` import fails on sight; + * `triton_heuristics.pointwise` is the AUTOTUNER. It picks XBLOCK at runtime + and derives `grid = cdiv(xnumel, XBLOCK)` from it. tnpu needs both to be + constants: the block size becomes a `tl.constexpr` in the ttir signature and + the grid is executed as a sequential loop by the generated C wrapper. + +So the decorator is stripped and XBLOCK is pinned as a constexpr instead. That is +not a workaround -- fixing the config at codegen time is what makes the kernel +statically describable, which is the whole premise of this route. +""" + +import math +import os +import re + +from torch._inductor.virtualized import V + +#: Triton signature token -> (torch dtype name, bytes). Only the dtypes +#: tnpu/wrapper.py can round-trip through .raw files. +_DTYPE = { + "*fp32": "float32", "*fp16": "float16", "*bf16": "bfloat16", + "*i64": "int64", "*i32": "int32", "*i8": "int8", "*i1": "bool", + "fp32": "float32", "i32": "int32", "i64": "int64", +} + + +class SpecIncomplete(RuntimeError): + """Metadata tnpu requires that this kernel did not provide. + + Raised with the missing field named, rather than writing a spec that fails + deeper in the pipeline where the cause is unrecoverable. + """ + + +# --------------------------------------------------------------------------- +# 1. codegen-time metadata capture +# --------------------------------------------------------------------------- +def _buffer_numel(name): + """Element count of an Inductor buffer, or None if it cannot be resolved.""" + try: + buf = V.graph.get_buffer(name) + if buf is None: + return None + size = buf.get_layout().size + n = 1 + for s in size: + n *= int(V.graph.sizevars.size_hint(s)) + return n + except Exception: # noqa: BLE001 - best effort; caller reports it as missing + return None + + +def _roles(kernel): + """arg name -> 'in' | 'out' | 'inout', from the kernel's buffer tables.""" + out = {} + for buf, arg in getattr(kernel.args, "input_buffers", {}).items(): + out[arg] = ("in", buf) + for buf, arg in getattr(kernel.args, "output_buffers", {}).items(): + out[arg] = ("out", buf) + for buf, arg in getattr(kernel.args, "inplace_buffers", {}).items(): + name = getattr(arg, "inner_name", arg) + out[name] = ("inout", buf) + return out + + +def collect_meta(kernel, kernel_name): + """Everything the compile step needs, as plain repr-able data. + + Must run while `V.graph` is live (i.e. inside define_kernel). + """ + triton_meta = dict(getattr(kernel, "triton_meta", None) or {}) + signature = dict(triton_meta.get("signature") or {}) + constants = dict(triton_meta.get("constants") or {}) + + roles = _roles(kernel) + arg_defs, _call_args, _precompile, _arg_types = kernel.args.python_argdefs() + + args = [] + for a in arg_defs: + name = getattr(a, "name", str(a)) + role, buf = roles.get(name, (None, None)) + if role is None: + continue # a numel / constexpr, not a tensor + args.append({ + "name": name, + "role": role, + "buffer": buf, + "dtype": _DTYPE.get(signature.get(name, ""), None), + "numel": _buffer_numel(buf) if buf else None, + }) + + # The numels Inductor appends to the call. They live in `kernel.numels`, + # keyed by iteration-space PREFIX ('x', 'y', 'r0', ...), not as xnumel/rnumel + # attributes (SIMDKernel.__init__ builds them from the tiling). These are + # what the grid is computed from. + numels = {} + for prefix, val in (getattr(kernel, "numels", None) or {}).items(): + try: + numels[f"{prefix}numel"] = int(V.graph.sizevars.size_hint(val)) + except Exception: # noqa: BLE001 - dynamic shape; reported by _grid + numels[f"{prefix}numel"] = None + + return { + "kernel_name": kernel_name, + "signature": {str(k): str(v) for k, v in signature.items()}, + "constants": {str(k): v for k, v in constants.items()}, + "args": args, + "numels": numels, + "inside_reduction": bool(getattr(kernel, "inside_reduction", False)), + "fixed_config": fixed_config_for(kernel), + } + + +def fixed_config_for(kernel): + """Block sizes pinned at codegen time. + + tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a + sequential loop, so there is no autotuner to choose XBLOCK later and no + runtime `grid=` callable. Fixing it here is what makes the launch shape + static. + + The lane count is the natural default: `bank_vectorize` distributes tile + dim 0 across the lanes, and a block equal to the lane count gives a per-lane + depth of 1 -- the case every tnpu baseline runs today. + """ + from PyTorchSimFrontend import extension_config + lanes = int(extension_config.vpu_num_lanes) + cfg = {"XBLOCK": lanes} + if getattr(kernel, "inside_reduction", False): + # A reduction block is NOT free to be the lane count: the reduced axis + # has to stay inside a lane (see triton-npu kernels/reduce.py). Left + # unset on purpose so the reduction path fails loudly rather than + # silently picking a layout the hardware cannot execute. + cfg["R0_BLOCK"] = None + return cfg + + +# --------------------------------------------------------------------------- +# 2. Triton source -> tnpu kernel file +# --------------------------------------------------------------------------- +_HEURISTIC_RE = re.compile(r"^@triton_heuristics\.") +_DROP_IMPORT_RE = re.compile( + r"^\s*(import torch|from torch\b|from __future__|import __main__)") +#: GPU-only runtime setup Inductor emits at module scope. Meaningless here (the +#: kernel is compiled ahead of time to a RISC-V ELF) and its import is dropped +#: above, so the call would be a NameError. +_DROP_CALL_RE = re.compile(r"^\s*triton_helpers\.set_driver_to_gpu\(\)") +#: Anything else from triton_helpers is a real dependency -- maximum/minimum/ +#: promote_to_tensor and friends, which reductions and clamps use constantly. +_HELPER_USE_RE = re.compile(r"\btriton_helpers\.(\w+)") + + +def strip_for_tnpu(src): + """Remove everything the torch-free tnpu venv cannot import. + + Drops torch/inductor imports and the `@triton_heuristics.*(...)` decorator + (keeping `@triton.jit`), then re-adds the two imports the kernel body needs. + + Raises SpecIncomplete if the kernel still calls into `triton_helpers`: that + module lives in torch, so it has to be vendored into the tnpu venv before + such a kernel can compile. Failing here names the missing helper; letting it + through fails as a bare NameError inside tnpu's stage-1 worker instead. + """ + lines = src.splitlines() + out, i = [], 0 + while i < len(lines): + line = lines[i] + if _HEURISTIC_RE.match(line.strip()) or _HEURISTIC_RE.match(line): + # skip the whole decorator call, up to (not including) @triton.jit + while i < len(lines) and lines[i].strip() != "@triton.jit": + i += 1 + continue + if _DROP_IMPORT_RE.match(line) or _DROP_CALL_RE.match(line): + i += 1 + continue + out.append(line) + i += 1 + body = "\n".join(out) + + used = sorted(set(_HELPER_USE_RE.findall(body))) + if used: + raise SpecIncomplete( + f"kernel uses triton_helpers.{{{','.join(used)}}}, which lives in " + f"torch and the tnpu venv has no torch. Vendor a minimal " + f"triton_helpers into the tnpu venv (or into TRITON_SRC) before this " + f"kernel can compile.") + + # The generated source already imports triton itself; only add what a + # stripped module might be missing. + prefix = "" + if "import triton.language as tl" not in body: + prefix = "import triton\nimport triton.language as tl\n\n" + return prefix + body + + +def _grid(meta): + """Sequential launch grid, from the numels and the pinned block sizes.""" + x = meta["numels"].get("xnumel") + xblock = (meta.get("fixed_config") or {}).get("XBLOCK") + if x is None or not xblock: + raise SpecIncomplete( + f"cannot compute the grid for {meta['kernel_name']}: " + f"xnumel={x!r}, XBLOCK={xblock!r}. Inductor defers the grid to " + f"triton_heuristics at runtime; this route needs it statically " + f"(see fixed_config_for).") + return (int(math.ceil(x / xblock)),) + + +SPEC_TEMPLATE = '''\ +"""Generated by PyTorchSimFrontend/triton_backend/kernel_spec.py -- do not edit. + +Inductor kernel {kernel_name!r}, rewritten for the tnpu pipeline: the +triton_heuristics autotuner decorator is stripped and its block sizes are pinned +as constexprs, so the launch shape is static. See kernel_spec.py for why. +""" +import importlib.util +import os +import sys + +sys.path.insert(0, {tnpu_dir!r}) +from tnpu.spec import KernelSpec, Arg # noqa: E402 + +#: The rewritten Triton source, beside this file. It must be a REAL file on +#: disk, not an exec'd string: triton's @jit reads the function back with +#: inspect.getsourcefile and rejects anything else ("@jit functions should be +#: defined in a Python file"). +TRITON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), + {triton_module!r}) + + +def kernel(): + spec = importlib.util.spec_from_file_location( + {kernel_name!r} + "_triton", TRITON_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, {kernel_name!r}) + + +def make_inputs(torch, seed=0): + g = torch.Generator().manual_seed(seed) + out = {{}} +{make_inputs_body} + return out + + +def reference(inputs): + # The Inductor route has no per-kernel torch reference: correctness is + # checked at the graph level by the test that ran torch.compile. tnpu's + # stage 7 is therefore not meaningful here and the pipeline is driven to + # stage 6 (spike) instead. + return {{}} + + +SPEC = KernelSpec( + name={kernel_name!r}, + kernel=kernel, + signature={signature!r}, + constexprs={constexprs!r}, + args=[ +{args_body} + ], + grid={grid!r}, + reference=reference, + make_inputs=make_inputs, + notes="generated from Inductor triton codegen", +) +''' + + +def write_spec_file(src_code, meta, path, tnpu_dir): + """Write a tnpu kernel file for this Inductor kernel. Returns `path`.""" + missing = [a["name"] for a in meta["args"] if not a["dtype"] or not a["numel"]] + if missing: + raise SpecIncomplete( + f"{meta['kernel_name']}: no dtype/numel for {missing} -- " + f"collect_meta could not resolve them from V.graph") + + signature = dict(meta["signature"]) + constexprs = dict(meta["constants"]) + for k, v in (meta.get("fixed_config") or {}).items(): + if v is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: block size {k} is unset " + f"(fixed_config_for leaves reduction blocks unset on purpose)") + constexprs[k] = v + signature[k] = "constexpr" + + args_body = "\n".join( + f" Arg({a['name']!r}, {a['role']!r}, {a['dtype']!r}, ({a['numel']},))," + for a in meta["args"]) + make_inputs_body = "\n".join( + f" out[{a['name']!r}] = torch.randn({a['numel']}, generator=g)" + f".to(torch.{a['dtype']})" + for a in meta["args"] if a["role"] in ("in", "inout")) or " pass" + + triton_module = f"{meta['kernel_name']}_triton.py" + with open(os.path.join(os.path.dirname(path), triton_module), "w") as f: + f.write(strip_for_tnpu(src_code)) + + text = SPEC_TEMPLATE.format( + kernel_name=meta["kernel_name"], + tnpu_dir=tnpu_dir, + triton_module=triton_module, + signature=signature, + constexprs=constexprs, + args_body=args_body, + make_inputs_body=make_inputs_body, + grid=_grid(meta), + ) + with open(path, "w") as f: + f.write(text) + return path diff --git a/PyTorchSimFrontend/triton_backend/scheduling.py b/PyTorchSimFrontend/triton_backend/scheduling.py new file mode 100644 index 00000000..07b86aec --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -0,0 +1,88 @@ +"""Inductor scheduling for the Triton route. + +`TritonNPUScheduling` keeps ALL of Inductor's Triton codegen -- fusion, index +expressions, masking, reductions, the kernel source itself -- and changes only +what happens to the generated source afterwards. Upstream hands it to +`async_compile.triton(...)`, which calls `triton.compile` for a GPU; we hand it +to `triton_npu_compile(...)`, which runs the triton-npu pipeline for this NPU. + +Two overrides, and nothing else: + + define_kernel emit our compile call into the wrapper instead of upstream's. + kernel_type a TritonKernel whose call site is a plain python call, because + the name is bound to our callable rather than to a triton + launcher with a `.run(grid=..., stream=...)` interface. +""" + +from torch._inductor.codegen.common import IndentedBuffer +from torch._inductor.codegen.triton import TritonKernel, TritonScheduling +from torch._inductor.utils import Placeholder, get_fused_kernel_name +from torch._inductor.virtualized import V + +from . import kernel_spec + + +class TritonNPUKernel(TritonKernel): + """A TritonKernel launched as a plain call. + + Upstream emits `kernel.run(a, b, xnumel, grid=grid(xnumel), stream=...)`, + where `grid` is resolved at RUNTIME by triton_heuristics from the autotuned + XBLOCK. There is no autotuner and no stream here: the kernel name is bound to + the callable `triton_npu_compile` returned, so the call is `kernel(a, b, n)`. + + That is also why the block sizes must be fixed at CODEGEN time -- see + kernel_spec.fixed_config_for. A grid that is only known after autotuning + cannot be written into a tnpu KernelSpec. + """ + + def call_kernel(self, name: str, node=None): + wrapper = V.graph.wrapper_code + _, call_args, _, arg_types = self.args.python_argdefs() + self.add_numel_to_call_args(name, call_args, arg_types) + # add_numel_to_call_args appends the numels as SYMPY values, which the + # triton path later renders through pexpr. ExtensionWrapperCodegen joins + # call args as plain strings (mlir_codegen_backend.py:241), so render + # them here instead of handing it a sympy Integer. + call_args = [a if isinstance(a, str) else str(a) for a in call_args] + # triton=False -> PythonWrapperCodegen emits `name(args...)`, the same + # shape the MLIR route uses (mlir_common.py:627). + wrapper.generate_kernel_call(name, call_args, triton=False) + + +class TritonNPUScheduling(TritonScheduling): + kernel_type = TritonNPUKernel + + count = 0 + + def define_kernel(self, src_code, node_schedule, kernel): + wrapper = V.graph.wrapper_code + if src_code in wrapper.src_to_kernel: + return wrapper.src_to_kernel[src_code] + + fused_name = get_fused_kernel_name(node_schedule, "original_aten") + kernel_name = "_".join( + x for x in ("triton_npu", fused_name, str(TritonNPUScheduling.count)) if x + ) + TritonNPUScheduling.count += 1 + wrapper.src_to_kernel[src_code] = kernel_name + + # Upstream substitutes these two placeholders inside define_kernel; the + # source still carries them here, and the tnpu side parses the source, so + # they have to be resolved before it leaves this function. + src_code = src_code.replace(str(Placeholder.DESCRIPTIVE_NAME), kernel_name) + src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name) + + meta = kernel_spec.collect_meta(kernel, kernel_name) + + compile_wrapper = IndentedBuffer() + compile_wrapper.writeline(f"triton_npu_compile('''{src_code}''',") + compile_wrapper.writeline(f" meta={meta!r},") + compile_wrapper.writeline(f" kernel_name={kernel_name!r})") + + origins = ", ".join( + sorted({str(o) for n in node_schedule + for o in getattr(getattr(n, "node", None), "origins", ()) or ()}) + ) + wrapper.define_kernel(kernel_name, compile_wrapper.getvalue(), + f"# origins: {origins}") + return kernel_name diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py new file mode 100644 index 00000000..5a49ccf5 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -0,0 +1,82 @@ +"""Run the triton-npu pipeline, out of process. + +WHY A SUBPROCESS +---------------- +tnpu's passes run on LLVM 23's MLIR python bindings; this process holds LLVM 20's +(TORCHSIM_LLVM_PATH). `mlir` ships without an `__init__.py`, so it is a NAMESPACE +package whose `__path__` is the union of every `mlir/` directory on sys.path -- +two LLVMs in one interpreter silently merge and fail later with an AttributeError +from a generated dialect module (tnpu/config.py:activate_bindings documents the +exact failure). They cannot share an interpreter, so tnpu gets its own. + +The seam between them is a FILE, which is measured to work: LLVM 23 prints the +IR, and LLVM 20's bindings parse it back without complaint (verified by feeding +tnpu's 04-custom.mlir to PyTorchSim's build_tog). That is what makes the split +viable rather than merely necessary. +""" + +import os +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + + +class TnpuError(RuntimeError): + def __init__(self, message, cmd=None, output=None): + super().__init__(message) + self.cmd = cmd + self.output = output + + +def tnpu_dir(): + d = extension_config.CONFIG_TNPU_DIR + if not os.path.isdir(d): + raise TnpuError( + f"triton-npu checkout not found at {d}. It is a separate repository " + f"and is not vendored; clone it there or set TNPU_DIR.") + return d + + +def doctor(): + """Return (ok, output) for tnpu's own toolchain check.""" + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, os.path.join(tnpu_dir(), "run.py"), "doctor"], + capture_output=True, text=True, cwd=tnpu_dir()) + return proc.returncode == 0, proc.stdout + proc.stderr + + +def run_pipeline(spec_path, workdir, to_stage="binary", from_stage="ttir", + verbose=False, timeout=1800): + """Drive tnpu's stages over `spec_path`, writing artifacts into `workdir`. + + Stops at `to_stage`. The default is `binary` (through the RISC-V ELF): + stage 6 (spike) needs the caller's real tensors as .raw files and stage 7 + compares against a per-kernel torch reference, neither of which exists on + the Inductor route -- correctness is a graph-level property there. + + Returns the workdir on success; raises TnpuError with tnpu's own stage + report (which names the failing command and its stderr) otherwise. + """ + cmd = [extension_config.CONFIG_TNPU_PYTHON, + os.path.join(tnpu_dir(), "run.py"), spec_path, + "--from", from_stage, "--to", to_stage, "--workdir", workdir] + if verbose: + cmd.append("-v") + + env = dict(os.environ) + # tnpu deliberately does not read TORCHSIM_LLVM_PATH (it would drag the + # backend back to LLVM 20 and break the textual seam), but a stale + # PYTHONPATH pointing at LLVM 20's mlir_core would still be picked up by the + # namespace package before tnpu's own activate_bindings() runs. + env.pop("PYTHONPATH", None) + + proc = subprocess.run(cmd, capture_output=True, text=True, + cwd=tnpu_dir(), env=env, timeout=timeout) + output = proc.stdout + proc.stderr + if proc.returncode != 0: + raise TnpuError(f"tnpu pipeline failed (exit {proc.returncode})", + cmd=" ".join(cmd), output=output) + logger.debug("[triton-npu] %s", output) + return workdir diff --git a/PyTorchSimFrontend/triton_backend/wrapper_codegen.py b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py new file mode 100644 index 00000000..cfc6490d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py @@ -0,0 +1,21 @@ +"""Wrapper codegen for the Triton route. + +Reuses `ExtensionWrapperCodegen` wholesale -- device guards, buffer allocation, +the TOGSimulator plumbing and the SRAM plan hooks are all route-independent -- +and adds the one import the generated module needs: `triton_npu_compile`, which +is to this route what `custom_async_compile` is to the MLIR one. +""" + +from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen + +from . import codecache + + +class TritonNPUWrapperCodegen(ExtensionWrapperCodegen): + def write_header(self): + super().write_header() + self.header.splice( + f""" + from {codecache.__name__} import triton_npu_compile + """ + ) diff --git a/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py new file mode 100644 index 00000000..a7b0af79 --- /dev/null +++ b/tests/system/test_triton_codegen.py @@ -0,0 +1,65 @@ +"""Drive the Triton codegen route as far as it currently goes. + +This route is WIP (see PyTorchSimFrontend/triton_backend/README.md). The test is +written to report WHERE it stops rather than to assert success: the value right +now is a reproducible statement of the next gap, not a pass/fail gate. Register +it in .github/workflows/pytorchsim_test.yml only once the route runs end to end. + + TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py +""" +import os +import sys +import traceback + +# Must be set before torch_openreg registers the Inductor backend for `npu`. +os.environ.setdefault("TORCHSIM_TRITON_CODEGEN", "1") + +import torch # noqa: E402 + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +N = 1024 + + +def build(): + def fn(x, y): + return x + y + + x = torch.randn(N) + y = torch.randn(N) + return fn, x, y + + +def main(): + from PyTorchSimFrontend import extension_config + from PyTorchSimFrontend.triton_backend import tnpu_bridge + + print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") + print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") + ok, _out = tnpu_bridge.doctor() + print(f"tnpu doctor = {'ok' if ok else 'FAILED (see run.py doctor)'}") + print() + + fn, x, y = build() + expected = fn(x, y) + + opt = torch.compile(fn, backend="inductor") + try: + got = opt(x.to("npu:0"), y.to("npu:0")) + except Exception as e: # noqa: BLE001 - the point is to report the stop + print(f"STOPPED AT: {type(e).__name__}") + print() + traceback.print_exc() + print() + print("The stage reached is what this test measures; see the traceback " + "above and README.md's gap list.") + return 1 + + err = (got.cpu() - expected).abs().max().item() + print(f"max_abs_err = {err}") + return 0 if err < 1e-4 else 1 + + +if __name__ == "__main__": + sys.exit(main()) From 568673573094c68f94f4a9a7731aaef0e8916c2e Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:04:14 +0900 Subject: [PATCH 02/35] [Build] Target torch 2.10 so the frontend and triton-npu agree on triton 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. --- CLAUDE.md | 2 +- Dockerfile.base | 14 +++- PyTorchSimFrontend/triton_backend/README.md | 16 ++-- PyTorchSimFrontend/triton_backend/__init__.py | 2 +- .../triton_backend/_triton_compat.py | 75 +++++++++++-------- .../triton_backend/scheduling.py | 6 +- 6 files changed, 71 insertions(+), 44 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index fb76c82d..bef6e374 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -137,7 +137,7 @@ Conan deps for TOGSim: `boost/1.79.0`, `robin-hood-hashing/3.11.5`, `spdlog/1.11 ## Gotchas / things I've already learned -- The repo expects `python` to be a Python 3.10+ binary with `torch==2.8.0`. The frontend extends the PyTorch 2 Inductor stack — pin to this version. +- The repo expects `python` to be a Python 3.10+ binary with `torch==2.10.0` (torchvision `0.25.0`, triton `3.6.0`). The frontend extends the PyTorch 2 Inductor stack — pin to this version. 2.10 specifically: it is the first release whose Inductor targets triton 3.6, the version triton-npu is built against. The pins live in `Dockerfile.base`, and editing that file changes the base-image tag automatically (the tag is `thirdparty-`, see `scripts/ci/thirdparty_base_pin.sh`). - The default Gem5 path is hard-coded to `/workspace/gem5/build/RISCV/gem5.opt`. Override with `GEM5_PATH` if you build elsewhere. - `_C.cpython-311-*.so` and `torch_openreg/lib/` are build artifacts — already in `.gitignore`, don't commit. - TOGSim creates a per-PID FIFO under `/tmp/togsim_fifo_` for command/event comm; if a previous run crashed and left stale FIFOs, they get cleaned up on the next start, but watch for orphaned processes if you Ctrl-C mid-run. diff --git a/Dockerfile.base b/Dockerfile.base index de023566..b4d58a7c 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -51,8 +51,20 @@ RUN apt-get -y update && \ rm -rf /var/lib/apt/lists/* # CPU PyTorch (no CUDA wheels). torchvision is required by the vision model tests. +# torch 2.10 is pinned for the Triton codegen route: it is the first release whose +# Inductor targets triton 3.6, which is the version triton-npu is built against +# (triton 3.6 pins LLVM 23, and both sides of triton-npu's textual IR seam must be +# the same LLVM). On 2.8 the frontend had to be shimmed onto a triton it did not +# expect; on 2.10 the versions simply agree. RUN python3.11 -m pip install --no-cache-dir \ - torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu + torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu + +# Triton, for the Inductor Triton codegen route (PyTorchSimFrontend/triton_backend). +# Inductor imports triton while GENERATING a kernel, so it is needed even though +# nothing here compiles or launches through triton's own runtime -- triton-npu +# compiles the kernel ahead of time to a RISC-V ELF using its own triton build. +# Not a dependency of the CPU torch wheels, hence installed explicitly. +RUN python3.11 -m pip install --no-cache-dir triton==3.6.0 # TorchSim Python dependencies (numpy pinned <2 for transformers/diffusers compat). RUN python3.11 -m pip install --no-cache-dir \ diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index ed286c0f..4f11466b 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -94,11 +94,11 @@ one interpreter silently merge. The seam between them is a file, and that is measured to work: LLVM 23 prints IR that LLVM 20's bindings parse without complaint. (`tnpu_bridge`) -**torch 2.8's Inductor is shimmed onto triton 3.6.** Inductor targets the -triton ~3.3 API; tnpu pins 3.6 because 3.6 pins LLVM 23, and both sides of the -IR seam must be the same LLVM. So the frontend bends: `triton_key` is injected -back into `triton.compiler.compiler`, and `triton_hash_with_backend` is -short-circuited because it asks a GPU driver for the current target and there is -no GPU. The alternative — a second, torch-compatible triton in the driver -interpreter purely for codegen — stays open if the API drift widens beyond this -one symbol. (`_triton_compat`) +**The torch pin is what makes triton 3.6 work.** 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 two +simply agree -- on 2.8 the frontend had to be shimmed onto a triton it did not +expect. What remains in `_triton_compat` is not a version shim: on a box with no +GPU, `triton_hash_with_backend()` raises "0 active drivers" because it asks the +triton runtime for the current target. We never launch through that runtime, so +the value is short-circuited to a deterministic cache key. diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py index f5cf1b60..460cf925 100644 --- a/PyTorchSimFrontend/triton_backend/__init__.py +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -30,7 +30,7 @@ from . import _triton_compat # Before anything imports Inductor's Triton codegen: it needs `triton` in THIS -# interpreter, and torch 2.8 expects a triton API older than the 3.6 tnpu pins. +# interpreter, and on a GPU-less box its backend hash cannot be computed. _triton_compat.install() from .scheduling import TritonNPUScheduling # noqa: E402,F401 diff --git a/PyTorchSimFrontend/triton_backend/_triton_compat.py b/PyTorchSimFrontend/triton_backend/_triton_compat.py index c97fe850..9c863b15 100644 --- a/PyTorchSimFrontend/triton_backend/_triton_compat.py +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -1,24 +1,24 @@ -"""Make torch 2.8's Inductor codegen work against tnpu's triton 3.6, CPU-only. - -TWO SKEWS, BOTH STRUCTURAL --------------------------- -1. VERSION. Inductor in torch 2.8 targets the triton ~3.3 API. tnpu pins triton - 3.6 and that pin is not negotiable: 3.6 is what pins LLVM 23, and both sides - of the textual IR seam must be the same LLVM (triton-npu/setup/versions.env). - So the frontend has to bend, not the backend. - -2. NO GPU. `triton_hash_with_backend()` asks the triton runtime driver for the - *current target*, which on a machine with no GPU has nothing to answer with. - We never launch through triton's runtime -- the kernel is compiled ahead of - time to a RISC-V ELF -- so the value is only a cache-key ingredient. - -Both are handled by replacing `torch.utils._triton.triton_hash_with_backend` -with a deterministic string. It is a monkeypatch, and it is the cheap half of a -real choice: the alternative is installing a torch-2.8-compatible triton (3.3.x) -in the driver interpreter purely for codegen, keeping 3.6 in the tnpu venv for -stage 1. That works because the two interpreters exchange only SOURCE TEXT and -never share objects -- but it means two tritons to keep straight, so it is worth -doing only if the API drift turns out to be wider than this one symbol. +"""Let Inductor's Triton codegen run on a machine with no GPU. + +ONE SKEW LEFT, AND IT IS NOT A VERSION SKEW +------------------------------------------- +`triton_hash_with_backend()` asks the triton runtime driver for the *current +target*, which on a box with no GPU raises "0 active drivers". We never launch +through triton's runtime -- triton-npu compiles the kernel ahead of time to a +RISC-V ELF -- so the value is only a cache-key ingredient, and a deterministic +string does the job. + +WHAT USED TO BE HERE, AND WHY IT IS GONE +---------------------------------------- +On torch 2.8 there was a second, larger skew: Inductor targeted the triton ~3.3 +API while triton-npu pins 3.6 (3.6 is what pins LLVM 23, and both sides of +triton-npu's textual IR seam must be the same LLVM), so `triton_key` had to be +injected back into `triton.compiler.compiler`. + +torch 2.10 pins triton 3.6.0 itself and reaches that symbol through its own +compat layer (`torch._inductor.runtime.triton_compat`), so on 2.10 the versions +simply agree and the injection is a no-op. It is kept, guarded, so the module +still works if someone runs an older torch -- `_torch_handles_triton()` decides. """ import functools @@ -87,13 +87,23 @@ def _stable_backend_hash(): return hashlib.sha256(key.encode("utf-8")).hexdigest().upper() -def _needs_hash_patch(): - """True when triton_hash_with_backend cannot work here.""" +def _torch_handles_triton(): + """True when this torch already knows how to reach triton's key itself. + + torch 2.10 routes it through torch._inductor.runtime.triton_compat, which + understands triton 3.6. Older torch imports `triton_key` straight out of + triton.compiler.compiler, where 3.6 no longer defines it. + """ + try: + from torch._inductor.runtime.triton_compat import triton_key # noqa: F401 + return True + except Exception: # noqa: BLE001 + pass try: mod = importlib.import_module("triton.compiler.compiler") except Exception: # noqa: BLE001 - return True - return not hasattr(mod, "triton_key") + return False + return hasattr(mod, "triton_key") def install(): @@ -109,16 +119,17 @@ def install(): if _installed: return notes - if _needs_hash_patch(): - # `triton_key` is imported from triton.compiler.compiler by SEVERAL torch - # call sites (codecache.CacheBase.get_system, _triton.triton_hash_with_ - # backend, ...), each with its own local import. Supplying the symbol on - # the triton side satisfies all of them at once instead of chasing every - # call site; it is a cache-key ingredient, so any stable string will do. + if not _torch_handles_triton(): + # Pre-2.10 torch: `triton_key` is imported from triton.compiler.compiler + # by SEVERAL call sites (codecache.CacheBase.get_system, + # _triton.triton_hash_with_backend, ...), each with its own local import. + # Supplying the symbol on the triton side satisfies all of them at once + # instead of chasing every call site; it is a cache-key ingredient, so + # any stable string will do. On torch 2.10 this branch does not run. mod = importlib.import_module("triton.compiler.compiler") mod.triton_key = _stable_backend_hash notes.append("injected triton.compiler.compiler.triton_key " - "(removed in triton 3.6; torch 2.8 still imports it)") + "(this torch predates the triton 3.6 compat layer)") # Separately: triton_hash_with_backend also asks the triton runtime driver # for the current target, which needs a GPU. We compile ahead of time to a diff --git a/PyTorchSimFrontend/triton_backend/scheduling.py b/PyTorchSimFrontend/triton_backend/scheduling.py index 07b86aec..56e7f1f8 100644 --- a/PyTorchSimFrontend/triton_backend/scheduling.py +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -35,7 +35,11 @@ class TritonNPUKernel(TritonKernel): cannot be written into a tnpu KernelSpec. """ - def call_kernel(self, name: str, node=None): + # **kwargs, not a fixed signature: Inductor keeps adding parameters here + # (2.10 added `deallocate_ws`). None of them apply to this route -- there is + # no triton launcher and no workspace to release -- so they are accepted and + # ignored rather than pinning us to one torch release. + def call_kernel(self, name: str, node=None, **kwargs): wrapper = V.graph.wrapper_code _, call_args, _, arg_types = self.args.python_argdefs() self.add_numel_to_call_args(name, call_args, arg_types) From ef76c24f732e190aab7e1c473faf08439e589102 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:12:19 +0900 Subject: [PATCH 03/35] [CI] Add the triton-npu toolchain image and its workflow 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. --- .github/workflows/triton_npu.yml | 234 ++++++++++++++++++++ CLAUDE.md | 2 + Dockerfile.tnpu | 102 +++++++++ PyTorchSimFrontend/triton_backend/README.md | 42 ++++ scripts/ci/tnpu_asset_env.sh | 64 ++++++ scripts/ci/tnpu_base_pin.sh | 8 + thirdparty/triton-npu.json | 22 ++ 7 files changed, 474 insertions(+) create mode 100644 .github/workflows/triton_npu.yml create mode 100644 Dockerfile.tnpu create mode 100644 scripts/ci/tnpu_asset_env.sh create mode 100755 scripts/ci/tnpu_base_pin.sh create mode 100644 thirdparty/triton-npu.json diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml new file mode 100644 index 00000000..6a915557 --- /dev/null +++ b/.github/workflows/triton_npu.yml @@ -0,0 +1,234 @@ +name: Triton codegen route (triton-npu) + +# Exercises the Triton codegen route: Inductor's own Triton backend produces the +# kernel and triton-npu lowers it to a RISC-V ELF. +# (PyTorchSimFrontend/triton_backend/README.md) +# +# Separate from the main CI on purpose. The route is WIP, and its toolchain layer +# is ~1.8 GiB that no other job needs, so it neither gates PRs nor slows them +# down. Promote the jobs into pytorchsim_test.yml once the route runs end to end. +# +# PREREQUISITES, both one-time and both outside this repo: +# 1. secrets.TNPU_TOKEN -- a PAT that can read PSAL-POSTECH/triton-npu. That +# repo is private, and the default Actions token is scoped to this +# repository only, so it cannot clone it or read its releases. +# 2. A release tagged `toolchain-llvm23` on PSAL-POSTECH/triton-npu carrying +# llvm23-install.tar.gz, spike-install.tar.gz and triton-runtime.tar.gz. +# As of writing that repo has no releases; the assets exist only on the +# upstream fork it was forked from. +# Both are checked by the preflight job below, which reports which one is missing +# instead of failing deep inside a docker build. + +on: + pull_request: + branches: [ "master", "develop" ] + paths: + - 'PyTorchSimFrontend/triton_backend/**' + - 'thirdparty/triton-npu.json' + - 'Dockerfile.tnpu' + - 'scripts/ci/tnpu_*.sh' + - '.github/workflows/triton_npu.yml' + workflow_dispatch: + +env: + BASE_IMAGE_REPO: ghcr.io/psal-postech/torchsim_base + TNPU_IMAGE_REPO: ghcr.io/psal-postech/torchsim_tnpu_base + APP_IMAGE_REPO: ghcr.io/psal-postech/torchsim_tnpu + SOURCE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + +jobs: + preflight: + name: Check tnpu access + runs-on: ubuntu-latest + outputs: + ready: ${{ steps.check.outputs.ready }} + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ env.SOURCE_SHA }} + persist-credentials: false + + - name: Token and release present + id: check + env: + TNPU_TOKEN: ${{ secrets.TNPU_TOKEN }} + run: | + if [ -z "${TNPU_TOKEN}" ]; then + echo "::error::secrets.TNPU_TOKEN is not set. PSAL-POSTECH/triton-npu is private and the default Actions token cannot read it." + echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1 + fi + REPO=$(jq -r '.triton_npu.repository' thirdparty/triton-npu.json) + TAG=$(jq -r '.triton_npu.release_tag' thirdparty/triton-npu.json) + if ! curl -fsS -H "Authorization: Bearer ${TNPU_TOKEN}" \ + "https://api.github.com/repos/${REPO}" -o /dev/null; then + echo "::error::TNPU_TOKEN cannot read ${REPO}." + echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1 + fi + if ! curl -fsS -H "Authorization: Bearer ${TNPU_TOKEN}" \ + "https://api.github.com/repos/${REPO}/releases/tags/${TAG}" -o /dev/null; then + echo "::error::${REPO} has no release tagged '${TAG}'. Mirror the toolchain assets there (see thirdparty/triton-npu.json)." + echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1 + fi + echo "ready=true" >> "$GITHUB_OUTPUT" + + ensure-tnpu-base: + name: Build tnpu toolchain image + needs: preflight + runs-on: ubuntu-latest + outputs: + tnpu_image: ${{ steps.pin.outputs.tnpu_image }} + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ env.SOURCE_SHA }} + submodules: recursive + persist-credentials: false + + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Pins + id: pin + run: | + BASE_PIN="$(bash scripts/ci/thirdparty_base_pin.sh)" + TNPU_PIN="$(bash scripts/ci/tnpu_base_pin.sh)" + echo "BASE_IMAGE=${BASE_IMAGE_REPO}:thirdparty-${BASE_PIN}" >> "$GITHUB_ENV" + # The tnpu layer sits on a specific base, so its tag carries both pins. + echo "TNPU_IMAGE=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_ENV" + echo "tnpu_image=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_OUTPUT" + + - name: Check tnpu image exists + id: exists + run: | + if docker manifest inspect "${TNPU_IMAGE}" > /dev/null 2>&1; then + echo "ok=true" >> "$GITHUB_OUTPUT" + else + echo "ok=false" >> "$GITHUB_OUTPUT" + fi + + - name: Resolve tnpu release asset IDs + if: steps.exists.outputs.ok != 'true' + env: + GITHUB_TOKEN: ${{ secrets.TNPU_TOKEN }} + run: bash scripts/ci/tnpu_asset_env.sh >> "$GITHUB_ENV" + + - name: Build and push tnpu toolchain image + if: steps.exists.outputs.ok != 'true' + uses: docker/build-push-action@v4 + with: + context: . + file: ./Dockerfile.tnpu + push: true + build-args: | + BASE_IMAGE=${{ env.BASE_IMAGE }} + TNPU_LLVM23_ASSET_ID=${{ env.TNPU_LLVM23_ASSET_ID }} + TNPU_SPIKE_ASSET_ID=${{ env.TNPU_SPIKE_ASSET_ID }} + TNPU_RUNTIME_ASSET_ID=${{ env.TNPU_RUNTIME_ASSET_ID }} + secrets: | + tnpu_token=${{ secrets.TNPU_TOKEN }} + tags: ${{ env.TNPU_IMAGE }} + + build-app: + name: Build app image on tnpu base + needs: ensure-tnpu-base + runs-on: ubuntu-latest + outputs: + app_image: ${{ steps.name.outputs.app_image }} + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ env.SOURCE_SHA }} + submodules: recursive + persist-credentials: false + + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Image name + id: name + run: echo "app_image=${APP_IMAGE_REPO}:${SOURCE_SHA}" >> "$GITHUB_OUTPUT" + + - name: Build and push + uses: docker/build-push-action@v4 + with: + context: . + file: ./Dockerfile + push: true + build-args: | + BASE_IMAGE=${{ needs.ensure-tnpu-base.outputs.tnpu_image }} + tags: ${{ steps.name.outputs.app_image }} + + tnpu-baselines: + name: triton-npu baselines + needs: build-app + runs-on: ubuntu-latest + steps: + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + # The harness's own kernels, end to end through Spike. This is the gate on + # the toolchain itself: if these regress, nothing downstream is meaningful. + # gemm/bmm need TNPU_VCIX_MATMUL=1 to reach the systolic array. + - name: doctor + add / mul / relu / gemm / bmm + run: | + docker run --rm -e TNPU_VCIX_MATMUL=1 \ + ${{ needs.build-app.outputs.app_image }} bash -lc ' + cd /workspace/triton-npu && + python3 run.py doctor && + for k in add mul relu gemm bmm; do + echo "=== $k ===" && python3 run.py kernels/$k.py || exit 1 + done' + + triton-route: + name: Inductor Triton route + needs: build-app + runs-on: ubuntu-latest + # WIP: the launch is deliberately unimplemented, so this reports how far the + # route gets rather than gating. Drop this once the launch lands. + continue-on-error: true + steps: + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: test_triton_codegen.py + run: | + docker run --rm -e TORCHSIM_TRITON_CODEGEN=1 \ + ${{ needs.build-app.outputs.app_image }} \ + python3 PyTorchSim/tests/system/test_triton_codegen.py + + mlir-route-regression: + name: MLIR route still passes + needs: build-app + runs-on: ubuntu-latest + steps: + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + # The tnpu layer adds a second LLVM and a second triton to the image. This + # is the check that it did not disturb the production path. + - name: test_add.py + run: | + docker run --rm \ + ${{ needs.build-app.outputs.app_image }} \ + python3 PyTorchSim/tests/ops/elementwise/test_add.py diff --git a/CLAUDE.md b/CLAUDE.md index bef6e374..e3424132 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,6 +47,8 @@ Run a model from `tests/models/Llama/`, `tests/models/DeepSeek/`, etc. similarly **CI coverage:** the GitHub Actions workflow `.github/workflows/pytorchsim_test.yml` runs an **explicit allowlist** of `tests/*.py` files (~40 jobs, one Docker container per test). Adding a new file under `tests/` does *not* automatically gate PRs — register it in `pytorchsim_test.yml` if you want CI to exercise it. Conversely, files like `tests/ops/attention/test_gqa.py`, `tests/ops/attention/test_gqa_decode.py`, and `tests/system/test_eager.py` exist in the repo but are *not* in CI, so local validation is the only safety net for them. +The Triton codegen route has its own workflow, `.github/workflows/triton_npu.yml`, kept separate because its toolchain layer is ~1.8 GiB that no other job needs. It builds `torchsim_tnpu_base` (pinned by `thirdparty/triton-npu.json` + `Dockerfile.tnpu`) and needs `secrets.TNPU_TOKEN` plus a toolchain release on the private `PSAL-POSTECH/triton-npu`; see `PyTorchSimFrontend/triton_backend/README.md`. + **For fast iteration** (skip functional check): ```bash export pytorchsim_functional_mode=False # skips Spike diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu new file mode 100644 index 00000000..fc7f74ce --- /dev/null +++ b/Dockerfile.tnpu @@ -0,0 +1,102 @@ +# syntax=docker/dockerfile:1.4 +# +# The triton-npu (tnpu) toolchain layer, on top of torchsim_base. +# +# Only the Triton codegen route needs this, and it is ~1.8 GiB compressed (LLVM 23 +# alone unpacks to ~4.3 GiB), so it is a SEPARATE image rather than part of +# torchsim_base -- every other CI job would otherwise pay for it. The app image for +# the Triton route is then the ordinary ./Dockerfile with BASE_IMAGE pointed here. +# +# Everything below mirrors triton-npu/setup/restore.sh --prebuilt. It is not simply +# invoked because that script has two problems in this context, both flagged where +# they are worked around: it clones a triton_shared fork that no longer exists, and +# the release tarballs omit the triton backend links the editable install needs. +# +# PSAL-POSTECH/triton-npu is private, so the clone and the asset downloads both need +# a token. It is passed as a BuildKit secret, never a build-arg: build-args are +# recorded in the image history. + +ARG BASE_IMAGE=ghcr.io/psal-postech/torchsim_base:latest +FROM ${BASE_IMAGE} + +ARG TNPU_LLVM23_ASSET_ID +ARG TNPU_SPIKE_ASSET_ID +ARG TNPU_RUNTIME_ASSET_ID +ARG TNPU_REPO=PSAL-POSTECH/triton-npu +ARG TNPU_REF=main +ARG TRITON_REPO=triton-lang/triton +ARG TRITON_COMMIT=df38505e451a1541555379bcf378be9e8c00545c + +WORKDIR /workspace + +# 1. The harness itself. Kept OUT of $TORCHSIM_DIR: ./Dockerfile copies the +# PyTorchSim checkout over that path afterwards, and a nested clone there would +# be shadowed or clobbered. extension_config reads TNPU_DIR, set below. +RUN --mount=type=secret,id=tnpu_token \ + TOKEN="$(cat /run/secrets/tnpu_token)" && \ + git clone --depth 1 --branch "${TNPU_REF}" \ + "https://x-access-token:${TOKEN}@github.com/${TNPU_REPO}.git" /workspace/triton-npu && \ + git -C /workspace/triton-npu rev-parse HEAD > /workspace/triton-npu/.checkout-sha + +# 2. The triton sources the prebuilt runtime was built against. The editable +# install in that tarball points back into this tree, so the checkout has to +# exist at exactly this revision -- and it is a main-branch commit, not the +# v3.6.0 tag: llvm-hash.txt moved from LLVM 22 to 23 after the tag, and the +# whole point of the pin is that both sides of the IR seam are LLVM 23. +RUN git clone --filter=blob:none --no-checkout \ + "https://github.com/${TRITON_REPO}.git" /workspace/hexagon-mlir/triton && \ + git -C /workspace/hexagon-mlir/triton fetch --depth 1 origin "${TRITON_COMMIT}" && \ + git -C /workspace/hexagon-mlir/triton checkout -q FETCH_HEAD && \ + for p in nvvm_reduction_kind_compatibility libdevice_sigmoid; do \ + git -C /workspace/hexagon-mlir/triton apply \ + "/workspace/triton-npu/setup/patches/triton/$p.patch"; \ + done +# nvvm_reduction_kind_compatibility is REQUIRED, not cosmetic: the NVVM +# ReductionKind API moved in this LLVM and triton does not compile without it. + +# 3. The prebuilt toolchain: LLVM 23 install, the CONFIG_DESC-capable spike, and +# the triton runtime (the mlir-env venv + triton-shared-opt + the built .so +# files). Each unpacks rooted at /workspace. A private repo's assets are only +# reachable through the API asset id with an octet-stream Accept header -- the +# /releases/download/ URL 404s even with a token. +RUN --mount=type=secret,id=tnpu_token \ + TOKEN="$(cat /run/secrets/tnpu_token)" && \ + for spec in "${TNPU_LLVM23_ASSET_ID}" "${TNPU_SPIKE_ASSET_ID}" "${TNPU_RUNTIME_ASSET_ID}"; do \ + curl -fL --retry 3 \ + -H "Authorization: Bearer ${TOKEN}" \ + -H "Accept: application/octet-stream" \ + "https://api.github.com/repos/${TNPU_REPO}/releases/assets/${spec}" \ + | tar -xz -C /workspace ; \ + done + +# 4. Re-create the triton backend links. `setup/package.sh` collects the triton +# tree with `find python -name '*.so'`, which misses +# python/triton/backends/{amd,nvidia,triton_shared} -- those are not source +# files but symlinks the build creates (setup.py:add_link_to_backends), so they +# are in neither the tarball nor git. Without them `import triton` dies in +# entry-point discovery with "No module named 'triton.backends.amd'". +# triton_shared is reconstructed from upstream microsoft/triton-shared: the fork +# the manifest used to name is gone, and tnpu never calls this backend -- it only +# has to import so discovery succeeds. +RUN T=/workspace/hexagon-mlir/triton && \ + for b in amd nvidia; do \ + ln -sfnT "$T/third_party/$b/backend" "$T/python/triton/backends/$b" && \ + if [ -d "$T/third_party/$b/language" ]; then \ + for x in "$T/third_party/$b/language"/*; do \ + ln -sfnT "$x" "$T/python/triton/language/extra/$(basename "$x")" ; \ + done ; \ + fi ; \ + done && \ + mkdir -p /workspace/hexagon-mlir/triton_shared/backend && \ + for f in compiler.py driver.py name.conf; do \ + curl -fsSL "https://raw.githubusercontent.com/microsoft/triton-shared/main/backend/$f" \ + -o "/workspace/hexagon-mlir/triton_shared/backend/$f" ; \ + done && \ + ln -sfnT /workspace/hexagon-mlir/triton_shared/backend "$T/python/triton/backends/triton_shared" + +# tnpu's own defaults already point at these paths (tnpu/config.py); TNPU_DIR is +# the one thing PyTorchSim needs told, since the checkout is not under TORCHSIM_DIR. +ENV TNPU_DIR=/workspace/triton-npu + +# Fail the build rather than the first CI job if the toolchain is not usable. +RUN python3 /workspace/triton-npu/run.py doctor diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 4f11466b..86b94da4 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -102,3 +102,45 @@ expect. What remains in `_triton_compat` is not a version shim: on a box with no GPU, `triton_hash_with_backend()` raises "0 active drivers" because it asks the triton runtime for the current target. We never launch through that runtime, so the value is short-circuited to a deterministic cache key. + +## CI + +`.github/workflows/triton_npu.yml`, separate from the main CI: this route is WIP, +and its toolchain layer is ~1.8 GiB that no other job needs. + +``` +preflight TNPU_TOKEN set? repo readable? release present? +ensure-tnpu-base torchsim_base + tnpu toolchain -> torchsim_tnpu_base: +build-app ./Dockerfile on that base +tnpu-baselines run.py doctor + add/mul/relu/gemm/bmm through Spike (gates) +triton-route tests/system/test_triton_codegen.py (reports, does not gate) +mlir-route-regression tests/ops/elementwise/test_add.py (gates) +``` + +The toolchain image is pinned the same way `torchsim_base` is — the tag carries +`sha256(thirdparty/triton-npu.json + Dockerfile.tnpu)`, so it is rebuilt only when +one of those moves, and its tag also carries the base pin it was built on. +`mlir-route-regression` is there because this layer adds a *second* LLVM and a +*second* triton to the image; it checks the production path did not notice. + +**Two prerequisites, both outside this repo.** `preflight` fails with which one is +missing rather than letting a docker build die deep: + +1. `secrets.TNPU_TOKEN` — a PAT that can read `PSAL-POSTECH/triton-npu`. That repo + is private, and the default Actions token is scoped to this repository, so it + can neither clone it 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, so they have to be mirrored across (or the repo made public and the + manifest pointed at whichever holds them). + +`Dockerfile.tnpu` mirrors `triton-npu/setup/restore.sh --prebuilt` rather than +calling it, because that script has two problems here, both worked around in place +with comments: it 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}` — those are symlinks the +build creates, so they are in neither the tarball nor git, and without them +`import triton` dies in entry-point discovery. Both are worth fixing upstream. + diff --git a/scripts/ci/tnpu_asset_env.sh b/scripts/ci/tnpu_asset_env.sh new file mode 100644 index 00000000..ecc56a63 --- /dev/null +++ b/scripts/ci/tnpu_asset_env.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# Emit TNPU_LLVM23_ASSET_ID / TNPU_SPIKE_ASSET_ID / TNPU_RUNTIME_ASSET_ID lines +# for appending to GITHUB_ENV. +# +# Same shape as thirdparty_github_asset_env.sh, with one difference that matters: +# PSAL-POSTECH/triton-npu is PRIVATE, so the default `secrets.GITHUB_TOKEN` (which +# is scoped to the PyTorchSim repo) cannot read its releases. Pass a PAT with read +# access instead -- the workflow puts secrets.TNPU_TOKEN here. +# +# Release assets of a private repo are also only reachable through the API asset +# id, not the /releases/download/ URL, which is why ids are resolved at all +# (thirdparty/github-releases.json documents the same constraint). +# +# Requires: jq, curl, GITHUB_TOKEN, repo root as cwd or GITHUB_WORKSPACE. +set -euo pipefail +ROOT="${GITHUB_WORKSPACE:-$(cd "$(dirname "$0")/../.." && pwd)}" +MANIFEST="${ROOT}/thirdparty/triton-npu.json" +if [ ! -f "$MANIFEST" ]; then + echo "Missing tnpu manifest: $MANIFEST" >&2 + exit 1 +fi +if [ -z "${GITHUB_TOKEN:-}" ]; then + echo "GITHUB_TOKEN is not set (needs a PAT that can read the private" >&2 + echo "PSAL-POSTECH/triton-npu; the default Actions token cannot)" >&2 + exit 1 +fi + +REPO=$(jq -r '.triton_npu.repository' "$MANIFEST") +TAG=$(jq -r '.triton_npu.release_tag' "$MANIFEST") +OWNER="${REPO%%/*}" +NAME="${REPO##*/}" + +if [ "$TAG" = "latest" ]; then + API_URL="https://api.github.com/repos/${OWNER}/${NAME}/releases/latest" +else + API_URL="https://api.github.com/repos/${OWNER}/${NAME}/releases/tags/${TAG}" +fi + +TMP=$(mktemp) +trap 'rm -f "$TMP"' EXIT +if ! curl -fsS -H "Authorization: Bearer ${GITHUB_TOKEN}" \ + -H "Accept: application/vnd.github+json" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + "$API_URL" -o "$TMP"; then + echo "Failed to read release ${TAG} of ${REPO}." >&2 + echo "Either the token cannot see the repo, or the release does not exist" >&2 + echo "yet -- as of writing, ${REPO} carries no releases and the toolchain" >&2 + echo "assets live only on the upstream fork. See thirdparty/triton-npu.json." >&2 + exit 1 +fi + +emit() { # emit + local id + id=$(jq -r --arg n "$1" '.assets[] | select(.name == $n) | .id' "$TMP" | head -n1) + if [ -z "$id" ] || [ "$id" = "null" ]; then + echo "Release ${TAG} of ${REPO} has no asset named '$1'" >&2 + exit 1 + fi + echo "$2=${id}" +} + +emit llvm23-install.tar.gz TNPU_LLVM23_ASSET_ID +emit spike-install.tar.gz TNPU_SPIKE_ASSET_ID +emit triton-runtime.tar.gz TNPU_RUNTIME_ASSET_ID diff --git a/scripts/ci/tnpu_base_pin.sh b/scripts/ci/tnpu_base_pin.sh new file mode 100755 index 00000000..59663836 --- /dev/null +++ b/scripts/ci/tnpu_base_pin.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Deterministic short pin for tagging torchsim_tnpu_base images. +# Mirrors thirdparty_base_pin.sh, over the tnpu manifest + its Dockerfile, so the +# ~1.8 GiB toolchain layer is rebuilt only when one of those two actually moves. +set -euo pipefail +ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +cd "$ROOT" +{ cat thirdparty/triton-npu.json; cat Dockerfile.tnpu; } | sha256sum | awk '{print substr($1,1,12)}' diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json new file mode 100644 index 00000000..408f4f6e --- /dev/null +++ b/thirdparty/triton-npu.json @@ -0,0 +1,22 @@ +{ + "description": "Pins for the triton-npu (tnpu) toolchain layer, used only by the Triton codegen route. Kept SEPARATE from github-releases.json on purpose: that file plus Dockerfile.base is hashed into the torchsim_base tag, so putting these here means the main base image is not rebuilt when the tnpu pins move, and the ~1.8 GiB toolchain is not forced onto every CI job. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing (pin = sha256 of this file plus Dockerfile.tnpu). The repository is PRIVATE and its release assets are not readable by the default Actions token, so the workflow needs a PAT in secrets.TNPU_TOKEN with read access to it.", + "triton_npu": { + "repository": "PSAL-POSTECH/triton-npu", + "ref": "main", + "release_tag": "toolchain-llvm23", + "assets": [ + "llvm23-install.tar.gz", + "spike-install.tar.gz", + "triton-runtime.tar.gz" + ] + }, + "triton": { + "description": "The triton checkout the prebuilt runtime was built against. Its editable install points back into this tree, so the sources must be present at exactly this revision for `import triton` to work in the tnpu venv. Mirrors triton-npu/setup/versions.env (TRITON_SHA); a main-branch commit, NOT the v3.6.0 tag, because llvm-hash.txt moved from LLVM 22 to 23 after the tag.", + "repository": "triton-lang/triton", + "commit": "df38505e451a1541555379bcf378be9e8c00545c", + "patches": [ + "setup/patches/triton/nvvm_reduction_kind_compatibility.patch", + "setup/patches/triton/libdevice_sigmoid.patch" + ] + } +} From 2173834dcf759e7afc67bead8daf8cd9c4821edd Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:24:03 +0900 Subject: [PATCH 04/35] [CI] Source triton-npu from PSAL-POSTECH and drop the local workarounds 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. --- .github/workflows/triton_npu.yml | 13 +-- Dockerfile.tnpu | 101 ++++---------------- PyTorchSimFrontend/triton_backend/README.md | 11 +-- scripts/ci/tnpu_asset_env.sh | 64 ------------- thirdparty/triton-npu.json | 20 +--- 5 files changed, 27 insertions(+), 182 deletions(-) delete mode 100644 scripts/ci/tnpu_asset_env.sh diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index 6a915557..92290f5f 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -26,7 +26,7 @@ on: - 'PyTorchSimFrontend/triton_backend/**' - 'thirdparty/triton-npu.json' - 'Dockerfile.tnpu' - - 'scripts/ci/tnpu_*.sh' + - 'scripts/ci/tnpu_base_pin.sh' - '.github/workflows/triton_npu.yml' workflow_dispatch: @@ -102,6 +102,7 @@ jobs: # The tnpu layer sits on a specific base, so its tag carries both pins. echo "TNPU_IMAGE=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_ENV" echo "tnpu_image=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_OUTPUT" + echo "TNPU_REF=$(jq -r '.triton_npu.ref' thirdparty/triton-npu.json)" >> "$GITHUB_ENV" - name: Check tnpu image exists id: exists @@ -112,12 +113,6 @@ jobs: echo "ok=false" >> "$GITHUB_OUTPUT" fi - - name: Resolve tnpu release asset IDs - if: steps.exists.outputs.ok != 'true' - env: - GITHUB_TOKEN: ${{ secrets.TNPU_TOKEN }} - run: bash scripts/ci/tnpu_asset_env.sh >> "$GITHUB_ENV" - - name: Build and push tnpu toolchain image if: steps.exists.outputs.ok != 'true' uses: docker/build-push-action@v4 @@ -127,9 +122,7 @@ jobs: push: true build-args: | BASE_IMAGE=${{ env.BASE_IMAGE }} - TNPU_LLVM23_ASSET_ID=${{ env.TNPU_LLVM23_ASSET_ID }} - TNPU_SPIKE_ASSET_ID=${{ env.TNPU_SPIKE_ASSET_ID }} - TNPU_RUNTIME_ASSET_ID=${{ env.TNPU_RUNTIME_ASSET_ID }} + TNPU_REF=${{ env.TNPU_REF }} secrets: | tnpu_token=${{ secrets.TNPU_TOKEN }} tags: ${{ env.TNPU_IMAGE }} diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu index fc7f74ce..c37f633a 100644 --- a/Dockerfile.tnpu +++ b/Dockerfile.tnpu @@ -1,102 +1,37 @@ # syntax=docker/dockerfile:1.4 # -# The triton-npu (tnpu) toolchain layer, on top of torchsim_base. +# triton-npu toolchain layer, for the Triton codegen route only. +# Separate from torchsim_base because it is ~1.8 GiB no other job needs. +# The app image for this route is ./Dockerfile with BASE_IMAGE pointed here. # -# Only the Triton codegen route needs this, and it is ~1.8 GiB compressed (LLVM 23 -# alone unpacks to ~4.3 GiB), so it is a SEPARATE image rather than part of -# torchsim_base -- every other CI job would otherwise pay for it. The app image for -# the Triton route is then the ordinary ./Dockerfile with BASE_IMAGE pointed here. -# -# Everything below mirrors triton-npu/setup/restore.sh --prebuilt. It is not simply -# invoked because that script has two problems in this context, both flagged where -# they are worked around: it clones a triton_shared fork that no longer exists, and -# the release tarballs omit the triton backend links the editable install needs. -# -# PSAL-POSTECH/triton-npu is private, so the clone and the asset downloads both need -# a token. It is passed as a BuildKit secret, never a build-arg: build-args are -# recorded in the image history. +# The repo is private, so the clone and the release downloads both need a token. +# It is a BuildKit secret, not a build-arg: build-args land in the image history. ARG BASE_IMAGE=ghcr.io/psal-postech/torchsim_base:latest FROM ${BASE_IMAGE} -ARG TNPU_LLVM23_ASSET_ID -ARG TNPU_SPIKE_ASSET_ID -ARG TNPU_RUNTIME_ASSET_ID ARG TNPU_REPO=PSAL-POSTECH/triton-npu ARG TNPU_REF=main -ARG TRITON_REPO=triton-lang/triton -ARG TRITON_COMMIT=df38505e451a1541555379bcf378be9e8c00545c WORKDIR /workspace -# 1. The harness itself. Kept OUT of $TORCHSIM_DIR: ./Dockerfile copies the -# PyTorchSim checkout over that path afterwards, and a nested clone there would -# be shadowed or clobbered. extension_config reads TNPU_DIR, set below. +# Not under $TORCHSIM_DIR: ./Dockerfile copies the PyTorchSim checkout over that +# path afterwards. extension_config reads TNPU_DIR, set below. RUN --mount=type=secret,id=tnpu_token \ TOKEN="$(cat /run/secrets/tnpu_token)" && \ - git clone --depth 1 --branch "${TNPU_REF}" \ - "https://x-access-token:${TOKEN}@github.com/${TNPU_REPO}.git" /workspace/triton-npu && \ - git -C /workspace/triton-npu rev-parse HEAD > /workspace/triton-npu/.checkout-sha - -# 2. The triton sources the prebuilt runtime was built against. The editable -# install in that tarball points back into this tree, so the checkout has to -# exist at exactly this revision -- and it is a main-branch commit, not the -# v3.6.0 tag: llvm-hash.txt moved from LLVM 22 to 23 after the tag, and the -# whole point of the pin is that both sides of the IR seam are LLVM 23. -RUN git clone --filter=blob:none --no-checkout \ - "https://github.com/${TRITON_REPO}.git" /workspace/hexagon-mlir/triton && \ - git -C /workspace/hexagon-mlir/triton fetch --depth 1 origin "${TRITON_COMMIT}" && \ - git -C /workspace/hexagon-mlir/triton checkout -q FETCH_HEAD && \ - for p in nvvm_reduction_kind_compatibility libdevice_sigmoid; do \ - git -C /workspace/hexagon-mlir/triton apply \ - "/workspace/triton-npu/setup/patches/triton/$p.patch"; \ - done -# nvvm_reduction_kind_compatibility is REQUIRED, not cosmetic: the NVVM -# ReductionKind API moved in this LLVM and triton does not compile without it. - -# 3. The prebuilt toolchain: LLVM 23 install, the CONFIG_DESC-capable spike, and -# the triton runtime (the mlir-env venv + triton-shared-opt + the built .so -# files). Each unpacks rooted at /workspace. A private repo's assets are only -# reachable through the API asset id with an octet-stream Accept header -- the -# /releases/download/ URL 404s even with a token. + git clone "https://x-access-token:${TOKEN}@github.com/${TNPU_REPO}.git" \ + /workspace/triton-npu && \ + git -C /workspace/triton-npu checkout -q "${TNPU_REF}" && \ + git -C /workspace/triton-npu remote set-url origin \ + "https://github.com/${TNPU_REPO}.git" + +# restore.sh owns every pin (setup/versions.env) and unpacks LLVM 23, spike and +# the triton runtime into /workspace. RUN --mount=type=secret,id=tnpu_token \ - TOKEN="$(cat /run/secrets/tnpu_token)" && \ - for spec in "${TNPU_LLVM23_ASSET_ID}" "${TNPU_SPIKE_ASSET_ID}" "${TNPU_RUNTIME_ASSET_ID}"; do \ - curl -fL --retry 3 \ - -H "Authorization: Bearer ${TOKEN}" \ - -H "Accept: application/octet-stream" \ - "https://api.github.com/repos/${TNPU_REPO}/releases/assets/${spec}" \ - | tar -xz -C /workspace ; \ - done - -# 4. Re-create the triton backend links. `setup/package.sh` collects the triton -# tree with `find python -name '*.so'`, which misses -# python/triton/backends/{amd,nvidia,triton_shared} -- those are not source -# files but symlinks the build creates (setup.py:add_link_to_backends), so they -# are in neither the tarball nor git. Without them `import triton` dies in -# entry-point discovery with "No module named 'triton.backends.amd'". -# triton_shared is reconstructed from upstream microsoft/triton-shared: the fork -# the manifest used to name is gone, and tnpu never calls this backend -- it only -# has to import so discovery succeeds. -RUN T=/workspace/hexagon-mlir/triton && \ - for b in amd nvidia; do \ - ln -sfnT "$T/third_party/$b/backend" "$T/python/triton/backends/$b" && \ - if [ -d "$T/third_party/$b/language" ]; then \ - for x in "$T/third_party/$b/language"/*; do \ - ln -sfnT "$x" "$T/python/triton/language/extra/$(basename "$x")" ; \ - done ; \ - fi ; \ - done && \ - mkdir -p /workspace/hexagon-mlir/triton_shared/backend && \ - for f in compiler.py driver.py name.conf; do \ - curl -fsSL "https://raw.githubusercontent.com/microsoft/triton-shared/main/backend/$f" \ - -o "/workspace/hexagon-mlir/triton_shared/backend/$f" ; \ - done && \ - ln -sfnT /workspace/hexagon-mlir/triton_shared/backend "$T/python/triton/backends/triton_shared" + GITHUB_TOKEN="$(cat /run/secrets/tnpu_token)" \ + /workspace/triton-npu/setup/restore.sh --prebuilt -# tnpu's own defaults already point at these paths (tnpu/config.py); TNPU_DIR is -# the one thing PyTorchSim needs told, since the checkout is not under TORCHSIM_DIR. ENV TNPU_DIR=/workspace/triton-npu -# Fail the build rather than the first CI job if the toolchain is not usable. +# Fail the build, not the first CI job. RUN python3 /workspace/triton-npu/run.py doctor diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 86b94da4..ddc55552 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -136,11 +136,6 @@ missing rather than letting a docker build die deep: came from, so they have to be mirrored across (or the repo made public and the manifest pointed at whichever holds them). -`Dockerfile.tnpu` mirrors `triton-npu/setup/restore.sh --prebuilt` rather than -calling it, because that script has two problems here, both worked around in place -with comments: it 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}` — those are symlinks the -build creates, so they are in neither the tarball nor git, and without them -`import triton` dies in entry-point discovery. Both are worth fixing upstream. - +`Dockerfile.tnpu` clones the harness and runs its own `setup/restore.sh +--prebuilt`; the pins all live in that repo's `setup/versions.env`. `ref` in the +manifest is a commit, so an upstream change there moves this image's tag too. diff --git a/scripts/ci/tnpu_asset_env.sh b/scripts/ci/tnpu_asset_env.sh deleted file mode 100644 index ecc56a63..00000000 --- a/scripts/ci/tnpu_asset_env.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# Emit TNPU_LLVM23_ASSET_ID / TNPU_SPIKE_ASSET_ID / TNPU_RUNTIME_ASSET_ID lines -# for appending to GITHUB_ENV. -# -# Same shape as thirdparty_github_asset_env.sh, with one difference that matters: -# PSAL-POSTECH/triton-npu is PRIVATE, so the default `secrets.GITHUB_TOKEN` (which -# is scoped to the PyTorchSim repo) cannot read its releases. Pass a PAT with read -# access instead -- the workflow puts secrets.TNPU_TOKEN here. -# -# Release assets of a private repo are also only reachable through the API asset -# id, not the /releases/download/ URL, which is why ids are resolved at all -# (thirdparty/github-releases.json documents the same constraint). -# -# Requires: jq, curl, GITHUB_TOKEN, repo root as cwd or GITHUB_WORKSPACE. -set -euo pipefail -ROOT="${GITHUB_WORKSPACE:-$(cd "$(dirname "$0")/../.." && pwd)}" -MANIFEST="${ROOT}/thirdparty/triton-npu.json" -if [ ! -f "$MANIFEST" ]; then - echo "Missing tnpu manifest: $MANIFEST" >&2 - exit 1 -fi -if [ -z "${GITHUB_TOKEN:-}" ]; then - echo "GITHUB_TOKEN is not set (needs a PAT that can read the private" >&2 - echo "PSAL-POSTECH/triton-npu; the default Actions token cannot)" >&2 - exit 1 -fi - -REPO=$(jq -r '.triton_npu.repository' "$MANIFEST") -TAG=$(jq -r '.triton_npu.release_tag' "$MANIFEST") -OWNER="${REPO%%/*}" -NAME="${REPO##*/}" - -if [ "$TAG" = "latest" ]; then - API_URL="https://api.github.com/repos/${OWNER}/${NAME}/releases/latest" -else - API_URL="https://api.github.com/repos/${OWNER}/${NAME}/releases/tags/${TAG}" -fi - -TMP=$(mktemp) -trap 'rm -f "$TMP"' EXIT -if ! curl -fsS -H "Authorization: Bearer ${GITHUB_TOKEN}" \ - -H "Accept: application/vnd.github+json" \ - -H "X-GitHub-Api-Version: 2022-11-28" \ - "$API_URL" -o "$TMP"; then - echo "Failed to read release ${TAG} of ${REPO}." >&2 - echo "Either the token cannot see the repo, or the release does not exist" >&2 - echo "yet -- as of writing, ${REPO} carries no releases and the toolchain" >&2 - echo "assets live only on the upstream fork. See thirdparty/triton-npu.json." >&2 - exit 1 -fi - -emit() { # emit - local id - id=$(jq -r --arg n "$1" '.assets[] | select(.name == $n) | .id' "$TMP" | head -n1) - if [ -z "$id" ] || [ "$id" = "null" ]; then - echo "Release ${TAG} of ${REPO} has no asset named '$1'" >&2 - exit 1 - fi - echo "$2=${id}" -} - -emit llvm23-install.tar.gz TNPU_LLVM23_ASSET_ID -emit spike-install.tar.gz TNPU_SPIKE_ASSET_ID -emit triton-runtime.tar.gz TNPU_RUNTIME_ASSET_ID diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 408f4f6e..9a744211 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -1,22 +1,8 @@ { - "description": "Pins for the triton-npu (tnpu) toolchain layer, used only by the Triton codegen route. Kept SEPARATE from github-releases.json on purpose: that file plus Dockerfile.base is hashed into the torchsim_base tag, so putting these here means the main base image is not rebuilt when the tnpu pins move, and the ~1.8 GiB toolchain is not forced onto every CI job. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing (pin = sha256 of this file plus Dockerfile.tnpu). The repository is PRIVATE and its release assets are not readable by the default Actions token, so the workflow needs a PAT in secrets.TNPU_TOKEN with read access to it.", + "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "main", - "release_tag": "toolchain-llvm23", - "assets": [ - "llvm23-install.tar.gz", - "spike-install.tar.gz", - "triton-runtime.tar.gz" - ] - }, - "triton": { - "description": "The triton checkout the prebuilt runtime was built against. Its editable install points back into this tree, so the sources must be present at exactly this revision for `import triton` to work in the tnpu venv. Mirrors triton-npu/setup/versions.env (TRITON_SHA); a main-branch commit, NOT the v3.6.0 tag, because llvm-hash.txt moved from LLVM 22 to 23 after the tag.", - "repository": "triton-lang/triton", - "commit": "df38505e451a1541555379bcf378be9e8c00545c", - "patches": [ - "setup/patches/triton/nvvm_reduction_kind_compatibility.patch", - "setup/patches/triton/libdevice_sigmoid.patch" - ] + "ref": "e686799e4681c9e01151e3d8237867c55ba285dd", + "release_tag": "toolchain-llvm23" } } From e4af91bc898f3277685ccdf22abb16d446d53441 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:32:04 +0900 Subject: [PATCH 05/35] [CI] Note only the token as a prerequisite The toolchain-llvm23 release now exists on PSAL-POSTECH/triton-npu, so secrets.TNPU_TOKEN is the only thing left to set up. --- .github/workflows/triton_npu.yml | 13 +++---------- PyTorchSimFrontend/triton_backend/README.md | 15 +++------------ 2 files changed, 6 insertions(+), 22 deletions(-) diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index 92290f5f..d38c6f19 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -8,16 +8,9 @@ name: Triton codegen route (triton-npu) # is ~1.8 GiB that no other job needs, so it neither gates PRs nor slows them # down. Promote the jobs into pytorchsim_test.yml once the route runs end to end. # -# PREREQUISITES, both one-time and both outside this repo: -# 1. secrets.TNPU_TOKEN -- a PAT that can read PSAL-POSTECH/triton-npu. That -# repo is private, and the default Actions token is scoped to this -# repository only, so it cannot clone it or read its releases. -# 2. A release tagged `toolchain-llvm23` on PSAL-POSTECH/triton-npu carrying -# llvm23-install.tar.gz, spike-install.tar.gz and triton-runtime.tar.gz. -# As of writing that repo has no releases; the assets exist only on the -# upstream fork it was forked from. -# Both are checked by the preflight job below, which reports which one is missing -# instead of failing deep inside a docker build. +# Needs secrets.TNPU_TOKEN: a PAT that can read PSAL-POSTECH/triton-npu and its +# toolchain-llvm23 release. The repo is private and the default Actions token is +# scoped to this repository. preflight checks it before the docker build. on: pull_request: diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index ddc55552..12b01b9d 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -123,18 +123,9 @@ one of those moves, and its tag also carries the base pin it was built on. `mlir-route-regression` is there because this layer adds a *second* LLVM and a *second* triton to the image; it checks the production path did not notice. -**Two prerequisites, both outside this repo.** `preflight` fails with which one is -missing rather than letting a docker build die deep: - -1. `secrets.TNPU_TOKEN` — a PAT that can read `PSAL-POSTECH/triton-npu`. That repo - is private, and the default Actions token is scoped to this repository, so it - can neither clone it 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, so they have to be mirrored across (or the repo made public and the - manifest pointed at whichever holds them). +**Needs `secrets.TNPU_TOKEN`** — a PAT that can read `PSAL-POSTECH/triton-npu` +and its `toolchain-llvm23` release. That repo is private and the default Actions +token is scoped to this repository. `preflight` checks it before the build. `Dockerfile.tnpu` clones the harness and runs its own `setup/restore.sh --prebuilt`; the pins all live in that repo's `setup/versions.env`. `ref` in the From 5f632d647316a06dac641c44440412aaf4192cfc Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:46:25 +0900 Subject: [PATCH 06/35] [CI] Bump the triton-npu pin to de2c767 --- thirdparty/triton-npu.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 9a744211..e8ae5a77 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "e686799e4681c9e01151e3d8237867c55ba285dd", + "ref": "de2c767800390e2a15bf83d2b48460859c5be9f7", "release_tag": "toolchain-llvm23" } } From 8faa1e509bd20973de2689e7e8a53b624aaef361 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 19:52:46 +0900 Subject: [PATCH 07/35] [CI] Bump the triton-npu pin to e113819 --- thirdparty/triton-npu.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index e8ae5a77..c5373872 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "de2c767800390e2a15bf83d2b48460859c5be9f7", + "ref": "e1138197961707aab495db0a43a344268700286f", "release_tag": "toolchain-llvm23" } } From b442aa06d3ef9aa84b2a57298c83a211d2de6e59 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 20:06:44 +0900 Subject: [PATCH 08/35] [CI] Bump the triton-npu pin to fd423b0 --- thirdparty/triton-npu.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index c5373872..0b2172bd 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "e1138197961707aab495db0a43a344268700286f", + "ref": "fd423b001ad562b4cbcd6405b03df2b0674656be", "release_tag": "toolchain-llvm23" } } From 600a6a93e0f99d5b027194b5ee581c55a6d96917 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:09:27 +0900 Subject: [PATCH 09/35] [Frontend] Rename print_operation to visit_operation 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. --- PyTorchSimFrontend/mlir/passes/build_tog.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index 5a40feec..fba9f5ea 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -568,7 +568,13 @@ def _process_dram_indices(self, value, loop_index_list, indirect_box): loop_index_list.append(("c" + str(c), c)) # ---- main recursion ---- - def print_operation(self, op, node): + def visit_operation(self, op, node): + """Walk `op` and attach the nodes it produces under `node`. + + Builds the graph; it does not print. (The C++ pass this is ported from + does both in one method, `printOperation` -- here `bfs`/`display` own + the printing.) + """ name = _op_name(op) if name in SKIP_OPS: return @@ -605,7 +611,7 @@ def bool_true(k): for region in oper.regions: for block in region.blocks: for inner in block.operations: - self.print_operation(inner, loop_node) + self.visit_operation(inner, loop_node) return if name == "togsim.transfer": @@ -1087,7 +1093,7 @@ def _build(module, builder): continue root = TOGNode("root") builder._reset_matmul_fsm() - builder.print_operation(op, root) + builder.visit_operation(op, root) root.bfs(out) return "".join(out) From c0df618b5ea03553970afe88c8ce95a2927f884c Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:09:27 +0900 Subject: [PATCH 10/35] [Frontend] Give the trace pipeline a kernel whose body is one work-item 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. --- .../mlir/passes/build_skeleton.py | 2 + PyTorchSimFrontend/mlir/passes/build_tog.py | 85 ++++++++-- .../mlir/passes/lower_to_emitc.py | 158 ++++++++++++++++-- 3 files changed, 218 insertions(+), 27 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index f4ed7d0d..ed52a56d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -366,6 +366,8 @@ def visit(n): for ln in builder.loop_nodes: visit(ln) + for dn in getattr(builder, "dma_nodes", ()): # DMAs outside any tile loop + visit(dn) return by_op diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index fba9f5ea..98590d88 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -414,6 +414,9 @@ def __init__(self): self.loop_var_name = {} # value-identity-key -> loop name self.compute_nodes = [] self.loop_nodes = [] + # `_collect_dma_nodes` descends from the loop nodes, so a DMA hanging off + # the root (no tile loop in the kernel) would be missed. + self.dma_nodes = [] self._reset_matmul_fsm() # ---- matmul FSM ---- @@ -825,9 +828,14 @@ def _handle_dma_start(self, op, node): loop_idx_list.append(key) loop_stride_list.append(reordered[key]) - # base address + # base address: which tensor this DMA touches. The operand is the block + # argument itself in PyTorchSim's codegen; when it is a view of one + # instead, only the producer knows which -- so it says so (`dram_arg`) + # rather than the consumer guessing its way back through view ops. address = "arg" - if _is_block_arg(dram_memref): + if "dram_arg" in oper.attributes: + address += str(ir.IntegerAttr(oper.attributes["dram_arg"]).value) + elif _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) # element size @@ -881,6 +889,7 @@ def _handle_dma_start(self, op, node): tag_stride_list, loop_idx_list, loop_stride_list, indirect_box[0]) dma_node.op = op + self.dma_nodes.append(dma_node) node.add_child(dma_node) dma_node.add_parent(node) @@ -924,7 +933,9 @@ def _handle_dma_wait(self, op, node): dram_memref = f["dst"] elif dst_space == 1 and src_space == 0: dram_memref = f["src"] - if dram_memref is not None and _is_block_arg(dram_memref): + if "dram_arg" in user.attributes: + address += str(ir.IntegerAttr(user.attributes["dram_arg"]).value) + elif dram_memref is not None and _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) if len(tag_stride_list) == 0: @@ -934,6 +945,7 @@ def _handle_dma_wait(self, op, node): wait_node = TOGDMAWaitNode("DMAWaitNode", tag_index_list, tag_stride_list, tag_divider_list, address) wait_node.op = op + self.dma_nodes.append(wait_node) node.add_child(wait_node) wait_node.add_parent(node) @@ -1070,12 +1082,47 @@ def _insert_compute_markers(builder): # Driver. # --------------------------------------------------------------------------- def _find_kernel(module): - for op in module.body.operations: - if op.operation.name != "func.func": - continue + """The kernel function: named `kernel` in PyTorchSim's codegen, else the + module's only func.func (triton-npu carries the Triton kernel's own name). + Declines when there is more than one -- the intent would be a guess.""" + funcs = [op for op in module.body.operations + if op.operation.name == "func.func"] + for op in funcs: if ir.StringAttr(op.operation.attributes["sym_name"]).value == "kernel": return op - return None + return funcs[0] if len(funcs) == 1 else None + + +#: The loop roles (sec 9.1). Without one, a loop is a micro-loop the compute FSM +#: folds into a single node, not a tile loop. +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") + + +def _has_loop_role(op): + attrs = op.operation.attributes + return any(k in attrs and ir.BoolAttr(attrs[k]).value for k in _LOOP_ROLE_ATTRS) + + +def _is_address_plumbing(op): + """Scalar index/integer math (DMA offsets, mask extents) and the terminator. + + Only consulted on the no-top-level-loop path. PyTorchSim's codegen puts this + math in `affine.apply`, which SKIP_OPS drops; triton-npu emits an + arith/index_cast chain that would otherwise count as vector compute. + + Keyed on result type: tile data here is always vector- or float-typed. A + top-level SCALAR arithmetic kernel would be misread, but no path emits one. + """ + name = _op_name(op) + if name in ("func.return", "memref.cast"): + return True + if not name.startswith("arith."): + return False + results = list(op.operation.results) + if not results: + return False + return all(ir.IndexType.isinstance(r.type) or ir.IntegerType.isinstance(r.type) + for r in results) def _build(module, builder): @@ -1088,13 +1135,29 @@ def _build(module, builder): block = func_op.regions[0].blocks[0] out = [] + # A root is a top-level TILE loop, identified by its role attribute (sec + # 9.1) -- not by being an affine.for: bank_vectorize leaves a bare one for + # the tile's vector work, and rooting there orphans every DMA. + roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] + if roots: + for op in roots: + root = TOGNode("root") + builder._reset_matmul_fsm() + builder.visit_operation(op, root) + root.bfs(out) + return "".join(out) + + # No top-level loop: the body is ONE work-item -- the shape a Triton kernel + # arrives in, its grid becoming the trace producer's dispatch loop (sec 9.3). + # PyTorchSim's codegen keeps the tile loops in the kernel and never lands here. + root = TOGNode("root") + builder._reset_matmul_fsm() for op in block.operations: - if op.operation.name != "affine.for": + if _is_address_plumbing(op): continue - root = TOGNode("root") - builder._reset_matmul_fsm() builder.visit_operation(op, root) - root.bfs(out) + root.bfs(out) return "".join(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 5633769a..bfb94e34 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -119,20 +119,130 @@ def _attr_bool(op, key): # --------------------------------------------------------------------------- # step 1: rewrite signature + togsim.* ops (the unregistered-op glue) # --------------------------------------------------------------------------- -def _strip_aux(module): - """Erase memref.global decls and every func except @kernel (the wrapper).""" +def _strip_aux(module, keep=None): + """Erase memref.global decls and every func except the kernel. + + `keep` is the kernel op: its name is `kernel` only in PyTorchSim's codegen, + so the caller passes what `_find_kernel` resolved. + """ + keep_op = keep.operation if keep is not None else None victims = [] for op in module.body.operations: name = op.operation.name if name == "memref.global": victims.append(op) elif name == "func.func": - if ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": + if keep_op is not None: + if op.operation != keep_op: + victims.append(op) + elif ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": victims.append(op) for op in victims: op.operation.erase() +class WorkItem: + """A kernel whose body is ONE work-item, plus the grid over it. + + A Triton kernel describes a single program instance; the grid lives outside + it. The trace producer already splits the same way (design sec 9.3), so only + the enumeration is missing. + + `parallel_args` are the argument positions holding the program ids + (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` + their extents. Both outermost-first. + """ + + def __init__(self, parallel_args, grid): + if len(parallel_args) != len(grid): + raise ValueError( + f"parallel_args {parallel_args} and grid {grid} must have the " + f"same length -- one program-id argument per grid axis") + self.parallel_args = list(parallel_args) + self.grid = [int(g) for g in grid] + + +def _materialize_grid_loop(kernel, work_item, ctx): + """Wrap the body in the grid loop the Triton kernel does not carry: + + func @k(..., %pid: i32) { + scf.for %p = 0 to G { index_cast %p> } {outer_loop} + } + + Downstream is then unchanged: `_parallel_loop_chain` finds the tagged loop, + the outliner threads its induction variable through `iv[]`, and the loop left + behind becomes the dispatch enumeration. `outer_loop` means "independent + work-item" (sec 9.1) -- exactly a Triton program id. + + MUST run before `_rewrite_signature`, which erases the arguments and first + asserts none are still used. + """ + from mlir.dialects import arith, scf + + block = kernel.regions[0].blocks[0] + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + + pid_args = [block.arguments[i] for i in work_item.parallel_args] + body_ops = [o for o in block.operations + if o.operation.name not in _LOOP_TERMINATORS] + terminator = [o for o in block.operations + if o.operation.name in _LOOP_TERMINATORS][0] + + # Every bound first, and all of them before the first loop: each is created + # just before the terminator, so one made after an outer loop would sit + # BELOW it in the block while an inner loop uses it -- which does not + # dominate, and the verifier rejects it (only reachable at rank >= 2). + with ir.InsertionPoint(terminator), loc: + c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result + c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e)).result + for e in work_item.grid] + + loops, inner = [], None + for ub in ubs: + # Nest inside the previous loop, BEFORE its yield: InsertionPoint on a + # block appends, and an scf.for body is already terminated. + ip = ir.InsertionPoint(terminator) if inner is None \ + else ir.InsertionPoint.at_block_terminator(inner.body) + with ip, loc: + loop = scf.ForOp(c0, ub, c1) + # ForOp leaves the body empty here; scf.for needs a terminator, and + # _outline_work_item inserts before it. + if len(loop.body.operations) == 0: + with ir.InsertionPoint(loop.body), loc: + scf.YieldOp([]) + loop.operation.attributes["outer_loop"] = ir.BoolAttr.get(True) + loops.append(loop) + inner = loop + + # Move the tile body inside the innermost loop, ahead of its yield. + inner_block = inner.body + inner_term = inner_block.operations[len(inner_block.operations) - 1] + for op in body_ops: + op.operation.move_before(inner_term) + + # Program ids are i32, induction variables index: cast once, at the top. + with ir.InsertionPoint(inner_block.operations[0]), loc: + casts = [] + for loop, pid in zip(loops, pid_args): + iv = loop.body.arguments[0] + casts.append(arith.IndexCastOp(pid.type, iv).result + if pid.type != idxty else iv) + + for pid, new in zip(pid_args, casts): + _replace_all_uses(pid, new) + + +def _replace_all_uses(old, new): + """The bindings expose no replaceAllUsesWith on a Value.""" + for use in list(old.uses): + owner = use.owner + for i in range(len(owner.operands)): + if owner.operands[i] == old: + owner.operands[i] = new + + def _rewrite_signature(kernel, ctx): """Replace @kernel's memref tensor args with the ABI args (EmitCtx*, int64_t* shape_args, int32_t n) and rename it to togsim_kernel. @@ -196,15 +306,22 @@ def _is_outer(forop): return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value +#: The role is carried by the `outer_loop` attribute, not the dialect: +#: PyTorchSim's codegen emits affine.for, _materialize_grid_loop scf.for. Both +#: keep the induction variable in block argument 0. +_LOOP_OPS = ("affine.for", "scf.for") +_LOOP_TERMINATORS = ("affine.yield", "scf.yield", "func.return") + + def _parallel_loop_chain(block): - """The nested chain of `affine.for {outer_loop}` from `block` inward (one + """The nested chain of `{outer_loop}` loops from `block` inward (one work-item's parallel indices). Empty if the kernel has no parallel loop.""" chain = [] cur = block while True: nxt = None for op in cur.operations: - if op.operation.name == "affine.for" and _is_outer(op): + if op.operation.name in _LOOP_OPS and _is_outer(op): nxt = op break if nxt is None: @@ -281,7 +398,7 @@ def _outline_work_item(ctx, kernel, ctx_val): # move the work-item body into the tile fn (terminators stay behind). for op in [o for o in Lbody.operations - if o.operation.name not in ("affine.yield", "func.return")]: + if o.operation.name not in _LOOP_TERMINATORS]: op.operation.move_before(tret) # remap captures (Value `==` is identity): ctx -> ctx2, each parallel IV -> @@ -337,7 +454,7 @@ def _remap(block): # --- the dispatcher: marshal the IVs and hand the tile fn to togsim_dispatch --- term = [o for o in Lbody.operations - if o.operation.name in ("affine.yield", "func.return")][0] + if o.operation.name in _LOOP_TERMINATORS][0] fn_ref = _opaque(ctx, ts.TILE_SYMBOL) # function name -> verbatim pointer in C with ir.InsertionPoint(term): if ivs: @@ -499,15 +616,22 @@ def _add_extern_c(module, ctx): # --------------------------------------------------------------------------- # driver # --------------------------------------------------------------------------- -def lower_to_emitc(skeleton_module): +def lower_to_emitc(skeleton_module, work_item=None): """Lower a skeleton+API module (in place) to an EmitC module with the - `togsim_kernel` entry function. Returns the same module.""" + `togsim_kernel` entry function. Returns the same module. + + `work_item` is for kernels whose body is one work-item with the grid outside + (Triton's shape); None keeps PyTorchSim's, where the tile loops are already + in the kernel. + """ ctx = skeleton_module.context kernel = _find_kernel(skeleton_module) if kernel is None: - raise ValueError("no @kernel found in skeleton module") + raise ValueError("no kernel function found in skeleton module") - _strip_aux(skeleton_module) + _strip_aux(skeleton_module, keep=kernel) + if work_item is not None: + _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch @@ -563,18 +687,20 @@ def _default_include_dir(): return os.path.join(root, "TOGSim", "include") -def skeleton_to_so(skeleton_module, so_path, include_dir=None): +def skeleton_to_so(skeleton_module, so_path, include_dir=None, work_item=None): """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the EmitC module text (for inspection / caching).""" - emitc = lower_to_emitc(skeleton_module) + emitc = lower_to_emitc(skeleton_module, work_item=work_item) inc = include_dir or _default_include_dir() cpp = emitc_to_cpp(emitc, include_dir=inc) compile_so(cpp, so_path, inc) return str(emitc) -def build_trace_so(postvcix_path, so_path, include_dir=None): - """Full P2 path from a post-vcix kernel .mlir to a trace `.so`.""" +def build_trace_so(postvcix_path, so_path, include_dir=None, work_item=None): + """Full P2 path from a post-vcix kernel .mlir to a trace `.so`. + + `work_item` -- see lower_to_emitc.""" from . import build_skeleton as bs ctx = ir.Context() @@ -582,7 +708,7 @@ def build_trace_so(postvcix_path, so_path, include_dir=None): with ctx: module = ir.Module.parse(open(postvcix_path).read(), ctx) bs.build_skeleton(module) - return skeleton_to_so(module, so_path, include_dir) + return skeleton_to_so(module, so_path, include_dir, work_item=work_item) def main(argv): From 15ce2a18465ee1f3e05e3244c1ff2aa4023d5445 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 11/35] [Frontend] Run the Triton route's kernel through TOGSim 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. --- .../triton_backend/codecache.py | 28 +++-- .../triton_backend/kernel_spec.py | 10 +- PyTorchSimFrontend/triton_backend/timing.py | 113 ++++++++++++++++++ tests/system/test_triton_codegen.py | 89 +++++++++++++- 4 files changed, 222 insertions(+), 18 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/timing.py diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index b4817cc6..69805a87 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -18,7 +18,7 @@ from torch._inductor.codecache import get_hash from PyTorchSimFrontend import extension_config -from . import kernel_spec, tnpu_bridge +from . import kernel_spec, timing, tnpu_bridge logger = extension_config.setup_logger() @@ -43,17 +43,20 @@ def __init__(self, kernel_name, workdir, meta): self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") def __call__(self, *args): - raise NotImplementedError( - f"{self.kernel_name}: compiled to {self.elf}, but the launch is not " - f"wired yet. Two pieces are missing and both are tracked in " - f"triton_backend/README.md:\n" - f" 1. functional -- marshal the caller's tensors into " - f"{self.workdir}/runtime/*.raw, run Spike on the ELF, read the " - f"outputs back into the caller's tensors;\n" - f" 2. timing -- emit trace.so + trace_cycles.tsv from the tnpu IR " - f"and hand them to TOGSim (needs the build_tog adapters).\n" - f"Compilation itself succeeded, so the codegen half of this route " - f"is exercised by getting this far.") + """One launch of the whole grid: simulate, return TOGSim's result. + + Does NOT write the caller's output tensors -- the functional launch is + not wired (README). Logged, so an undefined value cannot pass for a + computed one. + """ + if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): + timing.emit_trace(self.workdir, self.meta) + result = timing.run_togsim(self.workdir) + logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) + logger.warning( + "[Spike] %s: output tensors are NOT written; the functional launch " + "(tensors -> Spike -> tensors) is not wired yet", self.kernel_name) + return result def triton_npu_compile(src_code, meta, kernel_name): @@ -75,6 +78,7 @@ def triton_npu_compile(src_code, meta, kernel_name): tnpu_bridge.tnpu_dir()) with open(os.path.join(write_path, "kernel.py"), "w") as f: f.write(src_code) # the unmodified Inductor source, for diffing + timing.store_meta(write_path, meta) # lets the timing step run standalone tnpu_bridge.run_pipeline(spec_path, write_path, to_stage="binary") logger.info("[triton-npu] %s -> %s", kernel_name, write_path) return TritonNPULauncher(kernel_name, write_path, meta) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index ef779913..d74ea962 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -218,8 +218,12 @@ def strip_for_tnpu(src): return prefix + body -def _grid(meta): - """Sequential launch grid, from the numels and the pinned block sizes.""" +def grid_of(meta): + """Launch grid, from the numels and the pinned block sizes. + + Also read by the timing path, which needs the same extents to enumerate the + work-items -- so it lives here rather than being recomputed per consumer. + """ x = meta["numels"].get("xnumel") xblock = (meta.get("fixed_config") or {}).get("XBLOCK") if x is None or not xblock: @@ -330,7 +334,7 @@ def write_spec_file(src_code, meta, path, tnpu_dir): constexprs=constexprs, args_body=args_body, make_inputs_body=make_inputs_body, - grid=_grid(meta), + grid=grid_of(meta), ) with open(path, "w") as f: f.write(text) diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py new file mode 100644 index 00000000..9a045304 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -0,0 +1,113 @@ +"""The timing half of the Triton route: tnpu IR -> trace.so -> TOGSim. + +TOGSim simulates from a compiled trace producer (docs/design/togsim_cpp_trace.md). +PyTorchSim's codegen already emits one; this emits the same from a Triton-shaped +kernel, where the grid must be supplied -- see `lower_to_emitc.WorkItem`. + + emit_trace(workdir, meta) 04-custom.mlir -> trace.so + trace_cycles.tsv + run_togsim(workdir, ...) hand them to TOGSim, return its parsed result +""" + +import json +import os + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +#: Name TOGSim derives from the kernel directory (Simulator/simulator.py). +TRACE_SO = "trace.so" +CYCLE_TSV = "trace_cycles.tsv" +META_JSON = "meta.json" + +#: Stand-in per-tile cost until gem5 sampling lands. Deliberately not a +#: plausible-looking number: only an obvious non-measurement gets fixed. +PLACEHOLDER_CYCLE = 1 + + +def _runtime_arg_layout(meta): + """(n_tensor_args, n_scalar_args) of the lowered signature. + + triton-shared lays it out as pointers, user scalars, then its own six + (gridX,Y,Z / pidX,Y,Z). constexpr params never become arguments. + """ + sig = meta["signature"] + tensors = [k for k, v in sig.items() if v.startswith("*")] + scalars = [k for k, v in sig.items() + if not v.startswith("*") and v != "constexpr"] + return len(tensors), len(scalars) + + +def work_item_for(meta): + """The WorkItem describing this kernel's program-id args and grid extents.""" + from PyTorchSimFrontend.mlir.passes.lower_to_emitc import WorkItem + from . import kernel_spec + + n_tensor, n_scalar = _runtime_arg_layout(meta) + pid_x = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + grid = list(kernel_spec.grid_of(meta)) + return WorkItem(parallel_args=list(range(pid_x, pid_x + len(grid))), grid=grid) + + +def emit_trace(workdir, meta): + """Build `trace.so` + `trace_cycles.tsv` from tnpu's post-vcix IR. + + Returns the number of compute tiles the cycle table covers. + """ + from PyTorchSimFrontend.mlir.passes import build_skeleton as bs + from PyTorchSimFrontend.mlir.passes import cycle_table as ct + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + postvcix = os.path.join(workdir, "04-custom.mlir") + if not os.path.isfile(postvcix): + raise FileNotFoundError( + f"{postvcix} not found -- tnpu must have run at least to stage 4 " + f"(the post-vcix IR is what the trace is built from)") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(postvcix).read(), ctx) + bs.build_skeleton(module) + n_tiles = len(ct._compute_types(module)) + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), + work_item=work_item_for(meta)) + + # Until gem5 sampling lands, every tile costs PLACEHOLDER_CYCLE: TOGSim + # models the DMA and the dependency structure but NOT compute latency. + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) + logger.warning("[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d tiles); " + "compute latency is not modelled yet", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + return n_tiles + + +def run_togsim(workdir, attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict.""" + from Simulator.simulator import TOGSimulator + + so = os.path.join(workdir, TRACE_SO) + if not os.path.isfile(so): + raise FileNotFoundError(f"{so} not found -- call emit_trace first") + + # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its + # DIRECTORY, and reads the file itself only on the STONNE path. + handle = os.path.join(workdir, "tile_graph.onnx") + result_path = TOGSimulator.run_standalone( + handle, attribute_path or os.path.join(workdir, "attribute"), + timeout_sec=timeout_sec) + return TOGSimulator.get_result_from_file(result_path) + + +def store_meta(workdir, meta): + """Persist codegen metadata beside the artifacts, so the timing step can run + standalone.""" + with open(os.path.join(workdir, META_JSON), "w") as f: + json.dump(meta, f, indent=2) + + +def load_meta(workdir): + with open(os.path.join(workdir, META_JSON)) as f: + return json.load(f) diff --git a/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py index a7b0af79..b36715d6 100644 --- a/tests/system/test_triton_codegen.py +++ b/tests/system/test_triton_codegen.py @@ -31,10 +31,73 @@ def fn(x, y): return fn, x, y +def check_multi_axis_grid(): + """A 2-D grid must nest one loop per axis and hand both indices to iv[]. + + Guards the multi-axis path, which the add kernel does not reach: Inductor + only uses y/z when x would overflow, so a 1-D grid exercises just the first + iteration of the nest. + """ + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + src = """ + module { + func.func @k(%arg0: memref<*xf32>, %arg1: i32, %arg2: i32) { + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : i32 + %a = arith.muli %arg1, %c8 : i32 + %b = arith.addi %a, %arg2 : i32 + %o = arith.index_cast %b : i32 to index + "togsim.dma"(%o, %c0) {arg_id = 0 : i32, base = "arg0", dims = [128], + dir = 0 : i32, elem_bits = 32 : i32, is_async = false, read_bufs = [], + strides = [1], tag_id = 0 : i32, write_bufs = [0]} : (index, index) -> () + return + } + } + """ + problems = [] + # Verify the IR the pass itself produces: a bound created after an outer loop + # would not dominate an inner loop's use of it, which only shows at rank >= 2 + # and which the emitc lowering happens to paper over. + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + l2e._materialize_grid_loop( + l2e._find_kernel(module), + l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3]), ctx) + try: + module.operation.verify() + except Exception as e: # noqa: BLE001 + problems.append(f"materialized IR does not verify: {e}") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + emitc = l2e.lower_to_emitc( + module, work_item=l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3])) + cpp = l2e.emitc_to_cpp(emitc, include_dir=l2e._default_include_dir()) + + entry = cpp.split("togsim_kernel(EmitCtx*")[-1] + if entry.count("for (") != 2: + problems.append(f"expected 2 nested loops, found {entry.count('for (')}") + if "togsim_dispatch" not in entry: + problems.append("no togsim_dispatch call") + if ", 2);" not in entry: + problems.append("dispatch does not pass 2 indices") + for p in problems: + print(f" multi-axis grid: {p}") + return not problems + + def main(): from PyTorchSimFrontend import extension_config from PyTorchSimFrontend.triton_backend import tnpu_bridge + print(f"multi-axis grid = " + f"{'ok' if check_multi_axis_grid() else 'FAILED'}") print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") ok, _out = tnpu_bridge.doctor() @@ -56,9 +119,29 @@ def main(): "above and README.md's gap list.") return 1 - err = (got.cpu() - expected).abs().max().item() - print(f"max_abs_err = {err}") - return 0 if err < 1e-4 else 1 + # Values are NOT checked: the launch simulates the kernel but does not + # marshal tensors through Spike, so `got` is undefined. What is asserted is + # that the timing path ran end to end -- the two artifacts TOGSim consumes. + del got + import glob + + from PyTorchSimFrontend.triton_backend import timing + + dirs = glob.glob(os.path.join(extension_config.get_dump_path(), "triton_*")) + if not dirs: + print("no kernel directory was produced") + return 1 + workdir = max(dirs, key=os.path.getmtime) + for name in (timing.TRACE_SO, timing.CYCLE_TSV): + path = os.path.join(workdir, name) + if not os.path.isfile(path): + print(f"missing {name} in {workdir}") + return 1 + print(f" {name:18s} {os.path.getsize(path)} bytes") + print(f"\ntiming path OK ({workdir})") + print(f"values NOT verified -- torch would give {expected[:3].tolist()}...; " + f"the functional launch is not wired (triton_backend/README.md)") + return 0 if __name__ == "__main__": From 3c82969fe9ff4b96250e381f7b8ae3d35edeeeaa Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 12/35] [Frontend] Measure the Triton route's tile cost with gem5 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. --- PyTorchSimFrontend/triton_backend/README.md | 47 ++++++++++-- PyTorchSimFrontend/triton_backend/timing.py | 83 ++++++++++++++++++--- 2 files changed, 113 insertions(+), 17 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 12b01b9d..4b409158 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -35,11 +35,19 @@ torch.compile │ a tnpu kernel file (KernelSpec) kernel_spec.py ▼ run.py --to binary (subprocess) tnpu_bridge.py - │ 01-ttir → 02-ttshared → 03-adapted → 04-lowered → 05-*.elf + │ 01-ttir → 02-ttshared → 03-adapted → 04-custom → 05-*.elf ▼ - TritonNPULauncher.__call__ ← NOT WIRED YET + TritonNPULauncher.__call__ codecache.py + ├ timing 04-custom.mlir → trace.so + trace_cycles.tsv → TOGSim + │ cycles measured by gem5 on a one-tile binary + └ functional ← NOT WIRED YET ``` +The timing half reuses PyTorchSim's trace pipeline unchanged. The one structural +difference is that a Triton kernel body is a single program instance, so the grid +that enumerates instances is supplied by `lower_to_emitc.WorkItem` instead of +being read out of the kernel -- see "The grid is not in the kernel" below. + Artifacts land in one directory per source hash under the dump path (`outputs/triton_/`), alongside the unmodified Inductor source (`kernel.py`) so the rewrite is diffable. @@ -54,7 +62,18 @@ Artifacts land in one directory per source hash under the dump path - the lowering is correct in shape: `tl.load/store` became three `togsim.transfer` ops, and Inductor's `xmask` came through as a **masked DMA** (`masked_axes = [0]`, `masked_fill`), which tnpu already supports -- the run stops in `TritonNPULauncher.__call__`, by design +- the trace producer comes out in the shape the design calls for: a + `togsim_kernel_tile` computing `offset = iv[0]*128` around three `togsim_dma` + and one `togsim_compute`, and a `togsim_kernel` looping `p < 8` over + `togsim_dispatch` +- **TOGSim runs it: 650 cycles**, with channel-0 DRAM traffic of 16 reads x 32 B + x 16 channels = 8192 B, exactly the 8 work-items x 2 loads x 512 B the kernel + should move. The MLIR route on the same `x + y` reports 251 cycles -- the same + order, and higher here because tnpu emits synchronous DMA, so nothing overlaps + (gap 2) +- the tile's compute cost is a real measurement: gem5 samples **19 cycles** for + the vector-add tile, via `timing.measure_tile_cycles` +- values are NOT produced: the functional launch is still open (gap 1) ## Gap list, in order @@ -62,17 +81,18 @@ Artifacts land in one directory per source hash under the dump path `runtime/*.raw`, run Spike on the ELF, read outputs back. tnpu's stage 6 does this for its own kernels but generates inputs from the spec; here the tensors come from the caller. -2. **Launch (timing).** Emit `trace.so` + `trace_cycles.tsv` and hand them to - TOGSim. Blocked on the `build_tog` adapters — the tnpu IR is structurally - invisible to it today (no top-level `affine.for`, `scf.for` instead of - `affine.for`, vcix as LLVM intrinsics rather than dialect ops, DMA addresses - as `arith` chains rather than `affine.apply`, grid outside the IR). +2. **Double buffering.** tnpu emits synchronous DMA (`is_async=false`, no + `togsim.wait`), so load → compute → store serialize inside every work-item and + TOGSim has no overlap to model. This is the main remaining gap between the two + routes' cycle counts. 3. **`triton_helpers`.** Any kernel using `triton_helpers.*` (reductions, clamps, `maximum`/`minimum`) cannot compile: the module lives in torch and the tnpu venv has none. `strip_for_tnpu` raises and names the helper. Needs a minimal vendored copy. 4. **Reductions.** Independently blocked in tnpu itself — no lane-aware reduction path; see `triton-npu/kernels/reduce.py`. + Matmul is also still open on the timing side: `build_tog` finds compute nodes + by the `vcix.iv` op name, and tnpu emits `llvm.riscv.sf.vc.*` intrinsics. 5. **Block-size policy.** `fixed_config_for` pins `XBLOCK` to the lane count and deliberately leaves reduction blocks unset. Real tile selection (the MLIR route's autotuner / `codegen_mapping_strategy`) has no equivalent here yet. @@ -103,6 +123,17 @@ GPU, `triton_hash_with_backend()` raises "0 active drivers" because it asks the triton runtime for the current target. We never launch through that runtime, so the value is short-circuited to a deterministic cache key. +**The grid is not in the kernel.** PyTorchSim's codegen puts the tile loops +inside the kernel; a Triton kernel describes one program instance and leaves the +grid to the launch. The trace producer wants that same split already -- +`togsim_kernel_tile` per work-item, enumerated by `togsim_kernel` (design sec +9.3) -- so the models agree and only the enumeration was missing. +`_materialize_grid_loop` supplies it, on the trace artifact only: it wraps the +body in a loop tagged `outer_loop` with each program-id argument replaced by the +induction variable, and everything downstream is unchanged. It runs before +`_rewrite_signature`, which erases the kernel arguments and first asserts none +are still used -- that ordering is what decides where this can live. + ## CI `.github/workflows/triton_npu.yml`, separate from the main CI: this route is WIP, diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index 9a045304..f19423db 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -20,10 +20,56 @@ CYCLE_TSV = "trace_cycles.tsv" META_JSON = "meta.json" -#: Stand-in per-tile cost until gem5 sampling lands. Deliberately not a -#: plausible-looking number: only an obvious non-measurement gets fixed. +#: Used only when gem5 sampling fails. Deliberately not a plausible-looking +#: number: only an obvious non-measurement gets fixed. PLACEHOLDER_CYCLE = 1 +SAMPLE_MLIR = "04-sample.mlir" +CYCLE_BIN = "cycle_bin" + + +def measure_tile_cycles(workdir, meta): + """Per-compute-node cycle counts for ONE tile, measured under gem5. + + build_tog's sample mode marks each compute node and makes every loop a + single trip; tnpu lowers that to a binary (in ITS process -- the Gemmini/VCIX + lowering and its LLVM live there); gem5 runs it. Returns None on any failure, + and the caller falls back to the placeholder table. + """ + from PyTorchSimFrontend.mlir.passes.build_tog import run_tog + + kernel_name = meta["kernel_name"] + spec = os.path.join(workdir, f"{kernel_name}_spec.py") + if not os.path.isfile(spec): + logger.warning("[Gem5] %s not found; cannot sample cycles", spec) + return None + + run_tog(os.path.join(workdir, "04-custom.mlir"), + os.path.join(workdir, "tog_sample.py"), + os.path.join(workdir, SAMPLE_MLIR), sample_mode=True) + + import subprocess + + from . import tnpu_bridge + env = dict(os.environ) + env.pop("PYTHONPATH", None) # keep tnpu on its own MLIR bindings + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, "-m", "tnpu.cycle", spec, workdir], + capture_output=True, text=True, cwd=tnpu_bridge.tnpu_dir(), env=env) + if proc.returncode != 0: + logger.warning("[Gem5] cycle binary build failed:\n%s", + (proc.stdout + proc.stderr)[-2000:]) + return None + + from Simulator.simulator import CycleSimulator + try: + return CycleSimulator().compile_and_simulate( + os.path.join(workdir, CYCLE_BIN), int(extension_config.vpu_num_lanes), + silent_mode=True) + except Exception as e: # noqa: BLE001 - fall back to the placeholder table + logger.warning("[Gem5] sampling failed: %s", e) + return None + def _runtime_arg_layout(meta): """(n_tensor_args, n_scalar_args) of the lowered signature. @@ -65,22 +111,41 @@ def emit_trace(workdir, meta): f"{postvcix} not found -- tnpu must have run at least to stage 4 " f"(the post-vcix IR is what the trace is built from)") + # Before build_skeleton: both read the post-vcix IR, which it rewrites in place. + cycles = measure_tile_cycles(workdir, meta) + ctx = ir.Context() ctx.allow_unregistered_dialects = True with ctx: module = ir.Module.parse(open(postvcix).read(), ctx) bs.build_skeleton(module) - n_tiles = len(ct._compute_types(module)) + compute_types = ct._compute_types(module) + n_tiles = len(compute_types) + + if cycles: + # One numCycles per compute node; pad/truncate as the MLIR route does. + cl = list(cycles) + if len(cl) != n_tiles: + logger.warning("[Gem5] returned %d cycle(s) for %d " + "tile(s); padding with the last", len(cl), n_tiles) + cl = (cl + [cl[-1]] * n_tiles)[:n_tiles] + # Systolic-array fill; only matmul tiles use it. + lanes = int(extension_config.vpu_num_lanes) + table = ct.build_cycle_table(module, cl, x_offset=lanes, w_offset=0) + else: + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + logger.warning( + "[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d " + "tiles): gem5 sampling did not produce a measurement, so " + "compute latency is NOT modelled", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), work_item=work_item_for(meta)) - # Until gem5 sampling lands, every tile costs PLACEHOLDER_CYCLE: TOGSim - # models the DMA and the dependency structure but NOT compute latency. - table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) - logger.warning("[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d tiles); " - "compute latency is not modelled yet", - CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + if cycles: + logger.info("[Gem5] tile cycles: %s", table) return n_tiles From 85c4e664b07a7d1ff540d0f0222e6cfdc4c72755 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 13/35] [Frontend] Let the route ask for a multi-axis grid 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. --- .../triton_backend/kernel_spec.py | 68 +++++++++++++++---- PyTorchSimFrontend/triton_backend/timing.py | 18 ++++- 2 files changed, 68 insertions(+), 18 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index d74ea962..fe147286 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -136,21 +136,50 @@ def collect_meta(kernel, kernel_name): } +#: Parallel iteration prefixes, OUTERMOST first. Inductor's `x` is the +#: contiguous axis, so it is innermost; `r*` prefixes are reductions, looped +#: inside the kernel rather than spread over the grid (prefix_is_reduction). +_PARALLEL_PREFIXES = ("z", "y", "x") + + +def _block_name(prefix): + return f"{prefix.upper()}BLOCK" + + +def parallel_axes(numels): + """Grid axes this kernel uses, outermost first.""" + return [p for p in _PARALLEL_PREFIXES if f"{p}numel" in numels] + + def fixed_config_for(kernel): """Block sizes pinned at codegen time. tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a - sequential loop, so there is no autotuner to choose XBLOCK later and no - runtime `grid=` callable. Fixing it here is what makes the launch shape + sequential loop, so there is no autotuner to choose the blocks later and no + runtime `grid=` callable. Fixing them here is what makes the launch shape static. - The lane count is the natural default: `bank_vectorize` distributes tile - dim 0 across the lanes, and a block equal to the lane count gives a per-lane - depth of 1 -- the case every tnpu baseline runs today. + Tile dim 0 is the one `bank_vectorize` spreads over the lanes, so the + OUTERMOST axis gets the lane count -- a per-lane depth of 1, the shape every + tnpu baseline runs. The remaining axes get 1, which leaves the tile exactly + that verified shape and lets the grid cover the rest. It is conservative + rather than fast; choosing real tile sizes is the block-size policy gap in + README, not something to guess at here. """ from PyTorchSimFrontend import extension_config lanes = int(extension_config.vpu_num_lanes) - cfg = {"XBLOCK": lanes} + + axes = parallel_axes(getattr(kernel, "numels", None) or {}) + cfg = {_block_name(p): (lanes if i == 0 else 1) for i, p in enumerate(axes)} + if len(axes) > 1: + # Loud, because the shape is correct but pathological: an inner block of + # 1 makes every work-item move a strided column. Fine for getting a + # multi-axis kernel through the route, misleading to benchmark. + extension_config.setup_logger().warning( + "[triton-npu] %s tiles over %s; inner blocks pinned to 1, which is " + "correct but not a tiling worth measuring", + getattr(kernel, "kernel_name", "kernel"), axes) + cfg.setdefault("XBLOCK", lanes) # a kernel with no tiling info still has x if getattr(kernel, "inside_reduction", False): # A reduction block is NOT free to be the lane count: the reduced axis # has to stay inside a lane (see triton-npu kernels/reduce.py). Left @@ -219,20 +248,29 @@ def strip_for_tnpu(src): def grid_of(meta): - """Launch grid, from the numels and the pinned block sizes. + """Launch grid, from the numels and the pinned block sizes, outermost first. Also read by the timing path, which needs the same extents to enumerate the work-items -- so it lives here rather than being recomputed per consumer. """ - x = meta["numels"].get("xnumel") - xblock = (meta.get("fixed_config") or {}).get("XBLOCK") - if x is None or not xblock: + numels = meta["numels"] + cfg = meta.get("fixed_config") or {} + axes = parallel_axes(numels) + if not axes: raise SpecIncomplete( - f"cannot compute the grid for {meta['kernel_name']}: " - f"xnumel={x!r}, XBLOCK={xblock!r}. Inductor defers the grid to " - f"triton_heuristics at runtime; this route needs it statically " - f"(see fixed_config_for).") - return (int(math.ceil(x / xblock)),) + f"{meta['kernel_name']} has no parallel iteration axis to grid over") + + grid = [] + for prefix in axes: + n, block = numels.get(f"{prefix}numel"), cfg.get(_block_name(prefix)) + if n is None or not block: + raise SpecIncomplete( + f"cannot compute the grid for {meta['kernel_name']} axis " + f"'{prefix}': {prefix}numel={n!r}, {_block_name(prefix)}={block!r}. " + f"Inductor defers the grid to triton_heuristics at runtime; this " + f"route needs it statically (see fixed_config_for).") + grid.append(int(math.ceil(n / block))) + return tuple(grid) SPEC_TEMPLATE = '''\ diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index f19423db..1785e519 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -84,15 +84,27 @@ def _runtime_arg_layout(meta): return len(tensors), len(scalars) +#: triton-shared appends pidX, pidY, pidZ in that order, whatever the tiling is. +_PID_SLOT = {"x": 0, "y": 1, "z": 2} + + def work_item_for(meta): - """The WorkItem describing this kernel's program-id args and grid extents.""" + """The WorkItem describing this kernel's program-id args and grid extents. + + `grid_of` orders axes OUTERMOST first (z, y, x -- x is Inductor's contiguous + one), while the program-id arguments are always laid out x, y, z. The two + are zipped downstream, so the argument list is built per axis rather than as + a range. + """ from PyTorchSimFrontend.mlir.passes.lower_to_emitc import WorkItem from . import kernel_spec n_tensor, n_scalar = _runtime_arg_layout(meta) - pid_x = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + pid_base = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + axes = kernel_spec.parallel_axes(meta["numels"]) grid = list(kernel_spec.grid_of(meta)) - return WorkItem(parallel_args=list(range(pid_x, pid_x + len(grid))), grid=grid) + return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], + grid=grid) def emit_trace(workdir, meta): From 9657d38d4c36a3158a1a3bafdc812565e0e6b0d4 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 22:27:26 +0900 Subject: [PATCH 14/35] [Frontend] Take the grid at run time, so one trace serves every shape 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. --- .../mlir/passes/lower_to_emitc.py | 52 +++++++++++++++++-- .../triton_backend/codecache.py | 2 +- PyTorchSimFrontend/triton_backend/timing.py | 52 +++++++++++++++++-- TOGSim/src/main.cc | 13 ++++- 4 files changed, 110 insertions(+), 9 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index bfb94e34..537c8ad0 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -151,6 +151,13 @@ class WorkItem: `parallel_args` are the argument positions holding the program ids (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` their extents. Both outermost-first. + + An extent may be None, meaning "read it from shape_args at run time". Only + the NUMBER of axes has to be known when the kernel is compiled -- how many + loops to nest and how many iv[] slots to fill; the trip counts are just + values, and the producer ABI already takes them + (togsim_kernel(ctx, shape_args, n)). That is what lets one compiled trace + serve every shape. """ def __init__(self, parallel_args, grid): @@ -159,7 +166,12 @@ def __init__(self, parallel_args, grid): f"parallel_args {parallel_args} and grid {grid} must have the " f"same length -- one program-id argument per grid axis") self.parallel_args = list(parallel_args) - self.grid = [int(g) for g in grid] + self.grid = [None if g is None else int(g) for g in grid] + + @property + def dynamic_axes(self): + """Indices into `grid` whose extent arrives at run time.""" + return [i for i, g in enumerate(self.grid) if g is None] def _materialize_grid_loop(kernel, work_item, ctx): @@ -196,7 +208,10 @@ def _materialize_grid_loop(kernel, work_item, ctx): with ir.InsertionPoint(terminator), loc: c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result - ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e)).result + # A runtime extent still needs SOMETHING here: shape_args does not exist + # until _rewrite_signature adds it. The placeholder is replaced by + # _bind_runtime_bounds once it does. + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e or 1)).result for e in work_item.grid] loops, inner = [], None @@ -233,6 +248,35 @@ def _materialize_grid_loop(kernel, work_item, ctx): for pid, new in zip(pid_args, casts): _replace_all_uses(pid, new) + return [(loops[i], ubs[i]) for i in work_item.dynamic_axes] + + +def _bind_runtime_bounds(pending, shape_arg, ctx): + """Point each runtime loop bound at `shape_args[k]`. + + Runs AFTER _rewrite_signature, which is what creates the shape_args + argument. The loops stay in the entry function (the outliner moves only + their bodies), so the read is in scope where the bound is used. + """ + if not pending: + return + from mlir.dialects import arith + + i64 = ir.IntegerType.get_signless(64) + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + for k, (loop, placeholder) in enumerate(pending): + with ir.InsertionPoint(placeholder.owner), loc: + kc = ir.Operation.create( + "emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}).results[0] + elem = ir.Operation.create( + "emitc.subscript", results=[i64], + operands=[shape_arg, kc]).results[0] + bound = arith.IndexCastOp(idxty, elem).result + _replace_all_uses(placeholder, bound) + placeholder.owner.erase() + def _replace_all_uses(old, new): """The bindings expose no replaceAllUsesWith on a Value.""" @@ -630,9 +674,11 @@ def lower_to_emitc(skeleton_module, work_item=None): raise ValueError("no kernel function found in skeleton module") _strip_aux(skeleton_module, keep=kernel) + pending = [] if work_item is not None: - _materialize_grid_loop(kernel, work_item, ctx) + pending = _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) + _bind_runtime_bounds(pending, kernel.regions[0].blocks[0].arguments[1], ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index 69805a87..76dc264c 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -51,7 +51,7 @@ def __call__(self, *args): """ if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): timing.emit_trace(self.workdir, self.meta) - result = timing.run_togsim(self.workdir) + result = timing.run_togsim(self.workdir, meta=self.meta, args=args) logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) logger.warning( "[Spike] %s: output tensors are NOT written; the functional launch " diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index 1785e519..be28060b 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -18,6 +18,7 @@ #: Name TOGSim derives from the kernel directory (Simulator/simulator.py). TRACE_SO = "trace.so" CYCLE_TSV = "trace_cycles.tsv" +SHAPE_TXT = "trace_shape.txt" META_JSON = "meta.json" #: Used only when gem5 sampling fails. Deliberately not a plausible-looking @@ -102,9 +103,46 @@ def work_item_for(meta): n_tensor, n_scalar = _runtime_arg_layout(meta) pid_base = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ axes = kernel_spec.parallel_axes(meta["numels"]) - grid = list(kernel_spec.grid_of(meta)) + # Extents are left to run time: only the axis COUNT has to be compiled in, + # and the launch knows the real numels. One trace then serves every shape. return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], - grid=grid) + grid=[None] * len(axes)) + + +def write_shape(workdir, meta, args=()): + """Write the grid extents the trace producer reads as shape_args. + + `args` is the launch's positional arguments; Inductor appends the numels + after the tensors, so the trailing values are them, in `meta["numels"]` + order. Falls back to the compile-time hint when they are absent. + """ + from . import kernel_spec + + numels = dict(meta["numels"]) + # Only the PARALLEL numels ride along on the call -- a reduction axis is + # looped inside the kernel, so it is not passed and must not consume one of + # the trailing values. They arrive in kernel order, which is the dict's. + passed = [k for k in numels if not k.startswith("r")] + trailing = [a for a in args if isinstance(a, int) and not isinstance(a, bool)] + if passed and len(trailing) >= len(passed): + for key, val in zip(passed, trailing[-len(passed):]): + numels[key] = val + + axes = kernel_spec.parallel_axes(numels) + + cfg = meta.get("fixed_config") or {} + grid = [] + for p in axes: + n, block = numels.get(f"{p}numel"), cfg.get(f"{p.upper()}BLOCK") + if n is None or not block: + raise ValueError(f"no extent for grid axis '{p}': {n!r} / {block!r}") + grid.append(-(-int(n) // int(block))) # ceil-div + + path = os.path.join(workdir, SHAPE_TXT) + with open(path, "w") as f: + f.write("\n".join(str(g) for g in grid) + "\n") + logger.info("[TOGSim] grid %s -> %s", grid, SHAPE_TXT) + return grid def emit_trace(workdir, meta): @@ -161,13 +199,19 @@ def emit_trace(workdir, meta): return n_tiles -def run_togsim(workdir, attribute_path=None, timeout_sec=None): - """Simulate the emitted trace. Returns TOGSimulator's parsed result dict.""" +def run_togsim(workdir, meta=None, args=(), attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict. + + `meta`/`args` supply the grid: the trace producer takes its loop bounds from + shape_args, so they are written out per launch rather than compiled in. + """ from Simulator.simulator import TOGSimulator so = os.path.join(workdir, TRACE_SO) if not os.path.isfile(so): raise FileNotFoundError(f"{so} not found -- call emit_trace first") + if meta is not None: + write_shape(workdir, meta, args) # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its # DIRECTORY, and reads the file itself only on the STONNE path. diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 0ef98eff..b98d3305 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -40,7 +40,18 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } } if (cyc.empty()) { cyc.assign(256, 128); ovl.assign(256, 0); } - return trace_to_tilegraph(trace_so_path.c_str(), nullptr, 0, + // Shape args: the producer's grid bounds, one per axis, when the trace was + // compiled without them baked in. Same sidecar convention as the cycle table + // -- absent means the producer carries its own constants. + std::vector shape; + { + std::ifstream sh(fs::path(trace_so_path).parent_path() / "trace_shape.txt"); + int64_t v; + while (sh >> v) shape.push_back(v); + } + return trace_to_tilegraph(trace_so_path.c_str(), + shape.empty() ? nullptr : shape.data(), + (int32_t)shape.size(), bases.data(), (int)bases.size(), cyc.data(), ovl.data(), (int)cyc.size(), partition_cores.data(), (int32_t)partition_cores.size(), From 8da4bbc0a9556c34830d51be1680dc4ed962f250 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 23:30:34 +0900 Subject: [PATCH 15/35] [Frontend] Pin only the blocks the kernel takes as parameters 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. --- .../triton_backend/kernel_spec.py | 6 +++++ tests/system/test_triton_codegen.py | 26 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index fe147286..9fc2fc46 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -345,6 +345,12 @@ def write_spec_file(src_code, meta, path, tnpu_dir): signature = dict(meta["signature"]) constexprs = dict(meta["constants"]) for k, v in (meta.get("fixed_config") or {}).items(): + if k not in signature: + # Inductor already fixed this one in the kernel BODY rather than + # taking it as a parameter -- a persistent reduction does that with + # R0_BLOCK. Passing it would not match the signature, and there is + # nothing left for us to choose. + continue if v is None: raise SpecIncomplete( f"{meta['kernel_name']}: block size {k} is unset " diff --git a/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py index b36715d6..d98264cf 100644 --- a/tests/system/test_triton_codegen.py +++ b/tests/system/test_triton_codegen.py @@ -92,12 +92,38 @@ def check_multi_axis_grid(): return not problems +def check_reduction_is_refused(): + """A reduction must fail LOUDLY, not compile into wrong numbers. + + tnpu has no lane-aware reduction: the scratchpad is lane-banked, so the + reduced axis has to live inside a lane, and triton-shared hands over a + linalg.reduce (plus a linalg.transpose) that no pass lowers that way. Until + one does, reaching the launcher would mean simulating a kernel whose compute + is not what the hardware would do. + + Passing this check means the attempt still stops. When the lane path lands, + this is the test to delete. + """ + x = torch.randn(128, 64) + try: + torch.compile(lambda t: t.sum(dim=1))(x.to("npu:0")) + except Exception as e: # noqa: BLE001 - any diagnosed stop is the point + first = (str(e).strip().splitlines() or [type(e).__name__])[0] + print(f" reduction stops at: {type(e).__name__}: {first[:74]}") + return True + print(" reduction COMPILED -- if the lane-aware path landed, drop this " + "check; otherwise the numbers it produces are wrong") + return False + + def main(): from PyTorchSimFrontend import extension_config from PyTorchSimFrontend.triton_backend import tnpu_bridge print(f"multi-axis grid = " f"{'ok' if check_multi_axis_grid() else 'FAILED'}") + print(f"reduction refused = " + f"{'ok' if check_reduction_is_refused() else 'FAILED'}") print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") ok, _out = tnpu_bridge.doctor() From cdd53c08eabaef7099c081a0e7a7f29a1886f72e Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 23:56:48 +0900 Subject: [PATCH 16/35] [Frontend] Pass the kernel's scalar arguments to the tnpu wrapper 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. --- .../triton_backend/kernel_spec.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index 9fc2fc46..4f9056a6 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -49,6 +49,10 @@ def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): } +#: Triton scalar token -> C type, for the wrapper's kernel declaration. +_C_TYPE = {"i32": "int32_t", "i64": "int64_t", "fp32": "float"} + + class SpecIncomplete(RuntimeError): """Metadata tnpu requires that this kernel did not provide. @@ -247,6 +251,31 @@ def strip_for_tnpu(src): return prefix + body +def scalar_args(meta): + """User scalar parameters, in kernel order, as [(name, c_type, value)]. + + triton-shared keeps these in the lowered signature ahead of its own six + grid/pid arguments, so the wrapper must pass them or every later argument + lands one slot early -- pidX then reads pidY and only program 0 runs. + """ + numels = meta["numels"] + out = [] + for name, token in meta["signature"].items(): + if token.startswith("*") or token == "constexpr": + continue + ctype = _C_TYPE.get(token) + if ctype is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: scalar '{name}' has type {token!r}, " + f"which has no C mapping in _C_TYPE") + if numels.get(name) is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: no value for scalar '{name}' -- " + f"collect_meta resolves these from kernel.numels") + out.append((name, ctype, int(numels[name]))) + return out + + def grid_of(meta): """Launch grid, from the numels and the pinned block sizes, outermost first. @@ -329,6 +358,8 @@ def reference(inputs): grid={grid!r}, reference=reference, make_inputs=make_inputs, + extra={{"scalar_args": {scalar_decls!r}, + "scalar_values": {scalar_values!r}}}, notes="generated from Inductor triton codegen", ) ''' @@ -370,6 +401,7 @@ def write_spec_file(src_code, meta, path, tnpu_dir): with open(os.path.join(os.path.dirname(path), triton_module), "w") as f: f.write(strip_for_tnpu(src_code)) + scalars = scalar_args(meta) text = SPEC_TEMPLATE.format( kernel_name=meta["kernel_name"], tnpu_dir=tnpu_dir, @@ -379,6 +411,8 @@ def write_spec_file(src_code, meta, path, tnpu_dir): args_body=args_body, make_inputs_body=make_inputs_body, grid=grid_of(meta), + scalar_decls=[(n, c) for n, c, _ in scalars], + scalar_values={n: v for n, _, v in scalars}, ) with open(path, "w") as f: f.write(text) From 5e77a5faa9884a59be73d568acc0181d381b1164 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 27 Jul 2026 23:57:05 +0900 Subject: [PATCH 17/35] [Frontend] Run the Triton route's kernel on the launch's own tensors 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/.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. --- PyTorchSimFrontend/triton_backend/README.md | 25 ++-- .../triton_backend/codecache.py | 20 +-- .../triton_backend/functional.py | 132 ++++++++++++++++++ tests/system/test_triton_codegen.py | 15 +- 4 files changed, 170 insertions(+), 22 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/functional.py diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 4b409158..6deca819 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -38,9 +38,9 @@ torch.compile │ 01-ttir → 02-ttshared → 03-adapted → 04-custom → 05-*.elf ▼ TritonNPULauncher.__call__ codecache.py - ├ timing 04-custom.mlir → trace.so + trace_cycles.tsv → TOGSim - │ cycles measured by gem5 on a one-tile binary - └ functional ← NOT WIRED YET + ├ functional tensors → runtime/*.raw → Spike → tensors functional.py + └ timing 04-custom.mlir → trace.so + trace_cycles.tsv → TOGSim + cycles measured by gem5 on a one-tile binary ``` The timing half reuses PyTorchSim's trace pipeline unchanged. The one structural @@ -73,14 +73,23 @@ Artifacts land in one directory per source hash under the dump path (gap 2) - the tile's compute cost is a real measurement: gem5 samples **19 cycles** for the vector-add tile, via `timing.measure_tile_cycles` -- values are NOT produced: the functional launch is still open (gap 1) +- **values are correct**: the launch writes the caller's tensors from Spike and + `torch.allclose` holds over all 1024 elements, for the fused + `(x + y) * 2 - x` kernel too + +## Shape specialisation + +The functional binary is compiled for ONE shape: the spec bakes the grid, the +scalar values and the memref extents in. A dynamic-shape graph reuses that ELF, +so `functional.ShapeMismatch` rejects the launch instead of running against the +wrong bounds. The timing path has no such limit -- it takes the grid at run time +-- so `pytorchsim_functional_mode: False` studies cycles across shapes. ## Gap list, in order -1. **Launch (functional).** Marshal the caller's tensors into - `runtime/*.raw`, run Spike on the ELF, read outputs back. tnpu's stage 6 does - this for its own kernels but generates inputs from the spec; here the tensors - come from the caller. +1. **Shape-specialised functional launch.** Recompile per launch shape, or teach + the tnpu wrapper to take the grid and the extents as arguments the way the + trace producer already does. 2. **Double buffering.** tnpu emits synchronous DMA (`is_async=false`, no `togsim.wait`), so load → compute → store serialize inside every work-item and TOGSim has no overlap to model. This is the main remaining gap between the two diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index 76dc264c..00cb3b1c 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -18,7 +18,7 @@ from torch._inductor.codecache import get_hash from PyTorchSimFrontend import extension_config -from . import kernel_spec, timing, tnpu_bridge +from . import functional, kernel_spec, timing, tnpu_bridge logger = extension_config.setup_logger() @@ -43,19 +43,23 @@ def __init__(self, kernel_name, workdir, meta): self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") def __call__(self, *args): - """One launch of the whole grid: simulate, return TOGSim's result. + """One launch of the whole grid: run it on Spike, then time it. - Does NOT write the caller's output tensors -- the functional launch is - not wired (README). Logged, so an undefined value cannot pass for a - computed one. + Spike runs first so the caller's output tensors hold real values even if + TOGSim fails -- the two halves are independent. """ + if extension_config.pytorchsim_functional_mode: + written = functional.run(self.workdir, self.meta, args) + logger.info("[Spike] %s wrote %s", self.kernel_name, written) + else: + logger.warning( + "[Spike] %s: functional mode is off, so the output tensors keep " + "whatever they held", self.kernel_name) + if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): timing.emit_trace(self.workdir, self.meta) result = timing.run_togsim(self.workdir, meta=self.meta, args=args) logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) - logger.warning( - "[Spike] %s: output tensors are NOT written; the functional launch " - "(tensors -> Spike -> tensors) is not wired yet", self.kernel_name) return result diff --git a/PyTorchSimFrontend/triton_backend/functional.py b/PyTorchSimFrontend/triton_backend/functional.py new file mode 100644 index 00000000..4dd15682 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -0,0 +1,132 @@ +"""The functional half of the Triton route: real tensors -> Spike -> real tensors. + +The timing half (timing.py) tells you how long the kernel takes; this tells you +whether it computed the right thing. tnpu's stage 6 already runs the ELF under +Spike, but on inputs it generates itself. Here the launch's own tensors are +written as the `.raw` files stage 6 reads, and the outputs are copied back: + + run(workdir, meta, args) args -> runtime/*.raw -> spike -> args + +The binary is shape-specialised -- the spec bakes the grid, the scalar values and +the memref extents in -- so a launch whose shapes differ from the compiled ones +is rejected rather than silently run against the wrong bounds. +""" + +import os +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +RUNTIME_DIR = "runtime" + + +class ShapeMismatch(RuntimeError): + """The launch does not match the shapes the binary was compiled for.""" + + +def _np_dtype(name): + import numpy as np + return np.dtype("bool" if name == "bool" else name) + + +def tensor_args(meta, args): + """[(arg_meta, tensor)] for the launch, paired by position. + + Inductor passes the tensors first and the numels after, in signature order, + so `meta["args"]` (tensors only) lines up with the leading arguments. + """ + import torch + + tensors = [a for a in args if isinstance(a, torch.Tensor)] + metas = meta["args"] + if len(tensors) != len(metas): + raise ShapeMismatch( + f"{meta['kernel_name']}: launch passed {len(tensors)} tensor(s), " + f"but the spec declares {len(metas)} ({[m['name'] for m in metas]})") + return list(zip(metas, tensors)) + + +def _check(meta, pairs): + for m, t in pairs: + if t.numel() != m["numel"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' has {t.numel()} " + f"element(s), but the binary was compiled for {m['numel']}. " + f"tnpu bakes the extents, the grid and the scalar values into " + f"the kernel, so a dynamic-shape graph reuses an ELF that does " + f"not fit. The timing path does handle this (it takes the grid " + f"at run time); set pytorchsim_functional_mode: False to study " + f"cycles alone, or keep shapes static to check values.") + if str(t.dtype).removeprefix("torch.") != m["dtype"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' is {t.dtype}, but the " + f"binary was compiled for {m['dtype']}") + + +def write_inputs(workdir, meta, args): + """Write every arg as runtime/.raw. Returns the runtime directory. + + Outputs are written too, as zeros: the wrapper loads and dumps by argv + position, so a missing file shifts every later one. + """ + import numpy as np + + pairs = tensor_args(meta, args) + _check(meta, pairs) + + runtime = os.path.join(workdir, RUNTIME_DIR) + os.makedirs(runtime, exist_ok=True) + for m, t in pairs: + path = os.path.join(runtime, f"{m['name']}.raw") + if m["role"] in ("in", "inout"): + t.detach().to("cpu").contiguous().numpy().tofile(path) + else: + np.zeros(m["numel"], dtype=_np_dtype(m["dtype"])).tofile(path) + return runtime + + +def read_outputs(workdir, meta, args): + """Copy the .raw files Spike wrote back into the launch's output tensors.""" + import numpy as np + import torch + + runtime = os.path.join(workdir, RUNTIME_DIR) + written = [] + for m, t in tensor_args(meta, args): + if m["role"] not in ("out", "inout"): + continue + path = os.path.join(runtime, f"{m['name']}.raw") + flat = np.fromfile(path, dtype=_np_dtype(m["dtype"])) + if flat.size != m["numel"]: + raise RuntimeError( + f"{path} holds {flat.size} element(s), expected {m['numel']} " + f"-- Spike did not write the whole tensor") + t.copy_(torch.from_numpy(flat).view_as(t).to(t.dtype)) + written.append(m["name"]) + return written + + +def run(workdir, meta, args, timeout_sec=None): + """Execute the kernel on the launch's tensors. Returns the names written.""" + from . import tnpu_bridge + + spec = os.path.join(workdir, f"{meta['kernel_name']}_spec.py") + if not os.path.isfile(spec): + raise FileNotFoundError(f"{spec} not found -- compile the kernel first") + + write_inputs(workdir, meta, args) + + env = dict(os.environ) + env.pop("PYTHONPATH", None) # keep tnpu on its own MLIR bindings + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, "-m", "tnpu.spike", spec, workdir], + capture_output=True, text=True, cwd=tnpu_bridge.tnpu_dir(), env=env, + timeout=timeout_sec) + if proc.returncode != 0: + raise RuntimeError( + f"[Spike] {meta['kernel_name']} failed:\n" + + (proc.stdout + proc.stderr)[-2000:]) + + return read_outputs(workdir, meta, args) diff --git a/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py index d98264cf..2af839bd 100644 --- a/tests/system/test_triton_codegen.py +++ b/tests/system/test_triton_codegen.py @@ -145,10 +145,15 @@ def main(): "above and README.md's gap list.") return 1 - # Values are NOT checked: the launch simulates the kernel but does not - # marshal tensors through Spike, so `got` is undefined. What is asserted is - # that the timing path ran end to end -- the two artifacts TOGSim consumes. - del got + ok = torch.allclose(got.cpu(), expected, rtol=1e-4, atol=1e-4) + if not ok: + bad = (~torch.isclose(got.cpu(), expected, rtol=1e-4, atol=1e-4)) + print(f"VALUES WRONG: {int(bad.sum())}/{expected.numel()} elements") + print(f" got {got.cpu()[:4].tolist()}") + print(f" expected {expected[:4].tolist()}") + return 1 + print(f"values ok ({expected.numel()} elements through Spike)") + import glob from PyTorchSimFrontend.triton_backend import timing @@ -165,8 +170,6 @@ def main(): return 1 print(f" {name:18s} {os.path.getsize(path)} bytes") print(f"\ntiming path OK ({workdir})") - print(f"values NOT verified -- torch would give {expected[:3].tolist()}...; " - f"the functional launch is not wired (triton_backend/README.md)") return 0 From 68a6aaeca831c4f0ea995d41ccaede34035456f4 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 00:01:27 +0900 Subject: [PATCH 18/35] [CI] Bump the triton-npu pin to 22df065 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. --- thirdparty/triton-npu.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 0b2172bd..5c01580c 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "fd423b001ad562b4cbcd6405b03df2b0674656be", + "ref": "22df065b46b1a61f75a9847ac73207de8cc72482", "release_tag": "toolchain-llvm23" } } From 5f52306faf0f0ed09ce9effe20ebf34c3a160399 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 13:04:28 +0900 Subject: [PATCH 19/35] [Docs] Report the Triton route against the MLIR one 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. --- PyTorchSimFrontend/triton-codegen-route.md | 242 ++++++++++++++++++++ PyTorchSimFrontend/triton_backend/README.md | 4 + 2 files changed, 246 insertions(+) create mode 100644 PyTorchSimFrontend/triton-codegen-route.md diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md new file mode 100644 index 00000000..f009f448 --- /dev/null +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -0,0 +1,242 @@ +# Inductor Triton 코드젠 경로를 PyTorchSim에 연결 + +`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **triton-npu(tnpu)**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. + +**functional과 timing 양쪽이 연결되어 있고, 동적 shape도 처리됩니다.** + +모듈별 동작과 사용법은 [`triton_backend/README.md`](triton_backend/README.md)에 있습니다. +이 문서는 기존 경로와의 대조, 설계 판단, 측정 결과를 다룹니다. + +| | | +|---|---| +| functional | `x + y`, `(x+y)*2 - x` 모두 **max abs error 0.0** (1024 elements, Spike) | +| timing | TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | +| 동적 shape | 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32, 재컴파일 없음 | +| 변경량 | PyTorchSim 18 commits / 23 files / +2209−39, tnpu 3 commits / +177−4 | +| CI | 전 잡 green (툴체인 빌드 + 값 검증 + 기존 경로 회귀 확인) | + +--- + +## 1. 파이프라인 + +``` +torch.compile + └ TritonNPUScheduling.define_kernel scheduling.py + │ Inductor 가 만든 triton 소스 텍스트 + 수집한 메타데이터 + ▼ + triton_npu_compile(src, meta, name) codecache.py + │ tnpu KernelSpec 생성 kernel_spec.py + │ - 블록 크기를 constexpr 로 고정 + │ - 인자 역할(in/out/inout) · dtype · numel + │ - grid, 사용자 스칼라 값 + ▼ + tnpu 파이프라인 (별도 인터프리터, subprocess) tnpu_bridge.py + │ 1 ttir triton 커널 → Triton IR + │ 2 ttshared → linalg / memref / scf.for (triton-shared) + │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) + │ 4 lower vcix → gemmini DMA → LLVM + │ 5 binary mlir-translate → llc → RISC-V ELF + ▼ + TritonNPULauncher.__call__ codecache.py + │ + ├ functional 텐서 → runtime/*.raw → Spike → 텐서 functional.py + │ tnpu stage 6 (tnpu.spike) 재사용 + │ + └ timing 04-custom.mlir timing.py + ├ build_tog sample → gem5 → 타일 cycle 실측 + └ build_skeleton → trace.so + trace_cycles.tsv + → TOGSim +``` + +**LLVM 이음매.** tnpu는 stock LLVM 23을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. + +--- + +## 2. 기존 MLIR 경로와의 차이 + +### 갈라지는 지점과 합쳐지는 지점 + +``` + torch.compile / Inductor 스케줄 + │ + ┌───────────────┴───────────────┐ + │ │ + [기존] MLIR 경로 [신규] Triton 경로 + │ │ + Inductor 스케줄 → 손으로 쓴 Inductor 의 Triton 코드젠 + op별 MLIR 템플릿 이 낸 커널 소스를 가로챔 + (gemm, conv, sdpa, sort, (op별 템플릿 없음) + cat, maxpool, bmm …) + │ │ + PyTorchSim mlir/ 패스 tnpu 패스 (subprocess) + PSAL LLVM 20 stock LLVM 23 + │ │ + └───────────────┬───────────────┘ + │ + ▼ 여기서 다시 합류 ▼ + trace.so + trace_cycles.tsv + → TOGSim + (트레이스 계약은 완전히 동일) +``` + +핵심은 **TOGSim이 두 경로를 구분하지 못한다**는 점입니다. 트레이스 생산자의 형태가 같으므로 하드웨어 모델·DRAM·NoC·L2는 손대지 않았습니다. + +### 항목별 대조 + +| | 기존 MLIR 경로 | 신규 Triton 경로 | +|---|---|---| +| 커널을 만드는 주체 | PyTorchSim의 op별 MLIR 템플릿 | Inductor의 Triton 코드젠 | +| **커널 하나의 의미** | **루프 네스트 전체** | **타일 하나** | +| grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | +| lowering | `PyTorchSimFrontend/mlir/` (in-process) | tnpu (subprocess, LLVM 23) | +| 융합 | 템플릿과 `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | +| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | +| functional | `FunctionalSimulator.run_spike` | tnpu stage 6 (`tnpu.spike`) | +| timing | `trace.so` + `trace_cycles.tsv` → TOGSim | **동일** | +| 타일 cycle 실측 | gem5 | **동일** (`build_tog` sample 모드 공유) | +| DMA | 비동기 + `togsim.wait` 배리어 | **동기만** (`togsim.wait` 0개) | +| 동적 shape | 트레이스 경로는 아직 미지원 (PR #269 진행 중) | timing 경로에서 동작 | + +### 이 대조가 말해주는 것 + +**Triton 경로가 앞선 곳** — 동적 shape. 기존 경로의 C++ 트레이스는 `trace_to_tilegraph(..., nullptr, 0)`으로 shape 인자를 아예 넘기지 않아 shape마다 트레이스를 다시 만들어야 하고, 그걸 푸는 작업이 PR #269로 아직 열려 있습니다. Triton 경로는 `shape_args`를 통해 **트레이스 하나가 모든 shape을 섬깁니다.** + +**기존 경로가 앞선 곳** — op 커버리지와 DMA 겹침. 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 아래 사이클 격차의 원인입니다. + +**바뀌지 않은 것** — TOGSim, 하드웨어 설정, gem5 샘플링 방식, 트레이스 계약. 두 경로는 같은 시뮬레이터를 먹입니다. + +--- + +## 3. 핵심 설계 문제: 커널 하나가 무엇을 뜻하는가 + +``` +MLIR 경로 커널 = 루프 네스트 전체. TOG 가 루프에서 work-item 을 읽어냄 +Triton 커널 = 타일 하나. grid 는 커널 밖, launch 가 쥐고 있음 +``` + +TOGSim의 트레이스 계약(`docs/design/togsim_cpp_trace.md` §9.1/§9.3)이 이미 이 둘을 구분합니다: + +- `togsim_kernel_tile(ctx, iv, n)` — work-item 하나 +- `togsim_kernel(ctx, shape_args, n)` — 병렬 영역의 열거 + +Triton 커널 본문은 전자에 대응하므로, **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. 그 합성이 `lower_to_emitc.WorkItem` + `_materialize_grid_loop`입니다. + +### 동적 shape이 여기서 나옵니다 + +`_materialize_grid_loop`은 축 **개수**만 컴파일에 박고, **범위**는 `shape_args`에서 읽습니다. + +``` +컴파일 시 축이 몇 개인지만 안다 → 루프 네스트 골격 생성 +런타임 실제 numel 로 grid 계산 → trace_shape.txt 로 전달 + TOGSim 이 build_trace_tilegraph 에서 읽어 shape_args 로 주입 +``` + +측정: `dynamic=True`로 n=1024 → grid 8, n=4096 → grid 32. 트레이스 재생성 없음. + +다차원 grid(Triton 제약상 최대 3D)도 지원합니다. 구현 중 두 번 틀렸고 둘 다 rank ≥ 2에서만 드러났습니다 — 종료자가 있는 블록 끝에 삽입하는 문제, 그리고 bound를 루프 뒤에 만들어 dominance를 깨는 문제. 그래서 테스트가 생성된 C++가 아니라 **MLIR 모듈 자체를 verify**합니다. + +--- + +## 4. 측정 결과 + +### functional + +| 커널 | 원소 | max abs error | +|---|---:|---:| +| `x + y` | 1024 | 0.0 | +| `(x + y) * 2 - x` (Inductor가 단일 커널로 융합) | 1024 | 0.0 | + +### timing + +| 항목 | 값 | 확인 내용 | +|---|---:|---| +| 타일 compute (gem5) | 19–21 | 마커 사이 `numCycles` 실측. placeholder 아님 | +| TOGSim 총계 | 650 | DRAM 트래픽 8192 B = 8 work-item × 2 load × 512 B, 정확히 일치 | +| 기존 MLIR 경로 (동일 연산) | 251 | 같은 자릿수 | + +650 대 251은 모델 오류가 아닙니다. **tnpu가 동기 DMA만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. + +--- + +## 5. 도중에 찾은 실제 버그 + +functional 배선은 배관 작업일 줄 알았는데, 첫 실행에서 **1024개 중 896개가 틀렸습니다.** `pid_x=0` 블록만 맞고 나머지 7개는 전부 0. + +``` +MLIR func.func @k(%arg0..2: memref<*xf32>, in_ptr0, in_ptr1, out_ptr0 + %arg3: i32 xnumel <- 사용자 스칼라 + %arg4,5,6: i32 gridX,Y,Z + %arg7,8,9: i32 pidX,Y,Z ) + +wrapper k(1,&d_in_ptr0, 1,&d_in_ptr1, 1,&d_out_ptr0, 8, 1, 1, pid_x, pid_y, pid_z); + +------ i32 6개뿐 ------+ + xnumel 누락 +``` + +triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. tnpu wrapper는 이를 `spec.extra["scalar_args"]`에서 읽는데, PyTorchSim이 생성하는 spec에는 `extra`가 아예 없었습니다. 인자가 한 칸씩 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. + +**틀린 값이 쓰레기가 아니라 0으로 나온 점**이 고약합니다. 쓰레기값이면 즉시 눈에 띄지만 0은 그럴듯해 보입니다. timing 경로는 인자 위치를 lowered MLIR 시그니처에서 직접 읽어 애초에 정확했고, 그래서 functional을 붙이기 전까지 드러나지 않았습니다. + +--- + +## 6. 일반성을 위해 되돌린 설계 둘 + +**DMA가 어느 인자에 속하는지 — 추론에서 선언으로.** 처음에는 TOG 빌더가 memref view 연산을 거꾸로 걸어 올라가 인자 인덱스를 추론했습니다. 아는 view 연산에 대해서만 맞는 방식이라, 생산자(tnpu)가 `dram_arg`를 직접 적어 내려보내도록 바꾸고 추론 코드를 삭제했습니다. + +**grid — 컴파일 타임 상수에서 런타임 인자로.** 위 3절. + +--- + +## 7. 현재 상태 + +| 기능 | 상태 | 내용 | +|---|---|---| +| elementwise + 융합 | 동작 | 값·사이클 모두 통과, CI 포함 | +| 동적 shape (timing) | 동작 | 트레이스 하나가 모든 shape | +| 다차원 grid | 동작 | 테스트가 IR과 dispatch 양쪽 검증 | +| 동적 shape (functional) | 제약 | 바이너리가 shape 특수화 → `ShapeMismatch`로 거부 | +| double buffering | 미착수 | tnpu가 동기 DMA만 발행. 251 vs 650의 주원인 | +| matmul timing | 미착수 | `build_tog`는 `vcix.iv` 이름으로 compute 노드를 찾는데 tnpu는 `llvm.riscv.sf.vc.*` 인트린식을 냄 | +| `triton_helpers` | 차단 | 모듈이 torch 안에 있고 tnpu venv에는 없음 | +| reduction | 차단 | tnpu 자체 문제 — 아래 | + +**동적 shape의 한 가지 단서.** timing은 완전히 동작합니다. functional 바이너리는 tnpu가 grid·스칼라 값·memref extent를 전부 구워 넣어 shape 특수화되어 있어서, shape이 다른 launch를 `ShapeMismatch`로 **거부합니다** — 틀린 경계로 실행하는 대신. 사이클만 볼 때는 `pytorchsim_functional_mode: False`로 모든 shape을 돌릴 수 있습니다. + +**reduction이 막힌 지점.** `tt.reduce(axis=1)`이 triton-shared를 지나면 `linalg.transpose permutation=[1,0]` + `linalg.reduce dimensions=[0]`가 됩니다. transpose는 `transpose-reduce-to-rank0` 여부와 무관하게 삽입됩니다(rank 2에서 동일함을 측정). stage 3의 다섯 패스는 통과하고 `bank_vectorize`가 거부합니다 — 스크래치패드가 **레인 뱅킹**되어 있어 축소되는 축이 레인 안에 머물러야 하는데, identity-elementwise가 아니고 스칼라 폴백은 뱅킹된 스크래치패드를 읽게 되기 때문입니다. + +테스트(`check_reduction_is_refused`)가 이 경계를 못박습니다. **reduction이 컴파일에 성공하면 테스트가 실패합니다** — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. + +--- + +## 8. PR과 검증 + +| PR | 범위 | 상태 | +|---|---|---| +| [PyTorchSim #305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305) | 18 commits · 23 files · +2209/−39 | draft, mergeable, CI green | +| [triton-npu #1](https://github.com/PSAL-POSTECH/triton-npu/pull/1) | 3 commits · 5 files · +177/−4 | open | + +CI(`.github/workflows/triton_npu.yml`)는 툴체인 레이어가 ~1.8 GiB라 본 CI와 분리: + +``` +Check tnpu access success +Build tnpu toolchain image success +Build app image on tnpu base success +Inductor Triton route success <- test_triton_codegen.py (값 검증 포함) +MLIR route still passes success <- 기존 경로 회귀 없음 +triton-npu baselines success <- doctor + add/mul/relu/gemm/bmm +``` + +**머지 순서.** `thirdparty/triton-npu.json`이 tnpu 커밋 `22df065`를 핀하는데, 이는 `feature/timing-form`에만 있고 `main`에는 없습니다. sha라 CI fetch는 되지만 #1이 리베이스 머지되면 뜹니다 — **#1 머지 → 핀을 main 커밋으로 재조정 → #305** 순서가 안전합니다. + +--- + +## 9. 다음 우선순위 + +1. **double buffering** — tnpu가 비동기 DMA + `togsim.wait`를 내도록. 두 경로의 사이클 격차를 실제로 좁히는 유일한 항목이고, 기존 경로에 이미 있는 기계를 tnpu 쪽에 만드는 일입니다. +2. **shape 특수화 해소** — launch shape마다 재컴파일하거나, tnpu wrapper도 트레이스 생산자처럼 grid와 extent를 인자로 받게. 후자가 근본적. +3. **matmul timing** — `build_tog`가 vcix 인트린식을 인식하도록. systolic array 경로가 열립니다. +4. **reduction 레인 경로** — tnpu의 `bank_vectorize`에 reduction 추가, transpose를 `vlane_split_axis`로 흡수. 가장 큰 작업. + +--- + +측정 환경: torch 2.10.0+cpu / triton 3.6.0, `systolic_ws_128x128_c1_simple_noc_tpuv3.yml`, `vpu_num_lanes` 128. 기존 MLIR 경로는 `tests/ops/elementwise/test_add.py` 통과로 회귀 없음 확인. diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 6deca819..de701e20 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -9,6 +9,10 @@ and stays the production path. TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py ``` +This file is the working reference for the modules here. For how the route +compares with the MLIR one and what the numbers are, see +[`../triton-codegen-route.md`](../triton-codegen-route.md). + ## Why The MLIR route does not just emit loops — it hand-implements the whole hardware From 05e4edc7232df141ee3ab6d6797a0713e389d662 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 13:13:56 +0900 Subject: [PATCH 20/35] [Docs] State where the port ends and the lowering pass begins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- PyTorchSimFrontend/triton-codegen-route.md | 68 +++++++++++++++------ PyTorchSimFrontend/triton_backend/README.md | 10 ++- 2 files changed, 55 insertions(+), 23 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index f009f448..c0cdc7cc 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -1,6 +1,8 @@ # Inductor Triton 코드젠 경로를 PyTorchSim에 연결 -`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **triton-npu(tnpu)**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. +`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **PyTorchSim lowering pass**(담당 이정민)가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. + +**이 문서가 보고하는 것은 그 lowering pass 를 기존 PyTorchSim 스택에 이식·연결한 작업입니다.** lowering pass 자체는 범위 밖입니다 — 아래 "작업 경계" 참고. **functional과 timing 양쪽이 연결되어 있고, 동적 shape도 처리됩니다.** @@ -15,6 +17,19 @@ | 변경량 | PyTorchSim 18 commits / 23 files / +2209−39, tnpu 3 commits / +177−4 | | CI | 전 잡 green (툴체인 빌드 + 값 검증 + 기존 경로 회귀 확인) | +### 작업 경계 + +이 경로는 두 부분으로 나뉘고, **이 문서가 보고하는 작업은 아래쪽입니다.** + +| 부분 | 하는 일 | 소관 | +|---|---|---| +| **PyTorchSim lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 이 문서의 범위 밖 | +| **기존 PyTorchSim으로의 이식** | 위 lowering pass 를 기존 시뮬레이션 스택에 얹는 일: Inductor Triton 코드젠 가로채기, KernelSpec 생성, grid 합성, 트레이스/사이클 산출, functional launch, TOGSim 연결 | 이 문서의 작업 | + +즉 lowering pass 자체는 만들지 않았습니다. **이미 있는 lowering pass 를 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim/gem5/Spike 스택에 물린 것**이 여기서 한 일입니다. 그 과정에서 lowering pass 쪽에 필요해진 최소한의 훅 3개(`tnpu.cycle`, `dram_arg`, `tnpu.spike`)는 별도 PR로 올렸고, 8절에 있습니다. + +문서에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있고, 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image` 등) — 문서를 따라 실제 코드를 찾아갈 수 있어야 하기 때문입니다. + --- ## 1. 파이프라인 @@ -30,7 +45,9 @@ torch.compile │ - 인자 역할(in/out/inout) · dtype · numel │ - grid, 사용자 스칼라 값 ▼ - tnpu 파이프라인 (별도 인터프리터, subprocess) tnpu_bridge.py + PyTorchSim lowering pass (별도 인터프리터, subprocess) tnpu_bridge.py + │ ┌─ 담당 이정민 / 이 문서의 범위 밖. + │ └─ 여기서 한 일은 이 단계를 "호출하고 결과를 스택에 물린" 부분. │ 1 ttir triton 커널 → Triton IR │ 2 ttshared → linalg / memref / scf.for (triton-shared) │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) @@ -48,7 +65,7 @@ torch.compile → TOGSim ``` -**LLVM 이음매.** tnpu는 stock LLVM 23을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. +**LLVM 이음매.** lowering pass 는 stock LLVM 23 을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. --- @@ -68,8 +85,11 @@ torch.compile (gemm, conv, sdpa, sort, (op별 템플릿 없음) cat, maxpool, bmm …) │ │ - PyTorchSim mlir/ 패스 tnpu 패스 (subprocess) - PSAL LLVM 20 stock LLVM 23 + PyTorchSim mlir/ 패스 PyTorchSim lowering pass + PSAL LLVM 20 (subprocess, stock LLVM 23) + 담당 이정민 — 범위 밖 + 여기서 한 일은 이 블록을 + 아래 합류점까지 잇는 배선 │ │ └───────────────┬───────────────┘ │ @@ -88,7 +108,7 @@ torch.compile | 커널을 만드는 주체 | PyTorchSim의 op별 MLIR 템플릿 | Inductor의 Triton 코드젠 | | **커널 하나의 의미** | **루프 네스트 전체** | **타일 하나** | | grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | -| lowering | `PyTorchSimFrontend/mlir/` (in-process) | tnpu (subprocess, LLVM 23) | +| lowering | `PyTorchSimFrontend/mlir/` (in-process) | PyTorchSim lowering pass — 담당 이정민 (subprocess, LLVM 23) | | 융합 | 템플릿과 `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | | op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | | functional | `FunctionalSimulator.run_spike` | tnpu stage 6 (`tnpu.spike`) | @@ -154,7 +174,7 @@ Triton 커널 본문은 전자에 대응하므로, **후자를 합성해서 씌 | TOGSim 총계 | 650 | DRAM 트래픽 8192 B = 8 work-item × 2 load × 512 B, 정확히 일치 | | 기존 MLIR 경로 (동일 연산) | 251 | 같은 자릿수 | -650 대 251은 모델 오류가 아닙니다. **tnpu가 동기 DMA만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. +650 대 251은 모델 오류가 아닙니다. **lowering pass 가 동기 DMA 만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. --- @@ -173,7 +193,7 @@ wrapper k(1,&d_in_ptr0, 1,&d_in_ptr1, 1,&d_out_ptr0, 8, 1, 1, pid_x, pid_y, pid xnumel 누락 ``` -triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. tnpu wrapper는 이를 `spec.extra["scalar_args"]`에서 읽는데, PyTorchSim이 생성하는 spec에는 `extra`가 아예 없었습니다. 인자가 한 칸씩 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. +triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. lowering pass 의 wrapper 는 이를 `spec.extra["scalar_args"]`에서 읽는데, PyTorchSim이 생성하는 spec에는 `extra`가 아예 없었습니다. 인자가 한 칸씩 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. **틀린 값이 쓰레기가 아니라 0으로 나온 점**이 고약합니다. 쓰레기값이면 즉시 눈에 띄지만 0은 그럴듯해 보입니다. timing 경로는 인자 위치를 lowered MLIR 시그니처에서 직접 읽어 애초에 정확했고, 그래서 functional을 붙이기 전까지 드러나지 않았습니다. @@ -181,7 +201,7 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 ## 6. 일반성을 위해 되돌린 설계 둘 -**DMA가 어느 인자에 속하는지 — 추론에서 선언으로.** 처음에는 TOG 빌더가 memref view 연산을 거꾸로 걸어 올라가 인자 인덱스를 추론했습니다. 아는 view 연산에 대해서만 맞는 방식이라, 생산자(tnpu)가 `dram_arg`를 직접 적어 내려보내도록 바꾸고 추론 코드를 삭제했습니다. +**DMA가 어느 인자에 속하는지 — 추론에서 선언으로.** 처음에는 TOG 빌더가 memref view 연산을 거꾸로 걸어 올라가 인자 인덱스를 추론했습니다. 아는 view 연산에 대해서만 맞는 방식이라, 생산자(lowering pass)가 `dram_arg`를 직접 적어 내려보내도록 바꾸고 추론 코드를 삭제했습니다. **grid — 컴파일 타임 상수에서 런타임 인자로.** 위 3절. @@ -195,12 +215,12 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 | 동적 shape (timing) | 동작 | 트레이스 하나가 모든 shape | | 다차원 grid | 동작 | 테스트가 IR과 dispatch 양쪽 검증 | | 동적 shape (functional) | 제약 | 바이너리가 shape 특수화 → `ShapeMismatch`로 거부 | -| double buffering | 미착수 | tnpu가 동기 DMA만 발행. 251 vs 650의 주원인 | -| matmul timing | 미착수 | `build_tog`는 `vcix.iv` 이름으로 compute 노드를 찾는데 tnpu는 `llvm.riscv.sf.vc.*` 인트린식을 냄 | +| double buffering | 미착수 | lowering pass 가 동기 DMA 만 발행. 251 vs 650의 주원인 | +| matmul timing | 미착수 | `build_tog`는 `vcix.iv` 이름으로 compute 노드를 찾는데 lowering pass 는 `llvm.riscv.sf.vc.*` 인트린식을 냄 | | `triton_helpers` | 차단 | 모듈이 torch 안에 있고 tnpu venv에는 없음 | -| reduction | 차단 | tnpu 자체 문제 — 아래 | +| reduction | 차단 | lowering pass 쪽 문제 — 아래 | -**동적 shape의 한 가지 단서.** timing은 완전히 동작합니다. functional 바이너리는 tnpu가 grid·스칼라 값·memref extent를 전부 구워 넣어 shape 특수화되어 있어서, shape이 다른 launch를 `ShapeMismatch`로 **거부합니다** — 틀린 경계로 실행하는 대신. 사이클만 볼 때는 `pytorchsim_functional_mode: False`로 모든 shape을 돌릴 수 있습니다. +**동적 shape의 한 가지 단서.** timing은 완전히 동작합니다. functional 바이너리는 lowering pass 가 grid·스칼라 값·memref extent 를 전부 구워 넣어 shape 특수화되어 있어서, shape이 다른 launch를 `ShapeMismatch`로 **거부합니다** — 틀린 경계로 실행하는 대신. 사이클만 볼 때는 `pytorchsim_functional_mode: False`로 모든 shape을 돌릴 수 있습니다. **reduction이 막힌 지점.** `tt.reduce(axis=1)`이 triton-shared를 지나면 `linalg.transpose permutation=[1,0]` + `linalg.reduce dimensions=[0]`가 됩니다. transpose는 `transpose-reduce-to-rank0` 여부와 무관하게 삽입됩니다(rank 2에서 동일함을 측정). stage 3의 다섯 패스는 통과하고 `bank_vectorize`가 거부합니다 — 스크래치패드가 **레인 뱅킹**되어 있어 축소되는 축이 레인 안에 머물러야 하는데, identity-elementwise가 아니고 스칼라 폴백은 뱅킹된 스크래치패드를 읽게 되기 때문입니다. @@ -210,10 +230,18 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 ## 8. PR과 검증 -| PR | 범위 | 상태 | -|---|---|---| -| [PyTorchSim #305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305) | 18 commits · 23 files · +2209/−39 | draft, mergeable, CI green | -| [triton-npu #1](https://github.com/PSAL-POSTECH/triton-npu/pull/1) | 3 commits · 5 files · +177/−4 | open | +| PR | 내용 | 범위 | 상태 | +|---|---|---|---| +| [PyTorchSim #305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305) | 이식·연결 작업 본체 | 18 commits · 23 files · +2209/−39 | draft, mergeable, CI green | +| [triton-npu #1](https://github.com/PSAL-POSTECH/triton-npu/pull/1) | 이식에 필요해진 훅 3개 | 3 commits · 5 files · +177/−4 | open | + +triton-npu #1 은 lowering pass 를 고치는 PR이 아니라, **PyTorchSim이 그것을 호출하려면 있어야 했던 진입점**을 여는 PR입니다. 패스 로직 자체는 건드리지 않았습니다. + +| 훅 | 왜 필요했나 | +|---|---| +| `tnpu.cycle` | 타일 하나만 gem5로 돌려 cycle을 재려면, DMA를 지운 1-program 바이너리가 필요 | +| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 (추론 대신 생산자가 선언) | +| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌 수 있어야 함 | CI(`.github/workflows/triton_npu.yml`)는 툴체인 레이어가 ~1.8 GiB라 본 CI와 분리: @@ -232,10 +260,10 @@ triton-npu baselines success <- doctor + add/mul/relu/gemm/bmm ## 9. 다음 우선순위 -1. **double buffering** — tnpu가 비동기 DMA + `togsim.wait`를 내도록. 두 경로의 사이클 격차를 실제로 좁히는 유일한 항목이고, 기존 경로에 이미 있는 기계를 tnpu 쪽에 만드는 일입니다. -2. **shape 특수화 해소** — launch shape마다 재컴파일하거나, tnpu wrapper도 트레이스 생산자처럼 grid와 extent를 인자로 받게. 후자가 근본적. +1. **double buffering** — lowering pass 가 비동기 DMA + `togsim.wait` 를 내도록. 두 경로의 사이클 격차를 실제로 좁히는 유일한 항목이고, 기존 경로에 이미 있는 기계를 lowering pass 쪽에 만드는 일입니다. +2. **shape 특수화 해소** — launch shape마다 재컴파일하거나, lowering pass 의 wrapper 도 트레이스 생산자처럼 grid와 extent를 인자로 받게. 후자가 근본적. 3. **matmul timing** — `build_tog`가 vcix 인트린식을 인식하도록. systolic array 경로가 열립니다. -4. **reduction 레인 경로** — tnpu의 `bank_vectorize`에 reduction 추가, transpose를 `vlane_split_axis`로 흡수. 가장 큰 작업. +4. **reduction 레인 경로** — lowering pass 의 `bank_vectorize` 에 reduction 추가, transpose를 `vlane_split_axis`로 흡수. 가장 큰 작업. --- diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index de701e20..4952f6cb 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -1,9 +1,13 @@ # Triton codegen route (WIP) Replaces the hand-written MLIR emission in `PyTorchSimFrontend/mlir/` with -**Inductor's own Triton codegen**, lowered to this NPU by the **triton-npu** -(`tnpu`) pass pipeline. Opt-in and off by default; the MLIR route is untouched -and stays the production path. +**Inductor's own Triton codegen**, lowered to this NPU by the **PyTorchSim +lowering pass** (owned by 이정민; the code lives in the `triton-npu` repo, so +paths and module names read `tnpu`). Opt-in and off by default; the MLIR route +is untouched and stays the production path. + +The modules here are the PORT: they drive that lowering pass and wire its output +into the existing TOGSim / gem5 / Spike stack. The pass itself is not ours. ```bash TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py From 6670e5f29a69619d745c4333727cdcb43d8ceda8 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 13:43:43 +0900 Subject: [PATCH 21/35] [Docs] Make the next step coverage parity, measured rather than guessed 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. --- PyTorchSimFrontend/triton-codegen-route.md | 169 ++++++++++----------- 1 file changed, 82 insertions(+), 87 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index c0cdc7cc..23b8a7e0 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -14,8 +14,8 @@ | functional | `x + y`, `(x+y)*2 - x` 모두 **max abs error 0.0** (1024 elements, Spike) | | timing | TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | | 동적 shape | 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32, 재컴파일 없음 | -| 변경량 | PyTorchSim 18 commits / 23 files / +2209−39, tnpu 3 commits / +177−4 | | CI | 전 잡 green (툴체인 빌드 + 값 검증 + 기존 경로 회귀 확인) | +| 커버리지 | **elementwise와 그 융합까지만 확인됨.** non-contiguous 입력은 값이 틀리고, `triton_helpers`를 쓰는 커널은 멈춥니다 — 6절 | ### 작업 경계 @@ -26,50 +26,23 @@ | **PyTorchSim lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 이 문서의 범위 밖 | | **기존 PyTorchSim으로의 이식** | 위 lowering pass 를 기존 시뮬레이션 스택에 얹는 일: Inductor Triton 코드젠 가로채기, KernelSpec 생성, grid 합성, 트레이스/사이클 산출, functional launch, TOGSim 연결 | 이 문서의 작업 | -즉 lowering pass 자체는 만들지 않았습니다. **이미 있는 lowering pass 를 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim/gem5/Spike 스택에 물린 것**이 여기서 한 일입니다. 그 과정에서 lowering pass 쪽에 필요해진 최소한의 훅 3개(`tnpu.cycle`, `dram_arg`, `tnpu.spike`)는 별도 PR로 올렸고, 8절에 있습니다. +즉 lowering pass 자체는 만들지 않았습니다. **이미 있는 lowering pass 를 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim/gem5/Spike 스택에 물린 것**이 여기서 한 일입니다. -문서에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있고, 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image` 등) — 문서를 따라 실제 코드를 찾아갈 수 있어야 하기 때문입니다. +그 과정에서 lowering pass 쪽에 진입점 3개가 필요했습니다. 패스 로직을 고치는 것이 아니라 **바깥에서 호출할 수 있게 여는** 변경입니다 ([triton-npu#1](https://github.com/PSAL-POSTECH/triton-npu/pull/1), 5 files · +177−4). ---- +| 훅 | 왜 필요했나 | +|---|---| +| `tnpu.cycle` | 타일 하나만 gem5로 돌려 cycle을 재려면, DMA를 지운 1-program 바이너리가 필요 | +| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 (추론 대신 생산자가 선언) | +| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌 수 있어야 함 | -## 1. 파이프라인 +이식 작업 본체는 [PyTorchSim#305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305)입니다 (23 files · +2209−39, CI green). -``` -torch.compile - └ TritonNPUScheduling.define_kernel scheduling.py - │ Inductor 가 만든 triton 소스 텍스트 + 수집한 메타데이터 - ▼ - triton_npu_compile(src, meta, name) codecache.py - │ tnpu KernelSpec 생성 kernel_spec.py - │ - 블록 크기를 constexpr 로 고정 - │ - 인자 역할(in/out/inout) · dtype · numel - │ - grid, 사용자 스칼라 값 - ▼ - PyTorchSim lowering pass (별도 인터프리터, subprocess) tnpu_bridge.py - │ ┌─ 담당 이정민 / 이 문서의 범위 밖. - │ └─ 여기서 한 일은 이 단계를 "호출하고 결과를 스택에 물린" 부분. - │ 1 ttir triton 커널 → Triton IR - │ 2 ttshared → linalg / memref / scf.for (triton-shared) - │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) - │ 4 lower vcix → gemmini DMA → LLVM - │ 5 binary mlir-translate → llc → RISC-V ELF - ▼ - TritonNPULauncher.__call__ codecache.py - │ - ├ functional 텐서 → runtime/*.raw → Spike → 텐서 functional.py - │ tnpu stage 6 (tnpu.spike) 재사용 - │ - └ timing 04-custom.mlir timing.py - ├ build_tog sample → gem5 → 타일 cycle 실측 - └ build_skeleton → trace.so + trace_cycles.tsv - → TOGSim -``` - -**LLVM 이음매.** lowering pass 는 stock LLVM 23 을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. +문서에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있고, 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image` 등) — 문서를 따라 실제 코드를 찾아갈 수 있어야 하기 때문입니다. --- -## 2. 기존 MLIR 경로와의 차이 +## 1. 기존 MLIR 경로와의 차이 ### 갈라지는 지점과 합쳐지는 지점 @@ -110,7 +83,7 @@ torch.compile | grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | | lowering | `PyTorchSimFrontend/mlir/` (in-process) | PyTorchSim lowering pass — 담당 이정민 (subprocess, LLVM 23) | | 융합 | 템플릿과 `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | -| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | +| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 (6절 실측) | | functional | `FunctionalSimulator.run_spike` | tnpu stage 6 (`tnpu.spike`) | | timing | `trace.so` + `trace_cycles.tsv` → TOGSim | **동일** | | 타일 cycle 실측 | gem5 | **동일** (`build_tog` sample 모드 공유) | @@ -121,12 +94,49 @@ torch.compile **Triton 경로가 앞선 곳** — 동적 shape. 기존 경로의 C++ 트레이스는 `trace_to_tilegraph(..., nullptr, 0)`으로 shape 인자를 아예 넘기지 않아 shape마다 트레이스를 다시 만들어야 하고, 그걸 푸는 작업이 PR #269로 아직 열려 있습니다. Triton 경로는 `shape_args`를 통해 **트레이스 하나가 모든 shape을 섬깁니다.** -**기존 경로가 앞선 곳** — op 커버리지와 DMA 겹침. 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 아래 사이클 격차의 원인입니다. +**기존 경로가 앞선 곳** — op 커버리지와 DMA 겹침. 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 4절 사이클 격차의 원인입니다. **바뀌지 않은 것** — TOGSim, 하드웨어 설정, gem5 샘플링 방식, 트레이스 계약. 두 경로는 같은 시뮬레이터를 먹입니다. --- +## 2. 파이프라인 + +``` +torch.compile + └ TritonNPUScheduling.define_kernel scheduling.py + │ Inductor 가 만든 triton 소스 텍스트 + 수집한 메타데이터 + ▼ + triton_npu_compile(src, meta, name) codecache.py + │ tnpu KernelSpec 생성 kernel_spec.py + │ - 블록 크기를 constexpr 로 고정 + │ - 인자 역할(in/out/inout) · dtype · numel + │ - grid, 사용자 스칼라 값 + ▼ + PyTorchSim lowering pass (별도 인터프리터, subprocess) tnpu_bridge.py + │ ┌─ 담당 이정민 / 이 문서의 범위 밖. + │ └─ 여기서 한 일은 이 단계를 "호출하고 결과를 스택에 물린" 부분. + │ 1 ttir triton 커널 → Triton IR + │ 2 ttshared → linalg / memref / scf.for (triton-shared) + │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) + │ 4 lower vcix → gemmini DMA → LLVM + │ 5 binary mlir-translate → llc → RISC-V ELF + ▼ + TritonNPULauncher.__call__ codecache.py + │ + ├ functional 텐서 → runtime/*.raw → Spike → 텐서 functional.py + │ tnpu stage 6 (tnpu.spike) 재사용 + │ + └ timing 04-custom.mlir timing.py + ├ build_tog sample → gem5 → 타일 cycle 실측 + └ build_skeleton → trace.so + trace_cycles.tsv + → TOGSim +``` + +**LLVM 이음매.** lowering pass 는 stock LLVM 23 을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. + +--- + ## 3. 핵심 설계 문제: 커널 하나가 무엇을 뜻하는가 ``` @@ -199,71 +209,56 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 --- -## 6. 일반성을 위해 되돌린 설계 둘 +## 6. 다음 작업 — 기존 경로 커버리지까지 검증 -**DMA가 어느 인자에 속하는지 — 추론에서 선언으로.** 처음에는 TOG 빌더가 memref view 연산을 거꾸로 걸어 올라가 인자 인덱스를 추론했습니다. 아는 view 연산에 대해서만 맞는 방식이라, 생산자(lowering pass)가 `dram_arg`를 직접 적어 내려보내도록 바꾸고 추론 코드를 삭제했습니다. +지금 확인된 것은 elementwise와 그 융합뿐입니다. 다음 작업은 기능을 더 얹는 것이 아니라, **기존 MLIR 경로를 지탱하는 op 테스트 스위트를 Triton 경로로 그대로 돌려 어디까지 가는지 확인하는 것**입니다. -**grid — 컴파일 타임 상수에서 런타임 인자로.** 위 3절. +대상은 `tests/ops/` 아래 이미 있는 것들입니다 — `elementwise`, `reduce`, `gemm`, `conv`, `attention`, `view`, `sort`, `fusion`, `misc`. MLIR 경로가 통과하는 범위가 곧 목표선입니다. ---- +### 1차 실측 -## 7. 현재 상태 +대표 케이스를 `TORCHSIM_TRITON_CODEGEN=1`로 돌린 결과입니다. **경로 진입** 열은 Triton 경로를 실제로 탔는지(작업 디렉터리 생성 여부)를 뜻합니다 — 이걸 보지 않으면 Inductor가 extern으로 뺀 것을 통과로 착각합니다. -| 기능 | 상태 | 내용 | +| 케이스 | 경로 진입 | 결과 | |---|---|---| -| elementwise + 융합 | 동작 | 값·사이클 모두 통과, CI 포함 | -| 동적 shape (timing) | 동작 | 트레이스 하나가 모든 shape | -| 다차원 grid | 동작 | 테스트가 IR과 dispatch 양쪽 검증 | -| 동적 shape (functional) | 제약 | 바이너리가 shape 특수화 → `ShapeMismatch`로 거부 | -| double buffering | 미착수 | lowering pass 가 동기 DMA 만 발행. 251 vs 650의 주원인 | -| matmul timing | 미착수 | `build_tog`는 `vcix.iv` 이름으로 compute 노드를 찾는데 lowering pass 는 `llvm.riscv.sf.vc.*` 인트린식을 냄 | -| `triton_helpers` | 차단 | 모듈이 torch 안에 있고 tnpu venv에는 없음 | -| reduction | 차단 | lowering pass 쪽 문제 — 아래 | - -**동적 shape의 한 가지 단서.** timing은 완전히 동작합니다. functional 바이너리는 lowering pass 가 grid·스칼라 값·memref extent 를 전부 구워 넣어 shape 특수화되어 있어서, shape이 다른 launch를 `ShapeMismatch`로 **거부합니다** — 틀린 경계로 실행하는 대신. 사이클만 볼 때는 `pytorchsim_functional_mode: False`로 모든 shape을 돌릴 수 있습니다. - -**reduction이 막힌 지점.** `tt.reduce(axis=1)`이 triton-shared를 지나면 `linalg.transpose permutation=[1,0]` + `linalg.reduce dimensions=[0]`가 됩니다. transpose는 `transpose-reduce-to-rank0` 여부와 무관하게 삽입됩니다(rank 2에서 동일함을 측정). stage 3의 다섯 패스는 통과하고 `bank_vectorize`가 거부합니다 — 스크래치패드가 **레인 뱅킹**되어 있어 축소되는 축이 레인 안에 머물러야 하는데, identity-elementwise가 아니고 스칼라 폴백은 뱅킹된 스크래치패드를 읽게 되기 때문입니다. +| `x + y` | 예 | 값 일치 | +| `(x+y)*2 - x` (융합) | 예 | 값 일치 | +| `x.t() + 1` | 예 | **값 틀림 — 4030/4096** | +| `relu` | 예 | 중단: `SpecIncomplete: triton_helpers.maximum` | +| `softmax` | 예 | 중단: `SpecIncomplete: triton_helpers.max2` | +| `exp` | 예 | 중단: lowering pass 실패 | +| `sum(dim=1)` | 예 | 중단: lowering pass 실패 (`bank_vectorize`) | +| `cat` | 예 | 중단: lowering pass 실패 | +| `a @ b` | **아니오** | Inductor가 `aten.mm` extern 으로 처리 — 경로에 도달하지 않음 | -테스트(`check_reduction_is_refused`)가 이 경계를 못박습니다. **reduction이 컴파일에 성공하면 테스트가 실패합니다** — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. +### 작업 목록 (우선순위 순) ---- +**1. non-contiguous 텐서 — 값이 조용히 틀리는 유일한 항목이라 최우선.** -## 8. PR과 검증 +`x.t() + 1`의 출력이 정확히 `x + 1`입니다. transpose가 통째로 무시됩니다. 원인은 lowering pass 가 아니라 **이식 쪽**입니다: -| PR | 내용 | 범위 | 상태 | -|---|---|---|---| -| [PyTorchSim #305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305) | 이식·연결 작업 본체 | 18 commits · 23 files · +2209/−39 | draft, mergeable, CI green | -| [triton-npu #1](https://github.com/PSAL-POSTECH/triton-npu/pull/1) | 이식에 필요해진 훅 3개 | 3 commits · 5 files · +177/−4 | open | +``` +Inductor 가 낸 커널: tmp0 = tl.load(in_ptr0 + x0) <- 인덱스가 항등 + 즉 transpose 를 인덱스 식이 아니라 + 출력 버퍼의 stride (1,64) 로 접었음 -triton-npu #1 은 lowering pass 를 고치는 PR이 아니라, **PyTorchSim이 그것을 호출하려면 있어야 했던 진입점**을 여는 PR입니다. 패스 로직 자체는 건드리지 않았습니다. +functional.py: write_inputs t.contiguous().numpy().tofile() <- 저장 순서를 재배열 + read_outputs t.copy_(flat.view_as(t)) <- 논리 순서로 되씀 +``` -| 훅 | 왜 필요했나 | -|---|---| -| `tnpu.cycle` | 타일 하나만 gem5로 돌려 cycle을 재려면, DMA를 지운 1-program 바이너리가 필요 | -| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 (추론 대신 생산자가 선언) | -| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌 수 있어야 함 | +둘 다 **논리 순서와 저장 순서가 같다**고 가정합니다. contiguous 텐서에서만 참이고, 그래서 elementwise는 통과하고 transpose는 틀립니다. 저장 순서 기준으로 읽고 쓰도록 고치고, 비-contiguous 케이스를 테스트에 넣어야 합니다. -CI(`.github/workflows/triton_npu.yml`)는 툴체인 레이어가 ~1.8 GiB라 본 CI와 분리: +**2. `triton_helpers` 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등 상당수가 여기서 막힙니다. 모듈이 torch 안에 있고 lowering pass 쪽 venv 에는 없습니다. `strip_for_tnpu`가 어떤 헬퍼인지 이름을 대고 멈추므로, 필요한 것만 최소로 벤더링하면 커버리지가 한 번에 크게 늘어납니다. -``` -Check tnpu access success -Build tnpu toolchain image success -Build app image on tnpu base success -Inductor Triton route success <- test_triton_codegen.py (값 검증 포함) -MLIR route still passes success <- 기존 경로 회귀 없음 -triton-npu baselines success <- doctor + add/mul/relu/gemm/bmm -``` +**3. lowering pass 쪽 실패 (`exp`, `cat`, reduction).** 담당(이정민)과 나눠야 할 부분입니다. reduction은 원인이 파악돼 있습니다 — `tt.reduce(axis=1)`이 `linalg.transpose` + `linalg.reduce`가 되고, 스크래치패드가 레인 뱅킹되어 있어 `bank_vectorize`가 거부합니다. `exp`와 `cat`은 아직 원인 미확인. -**머지 순서.** `thirdparty/triton-npu.json`이 tnpu 커밋 `22df065`를 핀하는데, 이는 `feature/timing-form`에만 있고 `main`에는 없습니다. sha라 CI fetch는 되지만 #1이 리베이스 머지되면 뜹니다 — **#1 머지 → 핀을 main 커밋으로 재조정 → #305** 순서가 안전합니다. +**4. matmul 경로 진입.** 지금은 Inductor가 `aten.mm` extern 으로 빼서 Triton 경로를 아예 타지 않습니다. 통과한 것처럼 보이지만 시뮬레이터를 거치지 않은 값입니다. Triton 템플릿을 쓰게 하려면 `max_autotune` 계열 설정이 필요하고, 그래야 systolic array 경로를 볼 수 있습니다. ---- +**5. double buffering.** 커버리지가 아니라 정확도 문제 — 4절의 251 vs 650 격차. lowering pass 가 비동기 DMA + `togsim.wait`를 내야 하고, 기존 경로에 이미 있는 기계를 옮기는 일입니다. -## 9. 다음 우선순위 +### 회귀 방지 -1. **double buffering** — lowering pass 가 비동기 DMA + `togsim.wait` 를 내도록. 두 경로의 사이클 격차를 실제로 좁히는 유일한 항목이고, 기존 경로에 이미 있는 기계를 lowering pass 쪽에 만드는 일입니다. -2. **shape 특수화 해소** — launch shape마다 재컴파일하거나, lowering pass 의 wrapper 도 트레이스 생산자처럼 grid와 extent를 인자로 받게. 후자가 근본적. -3. **matmul timing** — `build_tog`가 vcix 인트린식을 인식하도록. systolic array 경로가 열립니다. -4. **reduction 레인 경로** — lowering pass 의 `bank_vectorize` 에 reduction 추가, transpose를 `vlane_split_axis`로 흡수. 가장 큰 작업. +`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서, 컴파일에 성공하면 테스트가 실패합니다 — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. --- From c099add4ad20ca922b6a6d58c8b4a97a34cad69f Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 15:44:20 +0900 Subject: [PATCH 22/35] [Docs] Rewrite the report as one piece 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. --- PyTorchSimFrontend/triton-codegen-route.md | 184 ++++++++++----------- 1 file changed, 87 insertions(+), 97 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index 23b8a7e0..db674daa 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -1,50 +1,35 @@ -# Inductor Triton 코드젠 경로를 PyTorchSim에 연결 +# Triton 코드젠 경로를 PyTorchSim에 이식 -`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **PyTorchSim lowering pass**(담당 이정민)가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. +`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **PyTorchSim lowering pass**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. -**이 문서가 보고하는 것은 그 lowering pass 를 기존 PyTorchSim 스택에 이식·연결한 작업입니다.** lowering pass 자체는 범위 밖입니다 — 아래 "작업 경계" 참고. +## 작업 경계 -**functional과 timing 양쪽이 연결되어 있고, 동적 shape도 처리됩니다.** - -모듈별 동작과 사용법은 [`triton_backend/README.md`](triton_backend/README.md)에 있습니다. -이 문서는 기존 경로와의 대조, 설계 판단, 측정 결과를 다룹니다. - -| | | -|---|---| -| functional | `x + y`, `(x+y)*2 - x` 모두 **max abs error 0.0** (1024 elements, Spike) | -| timing | TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | -| 동적 shape | 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32, 재컴파일 없음 | -| CI | 전 잡 green (툴체인 빌드 + 값 검증 + 기존 경로 회귀 확인) | -| 커버리지 | **elementwise와 그 융합까지만 확인됨.** non-contiguous 입력은 값이 틀리고, `triton_helpers`를 쓰는 커널은 멈춥니다 — 6절 | - -### 작업 경계 - -이 경로는 두 부분으로 나뉘고, **이 문서가 보고하는 작업은 아래쪽입니다.** +이 경로는 두 부분으로 나뉩니다. **이 문서가 보고하는 것은 아래쪽입니다.** | 부분 | 하는 일 | 소관 | |---|---|---| -| **PyTorchSim lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 이 문서의 범위 밖 | -| **기존 PyTorchSim으로의 이식** | 위 lowering pass 를 기존 시뮬레이션 스택에 얹는 일: Inductor Triton 코드젠 가로채기, KernelSpec 생성, grid 합성, 트레이스/사이클 산출, functional launch, TOGSim 연결 | 이 문서의 작업 | +| **PyTorchSim lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 범위 밖 | +| **기존 PyTorchSim으로의 이식** | 그 lowering pass를 기존 시뮬레이션 스택에 얹는 일 | 이 문서 | -즉 lowering pass 자체는 만들지 않았습니다. **이미 있는 lowering pass 를 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim/gem5/Spike 스택에 물린 것**이 여기서 한 일입니다. +lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim / gem5 / Spike 스택에 물린 것**이 여기서 한 일입니다. -그 과정에서 lowering pass 쪽에 진입점 3개가 필요했습니다. 패스 로직을 고치는 것이 아니라 **바깥에서 호출할 수 있게 여는** 변경입니다 ([triton-npu#1](https://github.com/PSAL-POSTECH/triton-npu/pull/1), 5 files · +177−4). +용어: 산문에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있어서 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image`). 문서를 따라 코드를 찾아갈 수 있어야 하기 때문입니다. -| 훅 | 왜 필요했나 | -|---|---| -| `tnpu.cycle` | 타일 하나만 gem5로 돌려 cycle을 재려면, DMA를 지운 1-program 바이너리가 필요 | -| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 (추론 대신 생산자가 선언) | -| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌 수 있어야 함 | - -이식 작업 본체는 [PyTorchSim#305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305)입니다 (23 files · +2209−39, CI green). +## 현재 도달점 -문서에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있고, 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image` 등) — 문서를 따라 실제 코드를 찾아갈 수 있어야 하기 때문입니다. +| | | +|---|---| +| functional | 연결됨. `x + y`, `(x+y)*2 - x` **max abs error 0.0** (1024 elements, Spike) | +| timing | 연결됨. TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | +| 동적 shape | 처리됨. 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32 | +| **커버리지** | **elementwise와 그 융합까지.** 그 밖은 5절 실측 참고 | +| CI | 전 잡 green | --- ## 1. 기존 MLIR 경로와의 차이 -### 갈라지는 지점과 합쳐지는 지점 +### 갈라지는 곳과 합쳐지는 곳 ``` torch.compile / Inductor 스케줄 @@ -53,50 +38,50 @@ │ │ [기존] MLIR 경로 [신규] Triton 경로 │ │ - Inductor 스케줄 → 손으로 쓴 Inductor 의 Triton 코드젠 - op별 MLIR 템플릿 이 낸 커널 소스를 가로챔 + Inductor 스케줄 → 손으로 쓴 Inductor 의 Triton 코드젠이 + op별 MLIR 템플릿 낸 커널 소스를 가로챔 (gemm, conv, sdpa, sort, (op별 템플릿 없음) cat, maxpool, bmm …) │ │ PyTorchSim mlir/ 패스 PyTorchSim lowering pass - PSAL LLVM 20 (subprocess, stock LLVM 23) - 담당 이정민 — 범위 밖 - 여기서 한 일은 이 블록을 - 아래 합류점까지 잇는 배선 + PSAL LLVM 20 (in-process) stock LLVM 23 (subprocess) + └ 담당 이정민 — 범위 밖. + 여기서 한 일은 이 블록을 + 아래 합류점까지 잇는 배선 │ │ └───────────────┬───────────────┘ │ - ▼ 여기서 다시 합류 ▼ - trace.so + trace_cycles.tsv - → TOGSim - (트레이스 계약은 완전히 동일) + ▼ 여기서 다시 합류 ▼ + trace.so + trace_cycles.tsv + → TOGSim + (트레이스 계약은 완전히 동일) ``` -핵심은 **TOGSim이 두 경로를 구분하지 못한다**는 점입니다. 트레이스 생산자의 형태가 같으므로 하드웨어 모델·DRAM·NoC·L2는 손대지 않았습니다. +핵심은 **TOGSim이 두 경로를 구분하지 못한다**는 점입니다. 트레이스 생산자의 형태가 같으므로 하드웨어 모델·DRAM·NoC·L2는 한 줄도 손대지 않았습니다. ### 항목별 대조 | | 기존 MLIR 경로 | 신규 Triton 경로 | |---|---|---| -| 커널을 만드는 주체 | PyTorchSim의 op별 MLIR 템플릿 | Inductor의 Triton 코드젠 | +| 커널을 만드는 주체 | op별 MLIR 템플릿 (직접 작성) | Inductor의 Triton 코드젠 | | **커널 하나의 의미** | **루프 네스트 전체** | **타일 하나** | | grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | -| lowering | `PyTorchSimFrontend/mlir/` (in-process) | PyTorchSim lowering pass — 담당 이정민 (subprocess, LLVM 23) | -| 융합 | 템플릿과 `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | -| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 (6절 실측) | +| lowering | `PyTorchSimFrontend/mlir/` | PyTorchSim lowering pass (subprocess) | +| 융합 | 템플릿 + `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | +| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | | functional | `FunctionalSimulator.run_spike` | tnpu stage 6 (`tnpu.spike`) | | timing | `trace.so` + `trace_cycles.tsv` → TOGSim | **동일** | | 타일 cycle 실측 | gem5 | **동일** (`build_tog` sample 모드 공유) | -| DMA | 비동기 + `togsim.wait` 배리어 | **동기만** (`togsim.wait` 0개) | -| 동적 shape | 트레이스 경로는 아직 미지원 (PR #269 진행 중) | timing 경로에서 동작 | +| DMA | 비동기 + `togsim.wait` 배리어 | **동기만** | +| 동적 shape | 트레이스 경로 미지원 (PR #269 진행 중) | 동작 | ### 이 대조가 말해주는 것 -**Triton 경로가 앞선 곳** — 동적 shape. 기존 경로의 C++ 트레이스는 `trace_to_tilegraph(..., nullptr, 0)`으로 shape 인자를 아예 넘기지 않아 shape마다 트레이스를 다시 만들어야 하고, 그걸 푸는 작업이 PR #269로 아직 열려 있습니다. Triton 경로는 `shape_args`를 통해 **트레이스 하나가 모든 shape을 섬깁니다.** +**Triton 경로가 앞선 곳 — 동적 shape.** 기존 경로의 C++ 트레이스는 `trace_to_tilegraph(..., nullptr, 0)`으로 shape 인자를 아예 넘기지 않아 shape마다 트레이스를 다시 만들어야 하고, 그걸 푸는 작업이 PR #269로 아직 열려 있습니다. Triton 경로는 `shape_args`를 통해 트레이스 하나가 모든 shape을 섬깁니다. -**기존 경로가 앞선 곳** — op 커버리지와 DMA 겹침. 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 4절 사이클 격차의 원인입니다. +**기존 경로가 앞선 곳 — op 커버리지와 DMA 겹침.** 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 4절 사이클 격차의 원인입니다. -**바뀌지 않은 것** — TOGSim, 하드웨어 설정, gem5 샘플링 방식, 트레이스 계약. 두 경로는 같은 시뮬레이터를 먹입니다. +**바뀌지 않은 것 — TOGSim 전체.** 하드웨어 설정, gem5 샘플링 방식, 트레이스 계약. 두 경로는 같은 시뮬레이터를 먹입니다. --- @@ -108,17 +93,18 @@ torch.compile │ Inductor 가 만든 triton 소스 텍스트 + 수집한 메타데이터 ▼ triton_npu_compile(src, meta, name) codecache.py - │ tnpu KernelSpec 생성 kernel_spec.py + │ KernelSpec 생성 kernel_spec.py │ - 블록 크기를 constexpr 로 고정 │ - 인자 역할(in/out/inout) · dtype · numel │ - grid, 사용자 스칼라 값 ▼ - PyTorchSim lowering pass (별도 인터프리터, subprocess) tnpu_bridge.py - │ ┌─ 담당 이정민 / 이 문서의 범위 밖. - │ └─ 여기서 한 일은 이 단계를 "호출하고 결과를 스택에 물린" 부분. + PyTorchSim lowering pass (subprocess) tnpu_bridge.py + │ 담당 이정민 — 범위 밖. 여기서 한 일은 이 단계를 + │ 호출하고 그 결과를 아래 스택에 물린 부분. + │ │ 1 ttir triton 커널 → Triton IR - │ 2 ttshared → linalg / memref / scf.for (triton-shared) - │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) + │ 2 ttshared → linalg / memref / scf.for (triton-shared) + │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) │ 4 lower vcix → gemmini DMA → LLVM │ 5 binary mlir-translate → llc → RISC-V ELF ▼ @@ -133,7 +119,15 @@ torch.compile → TOGSim ``` -**LLVM 이음매.** lowering pass 는 stock LLVM 23 을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. +**LLVM 이음매.** lowering pass는 stock LLVM 23을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. + +**lowering pass에 필요했던 진입점 3개.** 패스 로직을 고치는 것이 아니라 바깥에서 호출할 수 있게 여는 변경입니다 ([triton-npu#1](https://github.com/PSAL-POSTECH/triton-npu/pull/1)). + +| 훅 | 왜 필요했나 | +|---|---| +| `tnpu.cycle` | 타일 하나만 gem5로 재려면 DMA를 지운 1-program 바이너리가 필요 | +| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 | +| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌아야 함 | --- @@ -144,26 +138,26 @@ MLIR 경로 커널 = 루프 네스트 전체. TOG 가 루프에서 work-item Triton 커널 = 타일 하나. grid 는 커널 밖, launch 가 쥐고 있음 ``` -TOGSim의 트레이스 계약(`docs/design/togsim_cpp_trace.md` §9.1/§9.3)이 이미 이 둘을 구분합니다: +이식의 본질적 어려움은 여기 하나로 모입니다. 그런데 TOGSim의 트레이스 계약(`docs/design/togsim_cpp_trace.md` §9.1/§9.3)이 이미 둘을 구분하고 있었습니다: - `togsim_kernel_tile(ctx, iv, n)` — work-item 하나 - `togsim_kernel(ctx, shape_args, n)` — 병렬 영역의 열거 -Triton 커널 본문은 전자에 대응하므로, **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. 그 합성이 `lower_to_emitc.WorkItem` + `_materialize_grid_loop`입니다. +Triton 커널 본문은 전자에 대응하므로 **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. 그 합성이 `lower_to_emitc.WorkItem` + `_materialize_grid_loop`입니다. ### 동적 shape이 여기서 나옵니다 -`_materialize_grid_loop`은 축 **개수**만 컴파일에 박고, **범위**는 `shape_args`에서 읽습니다. +`_materialize_grid_loop`은 축 **개수**만 컴파일에 박고 **범위**는 `shape_args`에서 읽습니다. ``` -컴파일 시 축이 몇 개인지만 안다 → 루프 네스트 골격 생성 +컴파일 시 축이 몇 개인지만 안다 → 루프 네스트 골격 생성 런타임 실제 numel 로 grid 계산 → trace_shape.txt 로 전달 TOGSim 이 build_trace_tilegraph 에서 읽어 shape_args 로 주입 ``` 측정: `dynamic=True`로 n=1024 → grid 8, n=4096 → grid 32. 트레이스 재생성 없음. -다차원 grid(Triton 제약상 최대 3D)도 지원합니다. 구현 중 두 번 틀렸고 둘 다 rank ≥ 2에서만 드러났습니다 — 종료자가 있는 블록 끝에 삽입하는 문제, 그리고 bound를 루프 뒤에 만들어 dominance를 깨는 문제. 그래서 테스트가 생성된 C++가 아니라 **MLIR 모듈 자체를 verify**합니다. +다차원 grid(Triton 제약상 최대 3D)도 지원합니다. 구현 중 두 번 틀렸고 둘 다 rank ≥ 2에서만 드러났습니다 — 종료자가 있는 블록 끝에 삽입하는 문제, bound를 루프 뒤에 만들어 dominance를 깨는 문제. 그래서 테스트가 생성된 C++가 아니라 **MLIR 모듈 자체를 verify**합니다. --- @@ -184,36 +178,32 @@ Triton 커널 본문은 전자에 대응하므로, **후자를 합성해서 씌 | TOGSim 총계 | 650 | DRAM 트래픽 8192 B = 8 work-item × 2 load × 512 B, 정확히 일치 | | 기존 MLIR 경로 (동일 연산) | 251 | 같은 자릿수 | -650 대 251은 모델 오류가 아닙니다. **lowering pass 가 동기 DMA 만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. +650 대 251은 모델 오류가 아닙니다. **lowering pass가 동기 DMA만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. ---- - -## 5. 도중에 찾은 실제 버그 +### 도중에 찾은 버그: 인자 한 칸 밀림 -functional 배선은 배관 작업일 줄 알았는데, 첫 실행에서 **1024개 중 896개가 틀렸습니다.** `pid_x=0` 블록만 맞고 나머지 7개는 전부 0. +functional 배선은 배관 작업일 줄 알았는데 첫 실행에서 **1024개 중 896개가 틀렸습니다.** `pid_x=0` 블록만 맞고 나머지 7개는 전부 0. ``` -MLIR func.func @k(%arg0..2: memref<*xf32>, in_ptr0, in_ptr1, out_ptr0 - %arg3: i32 xnumel <- 사용자 스칼라 - %arg4,5,6: i32 gridX,Y,Z - %arg7,8,9: i32 pidX,Y,Z ) - -wrapper k(1,&d_in_ptr0, 1,&d_in_ptr1, 1,&d_out_ptr0, 8, 1, 1, pid_x, pid_y, pid_z); - +------ i32 6개뿐 ------+ - xnumel 누락 +lowered MLIR @k(%arg0..2: memref<*xf32> in_ptr0, in_ptr1, out_ptr0 + %arg3: i32 xnumel <- 사용자 스칼라 + %arg4,5,6: i32 gridX,Y,Z + %arg7,8,9: i32 pidX,Y,Z ) + +wrapper 호출 k(1,&d_in_ptr0, 1,&d_in_ptr1, 1,&d_out_ptr0, 8,1,1, pid_x,pid_y,pid_z) + +---- i32 6개뿐 ----+ + xnumel 누락 ``` -triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. lowering pass 의 wrapper 는 이를 `spec.extra["scalar_args"]`에서 읽는데, PyTorchSim이 생성하는 spec에는 `extra`가 아예 없었습니다. 인자가 한 칸씩 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. +triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. wrapper는 이를 `spec.extra["scalar_args"]`에서 읽는데 우리가 생성하는 spec에는 `extra`가 없었습니다. 인자가 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. **틀린 값이 쓰레기가 아니라 0으로 나온 점**이 고약합니다. 쓰레기값이면 즉시 눈에 띄지만 0은 그럴듯해 보입니다. timing 경로는 인자 위치를 lowered MLIR 시그니처에서 직접 읽어 애초에 정확했고, 그래서 functional을 붙이기 전까지 드러나지 않았습니다. --- -## 6. 다음 작업 — 기존 경로 커버리지까지 검증 - -지금 확인된 것은 elementwise와 그 융합뿐입니다. 다음 작업은 기능을 더 얹는 것이 아니라, **기존 MLIR 경로를 지탱하는 op 테스트 스위트를 Triton 경로로 그대로 돌려 어디까지 가는지 확인하는 것**입니다. +## 5. 다음 작업 — 기존 경로 커버리지까지 검증 -대상은 `tests/ops/` 아래 이미 있는 것들입니다 — `elementwise`, `reduce`, `gemm`, `conv`, `attention`, `view`, `sort`, `fusion`, `misc`. MLIR 경로가 통과하는 범위가 곧 목표선입니다. +확인된 것은 elementwise와 그 융합뿐입니다. 다음 작업은 기능을 더 얹는 것이 아니라, **기존 MLIR 경로를 지탱하는 op 테스트 스위트를 Triton 경로로 그대로 돌려 어디까지 가는지 확인하는 것**입니다. 대상은 `tests/ops/` 아래 이미 있는 것들이고(`elementwise`, `reduce`, `gemm`, `conv`, `attention`, `view`, `sort`, `fusion`, `misc`), MLIR 경로가 통과하는 범위가 목표선입니다. ### 1차 실측 @@ -229,37 +219,37 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 | `exp` | 예 | 중단: lowering pass 실패 | | `sum(dim=1)` | 예 | 중단: lowering pass 실패 (`bank_vectorize`) | | `cat` | 예 | 중단: lowering pass 실패 | -| `a @ b` | **아니오** | Inductor가 `aten.mm` extern 으로 처리 — 경로에 도달하지 않음 | +| `a @ b` | **아니오** | Inductor가 `aten.mm` extern으로 처리 — 경로에 도달하지 않음 | -### 작업 목록 (우선순위 순) +### 우선순위 -**1. non-contiguous 텐서 — 값이 조용히 틀리는 유일한 항목이라 최우선.** - -`x.t() + 1`의 출력이 정확히 `x + 1`입니다. transpose가 통째로 무시됩니다. 원인은 lowering pass 가 아니라 **이식 쪽**입니다: +**1. non-contiguous 텐서.** 값이 조용히 틀리는 유일한 항목이라 최우선입니다. `x.t() + 1`의 출력이 정확히 `x + 1`이고 — transpose가 통째로 무시되고 — 아무 예외도 나지 않습니다. 원인은 lowering pass가 아니라 **이식 쪽**입니다: ``` -Inductor 가 낸 커널: tmp0 = tl.load(in_ptr0 + x0) <- 인덱스가 항등 - 즉 transpose 를 인덱스 식이 아니라 +Inductor 가 낸 커널 tmp0 = tl.load(in_ptr0 + x0) <- 인덱스가 항등. + transpose 를 인덱스 식이 아니라 출력 버퍼의 stride (1,64) 로 접었음 -functional.py: write_inputs t.contiguous().numpy().tofile() <- 저장 순서를 재배열 - read_outputs t.copy_(flat.view_as(t)) <- 논리 순서로 되씀 +functional.py write_inputs t.contiguous()...tofile() <- 저장 순서를 재배열 + read_outputs t.copy_(flat.view_as(t)) <- 논리 순서로 되씀 ``` -둘 다 **논리 순서와 저장 순서가 같다**고 가정합니다. contiguous 텐서에서만 참이고, 그래서 elementwise는 통과하고 transpose는 틀립니다. 저장 순서 기준으로 읽고 쓰도록 고치고, 비-contiguous 케이스를 테스트에 넣어야 합니다. +둘 다 **논리 순서와 저장 순서가 같다**고 가정합니다. contiguous 텐서에서만 참이고, 지금까지 통과한 것이 정확히 그 집합입니다. 저장 순서 기준으로 읽고 쓰도록 고치고 비-contiguous 케이스를 테스트에 넣어야 합니다. -**2. `triton_helpers` 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등 상당수가 여기서 막힙니다. 모듈이 torch 안에 있고 lowering pass 쪽 venv 에는 없습니다. `strip_for_tnpu`가 어떤 헬퍼인지 이름을 대고 멈추므로, 필요한 것만 최소로 벤더링하면 커버리지가 한 번에 크게 늘어납니다. +**2. `triton_helpers` 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등 상당수가 여기서 막힙니다. 모듈이 torch 안에 있고 lowering pass 쪽 venv에는 없습니다. `strip_for_tnpu`가 어떤 헬퍼인지 이름을 대고 멈추므로, 필요한 것만 최소로 벤더링하면 커버리지가 한 번에 크게 늘어납니다. -**3. lowering pass 쪽 실패 (`exp`, `cat`, reduction).** 담당(이정민)과 나눠야 할 부분입니다. reduction은 원인이 파악돼 있습니다 — `tt.reduce(axis=1)`이 `linalg.transpose` + `linalg.reduce`가 되고, 스크래치패드가 레인 뱅킹되어 있어 `bank_vectorize`가 거부합니다. `exp`와 `cat`은 아직 원인 미확인. +**3. lowering pass 쪽 실패 (`exp`, `cat`, reduction).** 담당(이정민)과 나눌 부분입니다. reduction만 원인이 파악돼 있습니다 — `tt.reduce(axis=1)`이 `linalg.transpose` + `linalg.reduce`가 되고, 스크래치패드가 레인 뱅킹되어 축소 축이 레인 안에 머물러야 하는데 그렇지 못해 `bank_vectorize`가 거부합니다. `exp`와 `cat`은 미확인. -**4. matmul 경로 진입.** 지금은 Inductor가 `aten.mm` extern 으로 빼서 Triton 경로를 아예 타지 않습니다. 통과한 것처럼 보이지만 시뮬레이터를 거치지 않은 값입니다. Triton 템플릿을 쓰게 하려면 `max_autotune` 계열 설정이 필요하고, 그래야 systolic array 경로를 볼 수 있습니다. +**4. matmul 경로 진입.** 지금은 Inductor가 `aten.mm` extern으로 빼서 Triton 경로를 타지 않습니다. 값은 맞게 나오지만 시뮬레이터를 거치지 않은 값입니다. Triton 템플릿을 쓰게 하려면 `max_autotune` 계열 설정이 필요하고, 그래야 systolic array 경로를 볼 수 있습니다. -**5. double buffering.** 커버리지가 아니라 정확도 문제 — 4절의 251 vs 650 격차. lowering pass 가 비동기 DMA + `togsim.wait`를 내야 하고, 기존 경로에 이미 있는 기계를 옮기는 일입니다. +**5. double buffering.** 커버리지가 아니라 정확도 항목 — 4절의 251 대 650 격차. lowering pass가 비동기 DMA와 `togsim.wait`를 내야 하고, 기존 경로에 이미 있는 기계를 옮기는 일입니다. ### 회귀 방지 -`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서, 컴파일에 성공하면 테스트가 실패합니다 — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. +`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서 컴파일에 성공하면 테스트가 실패합니다 — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. --- +모듈별 동작과 사용법은 [`triton_backend/README.md`](triton_backend/README.md)에, 이식 작업 본체는 [PyTorchSim#305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305)에 있습니다. + 측정 환경: torch 2.10.0+cpu / triton 3.6.0, `systolic_ws_128x128_c1_simple_noc_tpuv3.yml`, `vpu_num_lanes` 128. 기존 MLIR 경로는 `tests/ops/elementwise/test_add.py` 통과로 회귀 없음 확인. From e3672a83b2658f5fa4fe33b98b49cc3f69b8b500 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 15:48:15 +0900 Subject: [PATCH 23/35] [Docs] Say lowering pass in prose, keep tnpu only where it is a real 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. --- PyTorchSimFrontend/triton-codegen-route.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index db674daa..1e671205 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -13,7 +13,7 @@ lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim / gem5 / Spike 스택에 물린 것**이 여기서 한 일입니다. -용어: 산문에서 **PyTorchSim lowering pass**는 이 lowering 계층 전체를 가리킵니다. 코드는 `triton-npu` 저장소에 있어서 파일 경로·모듈 이름·CI 잡 이름 등 **실제 식별자는 `tnpu`를 그대로** 씁니다(`tnpu/passes/`, `tnpu.spike`, `Build tnpu toolchain image`). 문서를 따라 코드를 찾아갈 수 있어야 하기 때문입니다. +용어: 이 lowering 계층을 문서 전체에서 **PyTorchSim lowering pass**로 부릅니다. 다만 코드가 `triton-npu` 저장소에 있어서 **실제 식별자는 `tnpu`로 남아 있고**(`tnpu_bridge.py`, `tnpu/passes/`, `tnpu.spike`, `strip_for_tnpu`), 문서에서 코드를 찾아갈 수 있도록 그 이름들은 그대로 인용합니다. ## 현재 도달점 @@ -69,7 +69,7 @@ lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorc | lowering | `PyTorchSimFrontend/mlir/` | PyTorchSim lowering pass (subprocess) | | 융합 | 템플릿 + `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | | op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | -| functional | `FunctionalSimulator.run_spike` | tnpu stage 6 (`tnpu.spike`) | +| functional | `FunctionalSimulator.run_spike` | lowering pass 의 stage 6 (`tnpu.spike`) | | timing | `trace.so` + `trace_cycles.tsv` → TOGSim | **동일** | | 타일 cycle 실측 | gem5 | **동일** (`build_tog` sample 모드 공유) | | DMA | 비동기 + `togsim.wait` 배리어 | **동기만** | @@ -111,7 +111,7 @@ torch.compile TritonNPULauncher.__call__ codecache.py │ ├ functional 텐서 → runtime/*.raw → Spike → 텐서 functional.py - │ tnpu stage 6 (tnpu.spike) 재사용 + │ lowering pass 의 stage 6 (tnpu.spike) 재사용 │ └ timing 04-custom.mlir timing.py ├ build_tog sample → gem5 → 타일 cycle 실측 From 83bce51e23b26ec38d8a4d743a336c4ca92bcb0c Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 15:50:05 +0900 Subject: [PATCH 24/35] [Docs] Call the lowering layer the NPU lowering pass 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. --- PyTorchSimFrontend/triton-codegen-route.md | 12 ++++++------ PyTorchSimFrontend/triton_backend/README.md | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index 1e671205..f2af5897 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -1,6 +1,6 @@ # Triton 코드젠 경로를 PyTorchSim에 이식 -`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **PyTorchSim lowering pass**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. +`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **NPU lowering pass**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. ## 작업 경계 @@ -8,12 +8,12 @@ | 부분 | 하는 일 | 소관 | |---|---|---| -| **PyTorchSim lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 범위 밖 | +| **NPU lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 범위 밖 | | **기존 PyTorchSim으로의 이식** | 그 lowering pass를 기존 시뮬레이션 스택에 얹는 일 | 이 문서 | lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim / gem5 / Spike 스택에 물린 것**이 여기서 한 일입니다. -용어: 이 lowering 계층을 문서 전체에서 **PyTorchSim lowering pass**로 부릅니다. 다만 코드가 `triton-npu` 저장소에 있어서 **실제 식별자는 `tnpu`로 남아 있고**(`tnpu_bridge.py`, `tnpu/passes/`, `tnpu.spike`, `strip_for_tnpu`), 문서에서 코드를 찾아갈 수 있도록 그 이름들은 그대로 인용합니다. +용어: 이 lowering 계층을 문서 전체에서 **NPU lowering pass**로 부릅니다. 다만 코드가 `triton-npu` 저장소에 있어서 **실제 식별자는 `tnpu`로 남아 있고**(`tnpu_bridge.py`, `tnpu/passes/`, `tnpu.spike`, `strip_for_tnpu`), 문서에서 코드를 찾아갈 수 있도록 그 이름들은 그대로 인용합니다. ## 현재 도달점 @@ -43,7 +43,7 @@ lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorc (gemm, conv, sdpa, sort, (op별 템플릿 없음) cat, maxpool, bmm …) │ │ - PyTorchSim mlir/ 패스 PyTorchSim lowering pass + PyTorchSim mlir/ 패스 NPU lowering pass PSAL LLVM 20 (in-process) stock LLVM 23 (subprocess) └ 담당 이정민 — 범위 밖. 여기서 한 일은 이 블록을 @@ -66,7 +66,7 @@ lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorc | 커널을 만드는 주체 | op별 MLIR 템플릿 (직접 작성) | Inductor의 Triton 코드젠 | | **커널 하나의 의미** | **루프 네스트 전체** | **타일 하나** | | grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | -| lowering | `PyTorchSimFrontend/mlir/` | PyTorchSim lowering pass (subprocess) | +| lowering | `PyTorchSimFrontend/mlir/` | NPU lowering pass (subprocess) | | 융합 | 템플릿 + `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | | op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | | functional | `FunctionalSimulator.run_spike` | lowering pass 의 stage 6 (`tnpu.spike`) | @@ -98,7 +98,7 @@ torch.compile │ - 인자 역할(in/out/inout) · dtype · numel │ - grid, 사용자 스칼라 값 ▼ - PyTorchSim lowering pass (subprocess) tnpu_bridge.py + NPU lowering pass (subprocess) tnpu_bridge.py │ 담당 이정민 — 범위 밖. 여기서 한 일은 이 단계를 │ 호출하고 그 결과를 아래 스택에 물린 부분. │ diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 4952f6cb..f0692132 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -1,10 +1,10 @@ # Triton codegen route (WIP) Replaces the hand-written MLIR emission in `PyTorchSimFrontend/mlir/` with -**Inductor's own Triton codegen**, lowered to this NPU by the **PyTorchSim -lowering pass** (owned by 이정민; the code lives in the `triton-npu` repo, so -paths and module names read `tnpu`). Opt-in and off by default; the MLIR route -is untouched and stays the production path. +**Inductor's own Triton codegen**, lowered to this NPU by the **NPU lowering +pass** (owned by 이정민; the code lives in the `triton-npu` repo, so paths and +module names read `tnpu`). Opt-in and off by default; the MLIR route is +untouched and stays the production path. The modules here are the PORT: they drive that lowering pass and wire its output into the existing TOGSim / gem5 / Spike stack. The pass itself is not ours. From 6ddab175fea19655b247fb959679ceca67957068 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 15:55:27 +0900 Subject: [PATCH 25/35] [Docs] Show the grid being rebuilt as the outer loop 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. --- PyTorchSimFrontend/triton-codegen-route.md | 33 +++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index f2af5897..37885129 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -143,7 +143,38 @@ Triton 커널 = 타일 하나. grid 는 커널 밖, launch 가 쥐 - `togsim_kernel_tile(ctx, iv, n)` — work-item 하나 - `togsim_kernel(ctx, shape_args, n)` — 병렬 영역의 열거 -Triton 커널 본문은 전자에 대응하므로 **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. 그 합성이 `lower_to_emitc.WorkItem` + `_materialize_grid_loop`입니다. +Triton 커널 본문은 전자에 대응하므로 **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. + +### grid 를 outer loop 으로 되세우기 + +Triton 쪽에서 grid 는 커널이 아니라 **KernelSpec 에 붙어 있습니다.** `kernel_spec.grid_of(meta)` 가 Inductor 의 numel 과 고정한 블록 크기로부터 축별 ceil-div 를 계산해 `grid=(8,)` 같은 값을 spec 에 적고, 커널 본문은 그 중 자기 몫이 몇 번째인지를 `pidX/Y/Z` 인자로 받을 뿐입니다. + +기존 PyTorchSim 은 정반대를 기대합니다. `build_tog` 는 **역할 속성이 붙은 최상위 루프**를 TOG 의 루트로 잡습니다: + +```python +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") +roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] +``` + +즉 work-item 을 **루프에서 읽어냅니다.** 그런데 Triton 커널에는 그 루프가 없습니다 — grid 로 흩어져 있으니까요. 루프가 없으면 루트도 없고, TOG 가 비게 됩니다. + +그래서 spec 의 grid 를 다시 루프로 세웁니다. `_materialize_grid_loop` 이 하는 일입니다: + +``` +들어올 때 func @k(..., %pidX: i32, ...) <- 타일 하나. 루프 없음 + body(%pidX) + +나갈 때 scf.for %p = 0 to G { {outer_loop = true} + body(<%pidX 를 index_cast %p 로 치환>) + } +``` + +`WorkItem(parallel_args, grid)` 이 **어느 인자가 program id 인지**와 **축이 몇 개인지**를 들고 있습니다. 패스는 축마다 루프를 하나씩 중첩하고, 본문을 그 안으로 옮기고, pid 인자의 모든 사용처를 루프 유도변수로 바꿉니다. 마지막에 `outer_loop` 속성을 답니다 — **`build_tog` 가 찾는 바로 그 표식**이고, 이게 붙어야 합성한 루프가 TOG 의 루트가 됩니다. + +(축별 범위는 `WorkItem` 이 들고 있을 수도, 런타임으로 미룰 수도 있습니다. timing 경로는 후자를 씁니다 — 바로 아래.) + +결과적으로 Triton 이 grid 로 표현한 것과 PyTorchSim 이 outer loop 로 표현한 것이 같은 것을 가리키게 되고, 그 뒤 파이프라인(`build_skeleton` → `trace.so` → TOGSim)은 MLIR 경로와 한 글자도 다르지 않게 흘러갑니다. ### 동적 shape이 여기서 나옵니다 From f8f3d0adc695f8e4c21d94b609b6f79738f739f9 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 16:05:32 +0900 Subject: [PATCH 26/35] [Docs] Make model coverage the goal, op coverage the step to it 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. --- PyTorchSimFrontend/triton-codegen-route.md | 41 +++++++++++++++++++--- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index 37885129..7129128b 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -22,7 +22,7 @@ lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorc | functional | 연결됨. `x + y`, `(x+y)*2 - x` **max abs error 0.0** (1024 elements, Spike) | | timing | 연결됨. TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | | 동적 shape | 처리됨. 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32 | -| **커버리지** | **elementwise와 그 융합까지.** 그 밖은 5절 실측 참고 | +| **커버리지** | **elementwise와 그 융합까지.** 남은 일은 op 스위트 → 모델까지 넓히는 것 — 5절 | | CI | 전 잡 green | --- @@ -232,11 +232,29 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 --- -## 5. 다음 작업 — 기존 경로 커버리지까지 검증 +## 5. 다음 작업 — 모델 커버리지까지 고도화 -확인된 것은 elementwise와 그 융합뿐입니다. 다음 작업은 기능을 더 얹는 것이 아니라, **기존 MLIR 경로를 지탱하는 op 테스트 스위트를 Triton 경로로 그대로 돌려 어디까지 가는지 확인하는 것**입니다. 대상은 `tests/ops/` 아래 이미 있는 것들이고(`elementwise`, `reduce`, `gemm`, `conv`, `attention`, `view`, `sort`, `fusion`, `misc`), MLIR 경로가 통과하는 범위가 목표선입니다. +확인된 것은 elementwise와 그 융합뿐입니다. **최종 목표는 기존 MLIR 경로가 돌리는 모델들을 Triton 경로로도 돌리는 것**이고, 새 기능을 얹기보다 이미 있는 테스트를 그대로 돌려 어디서 멈추는지 고쳐 나가는 일입니다. -### 1차 실측 +목표선은 저장소에 이미 있습니다: + +``` +tests/ops/ elementwise · reduce · gemm · conv · attention + view · sort · fusion · misc <- 1단계 + +tests/models/ MLP · MobileNet · Llama · Mixtral8x7B · DeepSeek + MoE · Diffusion · Yolov5 + test_resnet · test_vit · test_transformer + test_clip · test_convnextv2 · test_swinv2 <- 2단계 +``` + +**1단계 (op).** 모델은 op 의 조합이라 op 이 막히면 모델은 첫 커널에서 멈춥니다. 그래서 op 스위트를 먼저 통과시켜야 하고, 아래 실측이 그 출발점입니다. + +**2단계 (모델).** op 이 뚫리면 모델 단위로 올라갑니다. 여기서는 op 단위에서 드러나지 않는 것들이 나옵니다 — 커널 수십~수백 개가 이어질 때의 컴파일 시간과 캐시 거동, 커널 사이 버퍼 재사용, 실제 shape 조합(op 테스트는 대개 잘 나뉘는 크기를 씁니다), 그리고 training 경로의 backward 커널. 난이도 순으로 MLP → MobileNet/ResNet → ViT/Transformer → Llama 순이 무난합니다. + +각 단계의 판정 기준은 같습니다 — **값이 torch와 일치하고, 사이클이 나오고, 경로에 실제로 진입할 것.** + +### 1단계 1차 실측 대표 케이스를 `TORCHSIM_TRITON_CODEGEN=1`로 돌린 결과입니다. **경로 진입** 열은 Triton 경로를 실제로 탔는지(작업 디렉터리 생성 여부)를 뜻합니다 — 이걸 보지 않으면 Inductor가 extern으로 뺀 것을 통과로 착각합니다. @@ -275,6 +293,21 @@ functional.py write_inputs t.contiguous()...tofile() <- 저장 순서 **5. double buffering.** 커버리지가 아니라 정확도 항목 — 4절의 251 대 650 격차. lowering pass가 비동기 DMA와 `togsim.wait`를 내야 하고, 기존 경로에 이미 있는 기계를 옮기는 일입니다. +1~4가 풀리면 op 스위트가 대체로 통과할 것으로 보입니다. 5는 값이 아니라 사이클의 정확도라 모델 단계와 병행해도 됩니다. + +### 2단계에서 새로 볼 것 + +op 단위에서는 드러나지 않다가 모델에서 처음 나오는 항목들입니다. 지금은 예상이고, 실제로 돌려봐야 확정됩니다. + +| 항목 | 왜 모델에서만 나오나 | +|---|---| +| 컴파일 시간 | 커널 하나당 lowering pass subprocess 가 한 번 뜹니다. 커널이 수백 개면 이게 지배적이 되고, 캐시(`outputs/triton_/`)가 실제로 먹는지 확인해야 합니다 | +| 커널 간 버퍼 | op 테스트는 커널 하나로 끝나서 중간 버퍼 재사용이 없습니다 | +| 실제 shape | op 테스트는 대개 잘 나뉘는 크기를 씁니다. 모델은 나눠떨어지지 않는 shape 과 마스킹을 만듭니다 | +| dtype | f16/bf16 혼합. 지금 확인된 것은 f32 뿐입니다 | +| backward | training 경로의 커널은 forward 와 형태가 다릅니다 (`tests/models/MLP` 가 forward + backward 를 함께 돌립니다) | +| 메모리 | 모델 규모에서 스크래치패드와 DRAM 사용량이 설정값을 넘는지 | + ### 회귀 방지 `tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서 컴파일에 성공하면 테스트가 실패합니다 — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. From ebbd11357723be4cd9f2478c203c590a9e758f3f Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Tue, 28 Jul 2026 16:08:11 +0900 Subject: [PATCH 27/35] [Docs] Keep the remaining work high level 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. --- PyTorchSimFrontend/triton-codegen-route.md | 90 ++++++++-------------- 1 file changed, 30 insertions(+), 60 deletions(-) diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md index 7129128b..5c009924 100644 --- a/PyTorchSimFrontend/triton-codegen-route.md +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -232,85 +232,55 @@ triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니 --- -## 5. 다음 작업 — 모델 커버리지까지 고도화 +## 5. 남은 일 — 모델 커버리지까지 -확인된 것은 elementwise와 그 융합뿐입니다. **최종 목표는 기존 MLIR 경로가 돌리는 모델들을 Triton 경로로도 돌리는 것**이고, 새 기능을 얹기보다 이미 있는 테스트를 그대로 돌려 어디서 멈추는지 고쳐 나가는 일입니다. +지금 통과하는 것은 elementwise와 그 융합입니다. **목표는 기존 MLIR 경로가 돌리는 모델들을 Triton 경로로도 돌리는 것**이고, 새 기능을 얹기보다 이미 있는 테스트를 그대로 돌려 막히는 곳을 고쳐 나가는 일입니다. -목표선은 저장소에 이미 있습니다: +목표선은 저장소에 이미 있습니다. -``` -tests/ops/ elementwise · reduce · gemm · conv · attention - view · sort · fusion · misc <- 1단계 - -tests/models/ MLP · MobileNet · Llama · Mixtral8x7B · DeepSeek - MoE · Diffusion · Yolov5 - test_resnet · test_vit · test_transformer - test_clip · test_convnextv2 · test_swinv2 <- 2단계 -``` - -**1단계 (op).** 모델은 op 의 조합이라 op 이 막히면 모델은 첫 커널에서 멈춥니다. 그래서 op 스위트를 먼저 통과시켜야 하고, 아래 실측이 그 출발점입니다. - -**2단계 (모델).** op 이 뚫리면 모델 단위로 올라갑니다. 여기서는 op 단위에서 드러나지 않는 것들이 나옵니다 — 커널 수십~수백 개가 이어질 때의 컴파일 시간과 캐시 거동, 커널 사이 버퍼 재사용, 실제 shape 조합(op 테스트는 대개 잘 나뉘는 크기를 씁니다), 그리고 training 경로의 backward 커널. 난이도 순으로 MLP → MobileNet/ResNet → ViT/Transformer → Llama 순이 무난합니다. +| 단계 | 대상 | 지금 | +|---|---|---| +| 1. op | `tests/ops/` — elementwise, reduce, gemm, conv, attention, view, sort, fusion, misc | elementwise만 | +| 2. 모델 | `tests/models/` — MLP, MobileNet, ResNet, ViT, Transformer, Llama, Mixtral, DeepSeek, MoE, Diffusion, Yolov5 | 미착수 | -각 단계의 판정 기준은 같습니다 — **값이 torch와 일치하고, 사이클이 나오고, 경로에 실제로 진입할 것.** +모델은 op의 조합이라 op 하나가 막히면 모델은 첫 커널에서 멈춥니다. 그래서 op 먼저이고, 모델은 난이도 순으로 MLP → MobileNet/ResNet → ViT/Transformer → Llama가 무난합니다. 판정 기준은 두 단계가 같습니다 — **값이 torch와 일치하고, 사이클이 나오고, 경로에 실제로 진입할 것.** -### 1단계 1차 실측 +### 1단계에서 막히는 지점 -대표 케이스를 `TORCHSIM_TRITON_CODEGEN=1`로 돌린 결과입니다. **경로 진입** 열은 Triton 경로를 실제로 탔는지(작업 디렉터리 생성 여부)를 뜻합니다 — 이걸 보지 않으면 Inductor가 extern으로 뺀 것을 통과로 착각합니다. +대표 op를 돌려 확인한 것입니다. **경로 진입** 열이 필요한 이유는, Inductor가 일부 연산을 자체 커널 대신 외부 구현으로 빼기 때문입니다 — 그 경우 값은 맞지만 시뮬레이터를 거치지 않습니다. | 케이스 | 경로 진입 | 결과 | |---|---|---| -| `x + y` | 예 | 값 일치 | -| `(x+y)*2 - x` (융합) | 예 | 값 일치 | -| `x.t() + 1` | 예 | **값 틀림 — 4030/4096** | -| `relu` | 예 | 중단: `SpecIncomplete: triton_helpers.maximum` | -| `softmax` | 예 | 중단: `SpecIncomplete: triton_helpers.max2` | -| `exp` | 예 | 중단: lowering pass 실패 | -| `sum(dim=1)` | 예 | 중단: lowering pass 실패 (`bank_vectorize`) | -| `cat` | 예 | 중단: lowering pass 실패 | -| `a @ b` | **아니오** | Inductor가 `aten.mm` extern으로 처리 — 경로에 도달하지 않음 | - -### 우선순위 - -**1. non-contiguous 텐서.** 값이 조용히 틀리는 유일한 항목이라 최우선입니다. `x.t() + 1`의 출력이 정확히 `x + 1`이고 — transpose가 통째로 무시되고 — 아무 예외도 나지 않습니다. 원인은 lowering pass가 아니라 **이식 쪽**입니다: +| `x + y`, `(x+y)*2 - x` | 예 | 값 일치 | +| `x.t() + 1` | 예 | **값 틀림** | +| `relu`, `softmax` | 예 | 중단 — 헬퍼 모듈 부재 | +| `exp`, `cat`, `sum(dim=1)` | 예 | 중단 — NPU lowering pass | +| `a @ b` | **아니오** | 외부 구현으로 처리됨 | -``` -Inductor 가 낸 커널 tmp0 = tl.load(in_ptr0 + x0) <- 인덱스가 항등. - transpose 를 인덱스 식이 아니라 - 출력 버퍼의 stride (1,64) 로 접었음 - -functional.py write_inputs t.contiguous()...tofile() <- 저장 순서를 재배열 - read_outputs t.copy_(flat.view_as(t)) <- 논리 순서로 되씀 -``` - -둘 다 **논리 순서와 저장 순서가 같다**고 가정합니다. contiguous 텐서에서만 참이고, 지금까지 통과한 것이 정확히 그 집합입니다. 저장 순서 기준으로 읽고 쓰도록 고치고 비-contiguous 케이스를 테스트에 넣어야 합니다. +### 할 일 -**2. `triton_helpers` 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등 상당수가 여기서 막힙니다. 모듈이 torch 안에 있고 lowering pass 쪽 venv에는 없습니다. `strip_for_tnpu`가 어떤 헬퍼인지 이름을 대고 멈추므로, 필요한 것만 최소로 벤더링하면 커버리지가 한 번에 크게 늘어납니다. +1. **비연속 텐서 처리.** `x.t() + 1`이 조용히 틀린 값을 냅니다. 값이 틀리면서 아무 신호도 없는 유일한 항목이라 최우선입니다. 원인은 파악됐고 이식 쪽입니다. +2. **헬퍼 모듈 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등이 torch 안의 헬퍼를 참조하는데 lowering pass 쪽 환경에는 없습니다. 필요한 것만 옮기면 커버리지가 한 번에 크게 늘어납니다. +3. **NPU lowering pass 쪽 실패** (`exp`, `cat`, reduction). 담당(이정민)과 나눌 부분입니다. +4. **matmul을 경로 안으로.** 지금은 외부 구현으로 빠져 시뮬레이터를 거치지 않습니다. 이게 뚫려야 systolic array 경로를 볼 수 있습니다. +5. **DMA 겹침.** 값이 아니라 사이클 정확도 항목입니다(4절의 251 대 650). 기존 경로에 이미 있는 기계를 옮기는 일이라 모델 단계와 병행 가능합니다. -**3. lowering pass 쪽 실패 (`exp`, `cat`, reduction).** 담당(이정민)과 나눌 부분입니다. reduction만 원인이 파악돼 있습니다 — `tt.reduce(axis=1)`이 `linalg.transpose` + `linalg.reduce`가 되고, 스크래치패드가 레인 뱅킹되어 축소 축이 레인 안에 머물러야 하는데 그렇지 못해 `bank_vectorize`가 거부합니다. `exp`와 `cat`은 미확인. - -**4. matmul 경로 진입.** 지금은 Inductor가 `aten.mm` extern으로 빼서 Triton 경로를 타지 않습니다. 값은 맞게 나오지만 시뮬레이터를 거치지 않은 값입니다. Triton 템플릿을 쓰게 하려면 `max_autotune` 계열 설정이 필요하고, 그래야 systolic array 경로를 볼 수 있습니다. - -**5. double buffering.** 커버리지가 아니라 정확도 항목 — 4절의 251 대 650 격차. lowering pass가 비동기 DMA와 `togsim.wait`를 내야 하고, 기존 경로에 이미 있는 기계를 옮기는 일입니다. - -1~4가 풀리면 op 스위트가 대체로 통과할 것으로 보입니다. 5는 값이 아니라 사이클의 정확도라 모델 단계와 병행해도 됩니다. +1~4가 풀리면 op 스위트는 대체로 통과할 것으로 봅니다. ### 2단계에서 새로 볼 것 -op 단위에서는 드러나지 않다가 모델에서 처음 나오는 항목들입니다. 지금은 예상이고, 실제로 돌려봐야 확정됩니다. +op 단위에서는 드러나지 않다가 모델에서 처음 나오는 것들입니다. **아직 돌려보지 않았으므로 측정이 아니라 예상입니다.** -| 항목 | 왜 모델에서만 나오나 | -|---|---| -| 컴파일 시간 | 커널 하나당 lowering pass subprocess 가 한 번 뜹니다. 커널이 수백 개면 이게 지배적이 되고, 캐시(`outputs/triton_/`)가 실제로 먹는지 확인해야 합니다 | -| 커널 간 버퍼 | op 테스트는 커널 하나로 끝나서 중간 버퍼 재사용이 없습니다 | -| 실제 shape | op 테스트는 대개 잘 나뉘는 크기를 씁니다. 모델은 나눠떨어지지 않는 shape 과 마스킹을 만듭니다 | -| dtype | f16/bf16 혼합. 지금 확인된 것은 f32 뿐입니다 | -| backward | training 경로의 커널은 forward 와 형태가 다릅니다 (`tests/models/MLP` 가 forward + backward 를 함께 돌립니다) | -| 메모리 | 모델 규모에서 스크래치패드와 DRAM 사용량이 설정값을 넘는지 | +- **컴파일 시간** — 커널이 수백 개가 될 때 캐시가 실제로 먹는지 +- **커널 사이 버퍼 재사용** — op 테스트는 커널 하나로 끝나 드러나지 않음 +- **실제 shape** — op 테스트는 대개 잘 나뉘는 크기를 씀 +- **f16 / bf16** — 지금 확인된 것은 f32뿐 +- **backward 커널** — training 경로는 forward와 형태가 다름 +- **메모리 사용량** — 모델 규모에서 설정값을 넘는지 ### 회귀 방지 -`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서 컴파일에 성공하면 테스트가 실패합니다 — 레인 경로가 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. +`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서, 컴파일에 성공하면 테스트가 실패합니다 — 지원이 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. --- From 31baf29cd3c71f2cb605c3f421dbfca39b0492f3 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 17:08:30 +0900 Subject: [PATCH 28/35] [CI] Bump the triton-npu pin to 4835b38, on the released spike 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. --- .github/workflows/triton_npu.yml | 6 ++++++ PyTorchSimFrontend/triton_backend/README.md | 5 +++++ thirdparty/triton-npu.json | 2 +- 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index d38c6f19..56b4095c 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -170,9 +170,13 @@ jobs: # The harness's own kernels, end to end through Spike. This is the gate on # the toolchain itself: if these regress, nothing downstream is meaningful. # gemm/bmm need TNPU_VCIX_MATMUL=1 to reach the systolic array. + # TNPU_SPIKE/ISA: the released spike has no zvfp8, and an unknown extension + # stops spike at startup. Drop once riscv-isa-sim#7 is in the release. - name: doctor + add / mul / relu / gemm / bmm run: | docker run --rm -e TNPU_VCIX_MATMUL=1 \ + -e TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike \ + -e TNPU_SPIKE_ISA=rv64gcv_zfh \ ${{ needs.build-app.outputs.app_image }} bash -lc ' cd /workspace/triton-npu && python3 run.py doctor && @@ -197,6 +201,8 @@ jobs: - name: test_triton_codegen.py run: | docker run --rm -e TORCHSIM_TRITON_CODEGEN=1 \ + -e TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike \ + -e TNPU_SPIKE_ISA=rv64gcv_zfh \ ${{ needs.build-app.outputs.app_image }} \ python3 PyTorchSim/tests/system/test_triton_codegen.py diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index f0692132..e919cba0 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -178,3 +178,8 @@ token is scoped to this repository. `preflight` checks it before the build. `Dockerfile.tnpu` clones the harness and runs its own `setup/restore.sh --prebuilt`; the pins all live in that repo's `setup/versions.env`. `ref` in the manifest is a commit, so an upstream change there moves this image's tag too. + +The jobs set `TNPU_SPIKE` and `TNPU_SPIKE_ISA=rv64gcv_zfh`: tnpu asks for +`zvfp8`, which the released spike lacks, and an unknown extension stops spike at +startup. Costs only `ops_fp8_roundtrip.py`, which CI does not run. Drop once +`PSAL-POSTECH/riscv-isa-sim#7` is in the release. diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 5c01580c..9ece4f64 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "22df065b46b1a61f75a9847ac73207de8cc72482", + "ref": "4835b38cd6bc5122359486b51eaab753c13b2653", "release_tag": "toolchain-llvm23" } } From d37d4719ee61554e375d094dbe94e472c1d98812 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 20:35:53 +0900 Subject: [PATCH 29/35] [CI] Bump the triton-npu pin to 5d84caf 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. --- thirdparty/triton-npu.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 9ece4f64..4bac8dfb 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -2,7 +2,7 @@ "description": "Pin for the triton-npu toolchain layer (Triton codegen route only). Kept out of github-releases.json so the main torchsim_base image is not rebuilt when this moves. CI builds ghcr.io/.../torchsim_tnpu_base:tnpu-<12 hex> when missing, pin = sha256 of this file plus Dockerfile.tnpu. `ref` is a commit, not a branch, so any upstream change -- including the toolchain release its setup/versions.env points at -- moves the pin. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "4835b38cd6bc5122359486b51eaab753c13b2653", + "ref": "5d84cafdd5292f77548a84a6c44a4d2d15c4dd4f", "release_tag": "toolchain-llvm23" } } From 65ada99d86506d83068c38c1588ee9de36a571bc Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 20:41:45 +0900 Subject: [PATCH 30/35] [CI] Set the spike override in the image, not in each job 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. --- .github/workflows/triton_npu.yml | 6 ------ Dockerfile.tnpu | 5 +++++ PyTorchSimFrontend/triton_backend/README.md | 7 ++++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index 56b4095c..d38c6f19 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -170,13 +170,9 @@ jobs: # The harness's own kernels, end to end through Spike. This is the gate on # the toolchain itself: if these regress, nothing downstream is meaningful. # gemm/bmm need TNPU_VCIX_MATMUL=1 to reach the systolic array. - # TNPU_SPIKE/ISA: the released spike has no zvfp8, and an unknown extension - # stops spike at startup. Drop once riscv-isa-sim#7 is in the release. - name: doctor + add / mul / relu / gemm / bmm run: | docker run --rm -e TNPU_VCIX_MATMUL=1 \ - -e TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike \ - -e TNPU_SPIKE_ISA=rv64gcv_zfh \ ${{ needs.build-app.outputs.app_image }} bash -lc ' cd /workspace/triton-npu && python3 run.py doctor && @@ -201,8 +197,6 @@ jobs: - name: test_triton_codegen.py run: | docker run --rm -e TORCHSIM_TRITON_CODEGEN=1 \ - -e TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike \ - -e TNPU_SPIKE_ISA=rv64gcv_zfh \ ${{ needs.build-app.outputs.app_image }} \ python3 PyTorchSim/tests/system/test_triton_codegen.py diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu index c37f633a..b9d939a7 100644 --- a/Dockerfile.tnpu +++ b/Dockerfile.tnpu @@ -32,6 +32,11 @@ RUN --mount=type=secret,id=tnpu_token \ /workspace/triton-npu/setup/restore.sh --prebuilt ENV TNPU_DIR=/workspace/triton-npu +# tnpu defaults to a separate fp8 spike and asks for zvfp8; the released spike +# has neither, and an unknown extension stops spike at startup. Drop once +# PSAL-POSTECH/riscv-isa-sim#7 is in the release. +ENV TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike +ENV TNPU_SPIKE_ISA=rv64gcv_zfh # Fail the build, not the first CI job. RUN python3 /workspace/triton-npu/run.py doctor diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index e919cba0..35ca3ed5 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -179,7 +179,8 @@ token is scoped to this repository. `preflight` checks it before the build. --prebuilt`; the pins all live in that repo's `setup/versions.env`. `ref` in the manifest is a commit, so an upstream change there moves this image's tag too. -The jobs set `TNPU_SPIKE` and `TNPU_SPIKE_ISA=rv64gcv_zfh`: tnpu asks for -`zvfp8`, which the released spike lacks, and an unknown extension stops spike at -startup. Costs only `ops_fp8_roundtrip.py`, which CI does not run. Drop once +`Dockerfile.tnpu` sets `TNPU_SPIKE` and `TNPU_SPIKE_ISA=rv64gcv_zfh`: tnpu asks +for `zvfp8`, which the released spike lacks, and an unknown extension stops +spike at startup — including the doctor run inside the image build. Costs only +`ops_fp8_roundtrip.py`, which CI does not run. Drop once `PSAL-POSTECH/riscv-isa-sim#7` is in the release. From f8e744d27afa5671440711918630d8cd6b41fa13 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 21:44:32 +0900 Subject: [PATCH 31/35] [CI] Run the whole test suite on the Triton route, and report what fails 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. --- .github/workflows/triton_npu.yml | 54 ++- PyTorchSimFrontend/triton_backend/README.md | 34 ++ .../triton_backend/codecache.py | 7 +- scripts/ci/triton_route_passing.txt | 16 + scripts/ci/triton_route_sweep.py | 331 ++++++++++++++++++ 5 files changed, 434 insertions(+), 8 deletions(-) create mode 100644 scripts/ci/triton_route_passing.txt create mode 100755 scripts/ci/triton_route_sweep.py diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index d38c6f19..41f60ab3 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -32,7 +32,7 @@ env: jobs: preflight: name: Check tnpu access - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, x86_64] outputs: ready: ${{ steps.check.outputs.ready }} steps: @@ -67,7 +67,7 @@ jobs: ensure-tnpu-base: name: Build tnpu toolchain image needs: preflight - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, big, x86_64] outputs: tnpu_image: ${{ steps.pin.outputs.tnpu_image }} permissions: @@ -123,7 +123,7 @@ jobs: build-app: name: Build app image on tnpu base needs: ensure-tnpu-base - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, big, x86_64] outputs: app_image: ${{ steps.name.outputs.app_image }} permissions: @@ -159,7 +159,7 @@ jobs: tnpu-baselines: name: triton-npu baselines needs: build-app - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, x86_64] steps: - uses: docker/login-action@v3 with: @@ -183,7 +183,7 @@ jobs: triton-route: name: Inductor Triton route needs: build-app - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, x86_64] # WIP: the launch is deliberately unimplemented, so this reports how far the # route gets rather than gating. Drop this once the launch lands. continue-on-error: true @@ -200,10 +200,52 @@ jobs: ${{ needs.build-app.outputs.app_image }} \ python3 PyTorchSim/tests/system/test_triton_codegen.py + triton-route-suite: + name: Test suite on the Triton route + needs: build-app + runs-on: [self-hosted, slurm, big, x86_64] + steps: + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + # Gates on scripts/ci/triton_route_passing.txt: the tests that pass on this + # route today. Coverage grows by regenerating that file, so it cannot + # silently shrink. + - name: Allowlisted tests + run: | + docker run --rm \ + ${{ needs.build-app.outputs.app_image }} \ + python3 PyTorchSim/scripts/ci/triton_route_sweep.py + + # Everything else, for the count and the bucket breakdown. Reports only: + # most of the suite is blocked on the gaps in triton_backend/README.md. + # Each failure leaves its Inductor kernel and last stage IR in the + # artifact, so a report needs no rerun. + - name: Full sweep (report) + continue-on-error: true + run: | + mkdir -p sweep && chmod 777 sweep + docker run --rm -v "$PWD/sweep:/sweep" \ + ${{ needs.build-app.outputs.app_image }} \ + python3 PyTorchSim/scripts/ci/triton_route_sweep.py --all \ + --timeout 900 --json /sweep/results.json \ + --markdown /sweep/coverage.md --artifacts /sweep/failures + cat sweep/coverage.md >> "$GITHUB_STEP_SUMMARY" + + - uses: actions/upload-artifact@v4 + if: always() + with: + name: triton-route-coverage + path: sweep/ + if-no-files-found: warn + mlir-route-regression: name: MLIR route still passes needs: build-app - runs-on: ubuntu-latest + runs-on: [self-hosted, slurm, x86_64] steps: - uses: docker/login-action@v3 with: diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 35ca3ed5..8477ea4f 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -151,6 +151,30 @@ induction variable, and everything downstream is unchanged. It runs before `_rewrite_signature`, which erases the kernel arguments and first asserts none are still used -- that ordering is what decides where this can live. +## Running the whole suite on this route + +`TORCHSIM_TRITON_CODEGEN` is read once, at device registration, so every test +under `tests/` is already a test of this route — no test file knows which one it +is on. `scripts/ci/triton_route_sweep.py` runs them that way: + +```bash +python scripts/ci/triton_route_sweep.py # allowlist, gating +python scripts/ci/triton_route_sweep.py --all \ + --markdown coverage.md --artifacts failures # measure + report +``` + +`scripts/ci/triton_route_passing.txt` is the gate: the tests that pass today. +Coverage grows by regenerating it (`--update-allowlist`), so it cannot silently +shrink. A test that passes **without emitting a kernel** — CPU-only, eager +fallback, or an op Inductor sends to an extern call — is deliberately kept out +of it, since it would gate nothing. + +Each failure leaves a directory under `--artifacts`: the Inductor Triton kernel +that was rejected, whatever stage IR it reached (`01-ttir` … `04-custom`), +`stage.log`, and the error. That is the whole bug report for whoever owns the +pass, without a rerun. The bucket names the owning layer, and the stage says how +far it got, so the two together route it. + ## CI `.github/workflows/triton_npu.yml`, separate from the main CI: this route is WIP, @@ -162,9 +186,19 @@ ensure-tnpu-base torchsim_base + tnpu toolchain -> torchsim_tnpu_base: %s", kernel_name, write_path) diff --git a/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt new file mode 100644 index 00000000..38e81f02 --- /dev/null +++ b/scripts/ci/triton_route_passing.txt @@ -0,0 +1,16 @@ +# Tests that pass THROUGH the Triton codegen route. +# Gated by scripts/ci/triton_route_sweep.py; regenerate with +# python scripts/ci/triton_route_sweep.py --all --update-allowlist +# A test that passes without emitting a kernel (CPU-only, eager +# fallback) is deliberately absent -- it would not gate anything. +tests/ops/elementwise/test_add.py +tests/ops/fusion/test_addmm_residual.py +tests/ops/fusion/test_matmul_scalar.py +tests/ops/fusion/test_matmul_vector.py +tests/ops/fusion/test_prologue_fusion.py +tests/ops/misc/test_expert_mask.py +tests/ops/reduce/test_batchnorm.py +tests/ops/view/test_view3D_2D.py +tests/system/test_eager.py +tests/system/test_stonne.py +tests/system/test_triton_codegen.py diff --git a/scripts/ci/triton_route_sweep.py b/scripts/ci/triton_route_sweep.py new file mode 100755 index 00000000..8079a43c --- /dev/null +++ b/scripts/ci/triton_route_sweep.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +"""Run the existing test suite through the Triton codegen route. + +The route is selected at device-registration time by TORCHSIM_TRITON_CODEGEN +(PyTorchSimDevice/torch_openreg/__init__.py), so the tests themselves need no +change -- the same file is the MLIR route's test with the variable unset and the +Triton route's with it set. + +Three things come out of a run: + + * a GATE. `triton_route_passing.txt` lists the tests that pass today. Any of + them failing is a regression and exits non-zero. That is what makes coverage + grow monotonically instead of drifting. + * a REPORT bucketed by how far each failure got, so the gap list in + triton_backend/README.md is a measurement rather than a guess. + * ARTIFACTS per failure, under --artifacts: the Inductor Triton kernel that + did not survive, the last tnpu stage IR it produced, and the error. That is + what makes a failure reportable to whoever owns the pass, without a rerun. + + python scripts/ci/triton_route_sweep.py # the allowlist, gating + python scripts/ci/triton_route_sweep.py --all # every test, reports + python scripts/ci/triton_route_sweep.py --all --artifacts triton-failures +""" + +import argparse +import glob +import json +import os +import re +import shutil +import subprocess +import sys +import time + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +HERE = os.path.dirname(os.path.abspath(__file__)) +PASSING = os.path.join(HERE, "triton_route_passing.txt") + +#: How far the kernel got, outermost first. The stage a failure did NOT reach is +#: the one that owns it, so this doubles as the routing table for a bug report. +STAGES = [ + ("01-ttir.mlir", "1 triton -> ttir"), + ("02-ttshared.mlir", "2 ttir -> tts/linalg (triton-shared)"), + ("03-adapted.mlir", "3 tnpu adapt"), + ("04-custom.mlir", "4 tnpu lower (DMA, lanes, spad)"), + ("trace.so", "5 trace producer"), +] + +#: Failure buckets, first match wins. Each names the layer that owns the fix, +#: so the counts say which gap to close next rather than just how many failed. +BUCKETS = [ + ("missing_dep", r"ModuleNotFoundError|No module named"), + ("device_op", r"\w+_overrideable not implemented|not implemented\. .*privateuse"), + ("triton_helpers", r"triton_helpers"), + ("wrapper_gap", r"'TritonNPUWrapperCodegen' object has no attribute"), + ("spec_incomplete", r"SpecIncomplete"), + ("tnpu_stage", r"TnpuError|tnpu pipeline failed|triton-shared-opt|" + r"\[stage\d\]|failed to legalize"), + ("reduction", r"lane-aware|linalg\.reduce|no reduction path"), + ("dynamic_shape", r"ShapeMismatch|dynamic shape|size_hint returned None"), + ("matmul_timing", r"vcix\.iv|sf\.vc\.|no compute node"), + ("togsim", r"TOGSim|trace\.so|SIGSEGV|Signals\.SIG|'vpu_num_lanes'"), + ("wrong_values", r"VALUES WRONG|allclose|Test Failed"), + ("timeout", r"^__timeout__$"), +] + +#: Lines torch prints alongside an error that are not the error. +NOISE = re.compile( + r"TORCHDYNAMO_VERBOSE|torch\._dynamo|You can suppress this|set TORCH_LOGS|" + r"^During handling|^The above exception|for more information|^\s*\^+\s*$") + + +def discover(): + out = [] + for dirpath, _dirs, files in os.walk(os.path.join(ROOT, "tests")): + for f in files: + if f.startswith("test_") and f.endswith(".py"): + out.append(os.path.relpath(os.path.join(dirpath, f), ROOT)) + return sorted(out) + + +def load_allowlist(): + if not os.path.isfile(PASSING): + return [] + with open(PASSING) as f: + return [l.strip() for l in f + if l.strip() and not l.lstrip().startswith("#")] + + +def classify(output, timed_out): + if timed_out: + return "timeout" + for name, pat in BUCKETS: + if re.search(pat, output, re.I | re.M): + return name + return "other" + + +def first_error(output): + """The exception line, skipping torch's boilerplate around it.""" + lines = [l.strip() for l in output.strip().splitlines() if l.strip()] + for l in reversed(lines): + if NOISE.search(l): + continue + if re.match(r"^\w*(Error|Exception|Failure)\b.*:", l) or \ + re.match(r"^(assert|AssertionError)", l): + return l[:200] + for l in reversed(lines): + if not NOISE.search(l): + return l[:200] + return "" + + +def reached_stage(dump_dir): + """(label, workdir) of the furthest tnpu stage any kernel produced. + + A workdir with only kernel.py is still the one to collect: the route got as + far as generating a Triton kernel and rejected it before stage 1. + """ + best, best_dir, fallback = None, None, None + for wd in glob.glob(os.path.join(dump_dir, "triton_*")): + if os.path.isfile(os.path.join(wd, "kernel.py")): + fallback = wd + for i, (fname, label) in enumerate(STAGES): + if os.path.isfile(os.path.join(wd, fname)): + if best is None or i > best[0]: + best, best_dir = (i, label), wd + if best: + return best[1], best_dir + return ("0 kernel generated, not accepted" if fallback + else "0 nothing emitted"), fallback + + +def collect(test, dump_dir, out_root, output, bucket, stage, workdir): + """One directory per failing test: the kernel, the last IR, the error.""" + dest = os.path.join(out_root, test.replace("/", "_").removesuffix(".py")) + os.makedirs(dest, exist_ok=True) + with open(os.path.join(dest, "error.txt"), "w") as f: + f.write(f"test: {test}\nbucket: {bucket}\nstage: {stage}\n\n") + f.write("\n".join(output.strip().splitlines()[-60:])) + if workdir: + # The Inductor kernel is the thing to hand to whoever owns the pass; + # the stage IRs say where it stopped being representable. + for name in ("kernel.py", "stage.log", *(s[0] for s in STAGES[:-1])): + src = os.path.join(workdir, name) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(dest, name)) + return dest + + +def run_one(test, timeout, artifacts, scratch): + # A private dump dir per test: artifacts must be attributable, and a shared + # one lets a cached kernel from an earlier test answer for this one. + dump = os.path.join(scratch, test.replace("/", "_").removesuffix(".py")) + shutil.rmtree(dump, ignore_errors=True) + os.makedirs(dump, exist_ok=True) + env = dict(os.environ, TORCHSIM_TRITON_CODEGEN="1", TORCHSIM_DUMP_PATH=dump) + + t0, timed_out = time.time(), False + try: + p = subprocess.run([sys.executable, test], cwd=ROOT, env=env, + capture_output=True, text=True, timeout=timeout) + out, code = p.stdout + p.stderr, p.returncode + except subprocess.TimeoutExpired as e: + pre = (e.stdout or "") if isinstance(e.stdout, str) else "" + out, code, timed_out = pre + "\n__timeout__", 124, True + + ok = code == 0 + stage, workdir = reached_stage(dump) + r = {"test": test, "ok": ok, "returncode": code, + "seconds": round(time.time() - t0, 1), + "bucket": None if ok else classify(out, timed_out), + "stage": stage, + # A pass that emitted no kernel never used the route: a CPU-only test, + # an eager fallback, or a path that bypasses Inductor. Counting those as + # coverage would overstate it. + "exercised": workdir is not None, + "error": "" if ok else first_error(out)} + if not ok and artifacts: + r["artifacts"] = os.path.relpath( + collect(test, dump, artifacts, out, r["bucket"], stage, workdir), ROOT) + shutil.rmtree(dump, ignore_errors=True) + return r + + +def write_markdown(results, path): + """The report a human reads: counts by cause, then every failure.""" + passed = [r for r in results if r["ok"]] + failed = [r for r in results if not r["ok"]] + by = {} + for r in failed: + by.setdefault(r["bucket"], []).append(r) + + real = [r for r in passed if r["exercised"]] + L = ["# Triton route coverage", "", + f"**{len(real)}/{len(results)} pass through the Triton route.** " + f"({len(passed)-len(real)} more pass without exercising it -- CPU-only, " + f"eager fallback, or a path that bypasses Inductor.)", ""] + if failed: + L += ["| cause | count | owner |", "|---|---|---|"] + OWNER = { + "device_op": "PyTorchSimDevice -- op not registered for npu", + "triton_helpers": "triton_backend -- needs a vendored copy", + "wrapper_gap": "triton_backend -- TritonNPUWrapperCodegen incomplete", + "spec_incomplete": "triton_backend -- kernel_spec cannot describe it", + "tnpu_stage": "tnpu lowering passes", + "reduction": "tnpu -- no lane-aware reduction", + "dynamic_shape": "triton_backend -- shape-specialised launch", + "matmul_timing": "build_tog -- compute node lookup", + "togsim": "TOGSim / trace producer", + "wrong_values": "numerics -- investigate", + "missing_dep": "test environment (present in the CI image)", + "timeout": "too slow, or hung", + "other": "unclassified", + } + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + L.append(f"| {b} | {len(rs)} | {OWNER.get(b, '')} |") + L += ["", "## Failures", ""] + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + L.append(f"### {b} ({len(rs)})") + L.append("") + for r in sorted(rs, key=lambda r: r["test"]): + L.append(f"- `{r['test']}` — reached **{r['stage']}**") + if r["error"]: + L.append(f" - `{r['error'][:160]}`") + if r.get("artifacts"): + L.append(f" - artifacts: `{r['artifacts']}`") + L.append("") + if real: + L += ["## Passing through the route", ""] + L += [f"- `{r['test']}`" for r in real] + [""] + other = [r for r in passed if not r["exercised"]] + if other: + L += ["## Passing without exercising the route", ""] + L += [f"- `{r['test']}`" for r in other] + [""] + with open(path, "w") as f: + f.write("\n".join(L)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--all", action="store_true", + help="run every test, not just the passing allowlist") + ap.add_argument("--timeout", type=int, default=1800) + ap.add_argument("--json", help="write the full result list here") + ap.add_argument("--artifacts", metavar="DIR", + help="per-failure kernel + stage IR + error, for reporting") + ap.add_argument("--markdown", help="write the human-readable report here") + ap.add_argument("--update-allowlist", action="store_true", + help="rewrite the allowlist from what passed (use with --all)") + args = ap.parse_args() + + allow = load_allowlist() + tests = discover() if args.all else allow + if not tests: + print("no tests selected; the allowlist is empty and --all was not given") + return 1 + + scratch = os.path.join(ROOT, ".triton_sweep") + shutil.rmtree(scratch, ignore_errors=True) + os.makedirs(scratch, exist_ok=True) + if args.artifacts: + shutil.rmtree(args.artifacts, ignore_errors=True) + os.makedirs(args.artifacts, exist_ok=True) + + print(f"Triton route sweep: {len(tests)} tests" + f"{'' if args.all else ' (allowlist)'}\n") + results = [] + for i, t in enumerate(tests, 1): + r = run_one(t, args.timeout, args.artifacts, scratch) + results.append(r) + mark = ("ok " if r["exercised"] else "ok- ") if r["ok"] else "FAIL" + extra = ("" if r["exercised"] else " (route not exercised)") if r["ok"] \ + else f" [{r['bucket']}] @{r['stage']} {r['error'][:70]}" + print(f" {i:3d}/{len(tests)} {mark} {r['seconds']:7.1f}s " + f"{r['test']}{extra}", flush=True) + shutil.rmtree(scratch, ignore_errors=True) + + passed = [r for r in results if r["ok"]] + failed = [r for r in results if not r["ok"]] + real = [r for r in passed if r["exercised"]] + + print(f"\n{'='*72}\npassed {len(passed)}/{len(results)}" + f" ({len(real)} through the Triton route, " + f"{len(passed)-len(real)} without exercising it)") + if failed: + by = {} + for r in failed: + by.setdefault(r["bucket"], []).append(r) + print("\nfailures by cause:") + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + print(f" {b:16s} {len(rs):3d}") + print("\nhow far they got:") + st = {} + for r in failed: + st[r["stage"]] = st.get(r["stage"], 0) + 1 + for s, n in sorted(st.items()): + print(f" {s:40s} {n:3d}") + + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=2) + print(f"\nwrote {args.json}") + if args.markdown: + write_markdown(results, args.markdown) + print(f"wrote {args.markdown}") + if args.artifacts and failed: + print(f"wrote {args.artifacts}/ ({len(failed)} failure dirs)") + + if args.update_allowlist: + with open(PASSING, "w") as f: + f.write("# Tests that pass through the Triton codegen route.\n" + "# Gated by scripts/ci/triton_route_sweep.py; regenerate with\n" + "# python scripts/ci/triton_route_sweep.py --all " + "--update-allowlist\n") + for r in real: + f.write(r["test"] + "\n") + print(f"wrote {PASSING} ({len(real)} tests)") + return 0 + + regressed = [r for r in failed if r["test"] in allow] + if regressed: + print(f"\nREGRESSION: {len(regressed)} allowlisted test(s) failed") + for r in regressed: + print(f" {r['test']} [{r['bucket']}] {r['error']}") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From d356e3522b415a2267da87fe670e628ae3509720 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 22:22:16 +0900 Subject: [PATCH 32/35] [Frontend] Say why a kernel was rejected, and run the sweep in parallel 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. --- .github/workflows/triton_npu.yml | 9 +-- .../triton_backend/codecache.py | 5 +- .../triton_backend/kernel_spec.py | 14 ++++ .../triton_backend/tnpu_bridge.py | 27 ++++++- scripts/ci/triton_route_passing.txt | 3 +- scripts/ci/triton_route_sweep.py | 75 ++++++++++--------- 6 files changed, 86 insertions(+), 47 deletions(-) diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index 41f60ab3..05c6f893 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -211,19 +211,14 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - # Gates on scripts/ci/triton_route_passing.txt: the tests that pass on this - # route today. Coverage grows by regenerating that file, so it cannot - # silently shrink. + # Gates on the tests that pass today; coverage cannot silently shrink. - name: Allowlisted tests run: | docker run --rm \ ${{ needs.build-app.outputs.app_image }} \ python3 PyTorchSim/scripts/ci/triton_route_sweep.py - # Everything else, for the count and the bucket breakdown. Reports only: - # most of the suite is blocked on the gaps in triton_backend/README.md. - # Each failure leaves its Inductor kernel and last stage IR in the - # artifact, so a report needs no rerun. + # Reports the rest. Each failure leaves its kernel and stage IR behind. - name: Full sweep (report) continue-on-error: true run: | diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index c94cd124..d62de72d 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -78,9 +78,8 @@ def triton_npu_compile(src_code, meta, kernel_name): spec_path = os.path.join(write_path, f"{kernel_name}_spec.py") elf = os.path.join(write_path, f"05-{kernel_name}.elf") if not os.path.isfile(elf): - # Before anything that can raise: a kernel this route cannot compile - # is the one whose source is worth having, and write_spec_file - # rejects several (triton_helpers, SpecIncomplete). + # Before write_spec_file, which rejects exactly the kernels whose + # source is worth keeping. with open(os.path.join(write_path, "kernel.py"), "w") as f: f.write(src_code) # the unmodified Inductor source kernel_spec.write_spec_file(src_code, meta, spec_path, diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index 4f9056a6..e5a34400 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -248,6 +248,20 @@ def strip_for_tnpu(src): prefix = "" if "import triton.language as tl" not in body: prefix = "import triton\nimport triton.language as tl\n\n" + # tl_math is triton's own, re-exported through triton_helpers; the dropped + # torch import took it with it. + if re.search(r"\btl_math\.", body): + prefix += "from triton.language import math as tl_math\n" + + # libdevice members are @core.extern: no triton_shared implementation, so a + # call returns None and fails obscurely in stage 1. Name it here instead. + ext = sorted(set(re.findall(r"\blibdevice\.(\w+)", body))) + if ext: + raise SpecIncomplete( + f"kernel calls libdevice.{{{','.join(ext)}}}: those are extern math " + f"intrinsics with no implementation on the triton_shared backend. " + f"They need lowering to a VPU op (or a scalar fallback) before this " + f"kernel can compile.") return prefix + body diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py index 5a49ccf5..a30593d1 100644 --- a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -16,6 +16,7 @@ """ import os +import re import subprocess from PyTorchSimFrontend import extension_config @@ -24,10 +25,28 @@ class TnpuError(RuntimeError): + """A tnpu stage failed. Inductor reports only str(exc), so the stage's own + diagnostic has to travel in the message.""" + + #: How a failing stage names itself: MLIR diagnostics and exception lines. + _SIGNAL = re.compile( + r"^(?!\s|Traceback|During handling|The above)" + r"(.*\berror:\s.*|.*failed to legalize.*|" + r"[\w.]*(?:Error|Exception)\b.*|.*Assertion.*)$", re.M) + #: Frames and carets: context, not the diagnostic. + _FRAME = re.compile(r'^\s|^\s*File "|^\s*\^') + def __init__(self, message, cmd=None, output=None): - super().__init__(message) self.cmd = cmd self.output = output + if output: + hits = [h.strip() for h in self._SIGNAL.findall(output) + if not self._FRAME.match(h)] + if not hits: + hits = [l for l in output.strip().splitlines() + if l.strip() and not self._FRAME.match(l)] + message = message + "\n " + "\n ".join(l[:300] for l in hits[-3:]) + super().__init__(message) def tnpu_dir(): @@ -76,6 +95,12 @@ def run_pipeline(spec_path, workdir, to_stage="binary", from_stage="ttir", cwd=tnpu_dir(), env=env, timeout=timeout) output = proc.stdout + proc.stderr if proc.returncode != 0: + # run.py prints a stage table; the diagnostic itself only reaches + # stage.log. + log = os.path.join(workdir, "stage.log") + if os.path.isfile(log): + with open(log, errors="replace") as fh: + output += "\n" + fh.read() raise TnpuError(f"tnpu pipeline failed (exit {proc.returncode})", cmd=" ".join(cmd), output=output) logger.debug("[triton-npu] %s", output) diff --git a/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt index 38e81f02..213a8d58 100644 --- a/scripts/ci/triton_route_passing.txt +++ b/scripts/ci/triton_route_passing.txt @@ -1,8 +1,7 @@ # Tests that pass THROUGH the Triton codegen route. # Gated by scripts/ci/triton_route_sweep.py; regenerate with # python scripts/ci/triton_route_sweep.py --all --update-allowlist -# A test that passes without emitting a kernel (CPU-only, eager -# fallback) is deliberately absent -- it would not gate anything. +# A test that passes without emitting a kernel is deliberately absent. tests/ops/elementwise/test_add.py tests/ops/fusion/test_addmm_residual.py tests/ops/fusion/test_matmul_scalar.py diff --git a/scripts/ci/triton_route_sweep.py b/scripts/ci/triton_route_sweep.py index 8079a43c..df8660e3 100755 --- a/scripts/ci/triton_route_sweep.py +++ b/scripts/ci/triton_route_sweep.py @@ -1,21 +1,9 @@ #!/usr/bin/env python3 """Run the existing test suite through the Triton codegen route. -The route is selected at device-registration time by TORCHSIM_TRITON_CODEGEN -(PyTorchSimDevice/torch_openreg/__init__.py), so the tests themselves need no -change -- the same file is the MLIR route's test with the variable unset and the -Triton route's with it set. - -Three things come out of a run: - - * a GATE. `triton_route_passing.txt` lists the tests that pass today. Any of - them failing is a regression and exits non-zero. That is what makes coverage - grow monotonically instead of drifting. - * a REPORT bucketed by how far each failure got, so the gap list in - triton_backend/README.md is a measurement rather than a guess. - * ARTIFACTS per failure, under --artifacts: the Inductor Triton kernel that - did not survive, the last tnpu stage IR it produced, and the error. That is - what makes a failure reportable to whoever owns the pass, without a rerun. +TORCHSIM_TRITON_CODEGEN is read at device registration, so no test needs to know +which route it is on. Produces a gate (triton_route_passing.txt), a report +bucketed by cause and stage, and per-failure artifacts for reporting upstream. python scripts/ci/triton_route_sweep.py # the allowlist, gating python scripts/ci/triton_route_sweep.py --all # every test, reports @@ -23,6 +11,7 @@ """ import argparse +import concurrent.futures as cf import glob import json import os @@ -36,8 +25,7 @@ HERE = os.path.dirname(os.path.abspath(__file__)) PASSING = os.path.join(HERE, "triton_route_passing.txt") -#: How far the kernel got, outermost first. The stage a failure did NOT reach is -#: the one that owns it, so this doubles as the routing table for a bug report. +#: How far the kernel got. The stage a failure did not reach owns it. STAGES = [ ("01-ttir.mlir", "1 triton -> ttir"), ("02-ttshared.mlir", "2 ttir -> tts/linalg (triton-shared)"), @@ -46,8 +34,7 @@ ("trace.so", "5 trace producer"), ] -#: Failure buckets, first match wins. Each names the layer that owns the fix, -#: so the counts say which gap to close next rather than just how many failed. +#: First match wins. Each bucket names the layer that owns the fix. BUCKETS = [ ("missing_dep", r"ModuleNotFoundError|No module named"), ("device_op", r"\w+_overrideable not implemented|not implemented\. .*privateuse"), @@ -114,8 +101,7 @@ def first_error(output): def reached_stage(dump_dir): """(label, workdir) of the furthest tnpu stage any kernel produced. - A workdir with only kernel.py is still the one to collect: the route got as - far as generating a Triton kernel and rejected it before stage 1. + kernel.py alone still counts: a kernel was generated and rejected pre-stage-1. """ best, best_dir, fallback = None, None, None for wd in glob.glob(os.path.join(dump_dir, "triton_*")): @@ -139,8 +125,7 @@ def collect(test, dump_dir, out_root, output, bucket, stage, workdir): f.write(f"test: {test}\nbucket: {bucket}\nstage: {stage}\n\n") f.write("\n".join(output.strip().splitlines()[-60:])) if workdir: - # The Inductor kernel is the thing to hand to whoever owns the pass; - # the stage IRs say where it stopped being representable. + # The kernel to hand over, and the IR saying where it stopped. for name in ("kernel.py", "stage.log", *(s[0] for s in STAGES[:-1])): src = os.path.join(workdir, name) if os.path.isfile(src): @@ -149,8 +134,8 @@ def collect(test, dump_dir, out_root, output, bucket, stage, workdir): def run_one(test, timeout, artifacts, scratch): - # A private dump dir per test: artifacts must be attributable, and a shared - # one lets a cached kernel from an earlier test answer for this one. + # Private per test: a shared dump lets one test's cached kernel answer for + # another's. dump = os.path.join(scratch, test.replace("/", "_").removesuffix(".py")) shutil.rmtree(dump, ignore_errors=True) os.makedirs(dump, exist_ok=True) @@ -171,9 +156,8 @@ def run_one(test, timeout, artifacts, scratch): "seconds": round(time.time() - t0, 1), "bucket": None if ok else classify(out, timed_out), "stage": stage, - # A pass that emitted no kernel never used the route: a CPU-only test, - # an eager fallback, or a path that bypasses Inductor. Counting those as - # coverage would overstate it. + # No kernel emitted = the route was never used (CPU-only, eager + # fallback, extern call), so it is not coverage. "exercised": workdir is not None, "error": "" if ok else first_error(out)} if not ok and artifacts: @@ -242,6 +226,11 @@ def main(): ap.add_argument("--all", action="store_true", help="run every test, not just the passing allowlist") ap.add_argument("--timeout", type=int, default=1800) + ap.add_argument("-j", "--jobs", type=int, + default=max(1, min(8, (os.cpu_count() or 2) // 2)), + help="tests in flight at once; each may itself use several " + "cores (gem5, TOGSim), so this is half the box by " + "default") ap.add_argument("--json", help="write the full result list here") ap.add_argument("--artifacts", metavar="DIR", help="per-failure kernel + stage IR + error, for reporting") @@ -263,17 +252,35 @@ def main(): shutil.rmtree(args.artifacts, ignore_errors=True) os.makedirs(args.artifacts, exist_ok=True) - print(f"Triton route sweep: {len(tests)} tests" + print(f"Triton route sweep: {len(tests)} tests, {args.jobs} at a time" f"{'' if args.all else ' (allowlist)'}\n") - results = [] - for i, t in enumerate(tests, 1): - r = run_one(t, args.timeout, args.artifacts, scratch) - results.append(r) + results, done = [], 0 + + def report(r): + nonlocal done + done += 1 mark = ("ok " if r["exercised"] else "ok- ") if r["ok"] else "FAIL" extra = ("" if r["exercised"] else " (route not exercised)") if r["ok"] \ else f" [{r['bucket']}] @{r['stage']} {r['error'][:70]}" - print(f" {i:3d}/{len(tests)} {mark} {r['seconds']:7.1f}s " + print(f" {done:3d}/{len(tests)} {mark} {r['seconds']:7.1f}s " f"{r['test']}{extra}", flush=True) + + if args.jobs == 1: + for t in tests: + r = run_one(t, args.timeout, args.artifacts, scratch) + results.append(r) + report(r) + else: + # Threads: run_one only waits on a subprocess, and dump dir, Inductor + # cache and TOGSim FIFO are all already per-test. + with cf.ThreadPoolExecutor(max_workers=args.jobs) as pool: + futs = {pool.submit(run_one, t, args.timeout, args.artifacts, + scratch): t for t in tests} + for fut in cf.as_completed(futs): + r = fut.result() + results.append(r) + report(r) + results.sort(key=lambda r: r["test"]) shutil.rmtree(scratch, ignore_errors=True) passed = [r for r in results if r["ok"]] From 8e17519386f41dead0b837c0abf89b4f105a7ee1 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 22:32:32 +0900 Subject: [PATCH 33/35] [Docs] Record the first Triton route coverage measurement 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. --- docs/triton-route-coverage.md | 327 ++++++++++++++++++++++++++++++++++ 1 file changed, 327 insertions(+) create mode 100644 docs/triton-route-coverage.md diff --git a/docs/triton-route-coverage.md b/docs/triton-route-coverage.md new file mode 100644 index 00000000..87ddd2b1 --- /dev/null +++ b/docs/triton-route-coverage.md @@ -0,0 +1,327 @@ +# Triton codegen route — test suite coverage + +First measurement of PyTorchSim's existing test suite running through the Triton +codegen route (Inductor's Triton backend + the triton-npu lowering passes), +instead of the MLIR route. + +| | | +|---|---| +| Date | 2026-08-03 | +| Branch | `feature/triton-codegen` @ `d356e35` | +| tnpu pin | `5d84caf` | +| torch | 2.10.0, triton 3.6.0 | +| Tests | 69 (everything under `tests/`) | +| Runtime | 5 min at `-j 10` (~50 min serial) | + +Reproduce: + +```bash +python scripts/ci/triton_route_sweep.py --all -j 10 \ + --markdown coverage.md --artifacts failures +``` + +--- + +## 1. Headline + +``` +69 tests +├── 11 pass THROUGH the route ← this is the coverage number +├── 5 pass without using the route ← no kernel emitted at all +└── 53 fail + ├── 17 missing test deps (local venv only; present in the CI image) + └── 36 real blockers +``` + +**11/69, not 16/69.** Five tests pass while emitting no Triton kernel +whatsoever: `test_matmul`, `test_bmm`, `test_topk`, `test_moe_cpu`, +`test_mlir_bindings`. Inductor sends `mm`/`bmm` to an extern call rather than +generating a kernel, so those tests never exercise the thing under test. The +sweep records that separately (`exercised` in the JSON) and keeps them out of +the gate — otherwise the number would overstate coverage by 45%. + +### What passes + +| Test | Time | +|---|---| +| `tests/ops/elementwise/test_add.py` | 77.5s | +| `tests/ops/fusion/test_addmm_residual.py` | 33.0s | +| `tests/ops/fusion/test_matmul_scalar.py` | 11.5s | +| `tests/ops/fusion/test_matmul_vector.py` | 17.4s | +| `tests/ops/fusion/test_prologue_fusion.py` | 41.4s | +| `tests/ops/misc/test_expert_mask.py` | 11.0s | +| `tests/ops/reduce/test_batchnorm.py` | 37.7s | +| `tests/ops/view/test_view3D_2D.py` | 36.4s | +| `tests/system/test_eager.py` | 15.0s | +| `tests/system/test_stonne.py` | 9.7s | +| `tests/system/test_triton_codegen.py` | 10.0s | + +Four of the eleven are fusion tests. That is the encouraging part: Inductor's +fusion is the half of this migration we get for free, and it is already +producing kernels tnpu accepts. + +--- + +## 2. Where kernels stop + +Each test is placed at the furthest stage any of its kernels produced an +artifact for. The stage a kernel *fails to reach* is the one that owns the +failure. + +| Stage | Count | | +|---|---|---| +| — no kernel generated | 26 | died in torch/Inductor before codegen | +| 0 generated, rejected | 16 | `kernel_spec` refused to describe it | +| 1 triton → ttir | 1 | | +| 2 ttir → tts/linalg | 1 | triton-shared | +| 4 tnpu lower (DMA, lanes, spad) | 6 | | +| 5 trace producer | 3 | | + +**The lowering passes are not the bottleneck yet.** Only 2 of 53 failures are a +tnpu pass rejecting IR. The other 34 real blockers stop earlier — in the port +that feeds tnpu, or in torch itself. The next round of work is mostly on our +side of the seam, not upstream's. + +--- + +## 3. Failures by cause + +### `spec_incomplete` — 13 · owner: `triton_backend/kernel_spec.py` + +Three distinct sub-causes: + +**libdevice intrinsics (5).** `@core.extern` members with no triton_shared +implementation; a call returns `None`. + +| Test | Symbol | +|---|---| +| `ops/elementwise/test_exponent.py` | `libdevice.exp` | +| `ops/elementwise/test_pointwise.py` | `libdevice.isnan` | +| `ops/elementwise/test_transcendental.py` | `libdevice.tanh` | +| `ops/reduce/test_layernorm.py` | `libdevice.rsqrt` | +| `ops/view/test_floormod_axis_split.py` | `libdevice.rsqrt` | + +**Multi-axis grid (4).** `fixed_config_for` pins only the outermost axis, so +`YBLOCK` is `None` and the grid cannot be computed. This is the known +block-size policy gap. + +| Test | Detail | +|---|---| +| `ops/view/test_transpose2D.py` | axis `y`: ynumel=156, YBLOCK=None | +| `ops/view/test_transpose3D.py` | axis `y`: ynumel=2728, YBLOCK=None | +| `ops/fusion/test_conv_fusion.py` | axis `y`: ynumel=192, YBLOCK=None | +| `ops/conv/test_conv_view_input.py` | axis `y`: ynumel=512, YBLOCK=None | + +**Reduction blocks unset (2)** — `R0_BLOCK` is left unset on purpose: +`ops/fusion/test_bmm_reduction.py`, `ops/fusion/test_matmul_reduction.py`. + +**Genuine metadata hole (1)** — `ops/misc/test_widen_dtype.py`: no dtype/numel +for `out_ptr0`, `collect_meta` could not resolve it from `V.graph`. + +### `triton_helpers` — 7 · owner: `triton_backend` + +The module lives in torch; the tnpu venv deliberately has none. + +| Test | Helper | +|---|---| +| `ops/reduce/test_softmax.py` | `max2` | +| `ops/sort/test_sort.py` | `sort_with_index` | +| `ops/elementwise/test_activation.py` | `maximum` | +| `ops/conv/test_cnn.py` | `maximum` | +| `ops/fusion/test_matmul_activation.py` | `maximum` | +| `ops/sparsity/test_sparsity.py` | `maximum` | +| `models/test_mlp.py` | `maximum` | + +Four of seven want only `maximum`. Fix is a small vendored file, not a pass +change. + +### `wrapper_gap` — 6 · owner: `triton_backend` + +Every one: `'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak'`. + +`ops/attention/test_gqa.py`, `test_gqa_decode.py`, +`ops/fusion/test_attention_fusion.py`, `test_transformer_fusion.py`, +`models/Mixtral8x7B/test_attention.py`, `models/test_transformer.py` + +Every attention and transformer test in the suite, blocked on one unimplemented +method. + +### `device_op` — 3 · owner: `PyTorchSimDevice` + +Predates this route — the MLIR route intercepts these before the dispatcher. + +| Test | Error | +|---|---| +| `ops/conv/test_conv2d.py` | `convolution_overrideable not implemented` | +| `ops/conv/test_group_conv.py` | `convolution_overrideable not implemented` | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable not implemented` | + +### `tnpu_stage` — 2 · owner: triton-npu lowering passes + +The only failures that are genuinely a pass rejecting IR. + +| Test | Stage | Diagnostic | +|---|---|---| +| `ops/conv/test_pool.py` | 1 | ``Dialect `ttg' not found for custom op 'ttg.barrier'`` | +| `ops/reduce/test_reduce.py` | 2 | `'linalg.index' op expected dim (2) to be lower than the number of loops (2) of the enclosing LinalgOp` | + +Both artifacts carry the MLIR diagnostic and the offending `.mlir`, so they can +go upstream as-is. + +### `togsim` / `other` — 5 + +| Test | Stage | Detail | +|---|---|---| +| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed` | +| `ops/misc/test_masked_nondividing.py` | 4 | `[Spike] triton_npu_fused_constant_pad_nd_0 failed` | +| `ops/misc/test_indirect_access.py` | 5 | TOGSim returned `inf` cycles for `index_put` | +| `system/test_hetro.py` | — | `KeyError: 'vpu_num_lanes'` (hetero config lacks the key) | +| `ops/sparsity/test_sparse_core.py` | — | `TypeError: '>' between Tensor and torch.device` (test-side bug) | + +The two Spike failures are the most interesting in the sweep: the only cases +that compile all the way to a working RISC-V binary and then produce the wrong +thing. Everything else fails before it can be wrong. + +### `missing_dep` — 17 · not a route problem + +`transformers` (5), `torchvision` (4), `matplotlib` (4), `pytest` (2), +`diffusers`, `requests`, `sklearn`. Local venv only — these run for real in the +CI image, which is why the sweep belongs in CI. + +--- + +## 4. Infrastructure + +### The runner + +`scripts/ci/triton_route_sweep.py`. `TORCHSIM_TRITON_CODEGEN` is read once at +device registration (`PyTorchSimDevice/torch_openreg/__init__.py:30`), so **no +test file needed to change** — all 69 were already tests of this route. Only a +runner was missing. + +Three outputs: + +1. **Gate** — `scripts/ci/triton_route_passing.txt` lists what passes today. + CI fails if any regresses. Coverage grows by regenerating the file + (`--update-allowlist`), so it cannot silently shrink. +2. **Report** — bucketed by owning layer and by pipeline stage. +3. **Artifacts** — one directory per failing test. + +### Per-failure artifacts + +``` +failures/tests_ops_reduce_test_softmax/ + kernel.py the Inductor Triton kernel, unmodified + error.txt bucket, stage, last 60 lines + 01-ttir.mlir whatever stage IR it reached + stage.log +``` + +For `test_softmax` the kernel names its own blocker: + +```python +def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK: tl.constexpr): + ... + tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) + tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <- lives in torch + tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) + tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) +``` + +### Parallelism + +Tests are independent subprocesses with their own dump dir, Inductor cache +(`TORCHINDUCTOR_CACHE_DIR` follows `TORCHSIM_DUMP_PATH`) and TOGSim FIFO (keyed +by pid), so `-j` needs no coordination. Threads, not processes: `run_one` only +waits on a subprocess. **69 tests: ~50 min → 5 min at `-j 10`.** + +### CI + +`.github/workflows/triton_npu.yml`, job `triton-route-suite`: + +- **Allowlisted tests** — gates. +- **Full sweep** — `continue-on-error`, writes `coverage.md` to the step + summary and uploads `triton-route-coverage` (results.json + failures/). + +Jobs run on the PSAL Slurm runner farm (`PSAL-POSTECH/slurm-ghr`): `runs-on` +carries the `slurm` label, image builds and the sweep on `big` (16c/64G/2h), +the rest on small. Do not add `docker/setup-buildx-action` — the runner +registers its own builder. + +--- + +## 5. Three diagnostics fixed while measuring + +The reporting infrastructure could not be built until these were fixed, because +each one was destroying the evidence. + +**`kernel.py` was written after the check that rejects it.** +`write_spec_file` raises for exactly the kernels worth keeping +(`triton_helpers`, `SpecIncomplete`) and ran *before* the source was saved — so +the interesting sources were the ones being thrown away. Reordered; the dump +now exists for all 16 rejected kernels. + +**tnpu reported "exit 1" and nothing else.** `run.py` prints a stage table to +stdout and the real diagnostic only to `stage.log`. `TnpuError` now reads that +file and carries the failing line. That single change resolved six failures +into one bug: + +**`libdevice` and `tl_math` were collateral damage.** `strip_for_tnpu` drops +`from torch...`, and Inductor imports both names from +`torch._inductor.runtime.triton_helpers` — but they are re-exports of *triton's +own* symbols, not torch code. Six kernels died as a bare `NameError` inside +stage 1. + +- `tl_math` is now rebound from `triton.language`. `test_pointwise` 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 + explicitly, the way `triton_helpers` is. + +Net effect: `tnpu_stage` 8 → 2, `spec_incomplete` 7 → 13. The same 53 tests +fail; six of them now say something true. + +**Separately:** the local TOGSim build was from 07-20 and predated +`trace_shape.txt` support, so every Triton-route test died with SIGSEGV in +`trace_to_tilegraph`. A rebuild fixed it — not a code problem, and CI builds +from source so it was never affected. Worth knowing if anyone else has a stale +`TOGSim/build`. + +--- + +## 6. Next, ranked by tests unblocked per fix + +Counts are measured, not estimated — though a test unblocked at one stage may +simply fail at the next. + +| # | Fix | Unblocks | Owner | +|---|---|---|---| +| 1 | Implement `TritonNPUWrapperCodegen.estimate_peak` | 6 | triton_backend | +| 2 | Vendor a torch-free `triton_helpers` into the tnpu venv | 7 | triton_backend | +| 3 | Lower `libdevice` intrinsics (`exp`, `tanh`, `rsqrt`, `isnan`) to VPU ops | 5 | tnpu or triton_backend | +| 4 | Multi-axis block policy in `fixed_config_for` | 4 | triton_backend | +| 5 | Hand `ttg.barrier` + `linalg.index` rank error upstream | 2 | tnpu | +| 6 | Investigate the two Spike failures (`cat`, `constant_pad_nd`) | 2 | investigate | + +**1 is the cheapest by a wide margin** — one method, six tests, and it opens the +entire attention/transformer family. + +**3 needs a decision before work starts:** lower in a tnpu pass, or substitute a +triton-level polyfill in `strip_for_tnpu`. The former is correct; the latter is +cheap and would unblock measurement sooner. + +**6 is the most likely to be a real bug.** Everything else is a missing feature; +these two compile to a working binary and produce the wrong answer. + +--- + +## 7. Caveats + +- The 17 `missing_dep` failures are local-venv artifacts. In the CI image those + tests run for real and the buckets will shift — probably toward `wrapper_gap` + and `triton_helpers`, since most are transformer and CNN models. +- Unblocking a bucket moves its tests to the *next* failure, not necessarily to + passing. +- These numbers were taken with the section-5 fixes already applied, so they are + not comparable to a run from before them. From 6e3bd7e4459d684ac65f408011aa579e22d0275e Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 22:40:39 +0900 Subject: [PATCH 34/35] [Docs] Back the coverage report with the artifacts it came from 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. --- docs/triton-route-coverage.md | 293 +++++++++++++++++++++++++--------- 1 file changed, 214 insertions(+), 79 deletions(-) diff --git a/docs/triton-route-coverage.md b/docs/triton-route-coverage.md index 87ddd2b1..4e4a2590 100644 --- a/docs/triton-route-coverage.md +++ b/docs/triton-route-coverage.md @@ -1,13 +1,13 @@ # Triton codegen route — test suite coverage First measurement of PyTorchSim's existing test suite running through the Triton -codegen route (Inductor's Triton backend + the triton-npu lowering passes), +codegen route (Inductor's Triton backend + the triton-npu lowering passes) instead of the MLIR route. | | | |---|---| | Date | 2026-08-03 | -| Branch | `feature/triton-codegen` @ `d356e35` | +| Branch | `feature/triton-codegen` @ `8e17519` | | tnpu pin | `5d84caf` | | torch | 2.10.0, triton 3.6.0 | | Tests | 69 (everything under `tests/`) | @@ -20,6 +20,9 @@ python scripts/ci/triton_route_sweep.py --all -j 10 \ --markdown coverage.md --artifacts failures ``` +Every claim below is backed by a file in `failures/`. Paths are given so each +one can be checked. + --- ## 1. Headline @@ -33,12 +36,12 @@ python scripts/ci/triton_route_sweep.py --all -j 10 \ └── 36 real blockers ``` -**11/69, not 16/69.** Five tests pass while emitting no Triton kernel -whatsoever: `test_matmul`, `test_bmm`, `test_topk`, `test_moe_cpu`, -`test_mlir_bindings`. Inductor sends `mm`/`bmm` to an extern call rather than -generating a kernel, so those tests never exercise the thing under test. The -sweep records that separately (`exercised` in the JSON) and keeps them out of -the gate — otherwise the number would overstate coverage by 45%. +**11/69, not 16/69.** Five tests pass while emitting no Triton kernel at all: +`test_matmul`, `test_bmm`, `test_topk`, `test_moe_cpu`, `test_mlir_bindings`. +Inductor sends `mm`/`bmm` to an extern call rather than generating a kernel, so +those tests never exercise the thing under test. The sweep records that +separately (`exercised` in the JSON) and keeps them out of the gate — counting +them would overstate coverage by 45%. ### What passes @@ -56,17 +59,75 @@ the gate — otherwise the number would overstate coverage by 45%. | `tests/system/test_stonne.py` | 9.7s | | `tests/system/test_triton_codegen.py` | 10.0s | -Four of the eleven are fusion tests. That is the encouraging part: Inductor's -fusion is the half of this migration we get for free, and it is already -producing kernels tnpu accepts. +Four of the eleven are fusion tests. Inductor's fusion is the half of this +migration we get for free, and it is already producing kernels tnpu accepts. --- -## 2. Where kernels stop +## 2. Worked example — how one failure is diagnosed + +`tests/ops/reduce/test_softmax.py`. The sweep leaves this behind: + +``` +failures/tests_ops_reduce_test_softmax/ + kernel.py the Inductor Triton kernel, unmodified + error.txt bucket, stage, last 60 lines +``` + +`error.txt` opens with the routing header: + +``` +test: tests/ops/reduce/test_softmax.py +bucket: triton_helpers +stage: 0 kernel generated, not accepted +``` + +And `kernel.py` is the whole reason, in 28 lines: + +```python +@triton.jit +def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr): + xnumel = 64 + r0_numel = 128 + R0_BLOCK: tl.constexpr = 128 + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < xnumel + r0_index = tl.arange(0, R0_BLOCK)[None, :] + r0_1 = r0_index + x0 = xindex + tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) + tmp1 = tl.broadcast_to(tmp0, [XBLOCK, R0_BLOCK]) + tmp3 = tl.where(xmask, tmp1, float("-inf")) + tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <-- blocker 1 + tmp5 = tmp0 - tmp4 + tmp6 = libdevice.exp(tmp5) # <-- blocker 2 + tmp7 = tl.broadcast_to(tmp6, [XBLOCK, R0_BLOCK]) + tmp9 = tl.where(xmask, tmp7, 0) + tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) + tmp11 = (tmp6 / tmp10) + tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) +``` + +Two blockers, visible without running anything: + +- `triton_helpers.max2` — lives in `torch._inductor.runtime`, and the tnpu venv + deliberately has no torch. +- `libdevice.exp` — an `@core.extern` intrinsic with no triton_shared + implementation. + +Note what *is* fine: `tl.load` with a mask, `tl.where`, `tl.sum`, `tl.store`, +the 2-D `[XBLOCK, R0_BLOCK]` broadcast. Softmax is not blocked on anything +structural. It is blocked on two function calls. + +This is the whole bug report, and it needed no rerun to produce. + +--- + +## 3. Where kernels stop Each test is placed at the furthest stage any of its kernels produced an -artifact for. The stage a kernel *fails to reach* is the one that owns the -failure. +artifact for. The stage a kernel *fails to reach* owns the failure. | Stage | Count | | |---|---|---| @@ -80,16 +141,14 @@ failure. **The lowering passes are not the bottleneck yet.** Only 2 of 53 failures are a tnpu pass rejecting IR. The other 34 real blockers stop earlier — in the port that feeds tnpu, or in torch itself. The next round of work is mostly on our -side of the seam, not upstream's. +side of the seam. --- -## 3. Failures by cause +## 4. Failures by cause, with evidence ### `spec_incomplete` — 13 · owner: `triton_backend/kernel_spec.py` -Three distinct sub-causes: - **libdevice intrinsics (5).** `@core.extern` members with no triton_shared implementation; a call returns `None`. @@ -101,11 +160,21 @@ implementation; a call returns `None`. | `ops/reduce/test_layernorm.py` | `libdevice.rsqrt` | | `ops/view/test_floormod_axis_split.py` | `libdevice.rsqrt` | +Without the diagnostic added this session, these failed as a bare +`NameError('libdevice is not defined')` inside tnpu's stage-1 worker — which is +what made them look like six separate lowering bugs. They now say: + +``` +SpecIncomplete: kernel calls libdevice.{exp}: those are extern math intrinsics +with no implementation on the triton_shared backend. They need lowering to a +VPU op (or a scalar fallback) before this kernel can compile. +``` + **Multi-axis grid (4).** `fixed_config_for` pins only the outermost axis, so `YBLOCK` is `None` and the grid cannot be computed. This is the known block-size policy gap. -| Test | Detail | +| Test | Diagnostic | |---|---| | `ops/view/test_transpose2D.py` | axis `y`: ynumel=156, YBLOCK=None | | `ops/view/test_transpose3D.py` | axis `y`: ynumel=2728, YBLOCK=None | @@ -116,12 +185,10 @@ block-size policy gap. `ops/fusion/test_bmm_reduction.py`, `ops/fusion/test_matmul_reduction.py`. **Genuine metadata hole (1)** — `ops/misc/test_widen_dtype.py`: no dtype/numel -for `out_ptr0`, `collect_meta` could not resolve it from `V.graph`. +for `out_ptr0`; `collect_meta` could not resolve it from `V.graph`. ### `triton_helpers` — 7 · owner: `triton_backend` -The module lives in torch; the tnpu venv deliberately has none. - | Test | Helper | |---|---| | `ops/reduce/test_softmax.py` | `max2` | @@ -132,12 +199,16 @@ The module lives in torch; the tnpu venv deliberately has none. | `ops/sparsity/test_sparsity.py` | `maximum` | | `models/test_mlp.py` | `maximum` | -Four of seven want only `maximum`. Fix is a small vendored file, not a pass +Four of seven want only `maximum`. The fix is a small vendored file, not a pass change. ### `wrapper_gap` — 6 · owner: `triton_backend` -Every one: `'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak'`. +Every one, identically: + +``` +AttributeError: 'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak' +``` `ops/attention/test_gqa.py`, `test_gqa_decode.py`, `ops/fusion/test_attention_fusion.py`, `test_transformer_fusion.py`, @@ -158,29 +229,93 @@ Predates this route — the MLIR route intercepts these before the dispatcher. ### `tnpu_stage` — 2 · owner: triton-npu lowering passes -The only failures that are genuinely a pass rejecting IR. +The only failures that are genuinely a pass rejecting IR — and the artifacts +show the full chain from Python to the rejected op. -| Test | Stage | Diagnostic | -|---|---|---| -| `ops/conv/test_pool.py` | 1 | ``Dialect `ttg' not found for custom op 'ttg.barrier'`` | -| `ops/reduce/test_reduce.py` | 2 | `'linalg.index' op expected dim (2) to be lower than the number of loops (2) of the enclosing LinalgOp` | +**`ops/conv/test_pool.py`** — stage 1. -Both artifacts carry the MLIR diagnostic and the offending `.mlir`, so they can -go upstream as-is. +`kernel.py` ends with an innocuous line Inductor emits after a reduction: + +```python + tmp4 = tl.sum(tmp3, 1)[:, None].to(tl.float32) + tmp5 = 49.0 + tmp6 = (tmp4 / tmp5) + tl.debug_barrier() # <-- this +``` + +`01-ttir.mlir:46` is what that becomes: + +```mlir +%tmp4_17 = tt.expand_dims %tmp4 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> +%tmp6_18 = arith.divf %tmp4_17, %tmp6 : tensor<128x1xf32> +ttg.barrier all # <-- GPU dialect op +%0 = tt.splat %in_out_ptr0 : !tt.ptr -> tensor<128x1x!tt.ptr> +``` + +and `triton-shared-opt` cannot parse it: + +``` +01-ttir.mlir:46:5: error: Dialect `ttg' not found for custom op 'ttg.barrier' +``` + +`ttg` is the GPU dialect. A `tl.debug_barrier()` in a pointwise-after-reduction +kernel is meaningless on this target, but it is in the IR and the parser stops +on it. + +**`ops/reduce/test_reduce.py`** — stage 2. A plain `(a + b).sum(dim=1)`: + +```python +tmp0 = tl.load(in_ptr0 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp1 = tl.load(in_ptr1 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp2 = tmp0 + tmp1 +tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK]) +tmp5 = tl.where(r0_mask & xmask, tmp3, 0) +``` + +survives `01-ttir.mlir`, then fails converting to linalg: + +``` +error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than + the number of loops (2) of the enclosing LinalgOp +``` + +Both artifacts carry the diagnostic and the offending `.mlir`, so they can go +upstream as-is. ### `togsim` / `other` — 5 | Test | Stage | Detail | |---|---|---| -| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed` | +| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed`, exit 255 | | `ops/misc/test_masked_nondividing.py` | 4 | `[Spike] triton_npu_fused_constant_pad_nd_0 failed` | | `ops/misc/test_indirect_access.py` | 5 | TOGSim returned `inf` cycles for `index_put` | | `system/test_hetro.py` | — | `KeyError: 'vpu_num_lanes'` (hetero config lacks the key) | | `ops/sparsity/test_sparse_core.py` | — | `TypeError: '>' between Tensor and torch.device` (test-side bug) | The two Spike failures are the most interesting in the sweep: the only cases -that compile all the way to a working RISC-V binary and then produce the wrong -thing. Everything else fails before it can be wrong. +that compile all the way to a working RISC-V binary and then fail at run time. +`test_cat`'s `04-custom.mlir` shows the lowering did its job — + +```mlir +"togsim.transfer"(%reinterpret_cast_5, %c0, %2, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast_4, %c0, %1, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast, %9, %0, %c0, %7, %c0, %c3, %c1, %c0, %13) + {vlane_split_axis = 0, ...} +``` + +— two input DMAs and one output DMA, lane-split on axis 0, exactly the shape +`cat` should produce. It is the execution that goes wrong, not the lowering +structure. + +**Caveat on these two.** `tnpu.spike` reports only `StageError: command failed +with exit code 255`; spike's own stderr does not survive. Running the recorded +spike command by hand on the same workdir exits 0, because `write_inputs` +rewrites `runtime/*.raw` per launch and a by-hand run replays stale inputs. So +the failing input is not currently reproducible outside the pipeline. Surfacing +spike's stderr the way `TnpuError` now surfaces tnpu's is the prerequisite for +diagnosing these, and is not yet done. ### `missing_dep` — 17 · not a route problem @@ -190,7 +325,7 @@ CI image, which is why the sweep belongs in CI. --- -## 4. Infrastructure +## 5. Infrastructure ### The runner @@ -201,32 +336,15 @@ runner was missing. Three outputs: -1. **Gate** — `scripts/ci/triton_route_passing.txt` lists what passes today. - CI fails if any regresses. Coverage grows by regenerating the file +1. **Gate** — `scripts/ci/triton_route_passing.txt` lists what passes today; CI + fails if any regresses. Coverage grows by regenerating the file (`--update-allowlist`), so it cannot silently shrink. 2. **Report** — bucketed by owning layer and by pipeline stage. -3. **Artifacts** — one directory per failing test. - -### Per-failure artifacts - -``` -failures/tests_ops_reduce_test_softmax/ - kernel.py the Inductor Triton kernel, unmodified - error.txt bucket, stage, last 60 lines - 01-ttir.mlir whatever stage IR it reached - stage.log -``` +3. **Artifacts** — one directory per failing test (section 2). -For `test_softmax` the kernel names its own blocker: - -```python -def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK: tl.constexpr): - ... - tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) - tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <- lives in torch - tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) - tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) -``` +A deeper failure keeps more. `tests_ops_conv_test_cnn/` holds +`01-ttir.mlir 02-ttshared.mlir 03-adapted.mlir 04-custom.mlir kernel.py +stage.log error.txt` — the complete lowering chain up to the point it stopped. ### Parallelism @@ -240,7 +358,7 @@ waits on a subprocess. **69 tests: ~50 min → 5 min at `-j 10`.** `.github/workflows/triton_npu.yml`, job `triton-route-suite`: - **Allowlisted tests** — gates. -- **Full sweep** — `continue-on-error`, writes `coverage.md` to the step +- **Full sweep** — `continue-on-error`; writes `coverage.md` into the step summary and uploads `triton-route-coverage` (results.json + failures/). Jobs run on the PSAL Slurm runner farm (`PSAL-POSTECH/slurm-ghr`): `runs-on` @@ -250,21 +368,33 @@ registers its own builder. --- -## 5. Three diagnostics fixed while measuring +## 6. Three diagnostics fixed while measuring The reporting infrastructure could not be built until these were fixed, because -each one was destroying the evidence. +each was destroying the evidence. -**`kernel.py` was written after the check that rejects it.** -`write_spec_file` raises for exactly the kernels worth keeping -(`triton_helpers`, `SpecIncomplete`) and ran *before* the source was saved — so -the interesting sources were the ones being thrown away. Reordered; the dump -now exists for all 16 rejected kernels. +**`kernel.py` was written after the check that rejects it.** `write_spec_file` +raises for exactly the kernels worth keeping (`triton_helpers`, +`SpecIncomplete`) and ran *before* the source was saved — so the interesting +sources were the ones being thrown away. Reordered; the dump now exists for all +16 rejected kernels, including the softmax example in section 2. **tnpu reported "exit 1" and nothing else.** `run.py` prints a stage table to -stdout and the real diagnostic only to `stage.log`. `TnpuError` now reads that -file and carries the failing line. That single change resolved six failures -into one bug: +stdout and the real diagnostic only to `stage.log`. Before: + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) +``` + +After (`TnpuError` now reads `stage.log`): + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) + triton.compiler.errors.CompilationError: at 8:11: + NameError('tl_math is not defined') +``` + +That single change resolved six failures into one bug: **`libdevice` and `tl_math` were collateral damage.** `strip_for_tnpu` drops `from torch...`, and Inductor imports both names from @@ -283,14 +413,14 @@ Net effect: `tnpu_stage` 8 → 2, `spec_incomplete` 7 → 13. The same 53 tests fail; six of them now say something true. **Separately:** the local TOGSim build was from 07-20 and predated -`trace_shape.txt` support, so every Triton-route test died with SIGSEGV in -`trace_to_tilegraph`. A rebuild fixed it — not a code problem, and CI builds -from source so it was never affected. Worth knowing if anyone else has a stale -`TOGSim/build`. +`trace_shape.txt` support, so `togsim_kernel` was called with `shape_args = +nullptr` and every Triton-route test died with SIGSEGV in `trace_to_tilegraph`. +A rebuild fixed it — not a code problem, and CI builds from source so it was +never affected. Worth knowing if anyone else has a stale `TOGSim/build`. --- -## 6. Next, ranked by tests unblocked per fix +## 7. Next, ranked by tests unblocked per fix Counts are measured, not estimated — though a test unblocked at one stage may simply fail at the next. @@ -302,26 +432,31 @@ simply fail at the next. | 3 | Lower `libdevice` intrinsics (`exp`, `tanh`, `rsqrt`, `isnan`) to VPU ops | 5 | tnpu or triton_backend | | 4 | Multi-axis block policy in `fixed_config_for` | 4 | triton_backend | | 5 | Hand `ttg.barrier` + `linalg.index` rank error upstream | 2 | tnpu | -| 6 | Investigate the two Spike failures (`cat`, `constant_pad_nd`) | 2 | investigate | +| 6 | Surface spike's stderr, then diagnose `cat` / `constant_pad_nd` | 2 | triton_backend, then investigate | **1 is the cheapest by a wide margin** — one method, six tests, and it opens the entire attention/transformer family. +**2 and 3 unblock softmax together.** Section 2 shows softmax needs both; either +alone leaves it failing on the other. + **3 needs a decision before work starts:** lower in a tnpu pass, or substitute a triton-level polyfill in `strip_for_tnpu`. The former is correct; the latter is cheap and would unblock measurement sooner. -**6 is the most likely to be a real bug.** Everything else is a missing feature; -these two compile to a working binary and produce the wrong answer. +**6 has a prerequisite.** These are the only wrong-answer failures in the suite +and the most likely to be a real lowering bug, but they cannot be diagnosed +until spike's stderr survives the subprocess — the same fix already applied to +`TnpuError` in section 6. --- -## 7. Caveats +## 8. Caveats - The 17 `missing_dep` failures are local-venv artifacts. In the CI image those tests run for real and the buckets will shift — probably toward `wrapper_gap` and `triton_helpers`, since most are transformer and CNN models. - Unblocking a bucket moves its tests to the *next* failure, not necessarily to passing. -- These numbers were taken with the section-5 fixes already applied, so they are +- These numbers were taken with the section-6 fixes already applied, so they are not comparable to a run from before them. From 33a222a50077c72c076864185338424d3607d66d Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Mon, 3 Aug 2026 22:54:57 +0900 Subject: [PATCH 35/35] [Docs] Add the Korean coverage report --- docs/triton-route-coverage.ko.md | 461 +++++++++++++++++++++++++++++++ docs/triton-route-coverage.md | 3 +- 2 files changed, 463 insertions(+), 1 deletion(-) create mode 100644 docs/triton-route-coverage.ko.md diff --git a/docs/triton-route-coverage.ko.md b/docs/triton-route-coverage.ko.md new file mode 100644 index 00000000..86cb2ed8 --- /dev/null +++ b/docs/triton-route-coverage.ko.md @@ -0,0 +1,461 @@ +# Triton codegen route 커버리지 측정 보고서 + +기존 PyTorchSim 테스트 스위트를 MLIR 경로가 아니라 **Triton 경로**(Inductor의 +Triton 백엔드 + triton-npu lowering pass)로 돌린 첫 측정 결과입니다. + +| | | +|---|---| +| 측정일 | 2026-08-03 | +| 브랜치 | `feature/triton-codegen` @ `6e3bd7e` | +| tnpu 핀 | `5d84caf` | +| 환경 | torch 2.10.0, triton 3.6.0 | +| 대상 | 69개 (`tests/` 전체) | +| 소요 | `-j 10` 기준 5분 (직렬 약 50분) | + +재현: + +```bash +python scripts/ci/triton_route_sweep.py --all -j 10 \ + --markdown coverage.md --artifacts failures +``` + +아래 모든 주장은 `failures/` 아래 실제 파일로 뒷받침됩니다. 확인할 수 있도록 +경로를 함께 적었습니다. + +--- + +## 1. 결론부터 + +``` +69개 테스트 +├── 11 경로를 타고 통과 ← 이것이 커버리지 수치 +├── 5 통과하지만 경로 미사용 ← 커널을 아예 안 만듦 +└── 53 실패 + ├── 17 로컬 venv 패키지 없음 (CI 이미지에는 있음) + └── 36 실제 블로커 +``` + +**16/69가 아니라 11/69입니다.** 다섯 개(`test_matmul`, `test_bmm`, `test_topk`, +`test_moe_cpu`, `test_mlir_bindings`)는 통과하지만 Triton 커널을 하나도 만들지 +않습니다. Inductor가 `mm`/`bmm`을 커널 생성 대신 extern call로 내리기 때문에, +이 테스트들은 정작 검증 대상을 한 번도 거치지 않습니다. 스윕은 이를 별도로 +기록(JSON의 `exercised`)하고 gate에서 제외합니다. 포함시키면 커버리지가 45% +부풀려집니다. + +### 통과한 11개 + +| 테스트 | 시간 | +|---|---| +| `tests/ops/elementwise/test_add.py` | 77.5s | +| `tests/ops/fusion/test_addmm_residual.py` | 33.0s | +| `tests/ops/fusion/test_matmul_scalar.py` | 11.5s | +| `tests/ops/fusion/test_matmul_vector.py` | 17.4s | +| `tests/ops/fusion/test_prologue_fusion.py` | 41.4s | +| `tests/ops/misc/test_expert_mask.py` | 11.0s | +| `tests/ops/reduce/test_batchnorm.py` | 37.7s | +| `tests/ops/view/test_view3D_2D.py` | 36.4s | +| `tests/system/test_eager.py` | 15.0s | +| `tests/system/test_stonne.py` | 9.7s | +| `tests/system/test_triton_codegen.py` | 10.0s | + +11개 중 **4개가 fusion 테스트**입니다. 이 마이그레이션에서 Inductor의 fusion은 +공짜로 얻는 절반인데, 이미 tnpu가 받아들이는 커널을 만들어내고 있다는 뜻입니다. + +--- + +## 2. 실패 하나가 어떻게 진단되는가 — softmax 전 과정 + +`tests/ops/reduce/test_softmax.py`. 스윕이 남기는 것: + +``` +failures/tests_ops_reduce_test_softmax/ + kernel.py Inductor 가 만든 Triton 커널 원본 + error.txt 버킷, 단계, 로그 마지막 60줄 +``` + +`error.txt` 첫 줄이 담당자를 지정합니다: + +``` +test: tests/ops/reduce/test_softmax.py +bucket: triton_helpers +stage: 0 kernel generated, not accepted +``` + +그리고 `kernel.py`가 이유 전부를 보여줍니다: + +```python +@triton.jit +def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr): + xnumel = 64 + r0_numel = 128 + R0_BLOCK: tl.constexpr = 128 + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < xnumel + r0_index = tl.arange(0, R0_BLOCK)[None, :] + r0_1 = r0_index + x0 = xindex + tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) + tmp1 = tl.broadcast_to(tmp0, [XBLOCK, R0_BLOCK]) + tmp3 = tl.where(xmask, tmp1, float("-inf")) + tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <-- 블로커 1 + tmp5 = tmp0 - tmp4 + tmp6 = libdevice.exp(tmp5) # <-- 블로커 2 + tmp7 = tl.broadcast_to(tmp6, [XBLOCK, R0_BLOCK]) + tmp9 = tl.where(xmask, tmp7, 0) + tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) + tmp11 = (tmp6 / tmp10) + tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) +``` + +막힌 곳은 두 군데뿐이고, 아무것도 실행하지 않고 눈으로 확인됩니다. + +- `triton_helpers.max2` — `torch._inductor.runtime`에 있는 모듈인데, tnpu venv는 + 의도적으로 torch가 없습니다. +- `libdevice.exp` — `@core.extern` 인트린식으로 triton_shared 백엔드에 구현이 + 없습니다. + +**주목할 점은 나머지가 전부 멀쩡하다는 것입니다.** 마스크가 붙은 `tl.load`, +`tl.where`, `tl.sum`, `tl.store`, 2차원 `[XBLOCK, R0_BLOCK]` 브로드캐스트가 모두 +문제없이 통과합니다. softmax는 구조적으로 막힌 게 아니라 **함수 호출 두 개**에 +막혔습니다. + +이것이 버그 리포트 전체이고, 재실행 없이 만들어집니다. + +--- + +## 3. 커널이 어디서 멈추는가 + +각 테스트를, 그 테스트의 커널이 산출물을 남긴 가장 깊은 단계에 배치했습니다. +커널이 **도달하지 못한** 단계가 그 실패의 담당자입니다. + +| 단계 | 수 | | +|---|---|---| +| — 커널 생성 전 | 26 | codegen 이전에 torch/Inductor에서 사망 | +| 0 생성 후 거절 | 16 | `kernel_spec`이 기술을 거부 | +| 1 triton → ttir | 1 | | +| 2 ttir → tts/linalg | 1 | triton-shared | +| 4 tnpu lower (DMA, lane, spad) | 6 | | +| 5 trace producer | 3 | | + +**lowering pass는 아직 병목이 아닙니다.** 53개 실패 중 tnpu pass가 IR을 거절한 +것은 **단 2건**입니다. 나머지 34개 실제 블로커는 그보다 앞 — tnpu에 넘겨주는 +우리 쪽 포트, 또는 torch 자체 — 에서 멈춥니다. 다음 작업은 대부분 seam의 우리 +쪽에 있습니다. + +--- + +## 4. 원인별 상세 (근거 포함) + +### `spec_incomplete` — 13개 · 담당: `triton_backend/kernel_spec.py` + +**libdevice 인트린식 (5개).** `@core.extern` 멤버로 triton_shared 구현이 없어 +호출하면 `None`이 반환됩니다. + +| 테스트 | 심볼 | +|---|---| +| `ops/elementwise/test_exponent.py` | `libdevice.exp` | +| `ops/elementwise/test_pointwise.py` | `libdevice.isnan` | +| `ops/elementwise/test_transcendental.py` | `libdevice.tanh` | +| `ops/reduce/test_layernorm.py` | `libdevice.rsqrt` | +| `ops/view/test_floormod_axis_split.py` | `libdevice.rsqrt` | + +이번 세션에서 진단을 고치기 전에는 이것들이 tnpu stage-1 워커 안에서 +`NameError('libdevice is not defined')` 로만 죽었습니다. 그래서 서로 다른 여섯 +개의 lowering 버그처럼 보였습니다. 지금은 이렇게 말합니다: + +``` +SpecIncomplete: kernel calls libdevice.{exp}: those are extern math intrinsics +with no implementation on the triton_shared backend. They need lowering to a +VPU op (or a scalar fallback) before this kernel can compile. +``` + +**다축 grid (4개).** `fixed_config_for`가 가장 바깥 축만 고정하기 때문에 +`YBLOCK`이 `None`이 되고 grid를 계산할 수 없습니다. 알려진 block-size 정책 +공백입니다. + +| 테스트 | 진단 | +|---|---| +| `ops/view/test_transpose2D.py` | axis `y`: ynumel=156, YBLOCK=None | +| `ops/view/test_transpose3D.py` | axis `y`: ynumel=2728, YBLOCK=None | +| `ops/fusion/test_conv_fusion.py` | axis `y`: ynumel=192, YBLOCK=None | +| `ops/conv/test_conv_view_input.py` | axis `y`: ynumel=512, YBLOCK=None | + +**reduction block 미설정 (2개)** — `R0_BLOCK`을 의도적으로 비워둡니다: +`ops/fusion/test_bmm_reduction.py`, `ops/fusion/test_matmul_reduction.py`. + +**진짜 메타데이터 구멍 (1개)** — `ops/misc/test_widen_dtype.py`: `out_ptr0`의 +dtype/numel을 `collect_meta`가 `V.graph`에서 해결하지 못했습니다. + +### `triton_helpers` — 7개 · 담당: `triton_backend` + +| 테스트 | 헬퍼 | +|---|---| +| `ops/reduce/test_softmax.py` | `max2` | +| `ops/sort/test_sort.py` | `sort_with_index` | +| `ops/elementwise/test_activation.py` | `maximum` | +| `ops/conv/test_cnn.py` | `maximum` | +| `ops/fusion/test_matmul_activation.py` | `maximum` | +| `ops/sparsity/test_sparsity.py` | `maximum` | +| `models/test_mlp.py` | `maximum` | + +7개 중 4개가 `maximum` 하나만 필요합니다. pass 수정이 아니라 작은 파일 하나를 +vendoring 하는 작업입니다. + +### `wrapper_gap` — 6개 · 담당: `triton_backend` + +전부 동일합니다: + +``` +AttributeError: 'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak' +``` + +`ops/attention/test_gqa.py`, `test_gqa_decode.py`, +`ops/fusion/test_attention_fusion.py`, `test_transformer_fusion.py`, +`models/Mixtral8x7B/test_attention.py`, `models/test_transformer.py` + +스위트의 **attention·transformer 테스트 전부**가 미구현 메서드 하나에 막혀 +있습니다. + +### `device_op` — 3개 · 담당: `PyTorchSimDevice` + +이 경로 이전부터 있던 문제입니다. MLIR 경로는 이들을 dispatcher 도달 전에 +가로챕니다. + +| 테스트 | 오류 | +|---|---| +| `ops/conv/test_conv2d.py` | `convolution_overrideable not implemented` | +| `ops/conv/test_group_conv.py` | `convolution_overrideable not implemented` | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable not implemented` | + +### `tnpu_stage` — 2개 · 담당: triton-npu lowering pass + +진짜로 pass가 IR을 거절한 유일한 두 건입니다. 산출물에 Python부터 거절된 op까지 +사슬 전체가 남아 있습니다. + +**`ops/conv/test_pool.py`** — stage 1. + +`kernel.py` 끝에, Inductor가 reduction 뒤에 붙이는 무해해 보이는 한 줄: + +```python + tmp4 = tl.sum(tmp3, 1)[:, None].to(tl.float32) + tmp5 = 49.0 + tmp6 = (tmp4 / tmp5) + tl.debug_barrier() # <-- 이것 +``` + +`01-ttir.mlir:46`에서 이렇게 됩니다: + +```mlir +%tmp4_17 = tt.expand_dims %tmp4 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> +%tmp6_18 = arith.divf %tmp4_17, %tmp6 : tensor<128x1xf32> +ttg.barrier all # <-- GPU 다이얼렉트 op +%0 = tt.splat %in_out_ptr0 : !tt.ptr -> tensor<128x1x!tt.ptr> +``` + +그리고 `triton-shared-opt`가 파싱하지 못합니다: + +``` +01-ttir.mlir:46:5: error: Dialect `ttg' not found for custom op 'ttg.barrier' +``` + +`ttg`는 GPU 다이얼렉트입니다. reduction 뒤 pointwise 커널에서 이 타깃에는 의미 +없는 배리어인데, IR에 남아 있어서 파서가 멈춥니다. + +**`ops/reduce/test_reduce.py`** — stage 2. 평범한 `(a + b).sum(dim=1)`: + +```python +tmp0 = tl.load(in_ptr0 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp1 = tl.load(in_ptr1 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp2 = tmp0 + tmp1 +tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK]) +tmp5 = tl.where(r0_mask & xmask, tmp3, 0) +``` + +`01-ttir.mlir`까지는 살아남고, linalg 변환에서 실패합니다: + +``` +error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than + the number of loops (2) of the enclosing LinalgOp +``` + +두 산출물 모두 진단과 문제의 `.mlir`을 함께 갖고 있어, 그대로 업스트림에 넘길 +수 있습니다. + +### `togsim` / `기타` — 5개 + +| 테스트 | 단계 | 내용 | +|---|---|---| +| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed`, exit 255 | +| `ops/misc/test_masked_nondividing.py` | 4 | `[Spike] triton_npu_fused_constant_pad_nd_0 failed` | +| `ops/misc/test_indirect_access.py` | 5 | TOGSim이 `index_put`에 `inf` 사이클 반환 | +| `system/test_hetro.py` | — | `KeyError: 'vpu_num_lanes'` (hetero config에 키 없음) | +| `ops/sparsity/test_sparse_core.py` | — | `TypeError: '>' between Tensor and torch.device` (테스트 쪽 버그) | + +Spike 실패 두 건이 이번 스윕에서 가장 흥미롭습니다. **동작하는 RISC-V 바이너리 +까지 컴파일된 뒤 런타임에 실패하는 유일한 사례**입니다. `test_cat`의 +`04-custom.mlir`을 보면 lowering은 제 역할을 했습니다: + +```mlir +"togsim.transfer"(%reinterpret_cast_5, %c0, %2, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast_4, %c0, %1, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast, %9, %0, %c0, %7, %c0, %c3, %c1, %c0, %13) + {vlane_split_axis = 0, ...} +``` + +입력 DMA 2개와 출력 DMA 1개, 축 0으로 lane split — `cat`이 내야 할 정확한 +모양입니다. **lowering 구조가 아니라 실행이 잘못됩니다.** + +**이 두 건에 대한 단서.** `tnpu.spike`는 `StageError: command failed with exit +code 255`만 보고하고 spike 자신의 stderr는 살아남지 못합니다. 기록된 spike +명령을 같은 workdir에서 손으로 돌리면 exit 0이 나오는데, `write_inputs`가 매 +launch마다 `runtime/*.raw`를 새로 쓰기 때문에 손으로 돌린 실행은 **옛 입력을 +재생**하기 때문입니다. 따라서 실패하는 입력은 현재 파이프라인 밖에서 재현할 수 +없습니다. `TnpuError`에 적용한 것과 같은 방식으로 spike의 stderr를 노출시키는 +것이 이 두 건 진단의 선결 조건이고, 아직 하지 않았습니다. + +### `missing_dep` — 17개 · 경로 문제 아님 + +`transformers`(5), `torchvision`(4), `matplotlib`(4), `pytest`(2), `diffusers`, +`requests`, `sklearn`. 로컬 venv에만 없는 것으로, CI 이미지에서는 실제로 +돌아갑니다. 스윕이 CI에 있어야 하는 이유이기도 합니다. + +--- + +## 5. 인프라 + +### 러너 + +`scripts/ci/triton_route_sweep.py`. `TORCHSIM_TRITON_CODEGEN`은 device 등록 +시점에 한 번 읽히므로(`PyTorchSimDevice/torch_openreg/__init__.py:30`), +**테스트 파일은 하나도 고칠 필요가 없었습니다.** 69개 전부가 이미 이 경로의 +테스트였고, 없던 것은 러너뿐이었습니다. + +세 가지 산출물: + +1. **Gate** — `scripts/ci/triton_route_passing.txt`에 현재 통과 목록. 하나라도 + 깨지면 CI 실패. 커버리지는 이 파일을 재생성해서만 늘어나므로 + (`--update-allowlist`) 조용히 줄어들 수 없습니다. +2. **Report** — 담당 레이어와 파이프라인 단계로 분류. +3. **Artifacts** — 실패 테스트당 디렉토리 하나 (2절 참고). + +더 깊이 간 실패는 더 많이 남깁니다. `tests_ops_conv_test_cnn/`에는 +`01-ttir.mlir 02-ttshared.mlir 03-adapted.mlir 04-custom.mlir kernel.py +stage.log error.txt`가 있습니다 — 멈춘 지점까지의 lowering 사슬 전체입니다. + +### 병렬화 + +테스트는 각자 독립 서브프로세스이고 자기 덤프 디렉토리, Inductor 캐시 +(`TORCHINDUCTOR_CACHE_DIR`가 `TORCHSIM_DUMP_PATH`를 따라감), TOGSim FIFO(pid +기준)를 갖습니다. 따라서 `-j`에 조율이 필요 없습니다. 프로세스가 아니라 +스레드입니다 — `run_one`은 서브프로세스를 기다리기만 합니다. +**69개 기준 약 50분 → `-j 10`에서 5분.** + +### CI + +`.github/workflows/triton_npu.yml`의 `triton-route-suite` 잡: + +- **Allowlisted tests** — gate 역할. +- **Full sweep** — `continue-on-error`. `coverage.md`를 step summary에 쓰고 + `triton-route-coverage`(results.json + failures/)를 업로드. + +잡은 PSAL Slurm 러너 팜(`PSAL-POSTECH/slurm-ghr`)에서 돕니다. `runs-on`에 +`slurm` 라벨이 있어야 하고, 이미지 빌드와 스윕은 `big`(16c/64G/2h), 나머지는 +small 버킷입니다. `docker/setup-buildx-action`은 추가하면 안 됩니다 — 러너가 +자체 빌더를 등록해 둡니다. + +--- + +## 6. 측정하면서 고친 진단 3가지 + +이 셋을 고치기 전에는 보고 인프라를 만들 수 없었습니다. 각각이 증거를 파괴하고 +있었기 때문입니다. + +**`kernel.py`가 그것을 거절하는 검사 뒤에 저장되고 있었습니다.** +`write_spec_file`은 정확히 보존할 가치가 있는 커널(`triton_helpers`, +`SpecIncomplete`)에서 예외를 던지는데, 소스 저장보다 **먼저** 실행됐습니다. +결국 흥미로운 소스일수록 버려지고 있었습니다. 순서를 뒤집었고, 이제 거절된 16개 +커널 전부의 덤프가 남습니다 — 2절의 softmax 예시가 그중 하나입니다. + +**tnpu가 "exit 1"만 보고했습니다.** `run.py`는 stage 표를 stdout에 찍고 진짜 +진단은 `stage.log`에만 씁니다. 이전: + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) +``` + +이후 (`TnpuError`가 `stage.log`를 읽음): + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) + triton.compiler.errors.CompilationError: at 8:11: + NameError('tl_math is not defined') +``` + +이 한 가지 변경으로 여섯 개 실패가 **하나의 버그**로 정리됐습니다: + +**`libdevice`와 `tl_math`가 유탄을 맞고 있었습니다.** `strip_for_tnpu`가 +`from torch...`를 지우는데, Inductor는 이 두 이름을 +`torch._inductor.runtime.triton_helpers`에서 import합니다. 그런데 이들은 torch +코드가 아니라 **triton 자체 심볼의 재수출**입니다. 커널 여섯 개가 stage 1 안에서 +맨 `NameError`로 죽고 있었습니다. + +- `tl_math`는 `triton.language`에서 다시 바인딩했습니다. `test_pointwise`가 첫 + op에서 죽던 것이 op 14개를 지나 trace producer까지 갑니다. +- `libdevice`는 재바인딩이 불가능합니다(멤버가 `@core.extern`이고 triton_shared + 구현이 없어 호출하면 `None`). `triton_helpers`와 같은 방식으로 명시적으로 + 이름을 밝히도록 했습니다. + +순효과: `tnpu_stage` 8 → 2, `spec_incomplete` 7 → 13. **같은 53개가 실패하지만 +그중 6개가 이제 참을 말합니다.** + +**별건으로**, 로컬 TOGSim 빌드가 07-20자여서 `trace_shape.txt` 지원 이전이었고, +그래서 `togsim_kernel`이 `shape_args = nullptr`로 호출되어 모든 Triton 경로 +테스트가 `trace_to_tilegraph`에서 SIGSEGV로 죽었습니다. 재빌드로 해결됐습니다 — +코드 문제가 아니고, CI는 소스에서 빌드하므로 영향이 없었습니다. 다른 사람이 낡은 +`TOGSim/build`를 갖고 있다면 알아둘 만합니다. + +--- + +## 7. 다음 작업 — 해제되는 테스트 수 기준 + +수치는 측정값이지 추정이 아닙니다. 다만 한 단계에서 풀린 테스트가 다음 단계에서 +그냥 다시 실패할 수는 있습니다. + +| # | 작업 | 해제 | 담당 | +|---|---|---|---| +| 1 | `TritonNPUWrapperCodegen.estimate_peak` 구현 | 6 | triton_backend | +| 2 | torch 없는 `triton_helpers`를 tnpu venv에 vendoring | 7 | triton_backend | +| 3 | `libdevice` 인트린식(`exp`, `tanh`, `rsqrt`, `isnan`)을 VPU op로 lowering | 5 | tnpu 또는 triton_backend | +| 4 | `fixed_config_for`에 다축 block 정책 | 4 | triton_backend | +| 5 | `ttg.barrier` + `linalg.index` rank 오류를 업스트림에 전달 | 2 | tnpu | +| 6 | spike stderr 노출 후 `cat` / `constant_pad_nd` 진단 | 2 | triton_backend → 조사 | + +**1번이 압도적으로 쌉니다** — 메서드 하나로 6개, attention/transformer 계열 +전체가 열립니다. + +**2번과 3번은 함께 해야 softmax가 열립니다.** 2절에서 봤듯 softmax는 둘 다 +필요하고, 하나만 고치면 다른 하나에서 계속 실패합니다. + +**3번은 착수 전 결정이 필요합니다**: tnpu pass에서 lowering할 것인가, +`strip_for_tnpu`에서 triton 레벨 polyfill로 대체할 것인가. 전자가 옳고 후자가 +싸며 측정을 더 빨리 풀어줍니다. + +**6번은 선결 조건이 있습니다.** 스위트에서 유일하게 "답이 틀리는" 실패이고 진짜 +lowering 버그일 가능성이 가장 높지만, spike의 stderr가 서브프로세스를 넘어오기 +전에는 진단할 수 없습니다 — 6절에서 `TnpuError`에 이미 적용한 것과 같은 +수정입니다. + +--- + +## 8. 이 측정이 말해주지 않는 것 + +- `missing_dep` 17개는 로컬 venv 사정입니다. CI 이미지에서는 실제로 돌기 때문에 + 버킷이 이동할 것입니다 — 대부분 transformer·CNN 모델이므로 아마 `wrapper_gap` + 과 `triton_helpers` 쪽으로 갑니다. +- 어떤 버킷을 풀면 그 테스트들은 **다음 실패**로 이동하는 것이지, 반드시 통과로 + 가는 것이 아닙니다. +- 이 수치는 6절의 수정이 이미 적용된 상태에서 측정한 것이라, 그 이전 실행과 + 직접 비교할 수 없습니다. diff --git a/docs/triton-route-coverage.md b/docs/triton-route-coverage.md index 4e4a2590..51ce81d3 100644 --- a/docs/triton-route-coverage.md +++ b/docs/triton-route-coverage.md @@ -2,7 +2,8 @@ First measurement of PyTorchSim's existing test suite running through the Triton codegen route (Inductor's Triton backend + the triton-npu lowering passes) -instead of the MLIR route. +instead of the MLIR route. Korean version: +[`triton-route-coverage.ko.md`](triton-route-coverage.ko.md). | | | |---|---|