diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml new file mode 100644 index 000000000..c1d2ffdea --- /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/CLAUDE.md b/CLAUDE.md index fb76c82d8..e34241329 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/Dockerfile.base b/Dockerfile.base index de023566b..b4d58a7c7 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 \ diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu new file mode 100644 index 000000000..b9d939a77 --- /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/PyTorchSimDevice/torch_openreg/__init__.py b/PyTorchSimDevice/torch_openreg/__init__.py index e81583918..2667ac703 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/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py b/PyTorchSimDevice/torch_openreg/openreg/extension_device_op_overrides.py index 27a473571..45b10c9f1 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/extension_config.py b/PyTorchSimFrontend/extension_config.py index 2d706bbcc..09e5168da 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/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index 71927cf61..bd0c6fcce 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 f4ed7d0d9..ed52a56de 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 5a40feecb..98590d882 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 5633769a1..537c8ad06 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/README.md b/PyTorchSimFrontend/triton_backend/README.md new file mode 100644 index 000000000..d752a634a --- /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. diff --git a/PyTorchSimFrontend/triton_backend/__init__.py b/PyTorchSimFrontend/triton_backend/__init__.py new file mode 100644 index 000000000..0f879f128 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -0,0 +1,30 @@ +"""The Triton codegen route: Inductor's Triton backend + the tnpu lowering passes. + +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 + 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. README.md has the measured coverage and the gap list. +""" + +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 000000000..63d29a146 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -0,0 +1,127 @@ +"""Let Inductor's Triton codegen run on a machine with no GPU. + +`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 +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(): + # 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)") + + # 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 " + "(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 000000000..d62de72d5 --- /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 000000000..3c142b340 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -0,0 +1,163 @@ +"""The functional half of the Triton route: real tensors -> Spike -> real tensors. + +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 + +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 000000000..5e88b0686 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/helpers_shim.py @@ -0,0 +1,74 @@ +"""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. + +`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 +import shutil + +PACKAGE = "tnpu_helpers" + +#: Stands in for torch._inductor.runtime.triton_compat, whose only job here is +#: to resolve these seven names. Mirrors the upstream fallbacks so a triton +#: version change lands the same way on both sides. +_COMPAT = '''\ +"""Generated by PyTorchSimFrontend/triton_backend/helpers_shim.py. Do not edit.""" + +import inspect +from typing import Any + +import triton +import triton.language as tl +from triton.runtime.jit import JITFunction # noqa: F401 + +try: + from triton.language.extra import libdevice # noqa: F401 + + libdevice = tl.extra.libdevice # noqa: F811 + math = tl.math +except ImportError: + if hasattr(tl.extra, "cuda") and hasattr(tl.extra.cuda, "libdevice"): + libdevice = tl.extra.cuda.libdevice + math = tl.math + elif hasattr(tl.extra, "intel") and hasattr(tl.extra.intel, "libdevice"): + libdevice = tl.extra.intel.libdevice + math = tl.math + else: + libdevice = tl.math + math = tl + +try: + from triton.language.standard import _log2 +except ImportError: + + def _log2(x: Any) -> Any: + raise NotImplementedError + +builtins_use_semantic_kwarg = ( + "_semantic" in inspect.signature(triton.language.core.view).parameters +) +''' + + +def _source_path(): + from torch._inductor.runtime import triton_helpers + return triton_helpers.__file__ + + +def write_package(dest_dir): + """Write the importable package next to a kernel. Returns the import line.""" + pkg = os.path.join(dest_dir, PACKAGE) + os.makedirs(pkg, exist_ok=True) + with open(os.path.join(pkg, "__init__.py"), "w"): + pass + with open(os.path.join(pkg, "triton_compat.py"), "w") as f: + f.write(_COMPAT) + shutil.copyfile(_source_path(), os.path.join(pkg, "triton_helpers.py")) + return f"from {PACKAGE} import triton_helpers\n" diff --git a/PyTorchSimFrontend/triton_backend/inductor_templates.py b/PyTorchSimFrontend/triton_backend/inductor_templates.py new file mode 100644 index 000000000..f6459f85a --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/inductor_templates.py @@ -0,0 +1,150 @@ +"""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, 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 one per + compiled kernel. + """ + 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(): + 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. + 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 + 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": + return + from torch._inductor import config + + _register_npu_as_gpu() + _claim_triton_present() + _register_template_heuristics() + _short_circuit_degenerate_gemms() + _install_selection() + + # Not max_autotune: that also turns on pointwise autotuning, which appends + # a benchmark harness to every kernel module. + config.max_autotune_gemm = True + # 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. + config.benchmark_epilogue_fusion = False + _installed = True diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py new file mode 100644 index 000000000..354e36dc7 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -0,0 +1,399 @@ +"""Inductor kernel -> tnpu KernelSpec. + +`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 +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` (torch/_inductor/codegen/triton_utils.py) can emit. +_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, named rather + than left to fail deeper in the pipeline. + """ + + +# --------------------------------------------------------------------------- +# 1. codegen-time metadata capture +# --------------------------------------------------------------------------- +def _buffer_layout(name): + """(numel, size, stride) of an Inductor buffer; Nones if unresolvable. + 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) + 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(): + # 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(): + 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, + }) + + # 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: + 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. `r*` are reductions, looped inside the kernel. +_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; tnpu reads spec.grid positionally. + """ + return list(reversed(parallel_axes(numels))) + + +def fixed_config_for(kernel): + """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: + 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 + + # 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: + # 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. + 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__)") +#: 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: 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 + 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 ahead of its own grid/pid arguments, so omitting + one shifts every later argument a slot early. + """ + 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}) +# 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. 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}) + + +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): + # 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 {{}} + + +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: + # 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( + 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 000000000..d22331b3f --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -0,0 +1,84 @@ +"""Inductor scheduling for the Triton route. + +`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 +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: 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 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). + 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 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) + + 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/timing.py b/PyTorchSimFrontend/triton_backend/timing.py new file mode 100644 index 000000000..94834017c --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -0,0 +1,231 @@ +"""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"]) + # 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)) + + +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 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): + 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. + 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/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py new file mode 100644 index 000000000..dc12b6d8f --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -0,0 +1,96 @@ +"""Run the triton-npu pipeline, out of process. + +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 +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) + # 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 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 000000000..cfc6490d9 --- /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 + """ + ) diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 0ef98eff6..b98d33057 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(), diff --git a/scripts/ci/tnpu_base_pin.sh b/scripts/ci/tnpu_base_pin.sh new file mode 100755 index 000000000..596638363 --- /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/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt new file mode 100644 index 000000000..b806db41f --- /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 000000000..ff5b27ebf --- /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()) diff --git a/tests/models/MLP/test_mlp_cpu.py b/tests/models/MLP/test_mlp_cpu.py index 620d6cd29..57e55e6d3 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 21cd93446..ddde6d5de 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 000000000..2af839bd5 --- /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()) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json new file mode 100644 index 000000000..d6d9f4ae3 --- /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" + ] + } +}