diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml index 05c6f893..c1d2ffde 100644 --- a/.github/workflows/triton_npu.yml +++ b/.github/workflows/triton_npu.yml @@ -8,9 +8,10 @@ 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. # -# 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. +# Needs secrets.TNPU_TOKEN: a PAT that can read PSAL-POSTECH/triton-npu (and its +# toolchain-llvm23 release) plus every repo in the manifest's `also_reads`. They +# are private and the default Actions token is scoped to this repository; +# preflight checks each before the docker build. on: pull_request: @@ -62,6 +63,15 @@ jobs: 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 + # restore.sh clones these too; without them the failure is deep in the + # image build instead of here. + for R in $(jq -r '.triton_npu.also_reads[]?' thirdparty/triton-npu.json); do + if ! curl -fsS -H "Authorization: Bearer ${TNPU_TOKEN}" \ + "https://api.github.com/repos/${R}" -o /dev/null; then + echo "::error::TNPU_TOKEN cannot read ${R}, which restore.sh clones." + echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1 + fi + done echo "ready=true" >> "$GITHUB_OUTPUT" ensure-tnpu-base: diff --git a/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py b/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py index 27a47357..45b10c9f 100644 --- a/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py +++ b/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py @@ -21,7 +21,8 @@ def synchronize(self) -> str: return "pass" def device_guard(self, device_idx: int) -> str: - return "pass" + # The caller writes `with {this}:`, so "pass" is a SyntaxError. + return "torch._ops.contextlib.nullcontext()" register_device_op_overrides("npu", ExtensionDeviceOpOverrides()) register_device_op_overrides("cpu", CpuDeviceOpOverrides()) \ No newline at end of file diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 71927cf6..bd0c6fcc 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -215,7 +215,9 @@ def _generate_kernel_call_helper( original_fxnode_name=None, ): device = device or V.graph.get_current_device_or_throw() - self.writeline(self.wrap_kernel_call(kernel_name, call_args)) + # A template kernel passes sympy constants; wrap_kernel_call joins strings. + self.writeline(self.wrap_kernel_call( + kernel_name, [str(a) for a in call_args])) return def generate(self, is_inference): @@ -225,7 +227,8 @@ def generate(self, is_inference): self._fverify_seen = set() with contextlib.ExitStack() as stack: stack.enter_context(self.wrapper_call.indent()) - self.memory_plan_reuse() + # Upstream entry point: picks the planner and sets the state it needs. + self.run_wrapper_ir_passes(is_inference) with self.set_writeline(self.wrapper_call.writeline): for line in self.lines: # Add buffer plan hook for dealloc @@ -238,7 +241,8 @@ def generate(self, is_inference): if isinstance(line, wrapper.MemoryPlanningLine): line.codegen(self.wrapper_call) elif isinstance(line, wrapper.KernelCallLine): - self.wrapper_call.writeline(self.wrap_kernel_call(line.kernel_name, line.call_args)) + self.wrapper_call.writeline(self.wrap_kernel_call( + line.kernel_name, [str(a) for a in line.call_args])) if _func_verify.enabled(): self._fverify_emit_checks(line.call_args) else: diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py index 460cf925..cb3b88cd 100644 --- a/PyTorchSimFrontend/triton_backend/__init__.py +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -27,11 +27,12 @@ still owes; see README.md for the gap list. Expect failures, not results. """ -from . import _triton_compat +from . import _triton_compat, inductor_templates # Before anything imports Inductor's Triton codegen: it needs `triton` in THIS # interpreter, and on a GPU-less box its backend hash cannot be computed. _triton_compat.install() +inductor_templates.install() from .scheduling import TritonNPUScheduling # noqa: E402,F401 from .wrapper_codegen import TritonNPUWrapperCodegen # noqa: E402,F401 diff --git a/PyTorchSimFrontend/triton_backend/functional.py b/PyTorchSimFrontend/triton_backend/functional.py index 4dd15682..3ba91daa 100644 --- a/PyTorchSimFrontend/triton_backend/functional.py +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -65,6 +65,24 @@ def _check(meta, pairs): f"binary was compiled for {m['dtype']}") +def _storage_view(t, m): + """`t`'s values laid out the way the kernel indexes them, as a flat tensor. + + Inductor allocates with empty_strided and indexes by that stride, so the + element the kernel calls `k` lives at storage position `k` -- which is not + logical order unless the layout is contiguous. + """ + import torch + + size, stride = m.get("size"), m.get("stride") + flat = torch.empty(m["numel"], dtype=t.dtype) + if size and stride: + flat.as_strided(size, stride).copy_(t) + else: + flat.copy_(t.reshape(-1)) + return flat + + def write_inputs(workdir, meta, args): """Write every arg as runtime/.raw. Returns the runtime directory. @@ -81,7 +99,7 @@ def write_inputs(workdir, meta, args): 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) + _storage_view(t.detach().to("cpu"), m).numpy().tofile(path) else: np.zeros(m["numel"], dtype=_np_dtype(m["dtype"])).tofile(path) return runtime @@ -103,11 +121,23 @@ def read_outputs(workdir, meta, args): 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)) + buf = torch.from_numpy(flat).to(t.dtype) + size, stride = m.get("size"), m.get("stride") + t.copy_(buf.as_strided(size, stride) if size and stride + else buf.view_as(t)) written.append(m["name"]) return written +def _stage_log(workdir): + """tnpu's per-stage log, where every subprocess it runs leaves its output.""" + path = os.path.join(workdir, "stage.log") + if not os.path.isfile(path): + return "" + with open(path, errors="replace") as f: + return f.read() + + def run(workdir, meta, args, timeout_sec=None): """Execute the kernel on the launch's tensors. Returns the names written.""" from . import tnpu_bridge @@ -125,8 +155,10 @@ def run(workdir, meta, args, timeout_sec=None): 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:]) + raise tnpu_bridge.TnpuError( + f"[Spike] {meta['kernel_name']} failed", + cmd=" ".join([extension_config.CONFIG_TNPU_PYTHON, "-m", + "tnpu.spike", spec, workdir]), + output=proc.stdout + proc.stderr + "\n" + _stage_log(workdir)) return read_outputs(workdir, meta, args) diff --git a/PyTorchSimFrontend/triton_backend/helpers_shim.py b/PyTorchSimFrontend/triton_backend/helpers_shim.py new file mode 100644 index 00000000..f5440f1d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/helpers_shim.py @@ -0,0 +1,79 @@ +"""Make torch's `triton_helpers` importable from the torch-free tnpu venv. + +Inductor's kernels call `triton_helpers.maximum`, `.max2`, `.sort_with_index` +and friends. The module is inside torch, and tnpu's venv deliberately has none, +so a kernel that uses one cannot compile there. + +It does not have to be rewritten. `triton_helpers.py` imports nothing from +torch -- only `.triton_compat`, and that module touches torch in three places +(`torch.version.hip` twice, `torch.autograd.profiler` once), none of which +`triton_helpers` needs. So the file is copied verbatim next to the kernel and +paired with the small `triton_compat` below, which resolves the same seven +symbols straight from triton. + +Copying the installed torch's file rather than vendoring a snapshot keeps the +helpers matched to the torch that generated the kernel. +""" + +import os +import shutil + +PACKAGE = "tnpu_helpers" + +#: Stands in for torch._inductor.runtime.triton_compat, whose only job here is +#: to resolve these seven names. Mirrors the upstream fallbacks so a triton +#: version change lands the same way on both sides. +_COMPAT = '''\ +"""Generated by PyTorchSimFrontend/triton_backend/helpers_shim.py. Do not edit.""" + +import inspect +from typing import Any + +import triton +import triton.language as tl +from triton.runtime.jit import JITFunction # noqa: F401 + +try: + from triton.language.extra import libdevice # noqa: F401 + + libdevice = tl.extra.libdevice # noqa: F811 + math = tl.math +except ImportError: + if hasattr(tl.extra, "cuda") and hasattr(tl.extra.cuda, "libdevice"): + libdevice = tl.extra.cuda.libdevice + math = tl.math + elif hasattr(tl.extra, "intel") and hasattr(tl.extra.intel, "libdevice"): + libdevice = tl.extra.intel.libdevice + math = tl.math + else: + libdevice = tl.math + math = tl + +try: + from triton.language.standard import _log2 +except ImportError: + + def _log2(x: Any) -> Any: + raise NotImplementedError + +builtins_use_semantic_kwarg = ( + "_semantic" in inspect.signature(triton.language.core.view).parameters +) +''' + + +def _source_path(): + from torch._inductor.runtime import triton_helpers + return triton_helpers.__file__ + + +def write_package(dest_dir): + """Write the importable package next to a kernel. Returns the import line.""" + pkg = os.path.join(dest_dir, PACKAGE) + os.makedirs(pkg, exist_ok=True) + with open(os.path.join(pkg, "__init__.py"), "w"): + pass + with open(os.path.join(pkg, "triton_compat.py"), "w") as f: + f.write(_COMPAT) + shutil.copyfile(_source_path(), os.path.join(pkg, "triton_helpers.py")) + return f"from {PACKAGE} import triton_helpers\n" diff --git a/PyTorchSimFrontend/triton_backend/inductor_templates.py b/PyTorchSimFrontend/triton_backend/inductor_templates.py new file mode 100644 index 00000000..eb1b9712 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/inductor_templates.py @@ -0,0 +1,112 @@ +"""Let Inductor's mm/conv Triton templates reach this backend. + +Without this they go to `extern_kernels.*`, which on npu either raises +`convolution_overrideable not implemented` or falls back to eager and simulates +nothing. The templates themselves are not GPU-specific -- torch ships one +`triton_mm.py.jinja` for cuda, xpu, mtia and cpu -- but `use_triton_template` +gates on `is_gpu`, and GPU_TYPES is a hardcoded list with no registration hook. +""" + +import os + +import torch + + +def _register_npu_as_gpu(): + import torch._inductor.utils as inductor_utils + + if "npu" not in inductor_utils.GPU_TYPES: + inductor_utils.GPU_TYPES.append("npu") + + +def _claim_triton_present(): + # has_triton() asks whether a supported *device* is available, not whether + # triton is installed. The missing piece is a driver we never use. + import torch.utils._triton as triton_utils + import torch._inductor.scheduler as scheduler + + triton_utils.has_triton = lambda: True + if hasattr(scheduler, "has_triton"): + scheduler.has_triton = lambda: True + + +def _register_template_heuristics(): + from torch._inductor.kernel.bmm import bmm_template + from torch._inductor.kernel.mm import mm_template + from torch._inductor.template_heuristics.registry import ( + register_template_heuristic) + from torch._inductor.template_heuristics.triton import ( + AddMMConfigMixin, BaseConfigHeuristic, MMTemplateConfigMixin) + + @register_template_heuristic(mm_template.uid, "npu") + @register_template_heuristic(bmm_template.uid, "npu") + class NPUMMTemplateConfigHeuristic(MMTemplateConfigMixin, BaseConfigHeuristic): + # TODO: size these from the hardware config (lanes, spad per lane) + # rather than taking the generic set. + def __init__(self): + super().__init__() + self.exhaustive_configs = self.mm_configs + + # addmm and baddbmm carry a bias as input_nodes[0]; without their own entry + # the mm heuristic is used with prefix_args=0 and def_kernel asserts. + @register_template_heuristic(mm_template.uid, "npu", op_name="addmm") + @register_template_heuristic(bmm_template.uid, "npu", op_name="baddbmm") + class NPUAddmmTemplateConfigHeuristic(AddMMConfigMixin, + NPUMMTemplateConfigHeuristic): + pass + + +def pick_config(choices): + """Stand in for benchmarking: there is no device to time on. + + TODO: rank by simulated cycles. `timing.run_togsim` already returns a cycle + count per compiled kernel; a real implementation drives each candidate + through tnpu and caches the result per (kernel, config). Until then the + offered order wins -- deterministic, and not a claim about speed. + """ + return {c: 1.0 + i * 1e-3 for i, c in enumerate(choices)} + + +def _install_selection(): + from torch._inductor.select_algorithm import AlgorithmSelectorCache + + def benchmark_choices(cls, choices, autotune_args, is_collective=False): + return pick_config(choices) + + # Precompiling builds every candidate for the current GPU. We need only the + # chosen kernel's source; tnpu compiles it ahead of time. + AlgorithmSelectorCache.benchmark_choices = classmethod(benchmark_choices) + AlgorithmSelectorCache.make_precompile_fn = lambda self, *a, **k: (lambda: None) + + +_installed = False + + +def install(): + """On by default; TORCHSIM_TRITON_TEMPLATES=0 opts out. + + Sending mm to aten is not a working state -- the op is not simulated at all + -- so the templates are the default and the tests that now stop at + tl.assume in tnpu (PSAL-POSTECH/triton-npu#2) say so. + """ + global _installed + if _installed or os.environ.get("TORCHSIM_TRITON_TEMPLATES", "1") == "0": + return + from torch._inductor import config + + _register_npu_as_gpu() + _claim_triton_present() + _register_template_heuristics() + _install_selection() + + # max_autotune_gemm, not max_autotune: the latter also turns on pointwise + # autotuning, which appends a benchmark harness to every kernel module and + # breaks the ones that were already working. + config.max_autotune_gemm = True + config.max_autotune_gemm_backends = "TRITON" + config.max_autotune_conv_backends = "TRITON" + config.triton.autotune_at_compile_time = False + # Epilogue-fusion benchmarking renders a benchmark-flavoured kernel whose + # harness imports land indented in the real module. + config.benchmark_epilogue_fusion = False + _installed = True diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index e5a34400..fdcfeb5f 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -40,6 +40,8 @@ def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): from torch._inductor.virtualized import V +from . import helpers_shim + #: Triton signature token -> (torch dtype name, bytes). Only the dtypes #: tnpu/wrapper.py can round-trip through .raw files. _DTYPE = { @@ -64,19 +66,27 @@ class SpecIncomplete(RuntimeError): # --------------------------------------------------------------------------- # 1. codegen-time metadata capture # --------------------------------------------------------------------------- -def _buffer_numel(name): - """Element count of an Inductor buffer, or None if it cannot be resolved.""" +def _buffer_layout(name): + """(numel, size, stride) of an Inductor buffer; Nones if unresolvable. + + The stride is load-bearing: Inductor allocates outputs `empty_strided` and + indexes them by it, so a launch that assumes contiguous writes the elements + to the wrong places. + """ try: buf = V.graph.get_buffer(name) if buf is None: - return None - size = buf.get_layout().size + return None, None, None + layout = buf.get_layout() + hint = V.graph.sizevars.size_hint + size = [int(hint(s)) for s in layout.size] + stride = [int(hint(s)) for s in layout.stride] n = 1 for s in size: - n *= int(V.graph.sizevars.size_hint(s)) - return n + n *= s + return n, size, stride except Exception: # noqa: BLE001 - best effort; caller reports it as missing - return None + return None, None, None def _roles(kernel): @@ -85,7 +95,10 @@ def _roles(kernel): 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) + # A graph input here is mutated, not produced: unwritten elements must + # survive, so it has to be seeded. + role = "inout" if buf in getattr(V.graph, "graph_inputs", {}) else "out" + out[arg] = (role, buf) for buf, arg in getattr(kernel.args, "inplace_buffers", {}).items(): name = getattr(arg, "inner_name", arg) out[name] = ("inout", buf) @@ -110,12 +123,15 @@ def collect_meta(kernel, kernel_name): role, buf = roles.get(name, (None, None)) if role is None: continue # a numel / constexpr, not a tensor + numel, size, stride = _buffer_layout(buf) if buf else (None, None, None) args.append({ "name": name, "role": role, "buffer": buf, "dtype": _DTYPE.get(signature.get(name, ""), None), - "numel": _buffer_numel(buf) if buf else None, + "numel": numel, + "size": size, + "stride": stride, }) # The numels Inductor appends to the call. They live in `kernel.numels`, @@ -151,10 +167,19 @@ def _block_name(prefix): def parallel_axes(numels): - """Grid axes this kernel uses, outermost first.""" + """Grid axes this kernel uses, outermost first. For tile-shape decisions.""" return [p for p in _PARALLEL_PREFIXES if f"{p}numel" in numels] +def pid_axes(numels): + """The same axes in program-id order: pid 0 is x, whatever the tiling. + + Every grid TUPLE is in this order -- triton-shared lays the pid arguments + out x, y, z and tnpu's wrapper reads spec.grid positionally as gridX/Y/Z. + """ + return list(reversed(parallel_axes(numels))) + + def fixed_config_for(kernel): """Block sizes pinned at codegen time. @@ -171,9 +196,18 @@ def fixed_config_for(kernel): README, not something to guess at here. """ from PyTorchSimFrontend import extension_config - lanes = int(extension_config.vpu_num_lanes) - - axes = parallel_axes(getattr(kernel, "numels", None) or {}) + try: + lanes = int(extension_config.vpu_num_lanes) + except KeyError: + raise SpecIncomplete( + f"{extension_config.CONFIG_TOGSIM_CONFIG} has no vpu_num_lanes. " + f"This route pins every block size to the lane count, so a config " + f"without a VPU cannot describe a launch shape.") from None + + # kernel.numels is keyed by prefix; parallel_axes wants collect_meta's + # "numel" keys, and passing the raw dict silently matched nothing. + axes = parallel_axes([f"{p}numel" + for p in (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 @@ -185,11 +219,13 @@ def fixed_config_for(kernel): 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 - # unset on purpose so the reduction path fails loudly rather than - # silently picking a layout the hardware cannot execute. - cfg["R0_BLOCK"] = None + # The whole reduced extent, so the kernel's r0 loop runs once. Whether + # that tile fits the lanes is the lowering pass's call, not ours. + r0 = (getattr(kernel, "numels", None) or {}).get("r0_") + try: + cfg["R0_BLOCK"] = int(V.graph.sizevars.size_hint(r0)) + except Exception: # noqa: BLE001 - dynamic; write_spec_file reports it + cfg["R0_BLOCK"] = None return cfg @@ -199,13 +235,13 @@ def fixed_config_for(kernel): _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+)") +#: Calls that only mean something on a GPU. set_driver_to_gpu picks a runtime we +#: never launch through; debug_barrier works around a triton warp-scheduling bug +#: (triton-lang/triton#1615) and tnpu replays one work-item at a time, so there +#: are no warps to order. It also reaches ttir as ttg.barrier, a GPU-dialect op +#: triton-shared-opt cannot parse. +_DROP_CALL_RE = re.compile( + r"^\s*(triton_helpers\.set_driver_to_gpu|tl\.debug_barrier)\(\)") def strip_for_tnpu(src): @@ -235,33 +271,13 @@ def strip_for_tnpu(src): 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" - # 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.") + prefix = ( + "import triton\n" + "import triton.language as tl\n" + "from triton.language import math as tl_math\n" + "from triton.language.extra import libdevice\n" + f"from {helpers_shim.PACKAGE} import triton_helpers\n\n" + ) return prefix + body @@ -291,14 +307,14 @@ def scalar_args(meta): def grid_of(meta): - """Launch grid, from the numels and the pinned block sizes, outermost first. + """Launch grid, from the numels and the pinned block sizes, in pid order. 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. """ numels = meta["numels"] cfg = meta.get("fixed_config") or {} - axes = parallel_axes(numels) + axes = pid_axes(numels) if not axes: raise SpecIncomplete( f"{meta['kernel_name']} has no parallel iteration axis to grid over") @@ -328,6 +344,9 @@ def grid_of(meta): import sys sys.path.insert(0, {tnpu_dir!r}) +# This directory too: the kernel is loaded by path, so a sibling package +# (tnpu_helpers) would not otherwise be importable from it. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from tnpu.spec import KernelSpec, Arg # noqa: E402 #: The rewritten Triton source, beside this file. It must be a REAL file on @@ -412,8 +431,11 @@ def write_spec_file(src_code, meta, path, tnpu_dir): for a in meta["args"] if a["role"] in ("in", "inout")) or " pass" triton_module = f"{meta['kernel_name']}_triton.py" + stripped = strip_for_tnpu(src_code) + if helpers_shim.PACKAGE in stripped: + helpers_shim.write_package(os.path.dirname(path)) with open(os.path.join(os.path.dirname(path), triton_module), "w") as f: - f.write(strip_for_tnpu(src_code)) + f.write(stripped) scalars = scalar_args(meta) text = SPEC_TEMPLATE.format( diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index be28060b..e6bc7c44 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -92,17 +92,15 @@ def _runtime_arg_layout(meta): def work_item_for(meta): """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. + Axes are in pid order, the same as every grid tuple, so extent i and + parallel_args[i] describe the same axis. """ from PyTorchSimFrontend.mlir.passes.lower_to_emitc import WorkItem from . import kernel_spec 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"]) + axes = kernel_spec.pid_axes(meta["numels"]) # 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], @@ -128,7 +126,7 @@ def write_shape(workdir, meta, args=()): for key, val in zip(passed, trailing[-len(passed):]): numels[key] = val - axes = kernel_spec.parallel_axes(numels) + axes = kernel_spec.pid_axes(numels) cfg = meta.get("fixed_config") or {} grid = [] diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py index a30593d1..9599423d 100644 --- a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -32,7 +32,8 @@ class TnpuError(RuntimeError): _SIGNAL = re.compile( r"^(?!\s|Traceback|During handling|The above)" r"(.*\berror:\s.*|.*failed to legalize.*|" - r"[\w.]*(?:Error|Exception)\b.*|.*Assertion.*)$", re.M) + r"[\w.]*(?:Error|Exception)\b.*|.*Assertion.*|" + r".*(?:segfault|Illegal instruction|trap_|bad --isa).*)$", re.M) #: Frames and carets: context, not the diagnostic. _FRAME = re.compile(r'^\s|^\s*File "|^\s*\^') diff --git a/docs/toolchain-image-proposal.md b/docs/toolchain-image-proposal.md new file mode 100644 index 00000000..cde3e5ce --- /dev/null +++ b/docs/toolchain-image-proposal.md @@ -0,0 +1,125 @@ +# Proposal: build the tnpu toolchain in the image, not into it + +Status: proposal. Nothing here is implemented. + +## The claim + +`torchsim_tnpu_base` is a Docker image whose build step downloads a 1.8 GiB +build cache from a private GitHub release and unpacks it to absolute paths. A +Docker image layer *is* a build cache. We are running a second, hand-maintained +cache inside the first one, and every problem in this area comes from that. + +## What exists today + +``` +PSAL-POSTECH/triton-npu + setup/versions.env 8 pins, maintained by hand + setup/restore.sh ~300 lines of bash + release toolchain-llvm23 llvm23-install.tar.gz 1.18 GiB + triton-runtime.tar.gz 0.34 GiB + spike-install.tar.gz 0.01 GiB + MANIFEST.txt what the assets were built from, + written by hand +PyTorchSim + Dockerfile.tnpu runs restore.sh --prebuilt inside the image build + thirdparty/triton-npu.json pins all of the above as one commit sha +``` + +The assets are produced by `setup/package.sh` on a developer machine and +uploaded to the release by hand. + +## Four costs, all observed + +**1. The cache is maintained by a human.** `versions.env` carries +`TRITON_SHARED_PREBUILT_SHA`, whose comment records that it "had gone stale +twice over" — it is a note about what the uploaded tarball contains, and +nothing checks it. `MANIFEST.txt` is the same shape of problem. + +**2. Absolute paths make unpacking a restore, not a copy.** The triton install +is editable, so `__editable__*_finder.py` holds `/workspace/triton-src/python/ +triton` and the venv shebangs hold `/workspace/mlir-env/bin/python`. A tarball +must land exactly where it was built. Renaming one directory (2026-08-03, +`/workspace/triton` → `/workspace/triton-src`, because `triton` shadowed the +package for anything run from `/workspace`) required patching the editable +finder, re-linking the backend symlinks, repackaging the 340 MB tarball and +re-uploading it, and keeping the old asset under a `.flat-triton` name. + +**3. `--prebuilt` turns off more than it needs to.** It sets `STEPS=(layout)`, +which skips both `tritonenv` (the real build, correctly skipped) and +`tritonshared` (a git clone: 1 second, 6 MB). To compensate, `restore.sh` grew +a special case that fetches three files — `backend/{compiler.py,driver.py, +name.conf}` — from a *different repository at a different commit* +(`facebookincubator/triton-shared` at `TRITON_SHARED_BACKEND_SHA`), because +`raw.githubusercontent` cannot serve our private fork. Twelve lines of comment +defend the assumption that "every fork commit touches only include/ and lib/". +On 2026-08-04 a fork commit touched `backend/compiler.py`; the assumption broke +and the prebuilt path silently kept using upstream's older file. + +**4. Pins multiply.** Three variables name the same repository — +`TRITON_SHARED_SHA`, `TRITON_SHARED_BACKEND_SHA`, `TRITON_SHARED_PREBUILT_SHA` +— and only the first is authoritative. + +## The proposal + +Build the toolchain in `Dockerfile.tnpu` from source, and let the image layer be +the cache. + +```dockerfile +ARG LLVM_SHA=... +ARG TRITON_SHA=df38505... +ARG TRITON_SHARED_SHA=... +ARG SPIKE_SHA=... + +RUN git clone --depth 1 ... && cmake ... && ninja install # llvm 23 +RUN TRITON_PLUGIN_DIRS=/src/triton_shared pip install -e /src/triton +RUN ... spike, riscv-pk +``` + +What disappears: + +| | | +|---|---| +| `setup/restore.sh` | ~300 lines | +| release assets + `package.sh` + manual upload | 1.8 GiB, and the step that produces it | +| `MANIFEST.txt`, `TRITON_SHARED_PREBUILT_SHA` | records of what a tarball holds | +| `TRITON_SHARED_BACKEND_SHA` + the three-file fetch | the special case in cost 3 | +| "restore to absolute paths" | the paths are only ever inside the image | +| `GITHUB_TOKEN` for release assets | still needed for the private fork clone | + +What it costs: the first build of a given pin set compiles LLVM 23. Every later +build with the same pins is a registry pull. `ensure-tnpu-base` already skips +the build when the tag exists (`docker manifest inspect`), so the trigger is +unchanged — only the miss is more expensive. + +## The open question: walltime + +Jobs run on the PSAL Slurm farm, where `big` is 16 cores with a 2 hour +`--time`. An LLVM 23 build with MLIR will not finish in that. Options, in the +order I would try them: + +1. **A dedicated bucket** with a longer `--time` for this one job. Config-only + change on the runner farm (`~/.ghr/config.toml`); nothing in this repo moves. +2. **Split the image**: one layer per component (`llvm23` / `triton` / `spike`), + each its own tag and its own job. Each fits in 2 hours, and a pin move + rebuilds only its own layer instead of all of them. +3. **Keep LLVM prebuilt, build the rest.** LLVM is the only genuinely long one + and its pin almost never moves; triton + triton_shared + spike are minutes. + This keeps one tarball and deletes the other two, plus the whole special case + in cost 3. + +Option 3 is the smallest step that removes most of the pain, and is worth +measuring before committing to 1 or 2. + +## What this proposal is not + +It does not change what the toolchain *is* — same LLVM, same triton pin, same +passes. It changes only how the environment is assembled, and it should be +judged on whether the four costs above are worth ~300 lines of bash and a manual +upload step. + +## Sequencing + +This should not block the libdevice fix. That fix needs the three-file special +case gone, which is a ~15 line change to `restore.sh` (re-enable the +`tritonshared` step, delete the fetch block). Doing that first makes the +proposal smaller, not larger: it removes cost 3 on its own. diff --git a/docs/triton-route-coverage.ko.md b/docs/triton-route-coverage.ko.md index 86cb2ed8..f11fec91 100644 --- a/docs/triton-route-coverage.ko.md +++ b/docs/triton-route-coverage.ko.md @@ -1,461 +1,183 @@ -# Triton codegen route 커버리지 측정 보고서 +# Triton codegen route 커버리지 -기존 PyTorchSim 테스트 스위트를 MLIR 경로가 아니라 **Triton 경로**(Inductor의 -Triton 백엔드 + triton-npu lowering pass)로 돌린 첫 측정 결과입니다. +기존 테스트 스위트를 MLIR 경로가 아니라 **Triton 경로**(Inductor의 Triton 백엔드 ++ triton-npu lowering pass)로 돌린 결과입니다. | | | |---|---| -| 측정일 | 2026-08-03 | -| 브랜치 | `feature/triton-codegen` @ `6e3bd7e` | -| tnpu 핀 | `5d84caf` | +| 측정일 | 2026-08-04 | +| 브랜치 | `feature/triton-helpers` @ `7899a17` | +| tnpu 핀 | `d46995f` | | 환경 | 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 \ +python scripts/ci/triton_route_sweep.py --all -j 8 \ --markdown coverage.md --artifacts failures ``` -아래 모든 주장은 `failures/` 아래 실제 파일로 뒷받침됩니다. 확인할 수 있도록 -경로를 함께 적었습니다. +아래 모든 주장은 `failures/` 아래 파일로 뒷받침됩니다. --- -## 1. 결론부터 +## 1. 결론 ``` -69개 테스트 -├── 11 경로를 타고 통과 ← 이것이 커버리지 수치 -├── 5 통과하지만 경로 미사용 ← 커널을 아예 안 만듦 -└── 53 실패 - ├── 17 로컬 venv 패키지 없음 (CI 이미지에는 있음) - └── 36 실제 블로커 +69개 +├── 13 경로를 타고 통과 ← 커버리지 수치 +├── 2 통과하지만 경로 미사용 +└── 54 실패 + ├── 17 로컬 패키지 부재 (CI 이미지에는 있음) + └── 37 실제 블로커 ``` -**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% -부풀려집니다. +**37개 중 26개가 tnpu 버그 두 건입니다.** 서른일곱 개의 문제가 아니라 두 개이고, +재현 커널과 함께 `PSAL-POSTECH/triton-npu#2`에 보고돼 있습니다. -### 통과한 11개 +### 통과한 13개 | 테스트 | 시간 | |---|---| -| `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가 받아들이는 커널을 만들어내고 있다는 뜻입니다. +| `ops/elementwise/test_activation.py` | 83.7s | +| `ops/elementwise/test_add.py` | 79.7s | +| `ops/elementwise/test_exponent.py` | 10.4s | +| `ops/elementwise/test_transcendental.py` | 33.7s | +| `ops/misc/test_expert_mask.py` | 10.0s | +| `ops/reduce/test_batchnorm.py` | 37.6s | +| `ops/sparsity/test_sparse_core.py` | 15.2s | +| `ops/view/test_transpose2D.py` | 32.3s | +| `ops/view/test_transpose3D.py` | 182.6s | +| `ops/view/test_view3D_2D.py` | 38.6s | +| `system/test_eager.py` | 15.8s | +| `system/test_stonne.py` | 10.0s | +| `system/test_triton_codegen.py` | 10.4s | + +`test_topk`과 `test_mlir_bindings`는 커널을 안 만들고 통과합니다. 따로 기록하고 +gate에서 뺐습니다 — 세면 아무것도 아닌 것을 세는 셈입니다. --- -## 2. 실패 하나가 어떻게 진단되는가 — softmax 전 과정 +## 2. 원인별 -`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 테스트 전부**가 미구현 메서드 하나에 막혀 -있습니다. +| `tnpu_stage` | 30 | triton-npu | +| `missing_dep` | 17 | 테스트 환경 (CI 이미지엔 있음) | +| `wrong_values` | 3 | 수치 — 조사 필요 | +| `spec_incomplete` | 2 | triton_backend | +| `device_op` | 1 | PyTorchSimDevice | +| `other` | 1 | 미분류 | -### `device_op` — 3개 · 담당: `PyTorchSimDevice` +### `tnpu_stage` 30건 — 그중 26건이 두 버그 -이 경로 이전부터 있던 문제입니다. MLIR 경로는 이들을 dispatcher 도달 전에 -가로챕니다. +**`tl.assume` → `llvm.intr.assume` (16건).** Inductor의 mm 템플릿이 +`tl.assume(pid_m >= 0)`를 힌트로 냅니다. 이게 `triton-shared-opt`가 로드하지 않는 +다이얼렉트의 op가 됩니다. ttir에서 그 줄만 지우면 같은 stage 2 명령이 성공하고 +`linalg.matmul`이 나옵니다. -| 테스트 | 오류 | -|---|---| -| `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> +models/test_mlp.py ops/fusion/test_matmul_activation.py +models/test_transformer.py ops/fusion/test_matmul_reduction.py +ops/attention/test_gqa.py ops/fusion/test_matmul_scalar.py +ops/conv/test_conv2d.py ops/fusion/test_matmul_vector.py +ops/fusion/test_addmm_residual.py ops/fusion/test_prologue_fusion.py +ops/fusion/test_attention_fusion.py ops/fusion/test_transformer_fusion.py +ops/fusion/test_bmm_reduction.py ops/gemm/test_bmm.py +ops/sparsity/test_sparsity.py ops/gemm/test_matmul.py ``` -그리고 `triton-shared-opt`가 파싱하지 못합니다: +**matmul 결과에 대한 `select_lane_axis` (10건).** ``` -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) +lane_axis.Fatal: the demand 'linalg.matmul' has to move for the lane axes to +agree, but it is about a value 'linalg.matmul' PRODUCES -- a relayout makes a +new value, so there is no edge to put one on ``` -`01-ttir.mlir`까지는 살아남고, linalg 변환에서 실패합니다: +그 패스 자체 문서의 case 4입니다. 앞의 것과 달리 흘러들어온 힌트가 아니라 +**미지원 데이터플로**로 보입니다. ``` -error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than - the number of loops (2) of the enclosing LinalgOp +models/Mixtral8x7B/test_attention.py ops/elementwise/test_pointwise.py +ops/attention/test_gqa_decode.py ops/fusion/test_conv_fusion.py +ops/conv/test_cnn.py ops/misc/test_indirect_access.py +ops/conv/test_conv_view_input.py ops/misc/test_masked_nondividing.py +ops/conv/test_group_conv.py ops/view/test_cat.py ``` -두 산출물 모두 진단과 문제의 `.mlir`을 함께 갖고 있어, 그대로 업스트림에 넘길 -수 있습니다. +**`linalg.index` rank (3건)** — `ops/conv/test_pool.py`, +`ops/reduce/test_reduce.py`, `system/test_vectorops.py`. 그 PR의 원 리포트입니다. -### `togsim` / `기타` — 5개 +**`linalg.generic` shape (1건)** — `ops/sort/test_sort.py`. -| 테스트 | 단계 | 내용 | -|---|---|---| -| `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 구조가 아니라 실행이 잘못됩니다.** +### `wrong_values` 3건 -**이 두 건에 대한 단서.** `tnpu.spike`는 `StageError: command failed with exit -code 255`만 보고하고 spike 자신의 stderr는 살아남지 못합니다. 기록된 spike -명령을 같은 workdir에서 손으로 돌리면 exit 0이 나오는데, `write_inputs`가 매 -launch마다 `runtime/*.raw`를 새로 쓰기 때문에 손으로 돌린 실행은 **옛 입력을 -재생**하기 때문입니다. 따라서 실패하는 입력은 현재 파이프라인 밖에서 재현할 수 -없습니다. `TnpuError`에 적용한 것과 같은 방식으로 spike의 stderr를 노출시키는 -것이 이 두 건 진단의 선결 조건이고, 아직 하지 않았습니다. +`ops/reduce/test_softmax.py`, `ops/reduce/test_layernorm.py`, +`ops/view/test_floormod_axis_split.py`. 컴파일되고 실행되는데 **답이 틀립니다.** +가장 작은 범주이면서 가장 위험합니다. 앞의 둘은 reduction입니다. -### `missing_dep` — 17개 · 경로 문제 아님 +### 나머지 4건 -`transformers`(5), `torchvision`(4), `matplotlib`(4), `pytest`(2), `diffusers`, -`requests`, `sklearn`. 로컬 venv에만 없는 것으로, CI 이미지에서는 실제로 -돌아갑니다. 스윕이 CI에 있어야 하는 이유이기도 합니다. +| 테스트 | | +|---|---| +| `ops/misc/test_widen_dtype.py` | `collect_meta`가 `out_ptr0`의 dtype/numel을 못 구함 | +| `system/test_hetro.py` | stonne config에 `vpu_num_lanes`가 없음 — 모든 블록 크기가 그걸 기준으로 함 | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable` 미등록 | +| `models/MoE/test_moe_cpu.py` | 미분류 | --- -## 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` 잡: +## 3. 커널이 어디서 멈추는가 -- **Allowlisted tests** — gate 역할. -- **Full sweep** — `continue-on-error`. `coverage.md`를 step summary에 쓰고 - `triton-route-coverage`(results.json + failures/)를 업로드. +| 단계 | 수 | +|---|---| +| — 커널 생성 전 | 20 | +| 0 생성 후 거절 | 1 | +| 1 triton → ttir | 15 | +| 2 ttir → tts/linalg | 6 | +| 4 tnpu lower | 5 | +| 5 trace producer | 7 | -잡은 PSAL Slurm 러너 팜(`PSAL-POSTECH/slurm-ghr`)에서 돕니다. `runs-on`에 -`slurm` 라벨이 있어야 하고, 이미지 빌드와 스윕은 `big`(16c/64G/2h), 나머지는 -small 버킷입니다. `docker/setup-buildx-action`은 추가하면 안 됩니다 — 러너가 -자체 빌더를 등록해 둡니다. +stage 1의 15건이 `tl.assume` 무리입니다. 커널이 만들어지고 Inductor를 통과한 뒤, +tnpu 첫 단계에서 힌트에 걸립니다. --- -## 6. 측정하면서 고친 진단 3가지 - -이 셋을 고치기 전에는 보고 인프라를 만들 수 없었습니다. 각각이 증거를 파괴하고 -있었기 때문입니다. - -**`kernel.py`가 그것을 거절하는 검사 뒤에 저장되고 있었습니다.** -`write_spec_file`은 정확히 보존할 가치가 있는 커널(`triton_helpers`, -`SpecIncomplete`)에서 예외를 던지는데, 소스 저장보다 **먼저** 실행됐습니다. -결국 흥미로운 소스일수록 버려지고 있었습니다. 순서를 뒤집었고, 이제 거절된 16개 -커널 전부의 덤프가 남습니다 — 2절의 softmax 예시가 그중 하나입니다. +## 4. 첫 측정 이후 무엇이 바뀌었나 -**tnpu가 "exit 1"만 보고했습니다.** `run.py`는 stage 표를 stdout에 찍고 진짜 -진단은 `stage.log`에만 씁니다. 이전: +이 스윕의 첫 실행은 11/69였습니다. 지금은 13이지만 **둘을 비교하면 안 됩니다.** +움직인 것 대부분이 "무엇이 되는가"가 아니라 "무엇을 셌는가"입니다. -``` -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') -``` +**커버리지 기준이 두 번 엄격해졌습니다.** 커널을 안 내는 테스트는 원래 뺐고, 이제는 +커널을 내면서도 **자기 이름의 연산**을 `extern_kernels`에 넘기는 것도 뺍니다. 그렇게 +6개가 빠졌습니다 — `test_matmul_scalar`은 `mul`만 만들고 `extern_kernels.mm`을 +불렀습니다. -이 한 가지 변경으로 여섯 개 실패가 **하나의 버그**로 정리됐습니다: +**그리고 경로가 그 비상구를 안 쓰게 됐습니다.** `inductor_templates`가 npu를 +`GPU_TYPES`에 넣어 `use_triton_template`이 고려하게 하고, +`mm`/`bmm`/`addmm`/`baddbmm` heuristic을 등록하고, autotune을 고정 선택으로 +대체합니다 — 잴 하드웨어가 없으니까요. 이제 `mm`과 `addmm`이 aten이 아니라 Inductor +템플릿을 탑니다. 16개가 새로 `tl.assume`에 닿는 이유가 이것입니다. **그 전에는 +통과한 게 아니라 시뮬레이션을 안 했던 것입니다.** -**`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`를 갖고 있다면 알아둘 만합니다. +그 과정에서 고친 것: `libdevice`/`tl_math`가 그 이름을 만들던 torch import와 함께 +지워지고 있었고, 다축 grid가 한 번도 매칭되지 않는 키로 만들어지고 순서도 뒤집혀 +있었으며, 텐서 stride가 launch로 전달되지 않았고, `device_guard`가 `"pass"`를 +반환해 호출부가 `with pass:`를 쓰고 있었습니다. --- -## 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`에 이미 적용한 것과 같은 -수정입니다. - ---- +## 5. 다음 순서 -## 8. 이 측정이 말해주지 않는 것 +1. **`tl.assume`** — 16개. 의미 없는 힌트이고, tnpu를 쓰는 프론트엔드마다 벗기는 + 것보다 `normalize_upstream`에서 한 번 버리는 게 맞습니다. +2. **matmul 결과의 `select_lane_axis`** — 10개. 흘러든 op가 아니라 미지원 + 데이터플로라 더 무겁습니다. +3. **오답 3건** — 가장 작고 가장 위험합니다. 둘 다 reduction이라 원인이 하나일 + 가능성이 높습니다. +4. `linalg.index`(3), `linalg.generic`(1) — 이미 그 PR에 있습니다. -- `missing_dep` 17개는 로컬 venv 사정입니다. CI 이미지에서는 실제로 돌기 때문에 - 버킷이 이동할 것입니다 — 대부분 transformer·CNN 모델이므로 아마 `wrapper_gap` - 과 `triton_helpers` 쪽으로 갑니다. -- 어떤 버킷을 풀면 그 테스트들은 **다음 실패**로 이동하는 것이지, 반드시 통과로 - 가는 것이 아닙니다. -- 이 수치는 6절의 수정이 이미 적용된 상태에서 측정한 것이라, 그 이전 실행과 - 직접 비교할 수 없습니다. +1번과 2번이 실제 블로커 37개 중 26개이고 둘 다 업스트림입니다. 우리 쪽 목록은 +네 개뿐입니다. diff --git a/docs/triton-route-coverage.md b/docs/triton-route-coverage.md index 51ce81d3..551f9c71 100644 --- a/docs/triton-route-coverage.md +++ b/docs/triton-route-coverage.md @@ -1,28 +1,25 @@ # 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. Korean version: -[`triton-route-coverage.ko.md`](triton-route-coverage.ko.md). +The existing test suite run through the Triton codegen route (Inductor's Triton +backend + the triton-npu lowering passes) instead of the MLIR route. Korean +version: [`triton-route-coverage.ko.md`](triton-route-coverage.ko.md). | | | |---|---| -| Date | 2026-08-03 | -| Branch | `feature/triton-codegen` @ `8e17519` | -| tnpu pin | `5d84caf` | +| Date | 2026-08-04 | +| Branch | `feature/triton-helpers` @ `7899a17` | +| tnpu pin | `d46995f` | | 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 \ +python scripts/ci/triton_route_sweep.py --all -j 8 \ --markdown coverage.md --artifacts failures ``` -Every claim below is backed by a file in `failures/`. Paths are given so each -one can be checked. +Every claim below is backed by a file in `failures/`. --- @@ -30,434 +27,161 @@ one can be checked. ``` 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 +├── 13 pass THROUGH the route ← this is the coverage number +├── 2 pass without using the route +└── 54 fail ├── 17 missing test deps (local venv only; present in the CI image) - └── 36 real blockers + └── 37 real blockers ``` -**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%. +**26 of the 37 are two tnpu bugs.** Not thirty-seven problems; two, reported with +reproducers at `PSAL-POSTECH/triton-npu#2`. ### 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. Inductor's fusion is the half of this -migration we get for free, and it is already producing kernels tnpu accepts. +| `ops/elementwise/test_activation.py` | 83.7s | +| `ops/elementwise/test_add.py` | 79.7s | +| `ops/elementwise/test_exponent.py` | 10.4s | +| `ops/elementwise/test_transcendental.py` | 33.7s | +| `ops/misc/test_expert_mask.py` | 10.0s | +| `ops/reduce/test_batchnorm.py` | 37.6s | +| `ops/sparsity/test_sparse_core.py` | 15.2s | +| `ops/view/test_transpose2D.py` | 32.3s | +| `ops/view/test_transpose3D.py` | 182.6s | +| `ops/view/test_view3D_2D.py` | 38.6s | +| `system/test_eager.py` | 15.8s | +| `system/test_stonne.py` | 10.0s | +| `system/test_triton_codegen.py` | 10.4s | + +`test_topk` and `test_mlir_bindings` pass without emitting a kernel; they are +recorded separately and kept out of the gate, since counting them would be +counting nothing. --- -## 2. Worked example — how one failure is diagnosed +## 2. Failures by cause -`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* owns the failure. - -| Stage | Count | | +| bucket | count | owner | |---|---|---| -| — 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. - ---- - -## 4. Failures by cause, with evidence - -### `spec_incomplete` — 13 · owner: `triton_backend/kernel_spec.py` - -**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` | - -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 | Diagnostic | -|---|---| -| `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` - -| 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`. The fix is a small vendored file, not a pass -change. - -### `wrapper_gap` — 6 · owner: `triton_backend` - -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`, -`models/Mixtral8x7B/test_attention.py`, `models/test_transformer.py` - -Every attention and transformer test in the suite, blocked on one unimplemented -method. +| `tnpu_stage` | 30 | triton-npu | +| `missing_dep` | 17 | test environment (CI image has them) | +| `wrong_values` | 3 | numerics — investigate | +| `spec_incomplete` | 2 | triton_backend | +| `device_op` | 1 | PyTorchSimDevice | +| `other` | 1 | unclassified | -### `device_op` — 3 · owner: `PyTorchSimDevice` +### `tnpu_stage` — 30, and 26 of them are two bugs -Predates this route — the MLIR route intercepts these before the dispatcher. +**`tl.assume` → `llvm.intr.assume` (16).** Inductor's mm template emits +`tl.assume(pid_m >= 0)` as a hint; it lowers to an op in a dialect +`triton-shared-opt` does not load. Deleting only those lines from the ttir makes +the same stage-2 command succeed and produce a `linalg.matmul`. -| 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 — and the artifacts -show the full chain from Python to the rejected op. - -**`ops/conv/test_pool.py`** — stage 1. - -`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> +models/test_mlp.py ops/fusion/test_matmul_activation.py +models/test_transformer.py ops/fusion/test_matmul_reduction.py +ops/attention/test_gqa.py ops/fusion/test_matmul_scalar.py +ops/conv/test_conv2d.py ops/fusion/test_matmul_vector.py +ops/fusion/test_addmm_residual.py ops/fusion/test_prologue_fusion.py +ops/fusion/test_attention_fusion.py ops/fusion/test_transformer_fusion.py +ops/fusion/test_bmm_reduction.py ops/gemm/test_bmm.py +ops/sparsity/test_sparsity.py ops/gemm/test_matmul.py ``` -and `triton-shared-opt` cannot parse it: +**`select_lane_axis` on a matmul result (10).** ``` -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) +lane_axis.Fatal: the demand 'linalg.matmul' has to move for the lane axes to +agree, but it is about a value 'linalg.matmul' PRODUCES -- a relayout makes a +new value, so there is no edge to put one on ``` -survives `01-ttir.mlir`, then fails converting to linalg: +Case 4 in that pass's own docstring. Unlike the first, this looks like an +unsupported dataflow rather than a stray hint. ``` -error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than - the number of loops (2) of the enclosing LinalgOp +models/Mixtral8x7B/test_attention.py ops/elementwise/test_pointwise.py +ops/attention/test_gqa_decode.py ops/fusion/test_conv_fusion.py +ops/conv/test_cnn.py ops/misc/test_indirect_access.py +ops/conv/test_conv_view_input.py ops/misc/test_masked_nondividing.py +ops/conv/test_group_conv.py ops/view/test_cat.py ``` -Both artifacts carry the diagnostic and the offending `.mlir`, so they can go -upstream as-is. +**`linalg.index` rank (3)** — `ops/conv/test_pool.py`, `ops/reduce/test_reduce.py`, +`system/test_vectorops.py`. The original report on that PR. -### `togsim` / `other` — 5 +**`linalg.generic` shape (1)** — `ops/sort/test_sort.py`. -| Test | Stage | Detail | -|---|---|---| -| `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 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. +### `wrong_values` — 3 -**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. +`ops/reduce/test_softmax.py`, `ops/reduce/test_layernorm.py`, +`ops/view/test_floormod_axis_split.py`. These compile, run, and return the wrong +answer, which makes them the most dangerous category here even though it is the +smallest. Both reductions. -### `missing_dep` — 17 · not a route problem +### The rest — 4 -`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. +| test | | +|---|---| +| `ops/misc/test_widen_dtype.py` | `collect_meta` cannot resolve dtype/numel for `out_ptr0` | +| `system/test_hetro.py` | the stonne config has no `vpu_num_lanes`, which every block size is pinned to | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable` not registered | +| `models/MoE/test_moe_cpu.py` | unclassified | --- -## 5. 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 (section 2). - -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 - -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`: +## 3. Where kernels stop -- **Allowlisted tests** — gates. -- **Full sweep** — `continue-on-error`; writes `coverage.md` into the step - summary and uploads `triton-route-coverage` (results.json + failures/). +| Stage | Count | +|---|---| +| — no kernel generated | 20 | +| 0 generated, rejected | 1 | +| 1 triton → ttir | 15 | +| 2 ttir → tts/linalg | 6 | +| 4 tnpu lower | 5 | +| 5 trace producer | 7 | -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. +Fifteen at stage 1 is the `tl.assume` group: the kernel is generated and +survives Inductor, and the first tnpu stage is where the hint bites. --- -## 6. Three diagnostics fixed while measuring - -The reporting infrastructure could not be built until these were fixed, because -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, including the softmax example in section 2. +## 4. What changed since the first measurement -**tnpu reported "exit 1" and nothing else.** `run.py` prints a stage table to -stdout and the real diagnostic only to `stage.log`. Before: +The first run of this sweep read 11/69. The number is 13 now, but the two are not +comparable — most of the movement is in what got measured, not what works. -``` -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') -``` +**Coverage got stricter twice.** A test that emits no kernel was already +excluded; now a test that emits a kernel and still hands *its own op* to +`extern_kernels` is excluded too. Six tests moved out that way — +`test_matmul_scalar` generated the `mul` and called `extern_kernels.mm`. -That single change resolved six failures into one bug: +**And the route stopped taking that exit.** `inductor_templates` puts npu in +`GPU_TYPES` so `use_triton_template` considers it, registers heuristics for +`mm`/`bmm`/`addmm`/`baddbmm`, and replaces autotuning with a fixed choice — +there is no device to benchmark on. `mm` and `addmm` now go through Inductor's +template instead of aten, which is why sixteen tests newly reach `tl.assume`: +they were not passing before, they were not simulating. -**`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 `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`. +Fixed along the way: `libdevice`/`tl_math` were being stripped with the torch +import that named them; the multi-axis grid was built from a key that never +matched and emitted in the wrong order; tensor strides were not carried into the +launch; `device_guard` returned `"pass"`, which the caller writes as +`with pass:`. --- -## 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. - -| # | 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 | 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 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. - ---- +## 5. What to do next -## 8. Caveats +1. **`tl.assume`** — 16 tests. A hint with no semantics, dropped in + `normalize_upstream` rather than in every frontend that feeds tnpu. +2. **`select_lane_axis` on a matmul result** — 10 tests. Heavier: an unsupported + dataflow, not a stray op. +3. **The three wrong answers** — smallest and most dangerous. Both reductions, + so likely one cause. +4. `linalg.index` (3), `linalg.generic` (1) — already on that PR. -- 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-6 fixes already applied, so they are - not comparable to a run from before them. +1 and 2 are 26 of the 37 real blockers and both sit upstream. On our side the +list is four tests long. diff --git a/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt index 213a8d58..b806db41 100644 --- a/scripts/ci/triton_route_passing.txt +++ b/scripts/ci/triton_route_passing.txt @@ -2,13 +2,25 @@ # 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 is deliberately absent. +# +# Held out until PSAL-POSTECH/triton-npu#2 lands -- their matmul now goes +# through Inductor's template and stops at tl.assume, where before it went +# to aten and was never simulated: +# tests/ops/attention/test_gqa.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/sparsity/test_sparsity.py +tests/ops/elementwise/test_activation.py 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/elementwise/test_exponent.py +tests/ops/elementwise/test_transcendental.py tests/ops/misc/test_expert_mask.py tests/ops/reduce/test_batchnorm.py +tests/ops/sparsity/test_sparse_core.py +tests/ops/view/test_transpose2D.py +tests/ops/view/test_transpose3D.py tests/ops/view/test_view3D_2D.py tests/system/test_eager.py tests/system/test_stonne.py diff --git a/scripts/ci/triton_route_sweep.py b/scripts/ci/triton_route_sweep.py index df8660e3..675c2d53 100755 --- a/scripts/ci/triton_route_sweep.py +++ b/scripts/ci/triton_route_sweep.py @@ -43,11 +43,17 @@ ("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"), + ("reduction", r"lane-aware|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"), + # Before togsim: a kernel that simulated fine and then compared wrong is a + # numerics failure, and the log is full of TOGSim lines by then either way. + ("wrong_values", r"VALUES WRONG|Max abs diff|Test Failed|allclose"), + ("missing_artifact", r"FileNotFoundError"), + # Only a TOGSim *failure*. Matching the name alone caught every INFO line it + # writes, so anything that got as far as simulating landed here. + ("togsim", r"TOGSim subprocess|SIGSEGV|Signals\.SIG|'vpu_num_lanes'|" + r"trace\.so not found|\[TOGSim\].*(?:failed|Error)"), ("timeout", r"^__timeout__$"), ] @@ -139,7 +145,11 @@ def run_one(test, timeout, artifacts, scratch): 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) + # TORCHINDUCTOR_CACHE_DIR too: extension_config only re-points it at + # codegen time, by which point Inductor has already put the first graph + # in the shared /tmp cache -- concurrent tests then collide there. + env = dict(os.environ, TORCHSIM_TRITON_CODEGEN="1", TORCHSIM_DUMP_PATH=dump, + TORCHINDUCTOR_CACHE_DIR=os.path.join(dump, ".torchinductor")) t0, timed_out = time.time(), False try: @@ -192,6 +202,7 @@ def write_markdown(results, path): "dynamic_shape": "triton_backend -- shape-specialised launch", "matmul_timing": "build_tog -- compute node lookup", "togsim": "TOGSim / trace producer", + "missing_artifact": "an expected artifact was not written", "wrong_values": "numerics -- investigate", "missing_dep": "test environment (present in the CI image)", "timeout": "too slow, or hung", diff --git a/tests/ops/sparsity/test_sparse_core.py b/tests/ops/sparsity/test_sparse_core.py index 21cd9344..ddde6d5d 100644 --- a/tests/ops/sparsity/test_sparse_core.py +++ b/tests/ops/sparsity/test_sparse_core.py @@ -57,7 +57,7 @@ def forward(self, x): def test_sparse_mlp(device, batch_size=32, input_size=128, hidden_size=128, output_size=128): torch.manual_seed(0) # mlp = MLP(input_size, hidden_size, output_size) - mlp = SparseMLP(input_size, hidden_size, output_size, device) + mlp = SparseMLP(input_size, hidden_size, output_size, device=device) mlp = mlp.to(device=device) input = torch.randn(batch_size, input_size) x1 = input.to(device=device) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json index 4bac8dfb..2c2888b9 100644 --- a/thirdparty/triton-npu.json +++ b/thirdparty/triton-npu.json @@ -1,8 +1,11 @@ { - "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.", + "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. Both repositories are private: the workflow needs a PAT in secrets.TNPU_TOKEN with read on each, since the default Actions token is scoped to this repository. `also_reads` is what restore.sh clones besides the release, and preflight checks the token against it.", "triton_npu": { "repository": "PSAL-POSTECH/triton-npu", - "ref": "5d84cafdd5292f77548a84a6c44a4d2d15c4dd4f", - "release_tag": "toolchain-llvm23" + "ref": "d46995ffb530766af47e244d4af4a58890d604ea", + "release_tag": "toolchain-llvm23", + "also_reads": [ + "PSAL-POSTECH/triton_shared" + ] } }