Triton Backed Supported (Beta)

.github/workflows/pytorchsim_test.yml

@@ -613,7 +613,7 @@ jobs:
           echo "Running test_mistral.py"
           docker run --rm \
             -e TOGSIM_CONFIG="${{ inputs.togsim_config }}" \
-            ${{ inputs.image_name }} python3 PyTorchSim/tests/models/Mixtral8x7B/test_attention.py
+            ${{ inputs.image_name }} python3 PyTorchSim/tests/models/Mixtral8x7B/test_mistral.py
 
   test_vit:
     name: Run test_vit

.gitignore

@@ -10,3 +10,15 @@ tutorial/session1/fused/
 tutorial/session1/togsim_results/
 tutorial/session2/togsim_results/
 tutorial/session2/fx_graph.svg
+
+# Per-worktree environment, documented as local and uncommitted in
+# docs/worktrees.md -- it names absolute paths that differ per worktree.
+.envrc
+
+# PyTorchSimDevice build output. CMake's tree and the extension modules it
+# links -- 105 of these were committed by accident in 870a00e and are removed
+# in the commit that adds these lines. They are regenerated by
+#   cd PyTorchSimDevice && python -m pip install --no-build-isolation -e .
+PyTorchSimDevice/build/
+PyTorchSimDevice/torch_openreg/lib/
+PyTorchSimDevice/torch_openreg/_C.cpython-*.so

CLAUDE.md

@@ -153,3 +153,49 @@ Conan deps for TOGSim: `boost/1.79.0`, `robin-hood-hashing/3.11.5`, `spdlog/1.11
 - Fork, branch (`feature/<name>`), PR against `develop`, not `main`.
 - Commit prefix style observed: `[Frontend] ...`, `[TOGSim] ...`, etc.
 - Commit messages: plain text only. No Markdown formatting (no backticks, bold, bullet lists, headings). Avoid Unicode where ASCII works (use `->` not arrows, `--` not em-dashes, straight quotes).
+
+## Ship it — commit and push without being asked
+
+Work that changes behaviour is done when it is **pushed**, not when it runs. Commit
+and push in the same turn that finishes it. Do not stop to ask permission for
+either; this rule is the permission.
+
+    implement  ->  verify  ->  commit  ->  push
+
+**WHERE IT GOES.** `origin` is `PSAL-POSTECH/PyTorchSim`, the upstream everyone
+shares, and it is **not** our destination. Push to the fork:
+
+```bash
+git remote add fork git@github.com:student-Jungmin/PyTorchSim.git   # once
+git push fork <branch>
+```
+
+Reading "push" as "push to origin" puts work on the shared upstream, so check the
+remote before pushing rather than trusting whatever `origin` happens to be. A
+change spanning repositories is one commit per repository, each to its own remote
+— never one commit carrying another repository's work. (`triton-npu` has the
+matching rule and the full destination table.)
+
+**VERIFY FIRST, AND SAY WHAT RAN.** "Verified" means the test was executed, not
+that the code looks right. For the Triton route that means the affected test
+plus the allowlist (`scripts/ci/triton_route_passing.txt`) — and **clear
+`outputs/triton_*` and `outputs/.torchinductor` between runs**, or a cached
+artifact replays and a fix appears to change nothing. That has already caused one
+wrong conclusion. A push whose verification was skipped is a push of an unknown
+state.
+
+**PIN `TNPU_DIR` WHEN THE ROUTE IS INVOLVED.** Stages 1-5 live in a separate repo
+that someone else may be editing right now, and a pass mid-refactor produces
+failures that look like ours (`no lane axis`, bare `NameError`s that vanish on
+re-run). Point `TNPU_DIR` at a worktree pinned to a known-good commit so a result
+means something.
+
+**WHAT DOES NOT COUNT.** Exploration, scratch files, a half-finished edit, or a
+change whose verification failed. Broken work is not pushed — say what failed.
+
+**IF THE PUSH FAILS,** report it and why. Until it lands the work is committed,
+not shipped; never present the one as the other.
+
+**A REBUCKETED FAILURE IS NOT A FIX.** Clearing a guard so a test fails deeper is
+progress worth committing, but say so in those words rather than counting it as
+passing.

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/extension_config.py

@@ -71,7 +71,22 @@ def __getattr__(name):
     if name == "pytorchsim_functional_mode":
         return config_yaml['pytorchsim_functional_mode']
     if name == "pytorchsim_timing_mode":
-        return config_yaml['pytorchsim_timing_mode']
+        # FORCED OFF ON THIS BRANCH, ON PURPOSE AND TEMPORARILY. Correctness is
+        # what is being worked on here and the timing half costs a gem5 sample
+        # and a TOGSim run per kernel, so every sweep pays for a number nobody
+        # is reading. The YAML still says what the machine is; this says what
+        # this branch is doing.
+        #
+        # IT IS A FORCE RATHER THAN A DEFAULT because the default is per-config
+        # and there are nineteen of them: eighteen say 1, and a run that does
+        # not set TORCHSIM_DIR reads a DIFFERENT CHECKOUT's copy -- which is how
+        # the sweeps on this branch ran with timing on while this repo's default
+        # config said 0. Set TORCHSIM_TIMING_MODE=1 to get it back, and delete
+        # this arm when correctness work moves on.
+        env = os.environ.get("TORCHSIM_TIMING_MODE")
+        if env is not None:
+            return int(env)
+        return 0
     # Sub-option of functional mode: compare every realized Spike buffer against a CPU
     # golden to localize the first kernel whose value diverges. Auto-disabled when
     # functional mode is off (there are no Spike values to verify).

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -70,6 +70,23 @@ def reduction_partial_combine_vec(reduction_type, vector_value, init_value):
         return ops.logical_or(vector_value, init_value)
     raise AssertionError(reduction_type)
 
+def _fverify_writes(kernel_name, position):
+    """Does `kernel_name` write the tensor argument at `position`?
+
+    The roles are recorded at define_kernel by the Triton backend, which is
+    the only route that has them; the MLIR route records nothing and every
+    argument stays checkable, which is what this did for both routes before.
+    Unknown kernel, unknown position, backend not imported -> True.
+    """
+    if kernel_name is None:
+        return True
+    try:
+        from PyTorchSimFrontend.triton_backend import kernel_spec
+    except Exception:  # noqa: BLE001 - the MLIR route need not have it
+        return True
+    return kernel_spec.writes_arg(kernel_name, position)
+
+
 class ExtensionWrapperCodegen(wrapper.PythonWrapperCodegen):
     def __init__(self):
         super().__init__()
@@ -223,8 +240,18 @@ def generate(self, is_inference):
         # result.splice(self.header)
 
         self._fverify_seen = set()
+        self._fverify_last = None
         with contextlib.ExitStack() as stack:
             stack.enter_context(self.wrapper_call.indent())
+            # memory_plan_reuse() reaches self.estimate_peak through
+            # AllocateLine.should_reuse_buffer, and upstream sets it in
+            # run_wrapper_ir_passes -- which this override replaces, so nothing
+            # else will. Missing it is not a planning miss but an AttributeError,
+            # and only on a graph with a reuse candidate far enough back to need
+            # the estimate: ResNet-18 hits it, add does not. Same guard upstream
+            # uses, so buffer reuse off means no estimate to build.
+            if torch._inductor.config.allow_buffer_reuse:
+                self.estimate_peak = wrapper.EfficientPeakEstimate()
             self.memory_plan_reuse()
             with self.set_writeline(self.wrapper_call.writeline):
                 for line in self.lines:
@@ -240,7 +267,8 @@ def generate(self, is_inference):
                     elif isinstance(line, wrapper.KernelCallLine):
                         self.wrapper_call.writeline(self.wrap_kernel_call(line.kernel_name, line.call_args))
                         if _func_verify.enabled():
-                            self._fverify_emit_checks(line.call_args)
+                            self._fverify_emit_checks(line.call_args, id(line),
+                                                      line.kernel_name)
                     else:
                         if isinstance(line, wrapper.WrapperLine):
                             line.codegen(self.wrapper_call)
@@ -270,21 +298,83 @@ def generate(self, is_inference):
             self.kernel_declarations.getvaluewithlinemap(),
         )
 
-    def _fverify_emit_checks(self, call_args):
+    def _fverify_last_writer(self):
+        """{buffer name: id of the LAST kernel call that names it}.
+
+        THE FIRST KERNEL TO NAME A BUFFER IS NOT ALWAYS THE ONE THAT FINISHES
+        IT. One fx op can be split across several kernels, and then the buffer
+        is only complete after the last of them -- checking after the first
+        compares a half-built buffer against a finished golden and reports a
+        divergence that is not one.
+
+            measured   DeepSeek-V3's MoE router. `aten.scatter.value` comes out
+                       as two kernels sharing one origin node:
+
+                         triton_npu_fused_scatter_zeros_like_38(buf9, 256)
+                         _fverify.verify_check(buf9, ...)        <- here
+                         triton_npu_fused_scatter_zeros_like_39(buf8, buf9, 128)
+
+                       38 writes the zeros and 39 scatters the ones, so the
+                       check saw an all-zero buffer and reported "128/256
+                       elements over tol, all npu=0 cpu=1" -- every one of the
+                       scattered ones "missing". Running kernel 39 standalone
+                       against a torch reference gives max_abs_err 0.
+
+        So the walk is done twice: once to find where each buffer is last
+        written, and once to emit. Same order, same buffers, one check each --
+        only the position moves.
+        """
+        last = {}
+        for line in self.lines:
+            if not isinstance(line, wrapper.KernelCallLine):
+                continue
+            for pos, a in enumerate(line.call_args):
+                if not (isinstance(a, str) and a.strip().isidentifier()):
+                    continue
+                if not _fverify_writes(line.kernel_name, pos):
+                    continue
+                last[a.strip()] = id(line)
+        return last
+
+    def _fverify_emit_checks(self, call_args, line_id=None, kernel_name=None):
         """Emit per-kernel CPU verify calls for this kernel's output buffers.
 
-        A buffer's value is produced by the first kernel that names it (producer
-        precedes consumers in topo order), so we check each bare-identifier buffer
-        arg the first time it is seen -- that occurrence is its output. The buffer
-        is mapped to its originating fx node (op) so the runtime check can compare
-        against the CPU golden keyed by that node.
+        Each bare-identifier buffer arg the kernel WRITES is checked once,
+        after the LAST kernel that writes it -- see _fverify_last_writer for
+        why not the first. The buffer is mapped to its originating fx node (op)
+        so the runtime check can compare against the CPU golden keyed by that
+        node.
+
+        WRITES, NOT NAMES. This used to check every bare-identifier argument,
+        inputs included, and the docstrings on both halves said "writes" while
+        the code said "names". The two part company under buffer REUSE: the
+        wrapper renames storage (`buf20 = buf9  # reuse`), so one buffer's
+        contents live under another buffer's name, and that name's
+        `origin_node` describes what Inductor MEANT to put there. Check it
+        after a kernel that only reads it and the comparison is against a value
+        nothing computed.
+
+            measured   Stable Diffusion v1.5's UNet. Two kernels take an
+                       `in_out_ptr0` and never store to it; they are called
+                       eight times between them, on buf20, buf88, buf104,
+                       buf170, buf189, buf208, buf230 and buf298 -- and those
+                       eight are EXACTLY the eight divergences the run
+                       reported, each against an `add_N` node that no kernel
+                       materialises. Nothing else in the model diverges; 217
+                       kernels run between them.
         """
-        for a in call_args:
+        if self._fverify_last is None:
+            self._fverify_last = self._fverify_last_writer()
+        for pos, a in enumerate(call_args):
             if not isinstance(a, str):
                 continue
             name = a.strip()
             if not name.isidentifier() or name in self._fverify_seen:
                 continue
+            if not _fverify_writes(kernel_name, pos):
+                continue          # this kernel only reads it
+            if line_id is not None and self._fverify_last.get(name) != line_id:
+                continue          # a later kernel still writes this buffer
             self._fverify_seen.add(name)
             if name in V.graph.graph_inputs:
                 continue  # placeholders: golden == input, nothing to verify