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

.github/workflows/triton_npu.yml

@@ -0,0 +1,267 @@
+name: Triton codegen route (triton-npu)
+
+# Exercises the Triton codegen route: Inductor's own Triton backend produces the
+# kernel and triton-npu lowers it to a RISC-V ELF.
+# (PyTorchSimFrontend/triton_backend/README.md)
+#
+# Separate from the main CI on purpose. The route is WIP, and its toolchain layer
+# is ~1.8 GiB that no other job needs, so it neither gates PRs nor slows them
+# down. Promote the jobs into pytorchsim_test.yml once the route runs end to end.
+#
+# 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:
+    branches: [ "master", "develop" ]
+    paths:
+      - 'PyTorchSimFrontend/triton_backend/**'
+      - 'thirdparty/triton-npu.json'
+      - 'Dockerfile.tnpu'
+      - 'scripts/ci/tnpu_base_pin.sh'
+      - '.github/workflows/triton_npu.yml'
+  workflow_dispatch:
+
+env:
+  BASE_IMAGE_REPO: ghcr.io/psal-postech/torchsim_base
+  TNPU_IMAGE_REPO: ghcr.io/psal-postech/torchsim_tnpu_base
+  APP_IMAGE_REPO: ghcr.io/psal-postech/torchsim_tnpu
+  SOURCE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+
+jobs:
+  preflight:
+    name: Check tnpu access
+    runs-on: [self-hosted, slurm, x86_64]
+    outputs:
+      ready: ${{ steps.check.outputs.ready }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ env.SOURCE_SHA }}
+          persist-credentials: false
+
+      - name: Token and release present
+        id: check
+        env:
+          TNPU_TOKEN: ${{ secrets.TNPU_TOKEN }}
+        run: |
+          if [ -z "${TNPU_TOKEN}" ]; then
+            echo "::error::secrets.TNPU_TOKEN is not set. PSAL-POSTECH/triton-npu is private and the default Actions token cannot read it."
+            echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1
+          fi
+          REPO=$(jq -r '.triton_npu.repository' thirdparty/triton-npu.json)
+          TAG=$(jq -r '.triton_npu.release_tag' thirdparty/triton-npu.json)
+          if ! curl -fsS -H "Authorization: Bearer ${TNPU_TOKEN}" \
+               "https://api.github.com/repos/${REPO}" -o /dev/null; then
+            echo "::error::TNPU_TOKEN cannot read ${REPO}."
+            echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1
+          fi
+          if ! curl -fsS -H "Authorization: Bearer ${TNPU_TOKEN}" \
+               "https://api.github.com/repos/${REPO}/releases/tags/${TAG}" -o /dev/null; then
+            echo "::error::${REPO} has no release tagged '${TAG}'. Mirror the toolchain assets there (see thirdparty/triton-npu.json)."
+            echo "ready=false" >> "$GITHUB_OUTPUT"; exit 1
+          fi
+          # 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:
+    name: Build tnpu toolchain image
+    needs: preflight
+    runs-on: [self-hosted, slurm, big, x86_64]
+    outputs:
+      tnpu_image: ${{ steps.pin.outputs.tnpu_image }}
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ env.SOURCE_SHA }}
+          submodules: recursive
+          persist-credentials: false
+
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Pins
+        id: pin
+        run: |
+          BASE_PIN="$(bash scripts/ci/thirdparty_base_pin.sh)"
+          TNPU_PIN="$(bash scripts/ci/tnpu_base_pin.sh)"
+          echo "BASE_IMAGE=${BASE_IMAGE_REPO}:thirdparty-${BASE_PIN}" >> "$GITHUB_ENV"
+          # The tnpu layer sits on a specific base, so its tag carries both pins.
+          echo "TNPU_IMAGE=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_ENV"
+          echo "tnpu_image=${TNPU_IMAGE_REPO}:tnpu-${TNPU_PIN}-base-${BASE_PIN}" >> "$GITHUB_OUTPUT"
+          echo "TNPU_REF=$(jq -r '.triton_npu.ref' thirdparty/triton-npu.json)" >> "$GITHUB_ENV"
+
+      - name: Check tnpu image exists
+        id: exists
+        run: |
+          if docker manifest inspect "${TNPU_IMAGE}" > /dev/null 2>&1; then
+            echo "ok=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "ok=false" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Build and push tnpu toolchain image
+        if: steps.exists.outputs.ok != 'true'
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: ./Dockerfile.tnpu
+          push: true
+          build-args: |
+            BASE_IMAGE=${{ env.BASE_IMAGE }}
+            TNPU_REF=${{ env.TNPU_REF }}
+          secrets: |
+            tnpu_token=${{ secrets.TNPU_TOKEN }}
+          tags: ${{ env.TNPU_IMAGE }}
+
+  build-app:
+    name: Build app image on tnpu base
+    needs: ensure-tnpu-base
+    runs-on: [self-hosted, slurm, big, x86_64]
+    outputs:
+      app_image: ${{ steps.name.outputs.app_image }}
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ env.SOURCE_SHA }}
+          submodules: recursive
+          persist-credentials: false
+
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Image name
+        id: name
+        run: echo "app_image=${APP_IMAGE_REPO}:${SOURCE_SHA}" >> "$GITHUB_OUTPUT"
+
+      - name: Build and push
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: ./Dockerfile
+          push: true
+          build-args: |
+            BASE_IMAGE=${{ needs.ensure-tnpu-base.outputs.tnpu_image }}
+          tags: ${{ steps.name.outputs.app_image }}
+
+  tnpu-baselines:
+    name: triton-npu baselines
+    needs: build-app
+    runs-on: [self-hosted, slurm, x86_64]
+    steps:
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # The harness's own kernels, end to end through Spike. This is the gate on
+      # the toolchain itself: if these regress, nothing downstream is meaningful.
+      # gemm/bmm need TNPU_VCIX_MATMUL=1 to reach the systolic array.
+      - name: doctor + add / mul / relu / gemm / bmm
+        run: |
+          docker run --rm -e TNPU_VCIX_MATMUL=1 \
+            ${{ needs.build-app.outputs.app_image }} bash -lc '
+              cd /workspace/triton-npu &&
+              python3 run.py doctor &&
+              for k in add mul relu gemm bmm; do
+                echo "=== $k ===" && python3 run.py kernels/$k.py || exit 1
+              done'
+
+  triton-route:
+    name: Inductor Triton route
+    needs: build-app
+    runs-on: [self-hosted, slurm, x86_64]
+    # WIP: the launch is deliberately unimplemented, so this reports how far the
+    # route gets rather than gating. Drop this once the launch lands.
+    continue-on-error: true
+    steps:
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: test_triton_codegen.py
+        run: |
+          docker run --rm -e TORCHSIM_TRITON_CODEGEN=1 \
+            ${{ needs.build-app.outputs.app_image }} \
+            python3 PyTorchSim/tests/system/test_triton_codegen.py
+
+  triton-route-suite:
+    name: Test suite on the Triton route
+    needs: build-app
+    runs-on: [self-hosted, slurm, big, x86_64]
+    steps:
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Gates on the tests that pass today; coverage cannot silently shrink.
+      - name: Allowlisted tests
+        run: |
+          docker run --rm \
+            ${{ needs.build-app.outputs.app_image }} \
+            python3 PyTorchSim/scripts/ci/triton_route_sweep.py
+
+      # Reports the rest. Each failure leaves its kernel and stage IR behind.
+      - name: Full sweep (report)
+        continue-on-error: true
+        run: |
+          mkdir -p sweep && chmod 777 sweep
+          docker run --rm -v "$PWD/sweep:/sweep" \
+            ${{ needs.build-app.outputs.app_image }} \
+            python3 PyTorchSim/scripts/ci/triton_route_sweep.py --all \
+              --timeout 900 --json /sweep/results.json \
+              --markdown /sweep/coverage.md --artifacts /sweep/failures
+          cat sweep/coverage.md >> "$GITHUB_STEP_SUMMARY"
+
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: triton-route-coverage
+          path: sweep/
+          if-no-files-found: warn
+
+  mlir-route-regression:
+    name: MLIR route still passes
+    needs: build-app
+    runs-on: [self-hosted, slurm, x86_64]
+    steps:
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # The tnpu layer adds a second LLVM and a second triton to the image. This
+      # is the check that it did not disturb the production path.
+      - name: test_add.py
+        run: |
+          docker run --rm \
+            ${{ needs.build-app.outputs.app_image }} \
+            python3 PyTorchSim/tests/ops/elementwise/test_add.py

CLAUDE.md

@@ -47,6 +47,8 @@ Run a model from `tests/models/Llama/`, `tests/models/DeepSeek/`, etc. similarly
 
 **CI coverage:** the GitHub Actions workflow `.github/workflows/pytorchsim_test.yml` runs an **explicit allowlist** of `tests/*.py` files (~40 jobs, one Docker container per test). Adding a new file under `tests/` does *not* automatically gate PRs — register it in `pytorchsim_test.yml` if you want CI to exercise it. Conversely, files like `tests/ops/attention/test_gqa.py`, `tests/ops/attention/test_gqa_decode.py`, and `tests/system/test_eager.py` exist in the repo but are *not* in CI, so local validation is the only safety net for them.
 
+The Triton codegen route has its own workflow, `.github/workflows/triton_npu.yml`, kept separate because its toolchain layer is ~1.8 GiB that no other job needs. It builds `torchsim_tnpu_base` (pinned by `thirdparty/triton-npu.json` + `Dockerfile.tnpu`) and needs `secrets.TNPU_TOKEN` plus a toolchain release on the private `PSAL-POSTECH/triton-npu`; see `PyTorchSimFrontend/triton_backend/README.md`.
+
 **For fast iteration** (skip functional check):
 ```bash
 export pytorchsim_functional_mode=False   # skips Spike
@@ -137,7 +139,7 @@ Conan deps for TOGSim: `boost/1.79.0`, `robin-hood-hashing/3.11.5`, `spdlog/1.11
 
 ## Gotchas / things I've already learned
 
-- The repo expects `python` to be a Python 3.10+ binary with `torch==2.8.0`. The frontend extends the PyTorch 2 Inductor stack — pin to this version.
+- The repo expects `python` to be a Python 3.10+ binary with `torch==2.10.0` (torchvision `0.25.0`, triton `3.6.0`). The frontend extends the PyTorch 2 Inductor stack — pin to this version. 2.10 specifically: it is the first release whose Inductor targets triton 3.6, the version triton-npu is built against. The pins live in `Dockerfile.base`, and editing that file changes the base-image tag automatically (the tag is `thirdparty-<sha256 of thirdparty/github-releases.json + Dockerfile.base>`, see `scripts/ci/thirdparty_base_pin.sh`).
 - The default Gem5 path is hard-coded to `/workspace/gem5/build/RISCV/gem5.opt`. Override with `GEM5_PATH` if you build elsewhere.
 - `_C.cpython-311-*.so` and `torch_openreg/lib/` are build artifacts — already in `.gitignore`, don't commit.
 - TOGSim creates a per-PID FIFO under `/tmp/togsim_fifo_<pid>` for command/event comm; if a previous run crashed and left stale FIFOs, they get cleaned up on the next start, but watch for orphaned processes if you Ctrl-C mid-run.

Dockerfile.base

@@ -51,8 +51,20 @@ RUN apt-get -y update && \
     rm -rf /var/lib/apt/lists/*
 
 # CPU PyTorch (no CUDA wheels). torchvision is required by the vision model tests.
+# torch 2.10 is pinned for the Triton codegen route: it is the first release whose
+# Inductor targets triton 3.6, which is the version triton-npu is built against
+# (triton 3.6 pins LLVM 23, and both sides of triton-npu's textual IR seam must be
+# the same LLVM). On 2.8 the frontend had to be shimmed onto a triton it did not
+# expect; on 2.10 the versions simply agree.
 RUN python3.11 -m pip install --no-cache-dir \
-    torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu
+    torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu
+
+# Triton, for the Inductor Triton codegen route (PyTorchSimFrontend/triton_backend).
+# Inductor imports triton while GENERATING a kernel, so it is needed even though
+# nothing here compiles or launches through triton's own runtime -- triton-npu
+# compiles the kernel ahead of time to a RISC-V ELF using its own triton build.
+# Not a dependency of the CPU torch wheels, hence installed explicitly.
+RUN python3.11 -m pip install --no-cache-dir triton==3.6.0
 
 # TorchSim Python dependencies (numpy pinned <2 for transformers/diffusers compat).
 RUN python3.11 -m pip install --no-cache-dir \

Dockerfile.tnpu

@@ -0,0 +1,42 @@
+# syntax=docker/dockerfile:1.4
+#
+# triton-npu toolchain layer, for the Triton codegen route only.
+# Separate from torchsim_base because it is ~1.8 GiB no other job needs.
+# The app image for this route is ./Dockerfile with BASE_IMAGE pointed here.
+#
+# The repo is private, so the clone and the release downloads both need a token.
+# It is a BuildKit secret, not a build-arg: build-args land in the image history.
+
+ARG BASE_IMAGE=ghcr.io/psal-postech/torchsim_base:latest
+FROM ${BASE_IMAGE}
+
+ARG TNPU_REPO=PSAL-POSTECH/triton-npu
+ARG TNPU_REF=main
+
+WORKDIR /workspace
+
+# Not under $TORCHSIM_DIR: ./Dockerfile copies the PyTorchSim checkout over that
+# path afterwards. extension_config reads TNPU_DIR, set below.
+RUN --mount=type=secret,id=tnpu_token \
+    TOKEN="$(cat /run/secrets/tnpu_token)" && \
+    git clone "https://x-access-token:${TOKEN}@github.com/${TNPU_REPO}.git" \
+      /workspace/triton-npu && \
+    git -C /workspace/triton-npu checkout -q "${TNPU_REF}" && \
+    git -C /workspace/triton-npu remote set-url origin \
+      "https://github.com/${TNPU_REPO}.git"
+
+# restore.sh owns every pin (setup/versions.env) and unpacks LLVM 23, spike and
+# the triton runtime into /workspace.
+RUN --mount=type=secret,id=tnpu_token \
+    GITHUB_TOKEN="$(cat /run/secrets/tnpu_token)" \
+      /workspace/triton-npu/setup/restore.sh --prebuilt
+
+ENV TNPU_DIR=/workspace/triton-npu
+# tnpu defaults to a separate fp8 spike and asks for zvfp8; the released spike
+# has neither, and an unknown extension stops spike at startup. Drop once
+# PSAL-POSTECH/riscv-isa-sim#7 is in the release.
+ENV TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike
+ENV TNPU_SPIKE_ISA=rv64gcv_zfh
+
+# Fail the build, not the first CI job.
+RUN python3 /workspace/triton-npu/run.py doctor

PyTorchSimDevice/torch_openreg/__init__.py

@@ -18,13 +18,30 @@
 
 sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 import PyTorchSimFrontend.extension_config  # noqa: F401
+from PyTorchSimFrontend import extension_config as _extension_config
 from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen
-from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling
-torch._inductor.codegen.common.register_backend_for_device(
-    "npu",
-    lambda scheduling: MLIRScheduling(scheduling),
-    ExtensionWrapperCodegen
-)
+
+# Two mutually exclusive codegen routes for `npu`, chosen here because Inductor
+# registers a backend per device, once.
+#   MLIR   (default)  hand-written MLIR emission, PyTorchSimFrontend/mlir
+#   Triton (opt-in)   Inductor's own Triton codegen + the triton-npu passes,
+#                     TORCHSIM_TRITON_CODEGEN=1. WIP; see
+#                     PyTorchSimFrontend/triton_backend/README.md
+if _extension_config.CONFIG_TRITON_CODEGEN:
+    from PyTorchSimFrontend.triton_backend import (
+        TritonNPUScheduling, TritonNPUWrapperCodegen)
+    torch._inductor.codegen.common.register_backend_for_device(
+        "npu",
+        lambda scheduling: TritonNPUScheduling(scheduling),
+        TritonNPUWrapperCodegen
+    )
+else:
+    from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling
+    torch._inductor.codegen.common.register_backend_for_device(
+        "npu",
+        lambda scheduling: MLIRScheduling(scheduling),
+        ExtensionWrapperCodegen
+    )
 
 torch_openreg.openreg.init()
 sys.modules['torch.npu'] = torch_openreg.openreg

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