[Frontend] Unblock buffer reuse and triton_helpers on the Triton route

.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:

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())

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:

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

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/<name>.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)

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"

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