From 3366674324acc82838b5a510c4dfed5470fbcfe8 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:47:39 +0900 Subject: [PATCH 01/10] [Build] Target torch 2.10 so the frontend and triton-npu agree on triton 3.6 triton-npu is built against triton 3.6. 2.10 is the first torch release whose Inductor targets that version, so anything older emits kernels its passes were not written for. The pins live here rather than in a requirements file because the base image tag is derived from this file's hash: editing it moves the tag, and nothing else has to be kept in step by hand. --- Dockerfile.base | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/Dockerfile.base b/Dockerfile.base index de023566..b4d58a7c 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -51,8 +51,20 @@ RUN apt-get -y update && \ rm -rf /var/lib/apt/lists/* # CPU PyTorch (no CUDA wheels). torchvision is required by the vision model tests. +# torch 2.10 is pinned for the Triton codegen route: it is the first release whose +# Inductor targets triton 3.6, which is the version triton-npu is built against +# (triton 3.6 pins LLVM 23, and both sides of triton-npu's textual IR seam must be +# the same LLVM). On 2.8 the frontend had to be shimmed onto a triton it did not +# expect; on 2.10 the versions simply agree. RUN python3.11 -m pip install --no-cache-dir \ - torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu + torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu + +# Triton, for the Inductor Triton codegen route (PyTorchSimFrontend/triton_backend). +# Inductor imports triton while GENERATING a kernel, so it is needed even though +# nothing here compiles or launches through triton's own runtime -- triton-npu +# compiles the kernel ahead of time to a RISC-V ELF using its own triton build. +# Not a dependency of the CPU torch wheels, hence installed explicitly. +RUN python3.11 -m pip install --no-cache-dir triton==3.6.0 # TorchSim Python dependencies (numpy pinned <2 for transformers/diffusers compat). RUN python3.11 -m pip install --no-cache-dir \ From 0d64bbaeae5b23fb85a8c26fe8e3841de52f1554 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:47:39 +0900 Subject: [PATCH 02/10] [CI] Add the triton-npu toolchain image and its workflow The toolchain layer is about 1.8 GiB -- LLVM, Spike, the triton runtime, the triton-shared binary -- which no other job needs, so it gets its own image and its own workflow instead of growing the one every test pulls. torchsim_tnpu_base is pinned by thirdparty/triton-npu.json plus Dockerfile.tnpu, the same derived-tag scheme the base image uses. Building it needs secrets.TNPU_TOKEN and a toolchain release on PSAL-POSTECH/triton-npu, both private; the token wants Contents: Read-only and must not be a push token, since this tree ships inside the image. The pin is 8ebe408, which drops llvm.intr.assume in normalize_upstream and moves triton_shared to 9017bd4 for the llvm dialect registration. Both halves are needed: without registration the op stops stage 2 at the parse. Inductor's mm template emits tl.assume(pid >= 0) unconditionally, so every matmul from that template carried it. --- .github/workflows/triton_npu.yml | 267 +++++++++++++++++++++++++++++++ Dockerfile.tnpu | 42 +++++ scripts/ci/tnpu_base_pin.sh | 8 + thirdparty/triton-npu.json | 11 ++ 4 files changed, 328 insertions(+) create mode 100644 .github/workflows/triton_npu.yml create mode 100644 Dockerfile.tnpu create mode 100755 scripts/ci/tnpu_base_pin.sh create mode 100644 thirdparty/triton-npu.json diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml new file mode 100644 index 00000000..c1d2ffde --- /dev/null +++ b/.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 diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu new file mode 100644 index 00000000..b9d939a7 --- /dev/null +++ b/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 diff --git a/scripts/ci/tnpu_base_pin.sh b/scripts/ci/tnpu_base_pin.sh new file mode 100755 index 00000000..59663836 --- /dev/null +++ b/scripts/ci/tnpu_base_pin.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Deterministic short pin for tagging torchsim_tnpu_base images. +# Mirrors thirdparty_base_pin.sh, over the tnpu manifest + its Dockerfile, so the +# ~1.8 GiB toolchain layer is rebuilt only when one of those two actually moves. +set -euo pipefail +ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +cd "$ROOT" +{ cat thirdparty/triton-npu.json; cat Dockerfile.tnpu; } | sha256sum | awk '{print substr($1,1,12)}' diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json new file mode 100644 index 00000000..d6d9f4ae --- /dev/null +++ b/thirdparty/triton-npu.json @@ -0,0 +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. 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": "8ebe408dbe52b84a825dcc01152d4a70d211e601", + "release_tag": "toolchain-llvm23", + "also_reads": [ + "PSAL-POSTECH/triton_shared" + ] + } +} From 210a50534b24a8e37edd7b6e48388cbe7e528c64 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:47:59 +0900 Subject: [PATCH 03/10] [Frontend] Add the Triton codegen route (Inductor triton + tnpu passes) A second path from an fx graph to a simulated kernel: Inductor's own Triton backend emits the kernel, and triton-npu lowers it, instead of the MLIR templates doing both. It is selected by TORCHSIM_TRITON_CODEGEN, read once at device registration, so no test knows which route it is on. kernel_spec captures what tnpu needs while V.graph is still live -- there is no second chance once codegen returns -- and turns it into a tnpu KernelSpec. Several things have to be carried that Inductor never writes down for a device that benchmarks its own kernels: The grid. kernel.numels is keyed by iteration-space prefix, and the axes go outermost-first for tile decisions but x-first for the pid order tnpu reads. Getting either wrong silently produces a launch that matches nothing. The strides. Inductor allocates outputs with empty_strided and indexes them by stride, so a launch assuming contiguous writes elements to the wrong places -- and the wrongness is quiet. The scalars, which reach the tnpu wrapper as typed C arguments. helpers_shim copies torch's triton_helpers into the tnpu venv, since the kernel Inductor writes imports it. _triton_compat strips what only a GPU can honour: num_warps, num_stages, tl.debug_barrier, set_driver_to_gpu. The strip has to keep libdevice and tl_math, which used to go out with the torch import that named them. functional writes and reads tensors in storage order via as_strided, so a non-contiguous input round-trips through a .raw file unchanged. A kernel that cannot be described is rejected here with the missing field named, rather than written out to fail somewhere the cause is unrecoverable. --- PyTorchSimDevice/torch_openreg/__init__.py | 29 +- PyTorchSimFrontend/extension_config.py | 15 + PyTorchSimFrontend/triton_backend/__init__.py | 38 ++ .../triton_backend/_triton_compat.py | 143 ++++++ .../triton_backend/codecache.py | 90 ++++ .../triton_backend/functional.py | 164 +++++++ .../triton_backend/helpers_shim.py | 79 +++ .../triton_backend/kernel_spec.py | 463 ++++++++++++++++++ .../triton_backend/scheduling.py | 92 ++++ .../triton_backend/tnpu_bridge.py | 108 ++++ .../triton_backend/wrapper_codegen.py | 21 + 11 files changed, 1236 insertions(+), 6 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/__init__.py create mode 100644 PyTorchSimFrontend/triton_backend/_triton_compat.py create mode 100644 PyTorchSimFrontend/triton_backend/codecache.py create mode 100644 PyTorchSimFrontend/triton_backend/functional.py create mode 100644 PyTorchSimFrontend/triton_backend/helpers_shim.py create mode 100644 PyTorchSimFrontend/triton_backend/kernel_spec.py create mode 100644 PyTorchSimFrontend/triton_backend/scheduling.py create mode 100644 PyTorchSimFrontend/triton_backend/tnpu_bridge.py create mode 100644 PyTorchSimFrontend/triton_backend/wrapper_codegen.py diff --git a/PyTorchSimDevice/torch_openreg/__init__.py b/PyTorchSimDevice/torch_openreg/__init__.py index e8158391..2667ac70 100644 --- a/PyTorchSimDevice/torch_openreg/__init__.py +++ b/PyTorchSimDevice/torch_openreg/__init__.py @@ -18,13 +18,30 @@ sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')) import PyTorchSimFrontend.extension_config # noqa: F401 +from PyTorchSimFrontend import extension_config as _extension_config from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen -from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling -torch._inductor.codegen.common.register_backend_for_device( - "npu", - lambda scheduling: MLIRScheduling(scheduling), - ExtensionWrapperCodegen -) + +# Two mutually exclusive codegen routes for `npu`, chosen here because Inductor +# registers a backend per device, once. +# MLIR (default) hand-written MLIR emission, PyTorchSimFrontend/mlir +# Triton (opt-in) Inductor's own Triton codegen + the triton-npu passes, +# TORCHSIM_TRITON_CODEGEN=1. WIP; see +# PyTorchSimFrontend/triton_backend/README.md +if _extension_config.CONFIG_TRITON_CODEGEN: + from PyTorchSimFrontend.triton_backend import ( + TritonNPUScheduling, TritonNPUWrapperCodegen) + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: TritonNPUScheduling(scheduling), + TritonNPUWrapperCodegen + ) +else: + from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: MLIRScheduling(scheduling), + ExtensionWrapperCodegen + ) torch_openreg.openreg.init() sys.modules['torch.npu'] = torch_openreg.openreg diff --git a/PyTorchSimFrontend/extension_config.py b/PyTorchSimFrontend/extension_config.py index 2d706bbc..09e5168d 100644 --- a/PyTorchSimFrontend/extension_config.py +++ b/PyTorchSimFrontend/extension_config.py @@ -11,6 +11,21 @@ CONFIG_TORCHSIM_DUMP_MLIR_IR = int(os.environ.get("TORCHSIM_DUMP_MLIR_IR", default=False)) CONFIG_TORCHSIM_DUMP_LLVM_IR = int(os.environ.get("TORCHSIM_DUMP_LLVM_IR", default=False)) +# --- Triton codegen route (WIP, opt-in) -------------------------------------- +# Replaces the hand-written MLIR emission in PyTorchSimFrontend/mlir with +# Inductor's own Triton codegen, lowered to the NPU by the triton-npu (tnpu) +# pass pipeline. OFF by default: the MLIR route stays the production path until +# this one is complete. See PyTorchSimFrontend/triton_backend/README.md. +CONFIG_TRITON_CODEGEN = bool(int(os.environ.get("TORCHSIM_TRITON_CODEGEN", default=0))) +# The triton-npu checkout that owns stages 1-5 (ttir -> ttshared -> tnpu passes +# -> RISC-V ELF). It is a SEPARATE repository, deliberately not vendored. +CONFIG_TNPU_DIR = os.environ.get( + "TNPU_DIR", default=os.path.join(CONFIG_TORCHSIM_DIR, "triton-npu")) +# tnpu runs in its own process: its passes need LLVM 23's MLIR bindings while +# this process holds LLVM 20's, and `mlir` is a namespace package, so the two +# cannot coexist in one interpreter (tnpu/config.py:activate_bindings). +CONFIG_TNPU_PYTHON = os.environ.get("TNPU_PYTHON", default=sys.executable) + def get_dump_path(): """Resolve TORCHSIM_DUMP_PATH and re-point Inductor's cache dir at it. diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py new file mode 100644 index 00000000..cb3b88cd --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -0,0 +1,38 @@ +"""The Triton codegen route: Inductor's Triton backend + the tnpu lowering passes. + +WHAT THIS REPLACES +------------------ +The production path emits MLIR by hand (`PyTorchSimFrontend/mlir/`, ~5,500 lines +of string emission that also decides tiling, vectorization, DMA and scratchpad +placement). This route keeps Inductor's OWN Triton codegen for "what to compute" +and hands the resulting Triton kernel to triton-npu for "how to map it onto the +NPU": + + Inductor -> TritonNPUScheduling.define_kernel (scheduling.py) + | triton kernel SOURCE TEXT + v + -> TritonNPUCodeCache.load (codecache.py) + | a tnpu KernelSpec file (kernel_spec.py) + v + -> triton-npu, in a subprocess (tnpu_bridge.py) + ttir -> ttshared -> tnpu passes -> RISC-V ELF + v + -> Spike (functional) / TOGSim (timing) + +The two routes are mutually exclusive and chosen at device-registration time by +`extension_config.CONFIG_TRITON_CODEGEN` (env `TORCHSIM_TRITON_CODEGEN=1`). +Default OFF -- nothing here is on the production path yet. + +STATUS: scaffolding. The seams are wired and each stage names precisely what it +still owes; see README.md for the gap list. Expect failures, not results. +""" + +from . import _triton_compat, 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/_triton_compat.py b/PyTorchSimFrontend/triton_backend/_triton_compat.py new file mode 100644 index 00000000..9c863b15 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -0,0 +1,143 @@ +"""Let Inductor's Triton codegen run on a machine with no GPU. + +ONE SKEW LEFT, AND IT IS NOT A VERSION SKEW +------------------------------------------- +`triton_hash_with_backend()` asks the triton runtime driver for the *current +target*, which on a box with no GPU raises "0 active drivers". We never launch +through triton's runtime -- triton-npu compiles the kernel ahead of time to a +RISC-V ELF -- so the value is only a cache-key ingredient, and a deterministic +string does the job. + +WHAT USED TO BE HERE, AND WHY IT IS GONE +---------------------------------------- +On torch 2.8 there was a second, larger skew: Inductor targeted the triton ~3.3 +API while triton-npu pins 3.6 (3.6 is what pins LLVM 23, and both sides of +triton-npu's textual IR seam must be the same LLVM), so `triton_key` had to be +injected back into `triton.compiler.compiler`. + +torch 2.10 pins triton 3.6.0 itself and reaches that symbol through its own +compat layer (`torch._inductor.runtime.triton_compat`), so on 2.10 the versions +simply agree and the injection is a no-op. It is kept, guarded, so the module +still works if someone runs an older torch -- `_torch_handles_triton()` decides. +""" + +import functools +import hashlib +import importlib +import os +import sys + +_installed = False + + +def triton_src_dir(): + """Where tnpu's triton checkout lives (its editable install points here). + + Read out of tnpu's own `setup/versions.env` rather than guessed, so the two + repos cannot drift: that file is the single place the checkout layout is + pinned (HEXAGON_MLIR_ROOT). + """ + from PyTorchSimFrontend import extension_config + override = os.environ.get("TNPU_TRITON_SRC") + if override: + return override + + root = "/workspace/hexagon-mlir" + versions = os.path.join(extension_config.CONFIG_TNPU_DIR, "setup", "versions.env") + try: + with open(versions) as f: + for line in f: + if line.startswith("HEXAGON_MLIR_ROOT="): + root = line.split("=", 1)[1].strip() + break + except OSError: + pass + return os.path.join(root, "triton", "python") + + +def ensure_triton_importable(): + """`import triton` in THIS interpreter, borrowing tnpu's checkout if needed. + + Inductor's Triton codegen imports triton at codegen time (for metadata and + hashing), so the driver needs it even though it never compiles with it. + """ + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + cand = triton_src_dir() + if os.path.isdir(os.path.join(cand, "triton")): + sys.path.insert(0, cand) + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + return False + + +def _stable_backend_hash(): + try: + import triton + version = triton.__version__ + except Exception: # noqa: BLE001 + version = "unknown" + key = f"pytorchsim-tnpu-{version}" + return hashlib.sha256(key.encode("utf-8")).hexdigest().upper() + + +def _torch_handles_triton(): + """True when this torch already knows how to reach triton's key itself. + + torch 2.10 routes it through torch._inductor.runtime.triton_compat, which + understands triton 3.6. Older torch imports `triton_key` straight out of + triton.compiler.compiler, where 3.6 no longer defines it. + """ + try: + from torch._inductor.runtime.triton_compat import triton_key # noqa: F401 + return True + except Exception: # noqa: BLE001 + pass + try: + mod = importlib.import_module("triton.compiler.compiler") + except Exception: # noqa: BLE001 + return False + return hasattr(mod, "triton_key") + + +def install(): + """Idempotently apply the shims. Returns a short report for logging.""" + global _installed + notes = [] + if not ensure_triton_importable(): + raise ModuleNotFoundError( + f"the Triton codegen route needs `triton` importable in this " + f"interpreter (Inductor imports it during codegen). Not found, and " + f"no checkout at {triton_src_dir()}. Set TNPU_TRITON_SRC, or install " + f"triton into this environment.") + if _installed: + return notes + + if not _torch_handles_triton(): + # Pre-2.10 torch: `triton_key` is imported from triton.compiler.compiler + # by SEVERAL call sites (codecache.CacheBase.get_system, + # _triton.triton_hash_with_backend, ...), each with its own local import. + # Supplying the symbol on the triton side satisfies all of them at once + # instead of chasing every call site; it is a cache-key ingredient, so + # any stable string will do. On torch 2.10 this branch does not run. + mod = importlib.import_module("triton.compiler.compiler") + mod.triton_key = _stable_backend_hash + notes.append("injected triton.compiler.compiler.triton_key " + "(this torch predates the triton 3.6 compat layer)") + + # Separately: triton_hash_with_backend also asks the triton runtime driver + # for the current target, which needs a GPU. We compile ahead of time to a + # RISC-V ELF and never use triton's runtime, so short-circuit it. + import torch.utils._triton as _t + _t.triton_hash_with_backend = functools.cache(_stable_backend_hash) + notes.append("patched torch.utils._triton.triton_hash_with_backend " + "(no GPU target to query)") + + _installed = True + return notes diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py new file mode 100644 index 00000000..d62de72d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -0,0 +1,90 @@ +"""Compile cache for the Triton route -- the counterpart of extension_codecache. + +`triton_npu_compile` is what the generated wrapper calls, exactly where the MLIR +route calls `custom_async_compile.mlir(...)`. It compiles the Triton kernel via +tnpu and returns the callable the wrapper then invokes per launch. + + define_kernel -> triton_npu_compile(src, meta, kernel_name) -> launcher + call site -> launcher(arg0, arg1, ..., xnumel) + +Layout mirrors the MLIR route so the two are comparable: one directory per source +hash under the dump path, holding the generated tnpu kernel file and every tnpu +artifact (01-ttir.mlir ... 05-*.elf). +""" + +import os + +from filelock import FileLock +from torch._inductor.codecache import get_hash + +from PyTorchSimFrontend import extension_config +from . import functional, kernel_spec, timing, tnpu_bridge + +logger = extension_config.setup_logger() + +LOCK_TIMEOUT = 600 + + +def _write_path(src_code): + return os.path.join(extension_config.get_dump_path(), + "triton_" + get_hash(src_code.strip())[1:12]) + + +class TritonNPULauncher: + """What a compiled kernel name is bound to in the generated wrapper. + + Holds the compile result; each call is one launch of the whole grid. + """ + + def __init__(self, kernel_name, workdir, meta): + self.kernel_name = kernel_name + self.workdir = workdir + self.meta = meta + self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") + + def __call__(self, *args): + """One launch of the whole grid: run it on Spike, then time it. + + Spike runs first so the caller's output tensors hold real values even if + TOGSim fails -- the two halves are independent. + """ + if extension_config.pytorchsim_functional_mode: + written = functional.run(self.workdir, self.meta, args) + logger.info("[Spike] %s wrote %s", self.kernel_name, written) + else: + logger.warning( + "[Spike] %s: functional mode is off, so the output tensors keep " + "whatever they held", self.kernel_name) + + if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): + timing.emit_trace(self.workdir, self.meta) + result = timing.run_togsim(self.workdir, meta=self.meta, args=args) + logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) + return result + + +def triton_npu_compile(src_code, meta, kernel_name): + """Compile one Inductor-generated Triton kernel through tnpu. + + Called from the generated wrapper at module import time (same point as + `custom_async_compile.mlir`). Synchronous for now: the MLIR route's thread + pool buys nothing until the pipeline itself is proven. + """ + write_path = _write_path(src_code) + os.makedirs(write_path, exist_ok=True) + + lock = FileLock(os.path.join(write_path, ".compile.lock"), timeout=LOCK_TIMEOUT) + with lock: + spec_path = os.path.join(write_path, f"{kernel_name}_spec.py") + elf = os.path.join(write_path, f"05-{kernel_name}.elf") + if not os.path.isfile(elf): + # Before write_spec_file, which rejects exactly the kernels whose + # source is worth keeping. + with open(os.path.join(write_path, "kernel.py"), "w") as f: + f.write(src_code) # the unmodified Inductor source + kernel_spec.write_spec_file(src_code, meta, spec_path, + tnpu_bridge.tnpu_dir()) + timing.store_meta(write_path, meta) # lets the timing step run standalone + tnpu_bridge.run_pipeline(spec_path, write_path, to_stage="binary") + logger.info("[triton-npu] %s -> %s", kernel_name, write_path) + return TritonNPULauncher(kernel_name, write_path, meta) diff --git a/PyTorchSimFrontend/triton_backend/functional.py b/PyTorchSimFrontend/triton_backend/functional.py new file mode 100644 index 00000000..3ba91daa --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -0,0 +1,164 @@ +"""The functional half of the Triton route: real tensors -> Spike -> real tensors. + +The timing half (timing.py) tells you how long the kernel takes; this tells you +whether it computed the right thing. tnpu's stage 6 already runs the ELF under +Spike, but on inputs it generates itself. Here the launch's own tensors are +written as the `.raw` files stage 6 reads, and the outputs are copied back: + + run(workdir, meta, args) args -> runtime/*.raw -> spike -> args + +The binary is shape-specialised -- the spec bakes the grid, the scalar values and +the memref extents in -- so a launch whose shapes differ from the compiled ones +is rejected rather than silently run against the wrong bounds. +""" + +import os +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +RUNTIME_DIR = "runtime" + + +class ShapeMismatch(RuntimeError): + """The launch does not match the shapes the binary was compiled for.""" + + +def _np_dtype(name): + import numpy as np + return np.dtype("bool" if name == "bool" else name) + + +def tensor_args(meta, args): + """[(arg_meta, tensor)] for the launch, paired by position. + + Inductor passes the tensors first and the numels after, in signature order, + so `meta["args"]` (tensors only) lines up with the leading arguments. + """ + import torch + + tensors = [a for a in args if isinstance(a, torch.Tensor)] + metas = meta["args"] + if len(tensors) != len(metas): + raise ShapeMismatch( + f"{meta['kernel_name']}: launch passed {len(tensors)} tensor(s), " + f"but the spec declares {len(metas)} ({[m['name'] for m in metas]})") + return list(zip(metas, tensors)) + + +def _check(meta, pairs): + for m, t in pairs: + if t.numel() != m["numel"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' has {t.numel()} " + f"element(s), but the binary was compiled for {m['numel']}. " + f"tnpu bakes the extents, the grid and the scalar values into " + f"the kernel, so a dynamic-shape graph reuses an ELF that does " + f"not fit. The timing path does handle this (it takes the grid " + f"at run time); set pytorchsim_functional_mode: False to study " + f"cycles alone, or keep shapes static to check values.") + if str(t.dtype).removeprefix("torch.") != m["dtype"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' is {t.dtype}, but the " + f"binary was compiled for {m['dtype']}") + + +def _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. + + Outputs are written too, as zeros: the wrapper loads and dumps by argv + position, so a missing file shifts every later one. + """ + import numpy as np + + pairs = tensor_args(meta, args) + _check(meta, pairs) + + runtime = os.path.join(workdir, RUNTIME_DIR) + os.makedirs(runtime, exist_ok=True) + for m, t in pairs: + path = os.path.join(runtime, f"{m['name']}.raw") + if m["role"] in ("in", "inout"): + _storage_view(t.detach().to("cpu"), m).numpy().tofile(path) + else: + np.zeros(m["numel"], dtype=_np_dtype(m["dtype"])).tofile(path) + return runtime + + +def read_outputs(workdir, meta, args): + """Copy the .raw files Spike wrote back into the launch's output tensors.""" + import numpy as np + import torch + + runtime = os.path.join(workdir, RUNTIME_DIR) + written = [] + for m, t in tensor_args(meta, args): + if m["role"] not in ("out", "inout"): + continue + path = os.path.join(runtime, f"{m['name']}.raw") + flat = np.fromfile(path, dtype=_np_dtype(m["dtype"])) + if flat.size != m["numel"]: + raise RuntimeError( + f"{path} holds {flat.size} element(s), expected {m['numel']} " + f"-- Spike did not write the whole tensor") + 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 + + spec = os.path.join(workdir, f"{meta['kernel_name']}_spec.py") + if not os.path.isfile(spec): + raise FileNotFoundError(f"{spec} not found -- compile the kernel first") + + write_inputs(workdir, meta, args) + + env = dict(os.environ) + env.pop("PYTHONPATH", None) # keep tnpu on its own MLIR bindings + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, "-m", "tnpu.spike", spec, workdir], + capture_output=True, text=True, cwd=tnpu_bridge.tnpu_dir(), env=env, + timeout=timeout_sec) + if proc.returncode != 0: + raise 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/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py new file mode 100644 index 00000000..f01881d1 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -0,0 +1,463 @@ +"""Inductor kernel -> tnpu KernelSpec. + +Two jobs, both of which exist because Inductor's Triton output is written for a +GPU launcher and tnpu's is written for a static, ahead-of-time pipeline: + +1. `collect_meta` -- pull everything tnpu needs out of the Inductor kernel while + we still have `V.graph`: argument names/roles/dtypes/sizes, the constexprs, + and the numels the grid is computed from. This runs at codegen time; by the + time the compile callable fires, `V.graph` is gone. + +2. `write_spec_file` -- turn the Triton source + that metadata into a kernel file + tnpu can load (`tnpu.spec.load_spec`). + +WHY THE SOURCE HAS TO BE REWRITTEN +---------------------------------- +Inductor emits, above the kernel: + + from torch._inductor.runtime import triton_heuristics + @triton_heuristics.pointwise(size_hints={'x': 1024}, ..., inductor_meta=...) + @triton.jit + def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): + +Neither line can survive into tnpu: + + * the tnpu triton venv has NO torch (deliberately -- tnpu/spec.py), so the + `torch._inductor.runtime` import fails on sight; + * `triton_heuristics.pointwise` is the AUTOTUNER. It picks XBLOCK at runtime + and derives `grid = cdiv(xnumel, XBLOCK)` from it. tnpu needs both to be + constants: the block size becomes a `tl.constexpr` in the ttir signature and + the grid is executed as a sequential loop by the generated C wrapper. + +So the decorator is stripped and XBLOCK is pinned as a constexpr instead. That is +not a workaround -- fixing the config at codegen time is what makes the kernel +statically describable, which is the whole premise of this route. +""" + +import math +import os +import re + +from torch._inductor.virtualized import V + +from . import helpers_shim + +#: Triton signature token -> torch dtype name. The full set Inductor's +#: `_type_of` can emit (torch/_inductor/codegen/triton_utils.py); a token that +#: is absent here reads as an unresolvable arg, so partial coverage shows up as +#: SpecIncomplete on the first kernel that widens past int8. +_DTYPE = { + "*fp64": "float64", "*fp32": "float32", "*fp16": "float16", + "*bf16": "bfloat16", + "*i64": "int64", "*i32": "int32", "*i16": "int16", "*i8": "int8", + "*i1": "bool", + "*u64": "uint64", "*u32": "uint32", "*u16": "uint16", "*u8": "uint8", + "fp64": "float64", "fp32": "float32", "fp16": "float16", + "bf16": "bfloat16", + "i64": "int64", "i32": "int32", "i16": "int16", "i8": "int8", "i1": "bool", + "u64": "uint64", "u32": "uint32", "u16": "uint16", "u8": "uint8", +} + + +#: Triton scalar token -> C type, for the wrapper's kernel declaration. +_C_TYPE = {"i32": "int32_t", "i64": "int64_t", "fp32": "float"} + + +class SpecIncomplete(RuntimeError): + """Metadata tnpu requires that this kernel did not provide. + + Raised with the missing field named, rather than writing a spec that fails + deeper in the pipeline where the cause is unrecoverable. + """ + + +# --------------------------------------------------------------------------- +# 1. codegen-time metadata capture +# --------------------------------------------------------------------------- +def _buffer_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, 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 *= s + return n, size, stride + except Exception: # noqa: BLE001 - best effort; caller reports it as missing + return None, None, None + + +def _roles(kernel): + """arg name -> 'in' | 'out' | 'inout', from the kernel's buffer tables.""" + out = {} + for buf, arg in getattr(kernel.args, "input_buffers", {}).items(): + out[arg] = ("in", buf) + for buf, arg in getattr(kernel.args, "output_buffers", {}).items(): + # 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) + return out + + +def collect_meta(kernel, kernel_name): + """Everything the compile step needs, as plain repr-able data. + + Must run while `V.graph` is live (i.e. inside define_kernel). + """ + triton_meta = dict(getattr(kernel, "triton_meta", None) or {}) + signature = dict(triton_meta.get("signature") or {}) + constants = dict(triton_meta.get("constants") or {}) + + roles = _roles(kernel) + arg_defs, _call_args, _precompile, _arg_types = kernel.args.python_argdefs() + + args = [] + for a in arg_defs: + name = getattr(a, "name", str(a)) + role, buf = roles.get(name, (None, None)) + if role is None: + continue # a numel / constexpr, not a tensor + 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": numel, + "size": size, + "stride": stride, + }) + + # The numels Inductor appends to the call. They live in `kernel.numels`, + # keyed by iteration-space PREFIX ('x', 'y', 'r0', ...), not as xnumel/rnumel + # attributes (SIMDKernel.__init__ builds them from the tiling). These are + # what the grid is computed from. + numels = {} + for prefix, val in (getattr(kernel, "numels", None) or {}).items(): + try: + numels[f"{prefix}numel"] = int(V.graph.sizevars.size_hint(val)) + except Exception: # noqa: BLE001 - dynamic shape; reported by _grid + numels[f"{prefix}numel"] = None + + return { + "kernel_name": kernel_name, + "signature": {str(k): str(v) for k, v in signature.items()}, + "constants": {str(k): v for k, v in constants.items()}, + "args": args, + "numels": numels, + "inside_reduction": bool(getattr(kernel, "inside_reduction", False)), + "fixed_config": fixed_config_for(kernel), + } + + +#: Parallel iteration prefixes, OUTERMOST first. Inductor's `x` is the +#: contiguous axis, so it is innermost; `r*` prefixes are reductions, looped +#: inside the kernel rather than spread over the grid (prefix_is_reduction). +_PARALLEL_PREFIXES = ("z", "y", "x") + + +def _block_name(prefix): + return f"{prefix.upper()}BLOCK" + + +def parallel_axes(numels): + """Grid axes this kernel uses, outermost first. 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. + + tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a + sequential loop, so there is no autotuner to choose the blocks later and no + runtime `grid=` callable. Fixing them here is what makes the launch shape + static. + + Tile dim 0 is the one `bank_vectorize` spreads over the lanes, so the + OUTERMOST axis gets the lane count -- a per-lane depth of 1, the shape every + tnpu baseline runs. The remaining axes get 1, which leaves the tile exactly + that verified shape and lets the grid cover the rest. It is conservative + rather than fast; choosing real tile sizes is the block-size policy gap in + README, not something to guess at here. + """ + from PyTorchSimFrontend import extension_config + 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 + # 1 makes every work-item move a strided column. Fine for getting a + # multi-axis kernel through the route, misleading to benchmark. + extension_config.setup_logger().warning( + "[triton-npu] %s tiles over %s; inner blocks pinned to 1, which is " + "correct but not a tiling worth measuring", + getattr(kernel, "kernel_name", "kernel"), axes) + cfg.setdefault("XBLOCK", lanes) # a kernel with no tiling info still has x + if getattr(kernel, "inside_reduction", False): + # 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 + + +# --------------------------------------------------------------------------- +# 2. Triton source -> tnpu kernel file +# --------------------------------------------------------------------------- +_HEURISTIC_RE = re.compile(r"^@triton_heuristics\.") +_DROP_IMPORT_RE = re.compile( + r"^\s*(import torch|from torch\b|from __future__|import __main__)") +#: 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): + """Remove everything the torch-free tnpu venv cannot import. + + Drops torch/inductor imports and the `@triton_heuristics.*(...)` decorator + (keeping `@triton.jit`), then re-adds the two imports the kernel body needs. + + Raises SpecIncomplete if the kernel still calls into `triton_helpers`: that + module lives in torch, so it has to be vendored into the tnpu venv before + such a kernel can compile. Failing here names the missing helper; letting it + through fails as a bare NameError inside tnpu's stage-1 worker instead. + """ + lines = src.splitlines() + out, i = [], 0 + while i < len(lines): + line = lines[i] + if _HEURISTIC_RE.match(line.strip()) or _HEURISTIC_RE.match(line): + # skip the whole decorator call, up to (not including) @triton.jit + while i < len(lines) and lines[i].strip() != "@triton.jit": + i += 1 + continue + if _DROP_IMPORT_RE.match(line) or _DROP_CALL_RE.match(line): + i += 1 + continue + out.append(line) + i += 1 + body = "\n".join(out) + + 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 + + +def scalar_args(meta): + """User scalar parameters, in kernel order, as [(name, c_type, value)]. + + triton-shared keeps these in the lowered signature ahead of its own six + grid/pid arguments, so the wrapper must pass them or every later argument + lands one slot early -- pidX then reads pidY and only program 0 runs. + """ + numels = meta["numels"] + out = [] + for name, token in meta["signature"].items(): + if token.startswith("*") or token == "constexpr": + continue + ctype = _C_TYPE.get(token) + if ctype is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: scalar '{name}' has type {token!r}, " + f"which has no C mapping in _C_TYPE") + if numels.get(name) is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: no value for scalar '{name}' -- " + f"collect_meta resolves these from kernel.numels") + out.append((name, ctype, int(numels[name]))) + return out + + +def grid_of(meta): + """Launch grid, from the numels and the pinned block sizes, 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 = pid_axes(numels) + if not axes: + raise SpecIncomplete( + f"{meta['kernel_name']} has no parallel iteration axis to grid over") + + grid = [] + for prefix in axes: + n, block = numels.get(f"{prefix}numel"), cfg.get(_block_name(prefix)) + if n is None or not block: + raise SpecIncomplete( + f"cannot compute the grid for {meta['kernel_name']} axis " + f"'{prefix}': {prefix}numel={n!r}, {_block_name(prefix)}={block!r}. " + f"Inductor defers the grid to triton_heuristics at runtime; this " + f"route needs it statically (see fixed_config_for).") + grid.append(int(math.ceil(n / block))) + return tuple(grid) + + +SPEC_TEMPLATE = '''\ +"""Generated by PyTorchSimFrontend/triton_backend/kernel_spec.py -- do not edit. + +Inductor kernel {kernel_name!r}, rewritten for the tnpu pipeline: the +triton_heuristics autotuner decorator is stripped and its block sizes are pinned +as constexprs, so the launch shape is static. See kernel_spec.py for why. +""" +import importlib.util +import os +import sys + +sys.path.insert(0, {tnpu_dir!r}) +# 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 +#: disk, not an exec'd string: triton's @jit reads the function back with +#: inspect.getsourcefile and rejects anything else ("@jit functions should be +#: defined in a Python file"). +TRITON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), + {triton_module!r}) + + +def kernel(): + spec = importlib.util.spec_from_file_location( + {kernel_name!r} + "_triton", TRITON_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, {kernel_name!r}) + + +def make_inputs(torch, seed=0): + g = torch.Generator().manual_seed(seed) + out = {{}} +{make_inputs_body} + return out + + +def reference(inputs): + # The Inductor route has no per-kernel torch reference: correctness is + # checked at the graph level by the test that ran torch.compile. tnpu's + # stage 7 is therefore not meaningful here and the pipeline is driven to + # stage 6 (spike) instead. + return {{}} + + +SPEC = KernelSpec( + name={kernel_name!r}, + kernel=kernel, + signature={signature!r}, + constexprs={constexprs!r}, + args=[ +{args_body} + ], + grid={grid!r}, + reference=reference, + make_inputs=make_inputs, + extra={{"scalar_args": {scalar_decls!r}, + "scalar_values": {scalar_values!r}}}, + notes="generated from Inductor triton codegen", +) +''' + + +def write_spec_file(src_code, meta, path, tnpu_dir): + """Write a tnpu kernel file for this Inductor kernel. Returns `path`.""" + missing = [a["name"] for a in meta["args"] if not a["dtype"] or not a["numel"]] + if missing: + raise SpecIncomplete( + f"{meta['kernel_name']}: no dtype/numel for {missing} -- " + f"collect_meta could not resolve them from V.graph") + + signature = dict(meta["signature"]) + constexprs = dict(meta["constants"]) + for k, v in (meta.get("fixed_config") or {}).items(): + if k not in signature: + # Inductor already fixed this one in the kernel BODY rather than + # taking it as a parameter -- a persistent reduction does that with + # R0_BLOCK. Passing it would not match the signature, and there is + # nothing left for us to choose. + continue + if v is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: block size {k} is unset " + f"(fixed_config_for leaves reduction blocks unset on purpose)") + constexprs[k] = v + signature[k] = "constexpr" + + args_body = "\n".join( + f" Arg({a['name']!r}, {a['role']!r}, {a['dtype']!r}, ({a['numel']},))," + for a in meta["args"]) + make_inputs_body = "\n".join( + f" out[{a['name']!r}] = torch.randn({a['numel']}, generator=g)" + f".to(torch.{a['dtype']})" + for a in meta["args"] if a["role"] in ("in", "inout")) or " pass" + + triton_module = f"{meta['kernel_name']}_triton.py" + 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(stripped) + + scalars = scalar_args(meta) + text = SPEC_TEMPLATE.format( + kernel_name=meta["kernel_name"], + tnpu_dir=tnpu_dir, + triton_module=triton_module, + signature=signature, + constexprs=constexprs, + args_body=args_body, + make_inputs_body=make_inputs_body, + grid=grid_of(meta), + scalar_decls=[(n, c) for n, c, _ in scalars], + scalar_values={n: v for n, _, v in scalars}, + ) + with open(path, "w") as f: + f.write(text) + return path diff --git a/PyTorchSimFrontend/triton_backend/scheduling.py b/PyTorchSimFrontend/triton_backend/scheduling.py new file mode 100644 index 00000000..56e7f1f8 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -0,0 +1,92 @@ +"""Inductor scheduling for the Triton route. + +`TritonNPUScheduling` keeps ALL of Inductor's Triton codegen -- fusion, index +expressions, masking, reductions, the kernel source itself -- and changes only +what happens to the generated source afterwards. Upstream hands it to +`async_compile.triton(...)`, which calls `triton.compile` for a GPU; we hand it +to `triton_npu_compile(...)`, which runs the triton-npu pipeline for this NPU. + +Two overrides, and nothing else: + + define_kernel emit our compile call into the wrapper instead of upstream's. + kernel_type a TritonKernel whose call site is a plain python call, because + the name is bound to our callable rather than to a triton + launcher with a `.run(grid=..., stream=...)` interface. +""" + +from torch._inductor.codegen.common import IndentedBuffer +from torch._inductor.codegen.triton import TritonKernel, TritonScheduling +from torch._inductor.utils import Placeholder, get_fused_kernel_name +from torch._inductor.virtualized import V + +from . import kernel_spec + + +class TritonNPUKernel(TritonKernel): + """A TritonKernel launched as a plain call. + + Upstream emits `kernel.run(a, b, xnumel, grid=grid(xnumel), stream=...)`, + where `grid` is resolved at RUNTIME by triton_heuristics from the autotuned + XBLOCK. There is no autotuner and no stream here: the kernel name is bound to + the callable `triton_npu_compile` returned, so the call is `kernel(a, b, n)`. + + That is also why the block sizes must be fixed at CODEGEN time -- see + kernel_spec.fixed_config_for. A grid that is only known after autotuning + cannot be written into a tnpu KernelSpec. + """ + + # **kwargs, not a fixed signature: Inductor keeps adding parameters here + # (2.10 added `deallocate_ws`). None of them apply to this route -- there is + # no triton launcher and no workspace to release -- so they are accepted and + # ignored rather than pinning us to one torch release. + def call_kernel(self, name: str, node=None, **kwargs): + wrapper = V.graph.wrapper_code + _, call_args, _, arg_types = self.args.python_argdefs() + self.add_numel_to_call_args(name, call_args, arg_types) + # add_numel_to_call_args appends the numels as SYMPY values, which the + # triton path later renders through pexpr. ExtensionWrapperCodegen joins + # call args as plain strings (mlir_codegen_backend.py:241), so render + # them here instead of handing it a sympy Integer. + call_args = [a if isinstance(a, str) else str(a) for a in call_args] + # triton=False -> PythonWrapperCodegen emits `name(args...)`, the same + # shape the MLIR route uses (mlir_common.py:627). + wrapper.generate_kernel_call(name, call_args, triton=False) + + +class TritonNPUScheduling(TritonScheduling): + kernel_type = TritonNPUKernel + + count = 0 + + def define_kernel(self, src_code, node_schedule, kernel): + wrapper = V.graph.wrapper_code + if src_code in wrapper.src_to_kernel: + return wrapper.src_to_kernel[src_code] + + fused_name = get_fused_kernel_name(node_schedule, "original_aten") + kernel_name = "_".join( + x for x in ("triton_npu", fused_name, str(TritonNPUScheduling.count)) if x + ) + TritonNPUScheduling.count += 1 + wrapper.src_to_kernel[src_code] = kernel_name + + # Upstream substitutes these two placeholders inside define_kernel; the + # source still carries them here, and the tnpu side parses the source, so + # they have to be resolved before it leaves this function. + src_code = src_code.replace(str(Placeholder.DESCRIPTIVE_NAME), kernel_name) + src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name) + + meta = kernel_spec.collect_meta(kernel, kernel_name) + + compile_wrapper = IndentedBuffer() + compile_wrapper.writeline(f"triton_npu_compile('''{src_code}''',") + compile_wrapper.writeline(f" meta={meta!r},") + compile_wrapper.writeline(f" kernel_name={kernel_name!r})") + + origins = ", ".join( + sorted({str(o) for n in node_schedule + for o in getattr(getattr(n, "node", None), "origins", ()) or ()}) + ) + wrapper.define_kernel(kernel_name, compile_wrapper.getvalue(), + f"# origins: {origins}") + return kernel_name diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py new file mode 100644 index 00000000..9599423d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -0,0 +1,108 @@ +"""Run the triton-npu pipeline, out of process. + +WHY A SUBPROCESS +---------------- +tnpu's passes run on LLVM 23's MLIR python bindings; this process holds LLVM 20's +(TORCHSIM_LLVM_PATH). `mlir` ships without an `__init__.py`, so it is a NAMESPACE +package whose `__path__` is the union of every `mlir/` directory on sys.path -- +two LLVMs in one interpreter silently merge and fail later with an AttributeError +from a generated dialect module (tnpu/config.py:activate_bindings documents the +exact failure). They cannot share an interpreter, so tnpu gets its own. + +The seam between them is a FILE, which is measured to work: LLVM 23 prints the +IR, and LLVM 20's bindings parse it back without complaint (verified by feeding +tnpu's 04-custom.mlir to PyTorchSim's build_tog). That is what makes the split +viable rather than merely necessary. +""" + +import os +import re +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + + +class TnpuError(RuntimeError): + """A tnpu stage failed. Inductor reports only str(exc), so the stage's own + diagnostic has to travel in the message.""" + + #: How a failing stage names itself: MLIR diagnostics and exception lines. + _SIGNAL = re.compile( + r"^(?!\s|Traceback|During handling|The above)" + r"(.*\berror:\s.*|.*failed to legalize.*|" + r"[\w.]*(?:Error|Exception)\b.*|.*Assertion.*|" + r".*(?:segfault|Illegal instruction|trap_|bad --isa).*)$", re.M) + #: Frames and carets: context, not the diagnostic. + _FRAME = re.compile(r'^\s|^\s*File "|^\s*\^') + + def __init__(self, message, cmd=None, output=None): + self.cmd = cmd + self.output = output + if output: + hits = [h.strip() for h in self._SIGNAL.findall(output) + if not self._FRAME.match(h)] + if not hits: + hits = [l for l in output.strip().splitlines() + if l.strip() and not self._FRAME.match(l)] + message = message + "\n " + "\n ".join(l[:300] for l in hits[-3:]) + super().__init__(message) + + +def tnpu_dir(): + d = extension_config.CONFIG_TNPU_DIR + if not os.path.isdir(d): + raise TnpuError( + f"triton-npu checkout not found at {d}. It is a separate repository " + f"and is not vendored; clone it there or set TNPU_DIR.") + return d + + +def doctor(): + """Return (ok, output) for tnpu's own toolchain check.""" + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, os.path.join(tnpu_dir(), "run.py"), "doctor"], + capture_output=True, text=True, cwd=tnpu_dir()) + return proc.returncode == 0, proc.stdout + proc.stderr + + +def run_pipeline(spec_path, workdir, to_stage="binary", from_stage="ttir", + verbose=False, timeout=1800): + """Drive tnpu's stages over `spec_path`, writing artifacts into `workdir`. + + Stops at `to_stage`. The default is `binary` (through the RISC-V ELF): + stage 6 (spike) needs the caller's real tensors as .raw files and stage 7 + compares against a per-kernel torch reference, neither of which exists on + the Inductor route -- correctness is a graph-level property there. + + Returns the workdir on success; raises TnpuError with tnpu's own stage + report (which names the failing command and its stderr) otherwise. + """ + cmd = [extension_config.CONFIG_TNPU_PYTHON, + os.path.join(tnpu_dir(), "run.py"), spec_path, + "--from", from_stage, "--to", to_stage, "--workdir", workdir] + if verbose: + cmd.append("-v") + + env = dict(os.environ) + # tnpu deliberately does not read TORCHSIM_LLVM_PATH (it would drag the + # backend back to LLVM 20 and break the textual seam), but a stale + # PYTHONPATH pointing at LLVM 20's mlir_core would still be picked up by the + # namespace package before tnpu's own activate_bindings() runs. + env.pop("PYTHONPATH", None) + + proc = subprocess.run(cmd, capture_output=True, text=True, + cwd=tnpu_dir(), env=env, timeout=timeout) + output = proc.stdout + proc.stderr + if proc.returncode != 0: + # run.py prints a stage table; the diagnostic itself only reaches + # stage.log. + log = os.path.join(workdir, "stage.log") + if os.path.isfile(log): + with open(log, errors="replace") as fh: + output += "\n" + fh.read() + raise TnpuError(f"tnpu pipeline failed (exit {proc.returncode})", + cmd=" ".join(cmd), output=output) + logger.debug("[triton-npu] %s", output) + return workdir diff --git a/PyTorchSimFrontend/triton_backend/wrapper_codegen.py b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py new file mode 100644 index 00000000..cfc6490d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py @@ -0,0 +1,21 @@ +"""Wrapper codegen for the Triton route. + +Reuses `ExtensionWrapperCodegen` wholesale -- device guards, buffer allocation, +the TOGSimulator plumbing and the SRAM plan hooks are all route-independent -- +and adds the one import the generated module needs: `triton_npu_compile`, which +is to this route what `custom_async_compile` is to the MLIR one. +""" + +from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen + +from . import codecache + + +class TritonNPUWrapperCodegen(ExtensionWrapperCodegen): + def write_header(self): + super().write_header() + self.header.splice( + f""" + from {codecache.__name__} import triton_npu_compile + """ + ) From 4412c25edfcc35942f73aba61f1ff65190623048 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:48:36 +0900 Subject: [PATCH 04/10] [Frontend] Run the route's kernels through gem5 and TOGSim The lowered kernel still has to become a tile-operation graph with a cycle cost per tile, the same as the MLIR route's output. timing drives gem5 over a body reduced to one work-item, then hands TOGSim the trace. The grid is taken at run time rather than baked in, so one trace serves every shape the kernel is launched with. Two fixes the shared passes needed. generate() went through memory_plan_reuse directly, which skips the wrapper IR passes estimate_peak depends on; it now enters through run_wrapper_ir_passes, the entry point upstream provides. And both wrap_kernel_call sites stringify their arguments, which were being formatted as objects. --- .../mlir/mlir_codegen_backend.py | 10 +- .../mlir/passes/build_skeleton.py | 2 + PyTorchSimFrontend/mlir/passes/build_tog.py | 97 ++++++-- .../mlir/passes/lower_to_emitc.py | 204 +++++++++++++-- PyTorchSimFrontend/triton_backend/timing.py | 232 ++++++++++++++++++ TOGSim/src/main.cc | 13 +- 6 files changed, 524 insertions(+), 34 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/timing.py 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/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index f4ed7d0d..ed52a56d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -366,6 +366,8 @@ def visit(n): for ln in builder.loop_nodes: visit(ln) + for dn in getattr(builder, "dma_nodes", ()): # DMAs outside any tile loop + visit(dn) return by_op diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index 5a40feec..98590d88 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -414,6 +414,9 @@ def __init__(self): self.loop_var_name = {} # value-identity-key -> loop name self.compute_nodes = [] self.loop_nodes = [] + # `_collect_dma_nodes` descends from the loop nodes, so a DMA hanging off + # the root (no tile loop in the kernel) would be missed. + self.dma_nodes = [] self._reset_matmul_fsm() # ---- matmul FSM ---- @@ -568,7 +571,13 @@ def _process_dram_indices(self, value, loop_index_list, indirect_box): loop_index_list.append(("c" + str(c), c)) # ---- main recursion ---- - def print_operation(self, op, node): + def visit_operation(self, op, node): + """Walk `op` and attach the nodes it produces under `node`. + + Builds the graph; it does not print. (The C++ pass this is ported from + does both in one method, `printOperation` -- here `bfs`/`display` own + the printing.) + """ name = _op_name(op) if name in SKIP_OPS: return @@ -605,7 +614,7 @@ def bool_true(k): for region in oper.regions: for block in region.blocks: for inner in block.operations: - self.print_operation(inner, loop_node) + self.visit_operation(inner, loop_node) return if name == "togsim.transfer": @@ -819,9 +828,14 @@ def _handle_dma_start(self, op, node): loop_idx_list.append(key) loop_stride_list.append(reordered[key]) - # base address + # base address: which tensor this DMA touches. The operand is the block + # argument itself in PyTorchSim's codegen; when it is a view of one + # instead, only the producer knows which -- so it says so (`dram_arg`) + # rather than the consumer guessing its way back through view ops. address = "arg" - if _is_block_arg(dram_memref): + if "dram_arg" in oper.attributes: + address += str(ir.IntegerAttr(oper.attributes["dram_arg"]).value) + elif _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) # element size @@ -875,6 +889,7 @@ def _handle_dma_start(self, op, node): tag_stride_list, loop_idx_list, loop_stride_list, indirect_box[0]) dma_node.op = op + self.dma_nodes.append(dma_node) node.add_child(dma_node) dma_node.add_parent(node) @@ -918,7 +933,9 @@ def _handle_dma_wait(self, op, node): dram_memref = f["dst"] elif dst_space == 1 and src_space == 0: dram_memref = f["src"] - if dram_memref is not None and _is_block_arg(dram_memref): + if "dram_arg" in user.attributes: + address += str(ir.IntegerAttr(user.attributes["dram_arg"]).value) + elif dram_memref is not None and _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) if len(tag_stride_list) == 0: @@ -928,6 +945,7 @@ def _handle_dma_wait(self, op, node): wait_node = TOGDMAWaitNode("DMAWaitNode", tag_index_list, tag_stride_list, tag_divider_list, address) wait_node.op = op + self.dma_nodes.append(wait_node) node.add_child(wait_node) wait_node.add_parent(node) @@ -1064,12 +1082,47 @@ def _insert_compute_markers(builder): # Driver. # --------------------------------------------------------------------------- def _find_kernel(module): - for op in module.body.operations: - if op.operation.name != "func.func": - continue + """The kernel function: named `kernel` in PyTorchSim's codegen, else the + module's only func.func (triton-npu carries the Triton kernel's own name). + Declines when there is more than one -- the intent would be a guess.""" + funcs = [op for op in module.body.operations + if op.operation.name == "func.func"] + for op in funcs: if ir.StringAttr(op.operation.attributes["sym_name"]).value == "kernel": return op - return None + return funcs[0] if len(funcs) == 1 else None + + +#: The loop roles (sec 9.1). Without one, a loop is a micro-loop the compute FSM +#: folds into a single node, not a tile loop. +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") + + +def _has_loop_role(op): + attrs = op.operation.attributes + return any(k in attrs and ir.BoolAttr(attrs[k]).value for k in _LOOP_ROLE_ATTRS) + + +def _is_address_plumbing(op): + """Scalar index/integer math (DMA offsets, mask extents) and the terminator. + + Only consulted on the no-top-level-loop path. PyTorchSim's codegen puts this + math in `affine.apply`, which SKIP_OPS drops; triton-npu emits an + arith/index_cast chain that would otherwise count as vector compute. + + Keyed on result type: tile data here is always vector- or float-typed. A + top-level SCALAR arithmetic kernel would be misread, but no path emits one. + """ + name = _op_name(op) + if name in ("func.return", "memref.cast"): + return True + if not name.startswith("arith."): + return False + results = list(op.operation.results) + if not results: + return False + return all(ir.IndexType.isinstance(r.type) or ir.IntegerType.isinstance(r.type) + for r in results) def _build(module, builder): @@ -1082,13 +1135,29 @@ def _build(module, builder): block = func_op.regions[0].blocks[0] out = [] + # A root is a top-level TILE loop, identified by its role attribute (sec + # 9.1) -- not by being an affine.for: bank_vectorize leaves a bare one for + # the tile's vector work, and rooting there orphans every DMA. + roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] + if roots: + for op in roots: + root = TOGNode("root") + builder._reset_matmul_fsm() + builder.visit_operation(op, root) + root.bfs(out) + return "".join(out) + + # No top-level loop: the body is ONE work-item -- the shape a Triton kernel + # arrives in, its grid becoming the trace producer's dispatch loop (sec 9.3). + # PyTorchSim's codegen keeps the tile loops in the kernel and never lands here. + root = TOGNode("root") + builder._reset_matmul_fsm() for op in block.operations: - if op.operation.name != "affine.for": + if _is_address_plumbing(op): continue - root = TOGNode("root") - builder._reset_matmul_fsm() - builder.print_operation(op, root) - root.bfs(out) + builder.visit_operation(op, root) + root.bfs(out) return "".join(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 5633769a..537c8ad0 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -119,20 +119,174 @@ def _attr_bool(op, key): # --------------------------------------------------------------------------- # step 1: rewrite signature + togsim.* ops (the unregistered-op glue) # --------------------------------------------------------------------------- -def _strip_aux(module): - """Erase memref.global decls and every func except @kernel (the wrapper).""" +def _strip_aux(module, keep=None): + """Erase memref.global decls and every func except the kernel. + + `keep` is the kernel op: its name is `kernel` only in PyTorchSim's codegen, + so the caller passes what `_find_kernel` resolved. + """ + keep_op = keep.operation if keep is not None else None victims = [] for op in module.body.operations: name = op.operation.name if name == "memref.global": victims.append(op) elif name == "func.func": - if ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": + if keep_op is not None: + if op.operation != keep_op: + victims.append(op) + elif ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": victims.append(op) for op in victims: op.operation.erase() +class WorkItem: + """A kernel whose body is ONE work-item, plus the grid over it. + + A Triton kernel describes a single program instance; the grid lives outside + it. The trace producer already splits the same way (design sec 9.3), so only + the enumeration is missing. + + `parallel_args` are the argument positions holding the program ids + (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` + their extents. Both outermost-first. + + An extent may be None, meaning "read it from shape_args at run time". Only + the NUMBER of axes has to be known when the kernel is compiled -- how many + loops to nest and how many iv[] slots to fill; the trip counts are just + values, and the producer ABI already takes them + (togsim_kernel(ctx, shape_args, n)). That is what lets one compiled trace + serve every shape. + """ + + def __init__(self, parallel_args, grid): + if len(parallel_args) != len(grid): + raise ValueError( + f"parallel_args {parallel_args} and grid {grid} must have the " + f"same length -- one program-id argument per grid axis") + self.parallel_args = list(parallel_args) + self.grid = [None if g is None else int(g) for g in grid] + + @property + def dynamic_axes(self): + """Indices into `grid` whose extent arrives at run time.""" + return [i for i, g in enumerate(self.grid) if g is None] + + +def _materialize_grid_loop(kernel, work_item, ctx): + """Wrap the body in the grid loop the Triton kernel does not carry: + + func @k(..., %pid: i32) { + scf.for %p = 0 to G { index_cast %p> } {outer_loop} + } + + Downstream is then unchanged: `_parallel_loop_chain` finds the tagged loop, + the outliner threads its induction variable through `iv[]`, and the loop left + behind becomes the dispatch enumeration. `outer_loop` means "independent + work-item" (sec 9.1) -- exactly a Triton program id. + + MUST run before `_rewrite_signature`, which erases the arguments and first + asserts none are still used. + """ + from mlir.dialects import arith, scf + + block = kernel.regions[0].blocks[0] + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + + pid_args = [block.arguments[i] for i in work_item.parallel_args] + body_ops = [o for o in block.operations + if o.operation.name not in _LOOP_TERMINATORS] + terminator = [o for o in block.operations + if o.operation.name in _LOOP_TERMINATORS][0] + + # Every bound first, and all of them before the first loop: each is created + # just before the terminator, so one made after an outer loop would sit + # BELOW it in the block while an inner loop uses it -- which does not + # dominate, and the verifier rejects it (only reachable at rank >= 2). + with ir.InsertionPoint(terminator), loc: + c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result + c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result + # A runtime extent still needs SOMETHING here: shape_args does not exist + # until _rewrite_signature adds it. The placeholder is replaced by + # _bind_runtime_bounds once it does. + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e or 1)).result + for e in work_item.grid] + + loops, inner = [], None + for ub in ubs: + # Nest inside the previous loop, BEFORE its yield: InsertionPoint on a + # block appends, and an scf.for body is already terminated. + ip = ir.InsertionPoint(terminator) if inner is None \ + else ir.InsertionPoint.at_block_terminator(inner.body) + with ip, loc: + loop = scf.ForOp(c0, ub, c1) + # ForOp leaves the body empty here; scf.for needs a terminator, and + # _outline_work_item inserts before it. + if len(loop.body.operations) == 0: + with ir.InsertionPoint(loop.body), loc: + scf.YieldOp([]) + loop.operation.attributes["outer_loop"] = ir.BoolAttr.get(True) + loops.append(loop) + inner = loop + + # Move the tile body inside the innermost loop, ahead of its yield. + inner_block = inner.body + inner_term = inner_block.operations[len(inner_block.operations) - 1] + for op in body_ops: + op.operation.move_before(inner_term) + + # Program ids are i32, induction variables index: cast once, at the top. + with ir.InsertionPoint(inner_block.operations[0]), loc: + casts = [] + for loop, pid in zip(loops, pid_args): + iv = loop.body.arguments[0] + casts.append(arith.IndexCastOp(pid.type, iv).result + if pid.type != idxty else iv) + + for pid, new in zip(pid_args, casts): + _replace_all_uses(pid, new) + + return [(loops[i], ubs[i]) for i in work_item.dynamic_axes] + + +def _bind_runtime_bounds(pending, shape_arg, ctx): + """Point each runtime loop bound at `shape_args[k]`. + + Runs AFTER _rewrite_signature, which is what creates the shape_args + argument. The loops stay in the entry function (the outliner moves only + their bodies), so the read is in scope where the bound is used. + """ + if not pending: + return + from mlir.dialects import arith + + i64 = ir.IntegerType.get_signless(64) + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + for k, (loop, placeholder) in enumerate(pending): + with ir.InsertionPoint(placeholder.owner), loc: + kc = ir.Operation.create( + "emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}).results[0] + elem = ir.Operation.create( + "emitc.subscript", results=[i64], + operands=[shape_arg, kc]).results[0] + bound = arith.IndexCastOp(idxty, elem).result + _replace_all_uses(placeholder, bound) + placeholder.owner.erase() + + +def _replace_all_uses(old, new): + """The bindings expose no replaceAllUsesWith on a Value.""" + for use in list(old.uses): + owner = use.owner + for i in range(len(owner.operands)): + if owner.operands[i] == old: + owner.operands[i] = new + + def _rewrite_signature(kernel, ctx): """Replace @kernel's memref tensor args with the ABI args (EmitCtx*, int64_t* shape_args, int32_t n) and rename it to togsim_kernel. @@ -196,15 +350,22 @@ def _is_outer(forop): return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value +#: The role is carried by the `outer_loop` attribute, not the dialect: +#: PyTorchSim's codegen emits affine.for, _materialize_grid_loop scf.for. Both +#: keep the induction variable in block argument 0. +_LOOP_OPS = ("affine.for", "scf.for") +_LOOP_TERMINATORS = ("affine.yield", "scf.yield", "func.return") + + def _parallel_loop_chain(block): - """The nested chain of `affine.for {outer_loop}` from `block` inward (one + """The nested chain of `{outer_loop}` loops from `block` inward (one work-item's parallel indices). Empty if the kernel has no parallel loop.""" chain = [] cur = block while True: nxt = None for op in cur.operations: - if op.operation.name == "affine.for" and _is_outer(op): + if op.operation.name in _LOOP_OPS and _is_outer(op): nxt = op break if nxt is None: @@ -281,7 +442,7 @@ def _outline_work_item(ctx, kernel, ctx_val): # move the work-item body into the tile fn (terminators stay behind). for op in [o for o in Lbody.operations - if o.operation.name not in ("affine.yield", "func.return")]: + if o.operation.name not in _LOOP_TERMINATORS]: op.operation.move_before(tret) # remap captures (Value `==` is identity): ctx -> ctx2, each parallel IV -> @@ -337,7 +498,7 @@ def _remap(block): # --- the dispatcher: marshal the IVs and hand the tile fn to togsim_dispatch --- term = [o for o in Lbody.operations - if o.operation.name in ("affine.yield", "func.return")][0] + if o.operation.name in _LOOP_TERMINATORS][0] fn_ref = _opaque(ctx, ts.TILE_SYMBOL) # function name -> verbatim pointer in C with ir.InsertionPoint(term): if ivs: @@ -499,16 +660,25 @@ def _add_extern_c(module, ctx): # --------------------------------------------------------------------------- # driver # --------------------------------------------------------------------------- -def lower_to_emitc(skeleton_module): +def lower_to_emitc(skeleton_module, work_item=None): """Lower a skeleton+API module (in place) to an EmitC module with the - `togsim_kernel` entry function. Returns the same module.""" + `togsim_kernel` entry function. Returns the same module. + + `work_item` is for kernels whose body is one work-item with the grid outside + (Triton's shape); None keeps PyTorchSim's, where the tile loops are already + in the kernel. + """ ctx = skeleton_module.context kernel = _find_kernel(skeleton_module) if kernel is None: - raise ValueError("no @kernel found in skeleton module") + raise ValueError("no kernel function found in skeleton module") - _strip_aux(skeleton_module) + _strip_aux(skeleton_module, keep=kernel) + pending = [] + if work_item is not None: + pending = _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) + _bind_runtime_bounds(pending, kernel.regions[0].blocks[0].arguments[1], ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch @@ -563,18 +733,20 @@ def _default_include_dir(): return os.path.join(root, "TOGSim", "include") -def skeleton_to_so(skeleton_module, so_path, include_dir=None): +def skeleton_to_so(skeleton_module, so_path, include_dir=None, work_item=None): """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the EmitC module text (for inspection / caching).""" - emitc = lower_to_emitc(skeleton_module) + emitc = lower_to_emitc(skeleton_module, work_item=work_item) inc = include_dir or _default_include_dir() cpp = emitc_to_cpp(emitc, include_dir=inc) compile_so(cpp, so_path, inc) return str(emitc) -def build_trace_so(postvcix_path, so_path, include_dir=None): - """Full P2 path from a post-vcix kernel .mlir to a trace `.so`.""" +def build_trace_so(postvcix_path, so_path, include_dir=None, work_item=None): + """Full P2 path from a post-vcix kernel .mlir to a trace `.so`. + + `work_item` -- see lower_to_emitc.""" from . import build_skeleton as bs ctx = ir.Context() @@ -582,7 +754,7 @@ def build_trace_so(postvcix_path, so_path, include_dir=None): with ctx: module = ir.Module.parse(open(postvcix_path).read(), ctx) bs.build_skeleton(module) - return skeleton_to_so(module, so_path, include_dir) + return skeleton_to_so(module, so_path, include_dir, work_item=work_item) def main(argv): diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py new file mode 100644 index 00000000..e6bc7c44 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -0,0 +1,232 @@ +"""The timing half of the Triton route: tnpu IR -> trace.so -> TOGSim. + +TOGSim simulates from a compiled trace producer (docs/design/togsim_cpp_trace.md). +PyTorchSim's codegen already emits one; this emits the same from a Triton-shaped +kernel, where the grid must be supplied -- see `lower_to_emitc.WorkItem`. + + emit_trace(workdir, meta) 04-custom.mlir -> trace.so + trace_cycles.tsv + run_togsim(workdir, ...) hand them to TOGSim, return its parsed result +""" + +import json +import os + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +#: Name TOGSim derives from the kernel directory (Simulator/simulator.py). +TRACE_SO = "trace.so" +CYCLE_TSV = "trace_cycles.tsv" +SHAPE_TXT = "trace_shape.txt" +META_JSON = "meta.json" + +#: Used only when gem5 sampling fails. Deliberately not a plausible-looking +#: number: only an obvious non-measurement gets fixed. +PLACEHOLDER_CYCLE = 1 + +SAMPLE_MLIR = "04-sample.mlir" +CYCLE_BIN = "cycle_bin" + + +def measure_tile_cycles(workdir, meta): + """Per-compute-node cycle counts for ONE tile, measured under gem5. + + build_tog's sample mode marks each compute node and makes every loop a + single trip; tnpu lowers that to a binary (in ITS process -- the Gemmini/VCIX + lowering and its LLVM live there); gem5 runs it. Returns None on any failure, + and the caller falls back to the placeholder table. + """ + from PyTorchSimFrontend.mlir.passes.build_tog import run_tog + + kernel_name = meta["kernel_name"] + spec = os.path.join(workdir, f"{kernel_name}_spec.py") + if not os.path.isfile(spec): + logger.warning("[Gem5] %s not found; cannot sample cycles", spec) + return None + + run_tog(os.path.join(workdir, "04-custom.mlir"), + os.path.join(workdir, "tog_sample.py"), + os.path.join(workdir, SAMPLE_MLIR), sample_mode=True) + + import subprocess + + from . import tnpu_bridge + env = dict(os.environ) + env.pop("PYTHONPATH", None) # keep tnpu on its own MLIR bindings + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, "-m", "tnpu.cycle", spec, workdir], + capture_output=True, text=True, cwd=tnpu_bridge.tnpu_dir(), env=env) + if proc.returncode != 0: + logger.warning("[Gem5] cycle binary build failed:\n%s", + (proc.stdout + proc.stderr)[-2000:]) + return None + + from Simulator.simulator import CycleSimulator + try: + return CycleSimulator().compile_and_simulate( + os.path.join(workdir, CYCLE_BIN), int(extension_config.vpu_num_lanes), + silent_mode=True) + except Exception as e: # noqa: BLE001 - fall back to the placeholder table + logger.warning("[Gem5] sampling failed: %s", e) + return None + + +def _runtime_arg_layout(meta): + """(n_tensor_args, n_scalar_args) of the lowered signature. + + triton-shared lays it out as pointers, user scalars, then its own six + (gridX,Y,Z / pidX,Y,Z). constexpr params never become arguments. + """ + sig = meta["signature"] + tensors = [k for k, v in sig.items() if v.startswith("*")] + scalars = [k for k, v in sig.items() + if not v.startswith("*") and v != "constexpr"] + return len(tensors), len(scalars) + + +#: triton-shared appends pidX, pidY, pidZ in that order, whatever the tiling is. +_PID_SLOT = {"x": 0, "y": 1, "z": 2} + + +def work_item_for(meta): + """The WorkItem describing this kernel's program-id args and grid extents. + + 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.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], + grid=[None] * len(axes)) + + +def write_shape(workdir, meta, args=()): + """Write the grid extents the trace producer reads as shape_args. + + `args` is the launch's positional arguments; Inductor appends the numels + after the tensors, so the trailing values are them, in `meta["numels"]` + order. Falls back to the compile-time hint when they are absent. + """ + from . import kernel_spec + + numels = dict(meta["numels"]) + # Only the PARALLEL numels ride along on the call -- a reduction axis is + # looped inside the kernel, so it is not passed and must not consume one of + # the trailing values. They arrive in kernel order, which is the dict's. + passed = [k for k in numels if not k.startswith("r")] + trailing = [a for a in args if isinstance(a, int) and not isinstance(a, bool)] + if passed and len(trailing) >= len(passed): + for key, val in zip(passed, trailing[-len(passed):]): + numels[key] = val + + axes = kernel_spec.pid_axes(numels) + + cfg = meta.get("fixed_config") or {} + grid = [] + for p in axes: + n, block = numels.get(f"{p}numel"), cfg.get(f"{p.upper()}BLOCK") + if n is None or not block: + raise ValueError(f"no extent for grid axis '{p}': {n!r} / {block!r}") + grid.append(-(-int(n) // int(block))) # ceil-div + + path = os.path.join(workdir, SHAPE_TXT) + with open(path, "w") as f: + f.write("\n".join(str(g) for g in grid) + "\n") + logger.info("[TOGSim] grid %s -> %s", grid, SHAPE_TXT) + return grid + + +def emit_trace(workdir, meta): + """Build `trace.so` + `trace_cycles.tsv` from tnpu's post-vcix IR. + + Returns the number of compute tiles the cycle table covers. + """ + from PyTorchSimFrontend.mlir.passes import build_skeleton as bs + from PyTorchSimFrontend.mlir.passes import cycle_table as ct + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + postvcix = os.path.join(workdir, "04-custom.mlir") + if not os.path.isfile(postvcix): + raise FileNotFoundError( + f"{postvcix} not found -- tnpu must have run at least to stage 4 " + f"(the post-vcix IR is what the trace is built from)") + + # Before build_skeleton: both read the post-vcix IR, which it rewrites in place. + cycles = measure_tile_cycles(workdir, meta) + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(postvcix).read(), ctx) + bs.build_skeleton(module) + compute_types = ct._compute_types(module) + n_tiles = len(compute_types) + + if cycles: + # One numCycles per compute node; pad/truncate as the MLIR route does. + cl = list(cycles) + if len(cl) != n_tiles: + logger.warning("[Gem5] returned %d cycle(s) for %d " + "tile(s); padding with the last", len(cl), n_tiles) + cl = (cl + [cl[-1]] * n_tiles)[:n_tiles] + # Systolic-array fill; only matmul tiles use it. + lanes = int(extension_config.vpu_num_lanes) + table = ct.build_cycle_table(module, cl, x_offset=lanes, w_offset=0) + else: + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + logger.warning( + "[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d " + "tiles): gem5 sampling did not produce a measurement, so " + "compute latency is NOT modelled", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), + work_item=work_item_for(meta)) + + ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) + if cycles: + logger.info("[Gem5] tile cycles: %s", table) + return n_tiles + + +def run_togsim(workdir, meta=None, args=(), attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict. + + `meta`/`args` supply the grid: the trace producer takes its loop bounds from + shape_args, so they are written out per launch rather than compiled in. + """ + from Simulator.simulator import TOGSimulator + + so = os.path.join(workdir, TRACE_SO) + if not os.path.isfile(so): + raise FileNotFoundError(f"{so} not found -- call emit_trace first") + if meta is not None: + write_shape(workdir, meta, args) + + # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its + # DIRECTORY, and reads the file itself only on the STONNE path. + handle = os.path.join(workdir, "tile_graph.onnx") + result_path = TOGSimulator.run_standalone( + handle, attribute_path or os.path.join(workdir, "attribute"), + timeout_sec=timeout_sec) + return TOGSimulator.get_result_from_file(result_path) + + +def store_meta(workdir, meta): + """Persist codegen metadata beside the artifacts, so the timing step can run + standalone.""" + with open(os.path.join(workdir, META_JSON), "w") as f: + json.dump(meta, f, indent=2) + + +def load_meta(workdir): + with open(os.path.join(workdir, META_JSON)) as f: + return json.load(f) diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 0ef98eff..b98d3305 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -40,7 +40,18 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } } if (cyc.empty()) { cyc.assign(256, 128); ovl.assign(256, 0); } - return trace_to_tilegraph(trace_so_path.c_str(), nullptr, 0, + // Shape args: the producer's grid bounds, one per axis, when the trace was + // compiled without them baked in. Same sidecar convention as the cycle table + // -- absent means the producer carries its own constants. + std::vector shape; + { + std::ifstream sh(fs::path(trace_so_path).parent_path() / "trace_shape.txt"); + int64_t v; + while (sh >> v) shape.push_back(v); + } + return trace_to_tilegraph(trace_so_path.c_str(), + shape.empty() ? nullptr : shape.data(), + (int32_t)shape.size(), bases.data(), (int)bases.size(), cyc.data(), ovl.data(), (int)cyc.size(), partition_cores.data(), (int32_t)partition_cores.size(), From 01190f36ada7b18a957359eb4f6877642f29c250 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:48:36 +0900 Subject: [PATCH 05/10] [Frontend] Reach Inductor's mm/conv templates instead of an extern kernel mm and conv went to extern_kernels, where on npu they either raise convolution_overrideable not implemented or fall back to eager and simulate nothing. Passing was not the same as working. The templates themselves are not GPU-specific: torch ships one triton_mm.py.jinja for cuda, xpu, mtia and cpu, and the per-backend difference is a config table. What gates them is use_triton_template asking is_gpu, and GPU_TYPES is a hardcoded list with no registration hook, so npu is appended to it here. Selection is a fixed choice, not a benchmark: there is no device to time on. The TODO is to rank by simulated cycles, which timing.run_togsim already returns per kernel. Until then the offered order wins, which is deterministic and not a claim about speed. max_autotune_gemm rather than max_autotune, and epilogue-fusion benchmarking off: both render a benchmark-flavoured kernel whose harness imports arrive indented in the real module, which broke kernels that already worked. addmm and baddbmm carry a bias as input_nodes[0] and need their own heuristic; without one the mm entry is used with prefix_args=0 and def_kernel asserts. device_guard returned "pass", which the caller writes as `with pass:`. --- .../openreg/extension_device_op_overrides.py | 3 +- .../triton_backend/inductor_templates.py | 112 ++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 PyTorchSimFrontend/triton_backend/inductor_templates.py 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/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 From c40b0c1781188aa5680b3e3a7e0f823a13715898 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:49:21 +0900 Subject: [PATCH 06/10] [CI] Run the whole test suite on the Triton route, and report what fails TORCHSIM_TRITON_CODEGEN is read at device registration, so the suite can be run on the other route without any test knowing. The sweep does that and produces three things: a gate (triton_route_passing.txt), a report bucketed by cause and by the tnpu stage a kernel last cleared, and per-failure artifacts -- the generated kernel, the stage IRs, the error -- so a problem can be reported upstream with the code rather than a description of it. Emitting a kernel is not enough to call the route exercised. A test can generate a kernel for part of its graph and still send the op it is named for to aten: test_matmul_scalar generated the mul and called extern_kernels.mm. Those are reported with the op that took the exit and kept out of the gate. Tests run in parallel, each with its own dump and TORCHINDUCTOR_CACHE_DIR; sharing the cache made results depend on what ran first. --- scripts/ci/triton_route_passing.txt | 27 ++ scripts/ci/triton_route_sweep.py | 380 ++++++++++++++++++++++++++++ 2 files changed, 407 insertions(+) create mode 100644 scripts/ci/triton_route_passing.txt create mode 100755 scripts/ci/triton_route_sweep.py diff --git a/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt new file mode 100644 index 00000000..b806db41 --- /dev/null +++ b/scripts/ci/triton_route_passing.txt @@ -0,0 +1,27 @@ +# Tests that pass THROUGH the Triton codegen route. +# Gated by scripts/ci/triton_route_sweep.py; regenerate with +# python scripts/ci/triton_route_sweep.py --all --update-allowlist +# A test that passes without emitting a kernel 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/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 +tests/system/test_triton_codegen.py diff --git a/scripts/ci/triton_route_sweep.py b/scripts/ci/triton_route_sweep.py new file mode 100755 index 00000000..ff5b27eb --- /dev/null +++ b/scripts/ci/triton_route_sweep.py @@ -0,0 +1,380 @@ +#!/usr/bin/env python3 +"""Run the existing test suite through the Triton codegen route. + +TORCHSIM_TRITON_CODEGEN is read at device registration, so no test needs to know +which route it is on. Produces a gate (triton_route_passing.txt), a report +bucketed by cause and stage, and per-failure artifacts for reporting upstream. + + python scripts/ci/triton_route_sweep.py # the allowlist, gating + python scripts/ci/triton_route_sweep.py --all # every test, reports + python scripts/ci/triton_route_sweep.py --all --artifacts triton-failures +""" + +import argparse +import concurrent.futures as cf +import glob +import json +import os +import re +import shutil +import subprocess +import sys +import time + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +HERE = os.path.dirname(os.path.abspath(__file__)) +PASSING = os.path.join(HERE, "triton_route_passing.txt") + +#: How far the kernel got. The stage a failure did not reach owns it. +STAGES = [ + ("01-ttir.mlir", "1 triton -> ttir"), + ("02-ttshared.mlir", "2 ttir -> tts/linalg (triton-shared)"), + ("03-adapted.mlir", "3 tnpu adapt"), + ("04-custom.mlir", "4 tnpu lower (DMA, lanes, spad)"), + ("trace.so", "5 trace producer"), +] + +#: First match wins. Each bucket names the layer that owns the fix. +BUCKETS = [ + ("missing_dep", r"ModuleNotFoundError|No module named"), + ("device_op", r"\w+_overrideable not implemented|not implemented\. .*privateuse"), + ("triton_helpers", r"triton_helpers"), + ("wrapper_gap", r"'TritonNPUWrapperCodegen' object has no attribute"), + ("spec_incomplete", r"SpecIncomplete"), + ("tnpu_stage", r"TnpuError|tnpu pipeline failed|triton-shared-opt|" + r"\[stage\d\]|failed to legalize"), + ("reduction", r"lane-aware|no reduction path"), + ("dynamic_shape", r"ShapeMismatch|dynamic shape|size_hint returned None"), + ("matmul_timing", r"vcix\.iv|sf\.vc\.|no compute node"), + # 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__$"), +] + +#: Lines torch prints alongside an error that are not the error. +NOISE = re.compile( + r"TORCHDYNAMO_VERBOSE|torch\._dynamo|You can suppress this|set TORCH_LOGS|" + r"^During handling|^The above exception|for more information|^\s*\^+\s*$") + + +def discover(): + out = [] + for dirpath, _dirs, files in os.walk(os.path.join(ROOT, "tests")): + for f in files: + if f.startswith("test_") and f.endswith(".py"): + out.append(os.path.relpath(os.path.join(dirpath, f), ROOT)) + return sorted(out) + + +def load_allowlist(): + if not os.path.isfile(PASSING): + return [] + with open(PASSING) as f: + return [l.strip() for l in f + if l.strip() and not l.lstrip().startswith("#")] + + +def classify(output, timed_out): + if timed_out: + return "timeout" + for name, pat in BUCKETS: + if re.search(pat, output, re.I | re.M): + return name + return "other" + + +def first_error(output): + """The exception line, skipping torch's boilerplate around it.""" + lines = [l.strip() for l in output.strip().splitlines() if l.strip()] + for l in reversed(lines): + if NOISE.search(l): + continue + if re.match(r"^\w*(Error|Exception|Failure)\b.*:", l) or \ + re.match(r"^(assert|AssertionError)", l): + return l[:200] + for l in reversed(lines): + if not NOISE.search(l): + return l[:200] + return "" + + +#: Inductor's own line naming the wrapper it wrote, and the call that sends an +#: op to a library instead of a generated kernel. +_WRAPPER_RE = re.compile(r"Wrapper Codegen Path = (\S+)") +_EXTERN_RE = re.compile(r"extern_kernels\.(\w+)") + + +def externs(output): + """Ops the graph sent to a library rather than through this route. + + A test can emit a Triton kernel for part of its graph and still hand the op + it is named for to aten -- test_matmul_scalar generated the mul and called + extern_kernels.mm. Counting that as coverage overstates it. + """ + found = set() + for path in _WRAPPER_RE.findall(output): + try: + with open(path, errors="replace") as f: + found.update(_EXTERN_RE.findall(f.read())) + except OSError: + continue + return sorted(found) + + +def reached_stage(dump_dir): + """(label, workdir) of the furthest tnpu stage any kernel produced. + + kernel.py alone still counts: a kernel was generated and rejected pre-stage-1. + """ + best, best_dir, fallback = None, None, None + for wd in glob.glob(os.path.join(dump_dir, "triton_*")): + if os.path.isfile(os.path.join(wd, "kernel.py")): + fallback = wd + for i, (fname, label) in enumerate(STAGES): + if os.path.isfile(os.path.join(wd, fname)): + if best is None or i > best[0]: + best, best_dir = (i, label), wd + if best: + return best[1], best_dir + return ("0 kernel generated, not accepted" if fallback + else "0 nothing emitted"), fallback + + +def collect(test, dump_dir, out_root, output, bucket, stage, workdir): + """One directory per failing test: the kernel, the last IR, the error.""" + dest = os.path.join(out_root, test.replace("/", "_").removesuffix(".py")) + os.makedirs(dest, exist_ok=True) + with open(os.path.join(dest, "error.txt"), "w") as f: + f.write(f"test: {test}\nbucket: {bucket}\nstage: {stage}\n\n") + f.write("\n".join(output.strip().splitlines()[-60:])) + if workdir: + # The kernel to hand over, and the IR saying where it stopped. + for name in ("kernel.py", "stage.log", *(s[0] for s in STAGES[:-1])): + src = os.path.join(workdir, name) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(dest, name)) + return dest + + +def run_one(test, timeout, artifacts, scratch): + # Private per test: a shared dump lets one test's cached kernel answer for + # another's. + dump = os.path.join(scratch, test.replace("/", "_").removesuffix(".py")) + shutil.rmtree(dump, ignore_errors=True) + os.makedirs(dump, exist_ok=True) + # 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: + p = subprocess.run([sys.executable, test], cwd=ROOT, env=env, + capture_output=True, text=True, timeout=timeout) + out, code = p.stdout + p.stderr, p.returncode + except subprocess.TimeoutExpired as e: + pre = (e.stdout or "") if isinstance(e.stdout, str) else "" + out, code, timed_out = pre + "\n__timeout__", 124, True + + ok = code == 0 + stage, workdir = reached_stage(dump) + ext = externs(out) + r = {"test": test, "ok": ok, "returncode": code, + "seconds": round(time.time() - t0, 1), + "bucket": None if ok else classify(out, timed_out), + "stage": stage, + # No kernel emitted = the route was never used (CPU-only, eager + # fallback, extern call), so it is not coverage. + # Emitting a kernel is not enough: an extern call means part of the + # graph bypassed the route entirely. + "exercised": workdir is not None and not ext, + "externs": ext, + "error": "" if ok else first_error(out)} + if not ok and artifacts: + r["artifacts"] = os.path.relpath( + collect(test, dump, artifacts, out, r["bucket"], stage, workdir), ROOT) + shutil.rmtree(dump, ignore_errors=True) + return r + + +def write_markdown(results, path): + """The report a human reads: counts by cause, then every failure.""" + passed = [r for r in results if r["ok"]] + failed = [r for r in results if not r["ok"]] + by = {} + for r in failed: + by.setdefault(r["bucket"], []).append(r) + + real = [r for r in passed if r["exercised"]] + L = ["# Triton route coverage", "", + f"**{len(real)}/{len(results)} pass through the Triton route.** " + f"({len(passed)-len(real)} more pass without exercising it -- CPU-only, " + f"eager fallback, or a path that bypasses Inductor.)", ""] + if failed: + L += ["| cause | count | owner |", "|---|---|---|"] + OWNER = { + "device_op": "PyTorchSimDevice -- op not registered for npu", + "triton_helpers": "triton_backend -- needs a vendored copy", + "wrapper_gap": "triton_backend -- TritonNPUWrapperCodegen incomplete", + "spec_incomplete": "triton_backend -- kernel_spec cannot describe it", + "tnpu_stage": "tnpu lowering passes", + "reduction": "tnpu -- no lane-aware reduction", + "dynamic_shape": "triton_backend -- shape-specialised launch", + "matmul_timing": "build_tog -- compute node lookup", + "togsim": "TOGSim / trace producer", + "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", + "other": "unclassified", + } + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + L.append(f"| {b} | {len(rs)} | {OWNER.get(b, '')} |") + L += ["", "## Failures", ""] + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + L.append(f"### {b} ({len(rs)})") + L.append("") + for r in sorted(rs, key=lambda r: r["test"]): + L.append(f"- `{r['test']}` — reached **{r['stage']}**") + if r["error"]: + L.append(f" - `{r['error'][:160]}`") + if r.get("artifacts"): + L.append(f" - artifacts: `{r['artifacts']}`") + L.append("") + if real: + L += ["## Passing through the route", ""] + L += [f"- `{r['test']}`" for r in real] + [""] + other = [r for r in passed if not r["exercised"]] + if other: + L += ["## Passing without exercising the route", ""] + L += [f"- `{r['test']}`" for r in other] + [""] + with open(path, "w") as f: + f.write("\n".join(L)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("tests", nargs="*", metavar="TEST", + help="run these instead of the allowlist") + ap.add_argument("--all", action="store_true", + help="run every test, not just the passing allowlist") + ap.add_argument("--timeout", type=int, default=1800) + ap.add_argument("-j", "--jobs", type=int, + default=max(1, min(8, (os.cpu_count() or 2) // 2)), + help="tests in flight at once; each may itself use several " + "cores (gem5, TOGSim), so this is half the box by " + "default") + ap.add_argument("--json", help="write the full result list here") + ap.add_argument("--artifacts", metavar="DIR", + help="per-failure kernel + stage IR + error, for reporting") + ap.add_argument("--markdown", help="write the human-readable report here") + ap.add_argument("--update-allowlist", action="store_true", + help="rewrite the allowlist from what passed (use with --all)") + args = ap.parse_args() + + allow = load_allowlist() + tests = args.tests or (discover() if args.all else allow) + if not tests: + print("no tests selected; the allowlist is empty and --all was not given") + return 1 + + scratch = os.path.join(ROOT, ".triton_sweep") + shutil.rmtree(scratch, ignore_errors=True) + os.makedirs(scratch, exist_ok=True) + if args.artifacts: + shutil.rmtree(args.artifacts, ignore_errors=True) + os.makedirs(args.artifacts, exist_ok=True) + + print(f"Triton route sweep: {len(tests)} tests, {args.jobs} at a time" + f"{'' if args.all else ' (allowlist)'}\n") + results, done = [], 0 + + def report(r): + nonlocal done + done += 1 + mark = ("ok " if r["exercised"] else "ok- ") if r["ok"] else "FAIL" + why = (" (extern: " + ",".join(r["externs"]) + ")") if r.get("externs") \ + else " (route not exercised)" + extra = ("" if r["exercised"] else why) if r["ok"] \ + else f" [{r['bucket']}] @{r['stage']} {r['error'][:70]}" + print(f" {done:3d}/{len(tests)} {mark} {r['seconds']:7.1f}s " + f"{r['test']}{extra}", flush=True) + + if args.jobs == 1: + for t in tests: + r = run_one(t, args.timeout, args.artifacts, scratch) + results.append(r) + report(r) + else: + # Threads: run_one only waits on a subprocess, and dump dir, Inductor + # cache and TOGSim FIFO are all already per-test. + with cf.ThreadPoolExecutor(max_workers=args.jobs) as pool: + futs = {pool.submit(run_one, t, args.timeout, args.artifacts, + scratch): t for t in tests} + for fut in cf.as_completed(futs): + r = fut.result() + results.append(r) + report(r) + results.sort(key=lambda r: r["test"]) + shutil.rmtree(scratch, ignore_errors=True) + + passed = [r for r in results if r["ok"]] + failed = [r for r in results if not r["ok"]] + real = [r for r in passed if r["exercised"]] + + print(f"\n{'='*72}\npassed {len(passed)}/{len(results)}" + f" ({len(real)} through the Triton route, " + f"{len(passed)-len(real)} without exercising it)") + if failed: + by = {} + for r in failed: + by.setdefault(r["bucket"], []).append(r) + print("\nfailures by cause:") + for b, rs in sorted(by.items(), key=lambda kv: -len(kv[1])): + print(f" {b:16s} {len(rs):3d}") + print("\nhow far they got:") + st = {} + for r in failed: + st[r["stage"]] = st.get(r["stage"], 0) + 1 + for s, n in sorted(st.items()): + print(f" {s:40s} {n:3d}") + + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=2) + print(f"\nwrote {args.json}") + if args.markdown: + write_markdown(results, args.markdown) + print(f"wrote {args.markdown}") + if args.artifacts and failed: + print(f"wrote {args.artifacts}/ ({len(failed)} failure dirs)") + + if args.update_allowlist: + with open(PASSING, "w") as f: + f.write("# Tests that pass through the Triton codegen route.\n" + "# Gated by scripts/ci/triton_route_sweep.py; regenerate with\n" + "# python scripts/ci/triton_route_sweep.py --all " + "--update-allowlist\n") + for r in real: + f.write(r["test"] + "\n") + print(f"wrote {PASSING} ({len(real)} tests)") + return 0 + + regressed = [r for r in failed if r["test"] in allow] + if regressed: + print(f"\nREGRESSION: {len(regressed)} allowlisted test(s) failed") + for r in regressed: + print(f" {r['test']} [{r['bucket']}] {r['error']}") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 2bfc472ab231b8c43250226d7d511e4756088746 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 17:49:21 +0900 Subject: [PATCH 07/10] [Tests] Add a route test, and unblock two that stopped before their first op test_triton_codegen exercises the route directly, so a break in it is not only visible as a change in the sweep. test_sparse_core passed device positionally into a helper that takes it by keyword. test_mlp_cpu imported torchsummary, matplotlib and numpy, none of which it uses: the one summary() call is commented out. Importing a package the repo does not depend on was enough to stop it. It then loaded ./128_784_32_10_cpu/best_model.pth, which is neither the name computed ten lines above nor anything in the repo -- a checkpoint from whoever ran it last. The test trains and saves its own further down, so the load is a warm start, not a requirement. --- tests/models/MLP/test_mlp_cpu.py | 12 +- tests/ops/sparsity/test_sparse_core.py | 2 +- tests/system/test_triton_codegen.py | 177 +++++++++++++++++++++++++ 3 files changed, 183 insertions(+), 8 deletions(-) create mode 100644 tests/system/test_triton_codegen.py diff --git a/tests/models/MLP/test_mlp_cpu.py b/tests/models/MLP/test_mlp_cpu.py index 620d6cd2..57e55e6d 100644 --- a/tests/models/MLP/test_mlp_cpu.py +++ b/tests/models/MLP/test_mlp_cpu.py @@ -5,9 +5,6 @@ import contextlib import unittest import copy -import numpy as np -import matplotlib.pyplot as plt -from torchsummary import summary import torch @@ -112,10 +109,11 @@ def train_mlp_mnist(device): # save initial model state # torch.save(model_device.state_dict(), f"{name}/initial_model.pth") - # load from path - load_path = "./128_784_32_10_cpu/best_model.pth" - model_device.load_state_dict(torch.load(load_path)) - print("model loaded..") + # Optional warm start; this run trains and saves its own best_model.pth. + load_path = f"{name}/best_model.pth" + if os.path.exists(load_path): + model_device.load_state_dict(torch.load(load_path)) + print("model loaded..") model_device.train() 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/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py new file mode 100644 index 00000000..2af839bd --- /dev/null +++ b/tests/system/test_triton_codegen.py @@ -0,0 +1,177 @@ +"""Drive the Triton codegen route as far as it currently goes. + +This route is WIP (see PyTorchSimFrontend/triton_backend/README.md). The test is +written to report WHERE it stops rather than to assert success: the value right +now is a reproducible statement of the next gap, not a pass/fail gate. Register +it in .github/workflows/pytorchsim_test.yml only once the route runs end to end. + + TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py +""" +import os +import sys +import traceback + +# Must be set before torch_openreg registers the Inductor backend for `npu`. +os.environ.setdefault("TORCHSIM_TRITON_CODEGEN", "1") + +import torch # noqa: E402 + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +N = 1024 + + +def build(): + def fn(x, y): + return x + y + + x = torch.randn(N) + y = torch.randn(N) + return fn, x, y + + +def check_multi_axis_grid(): + """A 2-D grid must nest one loop per axis and hand both indices to iv[]. + + Guards the multi-axis path, which the add kernel does not reach: Inductor + only uses y/z when x would overflow, so a 1-D grid exercises just the first + iteration of the nest. + """ + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + src = """ + module { + func.func @k(%arg0: memref<*xf32>, %arg1: i32, %arg2: i32) { + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : i32 + %a = arith.muli %arg1, %c8 : i32 + %b = arith.addi %a, %arg2 : i32 + %o = arith.index_cast %b : i32 to index + "togsim.dma"(%o, %c0) {arg_id = 0 : i32, base = "arg0", dims = [128], + dir = 0 : i32, elem_bits = 32 : i32, is_async = false, read_bufs = [], + strides = [1], tag_id = 0 : i32, write_bufs = [0]} : (index, index) -> () + return + } + } + """ + problems = [] + # Verify the IR the pass itself produces: a bound created after an outer loop + # would not dominate an inner loop's use of it, which only shows at rank >= 2 + # and which the emitc lowering happens to paper over. + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + l2e._materialize_grid_loop( + l2e._find_kernel(module), + l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3]), ctx) + try: + module.operation.verify() + except Exception as e: # noqa: BLE001 + problems.append(f"materialized IR does not verify: {e}") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + emitc = l2e.lower_to_emitc( + module, work_item=l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3])) + cpp = l2e.emitc_to_cpp(emitc, include_dir=l2e._default_include_dir()) + + entry = cpp.split("togsim_kernel(EmitCtx*")[-1] + if entry.count("for (") != 2: + problems.append(f"expected 2 nested loops, found {entry.count('for (')}") + if "togsim_dispatch" not in entry: + problems.append("no togsim_dispatch call") + if ", 2);" not in entry: + problems.append("dispatch does not pass 2 indices") + for p in problems: + print(f" multi-axis grid: {p}") + return not problems + + +def check_reduction_is_refused(): + """A reduction must fail LOUDLY, not compile into wrong numbers. + + tnpu has no lane-aware reduction: the scratchpad is lane-banked, so the + reduced axis has to live inside a lane, and triton-shared hands over a + linalg.reduce (plus a linalg.transpose) that no pass lowers that way. Until + one does, reaching the launcher would mean simulating a kernel whose compute + is not what the hardware would do. + + Passing this check means the attempt still stops. When the lane path lands, + this is the test to delete. + """ + x = torch.randn(128, 64) + try: + torch.compile(lambda t: t.sum(dim=1))(x.to("npu:0")) + except Exception as e: # noqa: BLE001 - any diagnosed stop is the point + first = (str(e).strip().splitlines() or [type(e).__name__])[0] + print(f" reduction stops at: {type(e).__name__}: {first[:74]}") + return True + print(" reduction COMPILED -- if the lane-aware path landed, drop this " + "check; otherwise the numbers it produces are wrong") + return False + + +def main(): + from PyTorchSimFrontend import extension_config + from PyTorchSimFrontend.triton_backend import tnpu_bridge + + print(f"multi-axis grid = " + f"{'ok' if check_multi_axis_grid() else 'FAILED'}") + print(f"reduction refused = " + f"{'ok' if check_reduction_is_refused() else 'FAILED'}") + print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") + print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") + ok, _out = tnpu_bridge.doctor() + print(f"tnpu doctor = {'ok' if ok else 'FAILED (see run.py doctor)'}") + print() + + fn, x, y = build() + expected = fn(x, y) + + opt = torch.compile(fn, backend="inductor") + try: + got = opt(x.to("npu:0"), y.to("npu:0")) + except Exception as e: # noqa: BLE001 - the point is to report the stop + print(f"STOPPED AT: {type(e).__name__}") + print() + traceback.print_exc() + print() + print("The stage reached is what this test measures; see the traceback " + "above and README.md's gap list.") + return 1 + + ok = torch.allclose(got.cpu(), expected, rtol=1e-4, atol=1e-4) + if not ok: + bad = (~torch.isclose(got.cpu(), expected, rtol=1e-4, atol=1e-4)) + print(f"VALUES WRONG: {int(bad.sum())}/{expected.numel()} elements") + print(f" got {got.cpu()[:4].tolist()}") + print(f" expected {expected[:4].tolist()}") + return 1 + print(f"values ok ({expected.numel()} elements through Spike)") + + import glob + + from PyTorchSimFrontend.triton_backend import timing + + dirs = glob.glob(os.path.join(extension_config.get_dump_path(), "triton_*")) + if not dirs: + print("no kernel directory was produced") + return 1 + workdir = max(dirs, key=os.path.getmtime) + for name in (timing.TRACE_SO, timing.CYCLE_TSV): + path = os.path.join(workdir, name) + if not os.path.isfile(path): + print(f"missing {name} in {workdir}") + return 1 + print(f" {name:18s} {os.path.getsize(path)} bytes") + print(f"\ntiming path OK ({workdir})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 4e6c5196df2fa0065ebc8308a9a2c7e0bfacb329 Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 19:46:12 +0900 Subject: [PATCH 08/10] [Docs] Document the triton_backend modules and how to set the route up README.md is the setup, the seven tnpu stages, and what each module here owes. It is referenced from the workflow, extension_config, the device registration and the route test. CLAUDE.md gains the triton_npu workflow and the torch 2.10 pin. --- CLAUDE.md | 4 +- PyTorchSimFrontend/triton_backend/README.md | 218 ++++++++++++++++++++ 2 files changed, 221 insertions(+), 1 deletion(-) create mode 100644 PyTorchSimFrontend/triton_backend/README.md diff --git a/CLAUDE.md b/CLAUDE.md index fb76c82d..e3424132 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,6 +47,8 @@ Run a model from `tests/models/Llama/`, `tests/models/DeepSeek/`, etc. similarly **CI coverage:** the GitHub Actions workflow `.github/workflows/pytorchsim_test.yml` runs an **explicit allowlist** of `tests/*.py` files (~40 jobs, one Docker container per test). Adding a new file under `tests/` does *not* automatically gate PRs — register it in `pytorchsim_test.yml` if you want CI to exercise it. Conversely, files like `tests/ops/attention/test_gqa.py`, `tests/ops/attention/test_gqa_decode.py`, and `tests/system/test_eager.py` exist in the repo but are *not* in CI, so local validation is the only safety net for them. +The Triton codegen route has its own workflow, `.github/workflows/triton_npu.yml`, kept separate because its toolchain layer is ~1.8 GiB that no other job needs. It builds `torchsim_tnpu_base` (pinned by `thirdparty/triton-npu.json` + `Dockerfile.tnpu`) and needs `secrets.TNPU_TOKEN` plus a toolchain release on the private `PSAL-POSTECH/triton-npu`; see `PyTorchSimFrontend/triton_backend/README.md`. + **For fast iteration** (skip functional check): ```bash export pytorchsim_functional_mode=False # skips Spike @@ -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-`, see `scripts/ci/thirdparty_base_pin.sh`). - The default Gem5 path is hard-coded to `/workspace/gem5/build/RISCV/gem5.opt`. Override with `GEM5_PATH` if you build elsewhere. - `_C.cpython-311-*.so` and `torch_openreg/lib/` are build artifacts — already in `.gitignore`, don't commit. - TOGSim creates a per-PID FIFO under `/tmp/togsim_fifo_` for command/event comm; if a previous run crashed and left stale FIFOs, they get cleaned up on the next start, but watch for orphaned processes if you Ctrl-C mid-run. diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md new file mode 100644 index 00000000..d752a634 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -0,0 +1,218 @@ +# Triton codegen route (WIP) + +Replaces the hand-written MLIR emission in `PyTorchSimFrontend/mlir/` with +**Inductor's own Triton codegen**, lowered to this NPU by the **NPU lowering +pass** (owned by 이정민; the code lives in the `triton-npu` repo, so paths and +module names read `tnpu`). Opt-in and off by default; the MLIR route is +untouched and stays the production path. + +The modules here are the PORT: they drive that lowering pass and wire its output +into the existing TOGSim / gem5 / Spike stack. The pass itself is not ours. + +```bash +TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py +``` + +This file is the working reference for the modules here. + +## Why + +The MLIR route does not just emit loops — it hand-implements the whole hardware +mapping (tiling, vectorization, DMA, scratchpad, lane distribution) as ~5,500 +lines of Python string emission, which entangles *what to compute* with *how to +map it*. See `docs/linalg-codegen-migration.md` for the long form. + +This route keeps Inductor for the first and triton-npu for the second: + +| | owns | +|---|---| +| Inductor (upstream) | fusion, index expressions, masking, reductions, the kernel source | +| triton-shared | Triton IR -> `linalg` / `tts` pointer descriptors | +| tnpu passes | `tts` -> `togsim.transfer` DMA, scratchpad, lane-banked vectors, systolic array | + +## Flow + +``` +torch.compile + └ TritonNPUScheduling.define_kernel scheduling.py + │ Inductor's triton kernel SOURCE TEXT + collected metadata + ▼ + triton_npu_compile(src, meta, name) codecache.py + │ a tnpu kernel file (KernelSpec) kernel_spec.py + ▼ + run.py --to binary (subprocess) tnpu_bridge.py + │ 01-ttir → 02-ttshared → 03-adapted → 04-custom → 05-*.elf + ▼ + TritonNPULauncher.__call__ codecache.py + ├ functional tensors → runtime/*.raw → Spike → tensors functional.py + └ timing 04-custom.mlir → trace.so + trace_cycles.tsv → TOGSim + cycles measured by gem5 on a one-tile binary +``` + +The timing half reuses PyTorchSim's trace pipeline unchanged. The one structural +difference is that a Triton kernel body is a single program instance, so the grid +that enumerates instances is supplied by `lower_to_emitc.WorkItem` instead of +being read out of the kernel -- see "The grid is not in the kernel" below. + +Artifacts land in one directory per source hash under the dump path +(`outputs/triton_/`), alongside the unmodified Inductor source +(`kernel.py`) so the rewrite is diffable. + +## What works today (measured) + +`x + y`, 1024 elements, on `npu:0`: + +- Inductor generates the Triton kernel and our `define_kernel` intercepts it +- `kernel_spec` pins `XBLOCK` = lane count, computes `grid = (8,)`, writes the spec +- tnpu runs stages 1–5 and links **`05-triton_npu_fused_add_0.elf`** (20 B/lane spad) +- the lowering is correct in shape: `tl.load/store` became three + `togsim.transfer` ops, and Inductor's `xmask` came through as a **masked DMA** + (`masked_axes = [0]`, `masked_fill`), which tnpu already supports +- the trace producer comes out in the shape the design calls for: a + `togsim_kernel_tile` computing `offset = iv[0]*128` around three `togsim_dma` + and one `togsim_compute`, and a `togsim_kernel` looping `p < 8` over + `togsim_dispatch` +- **TOGSim runs it: 650 cycles**, with channel-0 DRAM traffic of 16 reads x 32 B + x 16 channels = 8192 B, exactly the 8 work-items x 2 loads x 512 B the kernel + should move. The MLIR route on the same `x + y` reports 251 cycles -- the same + order, and higher here because tnpu emits synchronous DMA, so nothing overlaps + (gap 2) +- the tile's compute cost is a real measurement: gem5 samples **19 cycles** for + the vector-add tile, via `timing.measure_tile_cycles` +- **values are correct**: the launch writes the caller's tensors from Spike and + `torch.allclose` holds over all 1024 elements, for the fused + `(x + y) * 2 - x` kernel too + +## Shape specialisation + +The functional binary is compiled for ONE shape: the spec bakes the grid, the +scalar values and the memref extents in. A dynamic-shape graph reuses that ELF, +so `functional.ShapeMismatch` rejects the launch instead of running against the +wrong bounds. The timing path has no such limit -- it takes the grid at run time +-- so `pytorchsim_functional_mode: False` studies cycles across shapes. + +## Gap list, in order + +1. **Shape-specialised functional launch.** Recompile per launch shape, or teach + the tnpu wrapper to take the grid and the extents as arguments the way the + trace producer already does. +2. **Double buffering.** tnpu emits synchronous DMA (`is_async=false`, no + `togsim.wait`), so load → compute → store serialize inside every work-item and + TOGSim has no overlap to model. This is the main remaining gap between the two + routes' cycle counts. +3. **`triton_helpers`.** Any kernel using `triton_helpers.*` (reductions, + clamps, `maximum`/`minimum`) cannot compile: the module lives in torch and + the tnpu venv has none. `strip_for_tnpu` raises and names the helper. Needs a + minimal vendored copy. +4. **Reductions.** Independently blocked in tnpu itself — no lane-aware + reduction path; see `triton-npu/kernels/reduce.py`. + Matmul is also still open on the timing side: `build_tog` finds compute nodes + by the `vcix.iv` op name, and tnpu emits `llvm.riscv.sf.vc.*` intrinsics. +5. **Block-size policy.** `fixed_config_for` pins `XBLOCK` to the lane count and + deliberately leaves reduction blocks unset. Real tile selection (the MLIR + route's autotuner / `codegen_mapping_strategy`) has no equivalent here yet. +6. **Dynamic shapes.** `collect_meta` resolves numels through `size_hint`; a + genuinely dynamic dim gives `None` and `_grid` raises. + +## Three design decisions + +**Block sizes are fixed at codegen time.** Inductor defers the grid to +`triton_heuristics` at runtime (`grid = cdiv(xnumel, XBLOCK)` after autotuning). +tnpu compiles one binary ahead of time and walks the grid as a sequential loop in +generated C, so there is nothing to autotune later and no runtime `grid=` +callable. Pinning the config is what makes the launch shape statically +describable — the premise of this route, not a shortcut. (`kernel_spec.fixed_config_for`) + +**tnpu runs in its own process.** Its passes need LLVM 23's MLIR bindings while +this process holds LLVM 20's, and `mlir` is a namespace package, so two LLVMs in +one interpreter silently merge. The seam between them is a file, and that is +measured to work: LLVM 23 prints IR that LLVM 20's bindings parse without +complaint. (`tnpu_bridge`) + +**The torch pin is what makes triton 3.6 work.** triton-npu pins triton 3.6 +because 3.6 pins LLVM 23, and both sides of its textual IR seam must be the same +LLVM. torch 2.10 is the first release whose Inductor targets 3.6, so the two +simply agree -- on 2.8 the frontend had to be shimmed onto a triton it did not +expect. What remains in `_triton_compat` is not a version shim: on a box with no +GPU, `triton_hash_with_backend()` raises "0 active drivers" because it asks the +triton runtime for the current target. We never launch through that runtime, so +the value is short-circuited to a deterministic cache key. + +**The grid is not in the kernel.** PyTorchSim's codegen puts the tile loops +inside the kernel; a Triton kernel describes one program instance and leaves the +grid to the launch. The trace producer wants that same split already -- +`togsim_kernel_tile` per work-item, enumerated by `togsim_kernel` (design sec +9.3) -- so the models agree and only the enumeration was missing. +`_materialize_grid_loop` supplies it, on the trace artifact only: it wraps the +body in a loop tagged `outer_loop` with each program-id argument replaced by the +induction variable, and everything downstream is unchanged. It runs before +`_rewrite_signature`, which erases the kernel arguments and first asserts none +are still used -- that ordering is what decides where this can live. + +## Running the whole suite on this route + +`TORCHSIM_TRITON_CODEGEN` is read once, at device registration, so every test +under `tests/` is already a test of this route — no test file knows which one it +is on. `scripts/ci/triton_route_sweep.py` runs them that way: + +```bash +python scripts/ci/triton_route_sweep.py # allowlist, gating +python scripts/ci/triton_route_sweep.py --all \ + --markdown coverage.md --artifacts failures # measure + report +``` + +`scripts/ci/triton_route_passing.txt` is the gate: the tests that pass today. +Coverage grows by regenerating it (`--update-allowlist`), so it cannot silently +shrink. A test that passes **without emitting a kernel** — CPU-only, eager +fallback, or an op Inductor sends to an extern call — is deliberately kept out +of it, since it would gate nothing. + +Each failure leaves a directory under `--artifacts`: the Inductor Triton kernel +that was rejected, whatever stage IR it reached (`01-ttir` … `04-custom`), +`stage.log`, and the error. That is the whole bug report for whoever owns the +pass, without a rerun. The bucket names the owning layer, and the stage says how +far it got, so the two together route it. + +## CI + +`.github/workflows/triton_npu.yml`, separate from the main CI: this route is WIP, +and its toolchain layer is ~1.8 GiB that no other job needs. + +``` +preflight TNPU_TOKEN set? repo readable? release present? +ensure-tnpu-base torchsim_base + tnpu toolchain -> torchsim_tnpu_base: +build-app ./Dockerfile on that base +tnpu-baselines run.py doctor + add/mul/relu/gemm/bmm through Spike (gates) +triton-route tests/system/test_triton_codegen.py (reports, does not gate) +triton-route-suite the allowlist (gates) + the full sweep (reports) +mlir-route-regression tests/ops/elementwise/test_add.py (gates) +``` + +The sweep uploads `triton-route-coverage`: `coverage.md`, `results.json`, and a +`failures/` directory per failing test. + +Jobs run on the PSAL Slurm runner farm (`PSAL-POSTECH/slurm-ghr`), so +`runs-on:` must carry the `slurm` label or the job never gets a runner. Image +builds and the sweep take `big` (16c/64G/2h); the rest take the small bucket. +Do not add `docker/setup-buildx-action` — the runner registers its own builder +and that action's driver cannot start under its podman. + +The toolchain image is pinned the same way `torchsim_base` is — the tag carries +`sha256(thirdparty/triton-npu.json + Dockerfile.tnpu)`, so it is rebuilt only when +one of those moves, and its tag also carries the base pin it was built on. +`mlir-route-regression` is there because this layer adds a *second* LLVM and a +*second* triton to the image; it checks the production path did not notice. + +**Needs `secrets.TNPU_TOKEN`** — a PAT that can read `PSAL-POSTECH/triton-npu` +and its `toolchain-llvm23` release. That repo is private and the default Actions +token is scoped to this repository. `preflight` checks it before the build. + +`Dockerfile.tnpu` clones the harness and runs its own `setup/restore.sh +--prebuilt`; the pins all live in that repo's `setup/versions.env`. `ref` in the +manifest is a commit, so an upstream change there moves this image's tag too. + +`Dockerfile.tnpu` sets `TNPU_SPIKE` and `TNPU_SPIKE_ISA=rv64gcv_zfh`: tnpu asks +for `zvfp8`, which the released spike lacks, and an unknown extension stops +spike at startup — including the doctor run inside the image build. Costs only +`ops_fp8_roundtrip.py`, which CI does not run. Drop once +`PSAL-POSTECH/riscv-isa-sim#7` is in the release. From a857cb8e0120aa6baaefceefa67c019db13ae0cd Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 19:33:59 +0900 Subject: [PATCH 09/10] [Frontend] Give a gemm a choice when the template heuristics have none Two ways the choice list came back empty, both raising NoValidChoicesError before any kernel was emitted. A zero-length axis has no tile, so the heuristics offer no config at all. A MoE expert that routes no tokens produces exactly that: [0, K] @ [K, N]. Those are short-circuited to zeros without going near the selection machinery -- correct whether the output is empty (M or N zero) or a sum over nothing (K zero). addmm and baddbmm with K == 0 are beta * bias rather than zeros, so they keep the original lowering unless the output is empty. The second was ours. max_autotune_gemm_backends is global but the heuristics are registered for npu only, so an ordinary cpu gemm in the same graph -- test_moe_cpu runs [32, 1000] @ [1000, 10] -- found no backend it could use. ATEN goes back in the list and pick_config ranks it last, so npu still takes a template wherever one exists. Measured over the whole suite afterwards: no test reaches an extern kernel. tnpu delivery goes from 60 to 61 of 69. --- .../triton_backend/inductor_templates.py | 74 ++++++++++++++----- 1 file changed, 56 insertions(+), 18 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/inductor_templates.py b/PyTorchSimFrontend/triton_backend/inductor_templates.py index eb1b9712..f6459f85 100644 --- a/PyTorchSimFrontend/triton_backend/inductor_templates.py +++ b/PyTorchSimFrontend/triton_backend/inductor_templates.py @@ -57,14 +57,51 @@ class NPUAddmmTemplateConfigHeuristic(AddMMConfigMixin, def pick_config(choices): - """Stand in for benchmarking: there is no device to time on. + """Stand in for benchmarking: there is no device to time on, so the offered + order wins. Extern ranks last, present only so a device with no registered + heuristic (cpu) still has a choice. - 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. + TODO: rank by simulated cycles; timing.run_togsim already returns one per + compiled kernel. """ - return {c: 1.0 + i * 1e-3 for i, c in enumerate(choices)} + from torch._inductor.select_algorithm import ExternKernelCaller + + return {c: (1e3 if isinstance(c, ExternKernelCaller) else 1.0) + i * 1e-3 + for i, c in enumerate(choices)} + + +def _short_circuit_degenerate_gemms(): + """A zero-length axis has no tile, so the heuristics offer no config and the + empty choice list raises. A MoE expert routing no tokens gives [0, K] @ [K, N]. + """ + from torch._inductor.kernel.mm_common import mm_args + from torch._inductor.lowering import full, lowerings + from torch._inductor.virtualized import V + + def wrap(op, bias): + def wrapped(*args, _orig=lowerings[op], **kwargs): + try: + m, n, k, layout = mm_args(*args[bias:bias + 2], + layout=kwargs.get("layout"))[:4] + m, n, k = (int(V.graph.sizevars.size_hint(s)) for s in (m, n, k)) + except Exception: # noqa: BLE001 - dynamic shape; leave it to _orig + return _orig(*args, **kwargs) + # k == 0 sums nothing, so zeros -- except addmm/baddbmm, which are + # then beta * bias. + if m == 0 or n == 0 or (k == 0 and not bias): + return full(layout.size, 0, dtype=layout.dtype, + device=layout.device) + return _orig(*args, **kwargs) + + return wrapped + + aten = torch.ops.aten + for op, bias in ((aten.mm, 0), (aten.bmm, 0), + (aten.addmm, 1), (aten.baddbmm, 1)): + for name in op.overloads(): + o = getattr(op, name) + if o in lowerings: + lowerings[o] = wrap(o, bias) def _install_selection(): @@ -73,8 +110,8 @@ def _install_selection(): 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. + # Precompiling builds every candidate for the current GPU; we need only the + # chosen kernel's source. AlgorithmSelectorCache.benchmark_choices = classmethod(benchmark_choices) AlgorithmSelectorCache.make_precompile_fn = lambda self, *a, **k: (lambda: None) @@ -83,11 +120,9 @@ def benchmark_choices(cls, choices, autotune_args, is_collective=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. + """On by default; TORCHSIM_TRITON_TEMPLATES=0 opts out. Sending mm to aten + simulates nothing, so a test that stops inside tnpu says more than one that + passes without running the op. """ global _installed if _installed or os.environ.get("TORCHSIM_TRITON_TEMPLATES", "1") == "0": @@ -97,14 +132,17 @@ def install(): _register_npu_as_gpu() _claim_triton_present() _register_template_heuristics() + _short_circuit_degenerate_gemms() _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. + # Not max_autotune: that also turns on pointwise autotuning, which appends + # a benchmark harness to every kernel module. config.max_autotune_gemm = True - config.max_autotune_gemm_backends = "TRITON" - config.max_autotune_conv_backends = "TRITON" + # These are global but the heuristics are registered for npu only, so ATEN + # stays in the list to keep a cpu gemm in the same graph from having no + # choice at all. pick_config ranks it last. + config.max_autotune_gemm_backends = "ATEN,TRITON" + config.max_autotune_conv_backends = "ATEN,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. From fc192b39bd695491d23a5ac53dc8c8870d0a850a Mon Sep 17 00:00:00 2001 From: Wonhyuk Yang Date: Wed, 5 Aug 2026 19:34:10 +0900 Subject: [PATCH 10/10] [Frontend] Cut the comments in the triton backend down to what is true now The prose had grown past what it was explaining. Removed: the account of a torch 2.8 API skew that no longer exists, a comparison against the MLIR route's line count, section banners in capitals, and the sentences arguing that a decision was the right one rather than saying what it does. Two comments were also stale. __init__ called the route scaffolding and said to expect failures, not results; fourteen tests now pass through it end to end. And inductor_templates pointed at tl.assume as the thing tests stop on, which was fixed upstream in triton-npu 8ebe408. Kept the pipeline diagram, the file and line references, and the TODOs. No code changed. --- PyTorchSimFrontend/triton_backend/__init__.py | 16 +- .../triton_backend/_triton_compat.py | 40 ++--- .../triton_backend/functional.py | 7 +- .../triton_backend/helpers_shim.py | 13 +- .../triton_backend/kernel_spec.py | 152 +++++------------- .../triton_backend/scheduling.py | 38 ++--- PyTorchSimFrontend/triton_backend/timing.py | 11 +- .../triton_backend/tnpu_bridge.py | 26 +-- 8 files changed, 94 insertions(+), 209 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py index cb3b88cd..0f879f12 100644 --- a/PyTorchSimFrontend/triton_backend/__init__.py +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -1,12 +1,7 @@ """The Triton codegen route: Inductor's Triton backend + the tnpu lowering passes. -WHAT THIS REPLACES ------------------- -The production path emits MLIR by hand (`PyTorchSimFrontend/mlir/`, ~5,500 lines -of string emission that also decides tiling, vectorization, DMA and scratchpad -placement). This route keeps Inductor's OWN Triton codegen for "what to compute" -and hands the resulting Triton kernel to triton-npu for "how to map it onto the -NPU": +Inductor decides what to compute; triton-npu decides how it maps onto the NPU, +in place of the hand-emitted MLIR under `PyTorchSimFrontend/mlir/`. Inductor -> TritonNPUScheduling.define_kernel (scheduling.py) | triton kernel SOURCE TEXT @@ -20,11 +15,8 @@ -> Spike (functional) / TOGSim (timing) The two routes are mutually exclusive and chosen at device-registration time by -`extension_config.CONFIG_TRITON_CODEGEN` (env `TORCHSIM_TRITON_CODEGEN=1`). -Default OFF -- nothing here is on the production path yet. - -STATUS: scaffolding. The seams are wired and each stage names precisely what it -still owes; see README.md for the gap list. Expect failures, not results. +`extension_config.CONFIG_TRITON_CODEGEN` (env `TORCHSIM_TRITON_CODEGEN=1`), +default off. README.md has the measured coverage and the gap list. """ from . import _triton_compat, inductor_templates diff --git a/PyTorchSimFrontend/triton_backend/_triton_compat.py b/PyTorchSimFrontend/triton_backend/_triton_compat.py index 9c863b15..63d29a14 100644 --- a/PyTorchSimFrontend/triton_backend/_triton_compat.py +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -1,24 +1,13 @@ """Let Inductor's Triton codegen run on a machine with no GPU. -ONE SKEW LEFT, AND IT IS NOT A VERSION SKEW -------------------------------------------- -`triton_hash_with_backend()` asks the triton runtime driver for the *current -target*, which on a box with no GPU raises "0 active drivers". We never launch -through triton's runtime -- triton-npu compiles the kernel ahead of time to a -RISC-V ELF -- so the value is only a cache-key ingredient, and a deterministic -string does the job. - -WHAT USED TO BE HERE, AND WHY IT IS GONE ----------------------------------------- -On torch 2.8 there was a second, larger skew: Inductor targeted the triton ~3.3 -API while triton-npu pins 3.6 (3.6 is what pins LLVM 23, and both sides of -triton-npu's textual IR seam must be the same LLVM), so `triton_key` had to be -injected back into `triton.compiler.compiler`. - -torch 2.10 pins triton 3.6.0 itself and reaches that symbol through its own -compat layer (`torch._inductor.runtime.triton_compat`), so on 2.10 the versions -simply agree and the injection is a no-op. It is kept, guarded, so the module -still works if someone runs an older torch -- `_torch_handles_triton()` decides. +`triton_hash_with_backend()` asks the triton runtime driver for the current +target, which raises "0 active drivers" without a GPU. We compile ahead of time +to a RISC-V ELF and never launch through that runtime, so the value is only a +cache-key ingredient and a deterministic string does. + +Pre-2.10 torch also needs `triton_key` injected into `triton.compiler.compiler`; +2.10 pins triton 3.6 itself and reaches it through its own compat layer. +`_torch_handles_triton()` decides. """ import functools @@ -120,20 +109,15 @@ def install(): return notes if not _torch_handles_triton(): - # Pre-2.10 torch: `triton_key` is imported from triton.compiler.compiler - # by SEVERAL call sites (codecache.CacheBase.get_system, - # _triton.triton_hash_with_backend, ...), each with its own local import. - # Supplying the symbol on the triton side satisfies all of them at once - # instead of chasing every call site; it is a cache-key ingredient, so - # any stable string will do. On torch 2.10 this branch does not run. + # Several call sites import triton_key with their own local import; + # supplying it on the triton side satisfies all of them at once. mod = importlib.import_module("triton.compiler.compiler") mod.triton_key = _stable_backend_hash notes.append("injected triton.compiler.compiler.triton_key " "(this torch predates the triton 3.6 compat layer)") - # Separately: triton_hash_with_backend also asks the triton runtime driver - # for the current target, which needs a GPU. We compile ahead of time to a - # RISC-V ELF and never use triton's runtime, so short-circuit it. + # triton_hash_with_backend asks the runtime driver for the current target, + # which needs a GPU. Short-circuited; it is only a cache key. import torch.utils._triton as _t _t.triton_hash_with_backend = functools.cache(_stable_backend_hash) notes.append("patched torch.utils._triton.triton_hash_with_backend " diff --git a/PyTorchSimFrontend/triton_backend/functional.py b/PyTorchSimFrontend/triton_backend/functional.py index 3ba91daa..3c142b34 100644 --- a/PyTorchSimFrontend/triton_backend/functional.py +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -1,9 +1,8 @@ """The functional half of the Triton route: real tensors -> Spike -> real tensors. -The timing half (timing.py) tells you how long the kernel takes; this tells you -whether it computed the right thing. tnpu's stage 6 already runs the ELF under -Spike, but on inputs it generates itself. Here the launch's own tensors are -written as the `.raw` files stage 6 reads, and the outputs are copied back: +tnpu's stage 6 runs the ELF under Spike on inputs it generates itself; here the +launch's own tensors are written as the `.raw` files it reads, and the outputs +copied back: run(workdir, meta, args) args -> runtime/*.raw -> spike -> args diff --git a/PyTorchSimFrontend/triton_backend/helpers_shim.py b/PyTorchSimFrontend/triton_backend/helpers_shim.py index f5440f1d..5e88b068 100644 --- a/PyTorchSimFrontend/triton_backend/helpers_shim.py +++ b/PyTorchSimFrontend/triton_backend/helpers_shim.py @@ -4,15 +4,10 @@ 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. +`triton_helpers.py` itself imports nothing from torch, only `.triton_compat`, +so it is copied verbatim next to the kernel and paired with the small +`triton_compat` below. Copying the installed torch's file rather than a vendored +snapshot keeps the helpers matched to the torch that generated the kernel. """ import os diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index f01881d1..354e36dc 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -1,37 +1,12 @@ """Inductor kernel -> tnpu KernelSpec. -Two jobs, both of which exist because Inductor's Triton output is written for a -GPU launcher and tnpu's is written for a static, ahead-of-time pipeline: - -1. `collect_meta` -- pull everything tnpu needs out of the Inductor kernel while - we still have `V.graph`: argument names/roles/dtypes/sizes, the constexprs, - and the numels the grid is computed from. This runs at codegen time; by the - time the compile callable fires, `V.graph` is gone. - -2. `write_spec_file` -- turn the Triton source + that metadata into a kernel file - tnpu can load (`tnpu.spec.load_spec`). - -WHY THE SOURCE HAS TO BE REWRITTEN ----------------------------------- -Inductor emits, above the kernel: - - from torch._inductor.runtime import triton_heuristics - @triton_heuristics.pointwise(size_hints={'x': 1024}, ..., inductor_meta=...) - @triton.jit - def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): - -Neither line can survive into tnpu: - - * the tnpu triton venv has NO torch (deliberately -- tnpu/spec.py), so the - `torch._inductor.runtime` import fails on sight; - * `triton_heuristics.pointwise` is the AUTOTUNER. It picks XBLOCK at runtime - and derives `grid = cdiv(xnumel, XBLOCK)` from it. tnpu needs both to be - constants: the block size becomes a `tl.constexpr` in the ttir signature and - the grid is executed as a sequential loop by the generated C wrapper. - -So the decorator is stripped and XBLOCK is pinned as a constexpr instead. That is -not a workaround -- fixing the config at codegen time is what makes the kernel -statically describable, which is the whole premise of this route. +`collect_meta` runs at codegen time, while `V.graph` still exists, and +`write_spec_file` turns the Triton source plus that metadata into a file +`tnpu.spec.load_spec` can load. + +The source is rewritten because Inductor's `@triton_heuristics.pointwise` +imports torch, which the tnpu venv does not have, and picks the block size at +runtime, which tnpu needs as a constexpr. """ import math @@ -42,10 +17,8 @@ def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): from . import helpers_shim -#: Triton signature token -> torch dtype name. The full set Inductor's -#: `_type_of` can emit (torch/_inductor/codegen/triton_utils.py); a token that -#: is absent here reads as an unresolvable arg, so partial coverage shows up as -#: SpecIncomplete on the first kernel that widens past int8. +#: Triton signature token -> torch dtype name; the full set Inductor's +#: `_type_of` (torch/_inductor/codegen/triton_utils.py) can emit. _DTYPE = { "*fp64": "float64", "*fp32": "float32", "*fp16": "float16", "*bf16": "bfloat16", @@ -64,10 +37,8 @@ def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): class SpecIncomplete(RuntimeError): - """Metadata tnpu requires that this kernel did not provide. - - Raised with the missing field named, rather than writing a spec that fails - deeper in the pipeline where the cause is unrecoverable. + """Metadata tnpu requires that this kernel did not provide, named rather + than left to fail deeper in the pipeline. """ @@ -76,10 +47,8 @@ class SpecIncomplete(RuntimeError): # --------------------------------------------------------------------------- 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. + Inductor allocates outputs empty_strided and indexes them by stride, so a + launch assuming contiguous writes to the wrong places. """ try: buf = V.graph.get_buffer(name) @@ -103,8 +72,7 @@ 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(): - # A graph input here is mutated, not produced: unwritten elements must - # survive, so it has to be seeded. + # Mutated, not produced: unwritten elements must survive, so seed it. 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(): @@ -142,10 +110,8 @@ def collect_meta(kernel, kernel_name): "stride": stride, }) - # The numels Inductor appends to the call. They live in `kernel.numels`, - # keyed by iteration-space PREFIX ('x', 'y', 'r0', ...), not as xnumel/rnumel - # attributes (SIMDKernel.__init__ builds them from the tiling). These are - # what the grid is computed from. + # kernel.numels is keyed by iteration-space prefix ('x', 'y', 'r0'), not by + # xnumel/rnumel attributes. The grid is computed from these. numels = {} for prefix, val in (getattr(kernel, "numels", None) or {}).items(): try: @@ -164,9 +130,8 @@ def collect_meta(kernel, kernel_name): } -#: Parallel iteration prefixes, OUTERMOST first. Inductor's `x` is the -#: contiguous axis, so it is innermost; `r*` prefixes are reductions, looped -#: inside the kernel rather than spread over the grid (prefix_is_reduction). +#: Parallel iteration prefixes, outermost first -- Inductor's `x` is the +#: contiguous axis. `r*` are reductions, looped inside the kernel. _PARALLEL_PREFIXES = ("z", "y", "x") @@ -181,27 +146,16 @@ def parallel_axes(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. + Every grid tuple is in this order; tnpu reads spec.grid positionally. """ return list(reversed(parallel_axes(numels))) def fixed_config_for(kernel): - """Block sizes pinned at codegen time. - - tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a - sequential loop, so there is no autotuner to choose the blocks later and no - runtime `grid=` callable. Fixing them here is what makes the launch shape - static. - - Tile dim 0 is the one `bank_vectorize` spreads over the lanes, so the - OUTERMOST axis gets the lane count -- a per-lane depth of 1, the shape every - tnpu baseline runs. The remaining axes get 1, which leaves the tile exactly - that verified shape and lets the grid cover the rest. It is conservative - rather than fast; choosing real tile sizes is the block-size policy gap in - README, not something to guess at here. + """Block sizes pinned at codegen time; tnpu has no autotuner to pick them + later. Tile dim 0 is the one bank_vectorize spreads over the lanes, so the + outermost axis gets the lane count and the rest get 1 -- conservative, and + the block-size policy gap in README. """ from PyTorchSimFrontend import extension_config try: @@ -212,23 +166,20 @@ def fixed_config_for(kernel): 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. + # parallel_axes wants collect_meta's "numel" keys, not raw prefixes. 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 - # 1 makes every work-item move a strided column. Fine for getting a - # multi-axis kernel through the route, misleading to benchmark. + # Correct but pathological: an inner block of 1 moves a strided column + # per work-item. Fine for coverage, misleading to benchmark. extension_config.setup_logger().warning( "[triton-npu] %s tiles over %s; inner blocks pinned to 1, which is " "correct but not a tiling worth measuring", getattr(kernel, "kernel_name", "kernel"), axes) cfg.setdefault("XBLOCK", lanes) # a kernel with no tiling info still has x if getattr(kernel, "inside_reduction", False): - # 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. + # The whole reduced extent, so the kernel's r0 loop runs once. r0 = (getattr(kernel, "numels", None) or {}).get("r0_") try: cfg["R0_BLOCK"] = int(V.graph.sizevars.size_hint(r0)) @@ -243,25 +194,18 @@ 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__)") -#: 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. +#: GPU-only calls. set_driver_to_gpu picks a runtime we never launch through; +#: debug_barrier orders warps that do not exist here and reaches ttir as +#: ttg.barrier, which 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): - """Remove everything the torch-free tnpu venv cannot import. - - Drops torch/inductor imports and the `@triton_heuristics.*(...)` decorator - (keeping `@triton.jit`), then re-adds the two imports the kernel body needs. - - Raises SpecIncomplete if the kernel still calls into `triton_helpers`: that - module lives in torch, so it has to be vendored into the tnpu venv before - such a kernel can compile. Failing here names the missing helper; letting it - through fails as a bare NameError inside tnpu's stage-1 worker instead. + """Remove everything the torch-free tnpu venv cannot import: torch imports + and the @triton_heuristics decorator, keeping @triton.jit. Raises + SpecIncomplete naming any triton_helpers call not yet vendored, which would + otherwise be a bare NameError inside tnpu's stage-1 worker. """ lines = src.splitlines() out, i = [], 0 @@ -291,10 +235,8 @@ def strip_for_tnpu(src): def scalar_args(meta): """User scalar parameters, in kernel order, as [(name, c_type, value)]. - - triton-shared keeps these in the lowered signature ahead of its own six - grid/pid arguments, so the wrapper must pass them or every later argument - lands one slot early -- pidX then reads pidY and only program 0 runs. + triton-shared keeps these ahead of its own grid/pid arguments, so omitting + one shifts every later argument a slot early. """ numels = meta["numels"] out = [] @@ -352,15 +294,13 @@ 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. +# The kernel is loaded by path, so a sibling package is not otherwise +# 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 -#: disk, not an exec'd string: triton's @jit reads the function back with -#: inspect.getsourcefile and rejects anything else ("@jit functions should be -#: defined in a Python file"). +#: The rewritten Triton source, beside this file. A real file, not an exec'd +#: string: @jit reads the function back with inspect.getsourcefile. TRITON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), {triton_module!r}) @@ -381,10 +321,8 @@ def make_inputs(torch, seed=0): def reference(inputs): - # The Inductor route has no per-kernel torch reference: correctness is - # checked at the graph level by the test that ran torch.compile. tnpu's - # stage 7 is therefore not meaningful here and the pipeline is driven to - # stage 6 (spike) instead. + # No per-kernel torch reference here -- correctness is checked at the graph + # level -- so the pipeline stops at stage 6 rather than tnpu's verify. return {{}} @@ -418,10 +356,8 @@ def write_spec_file(src_code, meta, path, tnpu_dir): constexprs = dict(meta["constants"]) for k, v in (meta.get("fixed_config") or {}).items(): if k not in signature: - # Inductor already fixed this one in the kernel BODY rather than - # taking it as a parameter -- a persistent reduction does that with - # R0_BLOCK. Passing it would not match the signature, and there is - # nothing left for us to choose. + # Fixed in the kernel body rather than taken as a parameter (a + # persistent reduction does this with R0_BLOCK); nothing to choose. continue if v is None: raise SpecIncomplete( diff --git a/PyTorchSimFrontend/triton_backend/scheduling.py b/PyTorchSimFrontend/triton_backend/scheduling.py index 56e7f1f8..d22331b3 100644 --- a/PyTorchSimFrontend/triton_backend/scheduling.py +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -1,17 +1,13 @@ """Inductor scheduling for the Triton route. -`TritonNPUScheduling` keeps ALL of Inductor's Triton codegen -- fusion, index -expressions, masking, reductions, the kernel source itself -- and changes only -what happens to the generated source afterwards. Upstream hands it to -`async_compile.triton(...)`, which calls `triton.compile` for a GPU; we hand it -to `triton_npu_compile(...)`, which runs the triton-npu pipeline for this NPU. - -Two overrides, and nothing else: - - define_kernel emit our compile call into the wrapper instead of upstream's. - kernel_type a TritonKernel whose call site is a plain python call, because - the name is bound to our callable rather than to a triton - launcher with a `.run(grid=..., stream=...)` interface. +`TritonNPUScheduling` keeps all of Inductor's Triton codegen and changes only +what happens to the source afterwards: `triton_npu_compile(...)` instead of +`async_compile.triton(...)`. + + define_kernel emit our compile call into the wrapper. + kernel_type a TritonKernel whose call site is a plain python call, since + the name binds to our callable, not a `.run(grid=...)` + launcher. """ from torch._inductor.codegen.common import IndentedBuffer @@ -35,18 +31,15 @@ class TritonNPUKernel(TritonKernel): cannot be written into a tnpu KernelSpec. """ - # **kwargs, not a fixed signature: Inductor keeps adding parameters here - # (2.10 added `deallocate_ws`). None of them apply to this route -- there is - # no triton launcher and no workspace to release -- so they are accepted and - # ignored rather than pinning us to one torch release. + # **kwargs: Inductor keeps adding parameters here and none apply to this + # route, so they are accepted and ignored rather than pinning a torch + # release. def call_kernel(self, name: str, node=None, **kwargs): wrapper = V.graph.wrapper_code _, call_args, _, arg_types = self.args.python_argdefs() self.add_numel_to_call_args(name, call_args, arg_types) - # add_numel_to_call_args appends the numels as SYMPY values, which the - # triton path later renders through pexpr. ExtensionWrapperCodegen joins - # call args as plain strings (mlir_codegen_backend.py:241), so render - # them here instead of handing it a sympy Integer. + # add_numel_to_call_args appends sympy values; ExtensionWrapperCodegen + # joins call args as plain strings, so render them here. call_args = [a if isinstance(a, str) else str(a) for a in call_args] # triton=False -> PythonWrapperCodegen emits `name(args...)`, the same # shape the MLIR route uses (mlir_common.py:627). @@ -70,9 +63,8 @@ def define_kernel(self, src_code, node_schedule, kernel): TritonNPUScheduling.count += 1 wrapper.src_to_kernel[src_code] = kernel_name - # Upstream substitutes these two placeholders inside define_kernel; the - # source still carries them here, and the tnpu side parses the source, so - # they have to be resolved before it leaves this function. + # Upstream substitutes these inside define_kernel; the tnpu side parses + # the source, so they must be resolved before it leaves here. src_code = src_code.replace(str(Placeholder.DESCRIPTIVE_NAME), kernel_name) src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name) diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index e6bc7c44..94834017 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -101,8 +101,8 @@ def work_item_for(meta): n_tensor, n_scalar = _runtime_arg_layout(meta) pid_base = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ axes = kernel_spec.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. + # Only the axis count is compiled in; the launch knows the real numels, so + # one trace serves every shape. return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], grid=[None] * len(axes)) @@ -117,9 +117,8 @@ def write_shape(workdir, meta, args=()): from . import kernel_spec numels = dict(meta["numels"]) - # Only the PARALLEL numels ride along on the call -- a reduction axis is - # looped inside the kernel, so it is not passed and must not consume one of - # the trailing values. They arrive in kernel order, which is the dict's. + # Only parallel numels ride along: a reduction axis is looped inside the + # kernel, so it must not consume one of the trailing values. passed = [k for k in numels if not k.startswith("r")] trailing = [a for a in args if isinstance(a, int) and not isinstance(a, bool)] if passed and len(trailing) >= len(passed): @@ -212,7 +211,7 @@ def run_togsim(workdir, meta=None, args=(), attribute_path=None, timeout_sec=Non write_shape(workdir, meta, args) # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its - # DIRECTORY, and reads the file itself only on the STONNE path. + # directory. handle = os.path.join(workdir, "tile_graph.onnx") result_path = TOGSimulator.run_standalone( handle, attribute_path or os.path.join(workdir, "attribute"), diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py index 9599423d..dc12b6d8 100644 --- a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -1,18 +1,9 @@ """Run the triton-npu pipeline, out of process. -WHY A SUBPROCESS ----------------- -tnpu's passes run on LLVM 23's MLIR python bindings; this process holds LLVM 20's -(TORCHSIM_LLVM_PATH). `mlir` ships without an `__init__.py`, so it is a NAMESPACE -package whose `__path__` is the union of every `mlir/` directory on sys.path -- -two LLVMs in one interpreter silently merge and fail later with an AttributeError -from a generated dialect module (tnpu/config.py:activate_bindings documents the -exact failure). They cannot share an interpreter, so tnpu gets its own. - -The seam between them is a FILE, which is measured to work: LLVM 23 prints the -IR, and LLVM 20's bindings parse it back without complaint (verified by feeding -tnpu's 04-custom.mlir to PyTorchSim's build_tog). That is what makes the split -viable rather than merely necessary. +tnpu's passes run on LLVM 23's MLIR bindings; this process holds LLVM 20's. +`mlir` is a namespace package, so two LLVMs on one sys.path merge silently and +fail later inside a generated dialect module. They get separate interpreters, +with a printed IR file as the seam. """ import os @@ -86,18 +77,15 @@ def run_pipeline(spec_path, workdir, to_stage="binary", from_stage="ttir", cmd.append("-v") env = dict(os.environ) - # tnpu deliberately does not read TORCHSIM_LLVM_PATH (it would drag the - # backend back to LLVM 20 and break the textual seam), but a stale - # PYTHONPATH pointing at LLVM 20's mlir_core would still be picked up by the - # namespace package before tnpu's own activate_bindings() runs. + # A PYTHONPATH pointing at LLVM 20's mlir_core would be picked up by the + # namespace package before tnpu's activate_bindings() runs. env.pop("PYTHONPATH", None) proc = subprocess.run(cmd, capture_output=True, text=True, cwd=tnpu_dir(), env=env, timeout=timeout) output = proc.stdout + proc.stderr if proc.returncode != 0: - # run.py prints a stage table; the diagnostic itself only reaches - # stage.log. + # run.py prints a stage table; the diagnostic only reaches stage.log. log = os.path.join(workdir, "stage.log") if os.path.isfile(log): with open(log, errors="replace") as fh: