diff --git a/.github/workflows/triton_npu.yml b/.github/workflows/triton_npu.yml new file mode 100644 index 00000000..05c6f893 --- /dev/null +++ b/.github/workflows/triton_npu.yml @@ -0,0 +1,257 @@ +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. The repo is private and the default Actions token is +# scoped to this repository. preflight checks it 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 + 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 fb76c82d..e3424132 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,6 +47,8 @@ Run a model from `tests/models/Llama/`, `tests/models/DeepSeek/`, etc. similarly **CI coverage:** the GitHub Actions workflow `.github/workflows/pytorchsim_test.yml` runs an **explicit allowlist** of `tests/*.py` files (~40 jobs, one Docker container per test). Adding a new file under `tests/` does *not* automatically gate PRs — register it in `pytorchsim_test.yml` if you want CI to exercise it. Conversely, files like `tests/ops/attention/test_gqa.py`, `tests/ops/attention/test_gqa_decode.py`, and `tests/system/test_eager.py` exist in the repo but are *not* in CI, so local validation is the only safety net for them. +The Triton codegen route has its own workflow, `.github/workflows/triton_npu.yml`, kept separate because its toolchain layer is ~1.8 GiB that no other job needs. It builds `torchsim_tnpu_base` (pinned by `thirdparty/triton-npu.json` + `Dockerfile.tnpu`) and needs `secrets.TNPU_TOKEN` plus a toolchain release on the private `PSAL-POSTECH/triton-npu`; see `PyTorchSimFrontend/triton_backend/README.md`. + **For fast iteration** (skip functional check): ```bash export pytorchsim_functional_mode=False # skips Spike @@ -137,7 +139,7 @@ Conan deps for TOGSim: `boost/1.79.0`, `robin-hood-hashing/3.11.5`, `spdlog/1.11 ## Gotchas / things I've already learned -- The repo expects `python` to be a Python 3.10+ binary with `torch==2.8.0`. The frontend extends the PyTorch 2 Inductor stack — pin to this version. +- The repo expects `python` to be a Python 3.10+ binary with `torch==2.10.0` (torchvision `0.25.0`, triton `3.6.0`). The frontend extends the PyTorch 2 Inductor stack — pin to this version. 2.10 specifically: it is the first release whose Inductor targets triton 3.6, the version triton-npu is built against. The pins live in `Dockerfile.base`, and editing that file changes the base-image tag automatically (the tag is `thirdparty-`, see `scripts/ci/thirdparty_base_pin.sh`). - The default Gem5 path is hard-coded to `/workspace/gem5/build/RISCV/gem5.opt`. Override with `GEM5_PATH` if you build elsewhere. - `_C.cpython-311-*.so` and `torch_openreg/lib/` are build artifacts — already in `.gitignore`, don't commit. - TOGSim creates a per-PID FIFO under `/tmp/togsim_fifo_` for command/event comm; if a previous run crashed and left stale FIFOs, they get cleaned up on the next start, but watch for orphaned processes if you Ctrl-C mid-run. diff --git a/Dockerfile.base b/Dockerfile.base index de023566..b4d58a7c 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -51,8 +51,20 @@ RUN apt-get -y update && \ rm -rf /var/lib/apt/lists/* # CPU PyTorch (no CUDA wheels). torchvision is required by the vision model tests. +# torch 2.10 is pinned for the Triton codegen route: it is the first release whose +# Inductor targets triton 3.6, which is the version triton-npu is built against +# (triton 3.6 pins LLVM 23, and both sides of triton-npu's textual IR seam must be +# the same LLVM). On 2.8 the frontend had to be shimmed onto a triton it did not +# expect; on 2.10 the versions simply agree. RUN python3.11 -m pip install --no-cache-dir \ - torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu + torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu + +# Triton, for the Inductor Triton codegen route (PyTorchSimFrontend/triton_backend). +# Inductor imports triton while GENERATING a kernel, so it is needed even though +# nothing here compiles or launches through triton's own runtime -- triton-npu +# compiles the kernel ahead of time to a RISC-V ELF using its own triton build. +# Not a dependency of the CPU torch wheels, hence installed explicitly. +RUN python3.11 -m pip install --no-cache-dir triton==3.6.0 # TorchSim Python dependencies (numpy pinned <2 for transformers/diffusers compat). RUN python3.11 -m pip install --no-cache-dir \ diff --git a/Dockerfile.tnpu b/Dockerfile.tnpu new file mode 100644 index 00000000..b9d939a7 --- /dev/null +++ b/Dockerfile.tnpu @@ -0,0 +1,42 @@ +# syntax=docker/dockerfile:1.4 +# +# triton-npu toolchain layer, for the Triton codegen route only. +# Separate from torchsim_base because it is ~1.8 GiB no other job needs. +# The app image for this route is ./Dockerfile with BASE_IMAGE pointed here. +# +# The repo is private, so the clone and the release downloads both need a token. +# It is a BuildKit secret, not a build-arg: build-args land in the image history. + +ARG BASE_IMAGE=ghcr.io/psal-postech/torchsim_base:latest +FROM ${BASE_IMAGE} + +ARG TNPU_REPO=PSAL-POSTECH/triton-npu +ARG TNPU_REF=main + +WORKDIR /workspace + +# Not under $TORCHSIM_DIR: ./Dockerfile copies the PyTorchSim checkout over that +# path afterwards. extension_config reads TNPU_DIR, set below. +RUN --mount=type=secret,id=tnpu_token \ + TOKEN="$(cat /run/secrets/tnpu_token)" && \ + git clone "https://x-access-token:${TOKEN}@github.com/${TNPU_REPO}.git" \ + /workspace/triton-npu && \ + git -C /workspace/triton-npu checkout -q "${TNPU_REF}" && \ + git -C /workspace/triton-npu remote set-url origin \ + "https://github.com/${TNPU_REPO}.git" + +# restore.sh owns every pin (setup/versions.env) and unpacks LLVM 23, spike and +# the triton runtime into /workspace. +RUN --mount=type=secret,id=tnpu_token \ + GITHUB_TOKEN="$(cat /run/secrets/tnpu_token)" \ + /workspace/triton-npu/setup/restore.sh --prebuilt + +ENV TNPU_DIR=/workspace/triton-npu +# tnpu defaults to a separate fp8 spike and asks for zvfp8; the released spike +# has neither, and an unknown extension stops spike at startup. Drop once +# PSAL-POSTECH/riscv-isa-sim#7 is in the release. +ENV TNPU_SPIKE=/workspace/riscv-isa-sim/install/bin/spike +ENV TNPU_SPIKE_ISA=rv64gcv_zfh + +# Fail the build, not the first CI job. +RUN python3 /workspace/triton-npu/run.py doctor diff --git a/PyTorchSimDevice/torch_openreg/__init__.py b/PyTorchSimDevice/torch_openreg/__init__.py index e8158391..2667ac70 100644 --- a/PyTorchSimDevice/torch_openreg/__init__.py +++ b/PyTorchSimDevice/torch_openreg/__init__.py @@ -18,13 +18,30 @@ sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')) import PyTorchSimFrontend.extension_config # noqa: F401 +from PyTorchSimFrontend import extension_config as _extension_config from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen -from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling -torch._inductor.codegen.common.register_backend_for_device( - "npu", - lambda scheduling: MLIRScheduling(scheduling), - ExtensionWrapperCodegen -) + +# Two mutually exclusive codegen routes for `npu`, chosen here because Inductor +# registers a backend per device, once. +# MLIR (default) hand-written MLIR emission, PyTorchSimFrontend/mlir +# Triton (opt-in) Inductor's own Triton codegen + the triton-npu passes, +# TORCHSIM_TRITON_CODEGEN=1. WIP; see +# PyTorchSimFrontend/triton_backend/README.md +if _extension_config.CONFIG_TRITON_CODEGEN: + from PyTorchSimFrontend.triton_backend import ( + TritonNPUScheduling, TritonNPUWrapperCodegen) + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: TritonNPUScheduling(scheduling), + TritonNPUWrapperCodegen + ) +else: + from PyTorchSimFrontend.mlir.mlir_scheduling import MLIRScheduling + torch._inductor.codegen.common.register_backend_for_device( + "npu", + lambda scheduling: MLIRScheduling(scheduling), + ExtensionWrapperCodegen + ) torch_openreg.openreg.init() sys.modules['torch.npu'] = torch_openreg.openreg diff --git a/PyTorchSimFrontend/extension_config.py b/PyTorchSimFrontend/extension_config.py index 2d706bbc..09e5168d 100644 --- a/PyTorchSimFrontend/extension_config.py +++ b/PyTorchSimFrontend/extension_config.py @@ -11,6 +11,21 @@ CONFIG_TORCHSIM_DUMP_MLIR_IR = int(os.environ.get("TORCHSIM_DUMP_MLIR_IR", default=False)) CONFIG_TORCHSIM_DUMP_LLVM_IR = int(os.environ.get("TORCHSIM_DUMP_LLVM_IR", default=False)) +# --- Triton codegen route (WIP, opt-in) -------------------------------------- +# Replaces the hand-written MLIR emission in PyTorchSimFrontend/mlir with +# Inductor's own Triton codegen, lowered to the NPU by the triton-npu (tnpu) +# pass pipeline. OFF by default: the MLIR route stays the production path until +# this one is complete. See PyTorchSimFrontend/triton_backend/README.md. +CONFIG_TRITON_CODEGEN = bool(int(os.environ.get("TORCHSIM_TRITON_CODEGEN", default=0))) +# The triton-npu checkout that owns stages 1-5 (ttir -> ttshared -> tnpu passes +# -> RISC-V ELF). It is a SEPARATE repository, deliberately not vendored. +CONFIG_TNPU_DIR = os.environ.get( + "TNPU_DIR", default=os.path.join(CONFIG_TORCHSIM_DIR, "triton-npu")) +# tnpu runs in its own process: its passes need LLVM 23's MLIR bindings while +# this process holds LLVM 20's, and `mlir` is a namespace package, so the two +# cannot coexist in one interpreter (tnpu/config.py:activate_bindings). +CONFIG_TNPU_PYTHON = os.environ.get("TNPU_PYTHON", default=sys.executable) + def get_dump_path(): """Resolve TORCHSIM_DUMP_PATH and re-point Inductor's cache dir at it. diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index f4ed7d0d..ed52a56d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -366,6 +366,8 @@ def visit(n): for ln in builder.loop_nodes: visit(ln) + for dn in getattr(builder, "dma_nodes", ()): # DMAs outside any tile loop + visit(dn) return by_op diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index 5a40feec..98590d88 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -414,6 +414,9 @@ def __init__(self): self.loop_var_name = {} # value-identity-key -> loop name self.compute_nodes = [] self.loop_nodes = [] + # `_collect_dma_nodes` descends from the loop nodes, so a DMA hanging off + # the root (no tile loop in the kernel) would be missed. + self.dma_nodes = [] self._reset_matmul_fsm() # ---- matmul FSM ---- @@ -568,7 +571,13 @@ def _process_dram_indices(self, value, loop_index_list, indirect_box): loop_index_list.append(("c" + str(c), c)) # ---- main recursion ---- - def print_operation(self, op, node): + def visit_operation(self, op, node): + """Walk `op` and attach the nodes it produces under `node`. + + Builds the graph; it does not print. (The C++ pass this is ported from + does both in one method, `printOperation` -- here `bfs`/`display` own + the printing.) + """ name = _op_name(op) if name in SKIP_OPS: return @@ -605,7 +614,7 @@ def bool_true(k): for region in oper.regions: for block in region.blocks: for inner in block.operations: - self.print_operation(inner, loop_node) + self.visit_operation(inner, loop_node) return if name == "togsim.transfer": @@ -819,9 +828,14 @@ def _handle_dma_start(self, op, node): loop_idx_list.append(key) loop_stride_list.append(reordered[key]) - # base address + # base address: which tensor this DMA touches. The operand is the block + # argument itself in PyTorchSim's codegen; when it is a view of one + # instead, only the producer knows which -- so it says so (`dram_arg`) + # rather than the consumer guessing its way back through view ops. address = "arg" - if _is_block_arg(dram_memref): + if "dram_arg" in oper.attributes: + address += str(ir.IntegerAttr(oper.attributes["dram_arg"]).value) + elif _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) # element size @@ -875,6 +889,7 @@ def _handle_dma_start(self, op, node): tag_stride_list, loop_idx_list, loop_stride_list, indirect_box[0]) dma_node.op = op + self.dma_nodes.append(dma_node) node.add_child(dma_node) dma_node.add_parent(node) @@ -918,7 +933,9 @@ def _handle_dma_wait(self, op, node): dram_memref = f["dst"] elif dst_space == 1 and src_space == 0: dram_memref = f["src"] - if dram_memref is not None and _is_block_arg(dram_memref): + if "dram_arg" in user.attributes: + address += str(ir.IntegerAttr(user.attributes["dram_arg"]).value) + elif dram_memref is not None and _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) if len(tag_stride_list) == 0: @@ -928,6 +945,7 @@ def _handle_dma_wait(self, op, node): wait_node = TOGDMAWaitNode("DMAWaitNode", tag_index_list, tag_stride_list, tag_divider_list, address) wait_node.op = op + self.dma_nodes.append(wait_node) node.add_child(wait_node) wait_node.add_parent(node) @@ -1064,12 +1082,47 @@ def _insert_compute_markers(builder): # Driver. # --------------------------------------------------------------------------- def _find_kernel(module): - for op in module.body.operations: - if op.operation.name != "func.func": - continue + """The kernel function: named `kernel` in PyTorchSim's codegen, else the + module's only func.func (triton-npu carries the Triton kernel's own name). + Declines when there is more than one -- the intent would be a guess.""" + funcs = [op for op in module.body.operations + if op.operation.name == "func.func"] + for op in funcs: if ir.StringAttr(op.operation.attributes["sym_name"]).value == "kernel": return op - return None + return funcs[0] if len(funcs) == 1 else None + + +#: The loop roles (sec 9.1). Without one, a loop is a micro-loop the compute FSM +#: folds into a single node, not a tile loop. +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") + + +def _has_loop_role(op): + attrs = op.operation.attributes + return any(k in attrs and ir.BoolAttr(attrs[k]).value for k in _LOOP_ROLE_ATTRS) + + +def _is_address_plumbing(op): + """Scalar index/integer math (DMA offsets, mask extents) and the terminator. + + Only consulted on the no-top-level-loop path. PyTorchSim's codegen puts this + math in `affine.apply`, which SKIP_OPS drops; triton-npu emits an + arith/index_cast chain that would otherwise count as vector compute. + + Keyed on result type: tile data here is always vector- or float-typed. A + top-level SCALAR arithmetic kernel would be misread, but no path emits one. + """ + name = _op_name(op) + if name in ("func.return", "memref.cast"): + return True + if not name.startswith("arith."): + return False + results = list(op.operation.results) + if not results: + return False + return all(ir.IndexType.isinstance(r.type) or ir.IntegerType.isinstance(r.type) + for r in results) def _build(module, builder): @@ -1082,13 +1135,29 @@ def _build(module, builder): block = func_op.regions[0].blocks[0] out = [] + # A root is a top-level TILE loop, identified by its role attribute (sec + # 9.1) -- not by being an affine.for: bank_vectorize leaves a bare one for + # the tile's vector work, and rooting there orphans every DMA. + roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] + if roots: + for op in roots: + root = TOGNode("root") + builder._reset_matmul_fsm() + builder.visit_operation(op, root) + root.bfs(out) + return "".join(out) + + # No top-level loop: the body is ONE work-item -- the shape a Triton kernel + # arrives in, its grid becoming the trace producer's dispatch loop (sec 9.3). + # PyTorchSim's codegen keeps the tile loops in the kernel and never lands here. + root = TOGNode("root") + builder._reset_matmul_fsm() for op in block.operations: - if op.operation.name != "affine.for": + if _is_address_plumbing(op): continue - root = TOGNode("root") - builder._reset_matmul_fsm() - builder.print_operation(op, root) - root.bfs(out) + builder.visit_operation(op, root) + root.bfs(out) return "".join(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 5633769a..537c8ad0 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -119,20 +119,174 @@ def _attr_bool(op, key): # --------------------------------------------------------------------------- # step 1: rewrite signature + togsim.* ops (the unregistered-op glue) # --------------------------------------------------------------------------- -def _strip_aux(module): - """Erase memref.global decls and every func except @kernel (the wrapper).""" +def _strip_aux(module, keep=None): + """Erase memref.global decls and every func except the kernel. + + `keep` is the kernel op: its name is `kernel` only in PyTorchSim's codegen, + so the caller passes what `_find_kernel` resolved. + """ + keep_op = keep.operation if keep is not None else None victims = [] for op in module.body.operations: name = op.operation.name if name == "memref.global": victims.append(op) elif name == "func.func": - if ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": + if keep_op is not None: + if op.operation != keep_op: + victims.append(op) + elif ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": victims.append(op) for op in victims: op.operation.erase() +class WorkItem: + """A kernel whose body is ONE work-item, plus the grid over it. + + A Triton kernel describes a single program instance; the grid lives outside + it. The trace producer already splits the same way (design sec 9.3), so only + the enumeration is missing. + + `parallel_args` are the argument positions holding the program ids + (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` + their extents. Both outermost-first. + + An extent may be None, meaning "read it from shape_args at run time". Only + the NUMBER of axes has to be known when the kernel is compiled -- how many + loops to nest and how many iv[] slots to fill; the trip counts are just + values, and the producer ABI already takes them + (togsim_kernel(ctx, shape_args, n)). That is what lets one compiled trace + serve every shape. + """ + + def __init__(self, parallel_args, grid): + if len(parallel_args) != len(grid): + raise ValueError( + f"parallel_args {parallel_args} and grid {grid} must have the " + f"same length -- one program-id argument per grid axis") + self.parallel_args = list(parallel_args) + self.grid = [None if g is None else int(g) for g in grid] + + @property + def dynamic_axes(self): + """Indices into `grid` whose extent arrives at run time.""" + return [i for i, g in enumerate(self.grid) if g is None] + + +def _materialize_grid_loop(kernel, work_item, ctx): + """Wrap the body in the grid loop the Triton kernel does not carry: + + func @k(..., %pid: i32) { + scf.for %p = 0 to G { index_cast %p> } {outer_loop} + } + + Downstream is then unchanged: `_parallel_loop_chain` finds the tagged loop, + the outliner threads its induction variable through `iv[]`, and the loop left + behind becomes the dispatch enumeration. `outer_loop` means "independent + work-item" (sec 9.1) -- exactly a Triton program id. + + MUST run before `_rewrite_signature`, which erases the arguments and first + asserts none are still used. + """ + from mlir.dialects import arith, scf + + block = kernel.regions[0].blocks[0] + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + + pid_args = [block.arguments[i] for i in work_item.parallel_args] + body_ops = [o for o in block.operations + if o.operation.name not in _LOOP_TERMINATORS] + terminator = [o for o in block.operations + if o.operation.name in _LOOP_TERMINATORS][0] + + # Every bound first, and all of them before the first loop: each is created + # just before the terminator, so one made after an outer loop would sit + # BELOW it in the block while an inner loop uses it -- which does not + # dominate, and the verifier rejects it (only reachable at rank >= 2). + with ir.InsertionPoint(terminator), loc: + c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result + c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result + # A runtime extent still needs SOMETHING here: shape_args does not exist + # until _rewrite_signature adds it. The placeholder is replaced by + # _bind_runtime_bounds once it does. + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e or 1)).result + for e in work_item.grid] + + loops, inner = [], None + for ub in ubs: + # Nest inside the previous loop, BEFORE its yield: InsertionPoint on a + # block appends, and an scf.for body is already terminated. + ip = ir.InsertionPoint(terminator) if inner is None \ + else ir.InsertionPoint.at_block_terminator(inner.body) + with ip, loc: + loop = scf.ForOp(c0, ub, c1) + # ForOp leaves the body empty here; scf.for needs a terminator, and + # _outline_work_item inserts before it. + if len(loop.body.operations) == 0: + with ir.InsertionPoint(loop.body), loc: + scf.YieldOp([]) + loop.operation.attributes["outer_loop"] = ir.BoolAttr.get(True) + loops.append(loop) + inner = loop + + # Move the tile body inside the innermost loop, ahead of its yield. + inner_block = inner.body + inner_term = inner_block.operations[len(inner_block.operations) - 1] + for op in body_ops: + op.operation.move_before(inner_term) + + # Program ids are i32, induction variables index: cast once, at the top. + with ir.InsertionPoint(inner_block.operations[0]), loc: + casts = [] + for loop, pid in zip(loops, pid_args): + iv = loop.body.arguments[0] + casts.append(arith.IndexCastOp(pid.type, iv).result + if pid.type != idxty else iv) + + for pid, new in zip(pid_args, casts): + _replace_all_uses(pid, new) + + return [(loops[i], ubs[i]) for i in work_item.dynamic_axes] + + +def _bind_runtime_bounds(pending, shape_arg, ctx): + """Point each runtime loop bound at `shape_args[k]`. + + Runs AFTER _rewrite_signature, which is what creates the shape_args + argument. The loops stay in the entry function (the outliner moves only + their bodies), so the read is in scope where the bound is used. + """ + if not pending: + return + from mlir.dialects import arith + + i64 = ir.IntegerType.get_signless(64) + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + for k, (loop, placeholder) in enumerate(pending): + with ir.InsertionPoint(placeholder.owner), loc: + kc = ir.Operation.create( + "emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}).results[0] + elem = ir.Operation.create( + "emitc.subscript", results=[i64], + operands=[shape_arg, kc]).results[0] + bound = arith.IndexCastOp(idxty, elem).result + _replace_all_uses(placeholder, bound) + placeholder.owner.erase() + + +def _replace_all_uses(old, new): + """The bindings expose no replaceAllUsesWith on a Value.""" + for use in list(old.uses): + owner = use.owner + for i in range(len(owner.operands)): + if owner.operands[i] == old: + owner.operands[i] = new + + def _rewrite_signature(kernel, ctx): """Replace @kernel's memref tensor args with the ABI args (EmitCtx*, int64_t* shape_args, int32_t n) and rename it to togsim_kernel. @@ -196,15 +350,22 @@ def _is_outer(forop): return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value +#: The role is carried by the `outer_loop` attribute, not the dialect: +#: PyTorchSim's codegen emits affine.for, _materialize_grid_loop scf.for. Both +#: keep the induction variable in block argument 0. +_LOOP_OPS = ("affine.for", "scf.for") +_LOOP_TERMINATORS = ("affine.yield", "scf.yield", "func.return") + + def _parallel_loop_chain(block): - """The nested chain of `affine.for {outer_loop}` from `block` inward (one + """The nested chain of `{outer_loop}` loops from `block` inward (one work-item's parallel indices). Empty if the kernel has no parallel loop.""" chain = [] cur = block while True: nxt = None for op in cur.operations: - if op.operation.name == "affine.for" and _is_outer(op): + if op.operation.name in _LOOP_OPS and _is_outer(op): nxt = op break if nxt is None: @@ -281,7 +442,7 @@ def _outline_work_item(ctx, kernel, ctx_val): # move the work-item body into the tile fn (terminators stay behind). for op in [o for o in Lbody.operations - if o.operation.name not in ("affine.yield", "func.return")]: + if o.operation.name not in _LOOP_TERMINATORS]: op.operation.move_before(tret) # remap captures (Value `==` is identity): ctx -> ctx2, each parallel IV -> @@ -337,7 +498,7 @@ def _remap(block): # --- the dispatcher: marshal the IVs and hand the tile fn to togsim_dispatch --- term = [o for o in Lbody.operations - if o.operation.name in ("affine.yield", "func.return")][0] + if o.operation.name in _LOOP_TERMINATORS][0] fn_ref = _opaque(ctx, ts.TILE_SYMBOL) # function name -> verbatim pointer in C with ir.InsertionPoint(term): if ivs: @@ -499,16 +660,25 @@ def _add_extern_c(module, ctx): # --------------------------------------------------------------------------- # driver # --------------------------------------------------------------------------- -def lower_to_emitc(skeleton_module): +def lower_to_emitc(skeleton_module, work_item=None): """Lower a skeleton+API module (in place) to an EmitC module with the - `togsim_kernel` entry function. Returns the same module.""" + `togsim_kernel` entry function. Returns the same module. + + `work_item` is for kernels whose body is one work-item with the grid outside + (Triton's shape); None keeps PyTorchSim's, where the tile loops are already + in the kernel. + """ ctx = skeleton_module.context kernel = _find_kernel(skeleton_module) if kernel is None: - raise ValueError("no @kernel found in skeleton module") + raise ValueError("no kernel function found in skeleton module") - _strip_aux(skeleton_module) + _strip_aux(skeleton_module, keep=kernel) + pending = [] + if work_item is not None: + pending = _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) + _bind_runtime_bounds(pending, kernel.regions[0].blocks[0].arguments[1], ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch @@ -563,18 +733,20 @@ def _default_include_dir(): return os.path.join(root, "TOGSim", "include") -def skeleton_to_so(skeleton_module, so_path, include_dir=None): +def skeleton_to_so(skeleton_module, so_path, include_dir=None, work_item=None): """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the EmitC module text (for inspection / caching).""" - emitc = lower_to_emitc(skeleton_module) + emitc = lower_to_emitc(skeleton_module, work_item=work_item) inc = include_dir or _default_include_dir() cpp = emitc_to_cpp(emitc, include_dir=inc) compile_so(cpp, so_path, inc) return str(emitc) -def build_trace_so(postvcix_path, so_path, include_dir=None): - """Full P2 path from a post-vcix kernel .mlir to a trace `.so`.""" +def build_trace_so(postvcix_path, so_path, include_dir=None, work_item=None): + """Full P2 path from a post-vcix kernel .mlir to a trace `.so`. + + `work_item` -- see lower_to_emitc.""" from . import build_skeleton as bs ctx = ir.Context() @@ -582,7 +754,7 @@ def build_trace_so(postvcix_path, so_path, include_dir=None): with ctx: module = ir.Module.parse(open(postvcix_path).read(), ctx) bs.build_skeleton(module) - return skeleton_to_so(module, so_path, include_dir) + return skeleton_to_so(module, so_path, include_dir, work_item=work_item) def main(argv): diff --git a/PyTorchSimFrontend/triton-codegen-route.md b/PyTorchSimFrontend/triton-codegen-route.md new file mode 100644 index 00000000..5c009924 --- /dev/null +++ b/PyTorchSimFrontend/triton-codegen-route.md @@ -0,0 +1,289 @@ +# Triton 코드젠 경로를 PyTorchSim에 이식 + +`torch.compile`이 `npu:0`에서 PyTorchSim의 자체 MLIR 코드젠 대신 **Inductor의 Triton 백엔드**를 쓰고, **NPU lowering pass**가 이를 RISC-V로 낮추는 두 번째 코드젠 경로. + +## 작업 경계 + +이 경로는 두 부분으로 나뉩니다. **이 문서가 보고하는 것은 아래쪽입니다.** + +| 부분 | 하는 일 | 소관 | +|---|---|---| +| **NPU lowering pass** | Triton IR → linalg/memref → tts 레벨 백엔드 패스 → vcix/gemmini → RISC-V ELF | **이정민** — 범위 밖 | +| **기존 PyTorchSim으로의 이식** | 그 lowering pass를 기존 시뮬레이션 스택에 얹는 일 | 이 문서 | + +lowering pass 자체는 만들지 않았습니다. **이미 있는 것을 PyTorchSim이 쓸 수 있는 형태로 이식하고, 기존 TOGSim / gem5 / Spike 스택에 물린 것**이 여기서 한 일입니다. + +용어: 이 lowering 계층을 문서 전체에서 **NPU lowering pass**로 부릅니다. 다만 코드가 `triton-npu` 저장소에 있어서 **실제 식별자는 `tnpu`로 남아 있고**(`tnpu_bridge.py`, `tnpu/passes/`, `tnpu.spike`, `strip_for_tnpu`), 문서에서 코드를 찾아갈 수 있도록 그 이름들은 그대로 인용합니다. + +## 현재 도달점 + +| | | +|---|---| +| functional | 연결됨. `x + y`, `(x+y)*2 - x` **max abs error 0.0** (1024 elements, Spike) | +| timing | 연결됨. TOGSim **650 cycles**, 타일 compute는 gem5 **19 cycles 실측** | +| 동적 shape | 처리됨. 트레이스 하나가 모든 shape을 섬김 — n=1024 → grid 8, n=4096 → grid 32 | +| **커버리지** | **elementwise와 그 융합까지.** 남은 일은 op 스위트 → 모델까지 넓히는 것 — 5절 | +| CI | 전 잡 green | + +--- + +## 1. 기존 MLIR 경로와의 차이 + +### 갈라지는 곳과 합쳐지는 곳 + +``` + torch.compile / Inductor 스케줄 + │ + ┌───────────────┴───────────────┐ + │ │ + [기존] MLIR 경로 [신규] Triton 경로 + │ │ + Inductor 스케줄 → 손으로 쓴 Inductor 의 Triton 코드젠이 + op별 MLIR 템플릿 낸 커널 소스를 가로챔 + (gemm, conv, sdpa, sort, (op별 템플릿 없음) + cat, maxpool, bmm …) + │ │ + PyTorchSim mlir/ 패스 NPU lowering pass + PSAL LLVM 20 (in-process) stock LLVM 23 (subprocess) + └ 담당 이정민 — 범위 밖. + 여기서 한 일은 이 블록을 + 아래 합류점까지 잇는 배선 + │ │ + └───────────────┬───────────────┘ + │ + ▼ 여기서 다시 합류 ▼ + trace.so + trace_cycles.tsv + → TOGSim + (트레이스 계약은 완전히 동일) +``` + +핵심은 **TOGSim이 두 경로를 구분하지 못한다**는 점입니다. 트레이스 생산자의 형태가 같으므로 하드웨어 모델·DRAM·NoC·L2는 한 줄도 손대지 않았습니다. + +### 항목별 대조 + +| | 기존 MLIR 경로 | 신규 Triton 경로 | +|---|---|---| +| 커널을 만드는 주체 | op별 MLIR 템플릿 (직접 작성) | Inductor의 Triton 코드젠 | +| **커널 하나의 의미** | **루프 네스트 전체** | **타일 하나** | +| grid | 루프 네스트에서 읽어냄 | 커널 밖에 있음 → `WorkItem`이 합성 | +| lowering | `PyTorchSimFrontend/mlir/` | NPU lowering pass (subprocess) | +| 융합 | 템플릿 + `codegen_compiler_optimization` | Inductor가 이미 한 것을 물려받음 | +| op 커버리지 | gemm, conv×4, sdpa, sort, cat, maxpool, bmm | elementwise + 그 융합 | +| functional | `FunctionalSimulator.run_spike` | lowering pass 의 stage 6 (`tnpu.spike`) | +| timing | `trace.so` + `trace_cycles.tsv` → TOGSim | **동일** | +| 타일 cycle 실측 | gem5 | **동일** (`build_tog` sample 모드 공유) | +| DMA | 비동기 + `togsim.wait` 배리어 | **동기만** | +| 동적 shape | 트레이스 경로 미지원 (PR #269 진행 중) | 동작 | + +### 이 대조가 말해주는 것 + +**Triton 경로가 앞선 곳 — 동적 shape.** 기존 경로의 C++ 트레이스는 `trace_to_tilegraph(..., nullptr, 0)`으로 shape 인자를 아예 넘기지 않아 shape마다 트레이스를 다시 만들어야 하고, 그걸 푸는 작업이 PR #269로 아직 열려 있습니다. Triton 경로는 `shape_args`를 통해 트레이스 하나가 모든 shape을 섬깁니다. + +**기존 경로가 앞선 곳 — op 커버리지와 DMA 겹침.** 템플릿 9종 대 elementwise, 그리고 비동기 DMA 유무. 후자가 4절 사이클 격차의 원인입니다. + +**바뀌지 않은 것 — TOGSim 전체.** 하드웨어 설정, gem5 샘플링 방식, 트레이스 계약. 두 경로는 같은 시뮬레이터를 먹입니다. + +--- + +## 2. 파이프라인 + +``` +torch.compile + └ TritonNPUScheduling.define_kernel scheduling.py + │ Inductor 가 만든 triton 소스 텍스트 + 수집한 메타데이터 + ▼ + triton_npu_compile(src, meta, name) codecache.py + │ KernelSpec 생성 kernel_spec.py + │ - 블록 크기를 constexpr 로 고정 + │ - 인자 역할(in/out/inout) · dtype · numel + │ - grid, 사용자 스칼라 값 + ▼ + NPU lowering pass (subprocess) tnpu_bridge.py + │ 담당 이정민 — 범위 밖. 여기서 한 일은 이 단계를 + │ 호출하고 그 결과를 아래 스택에 물린 부분. + │ + │ 1 ttir triton 커널 → Triton IR + │ 2 ttshared → linalg / memref / scf.for (triton-shared) + │ 3 adapt tts 레벨 백엔드 6개 패스 (tnpu/passes/) + │ 4 lower vcix → gemmini DMA → LLVM + │ 5 binary mlir-translate → llc → RISC-V ELF + ▼ + TritonNPULauncher.__call__ codecache.py + │ + ├ functional 텐서 → runtime/*.raw → Spike → 텐서 functional.py + │ lowering pass 의 stage 6 (tnpu.spike) 재사용 + │ + └ timing 04-custom.mlir timing.py + ├ build_tog sample → gem5 → 타일 cycle 실측 + └ build_skeleton → trace.so + trace_cycles.tsv + → TOGSim +``` + +**LLVM 이음매.** lowering pass는 stock LLVM 23을, PyTorchSim은 PSAL LLVM 20을 씁니다. `mlir`이 namespace 패키지라 한 인터프리터에 공존할 수 없어, 두 쪽은 **텍스트 MLIR을 주고받는 subprocess**로 갈라져 있습니다. + +**lowering pass에 필요했던 진입점 3개.** 패스 로직을 고치는 것이 아니라 바깥에서 호출할 수 있게 여는 변경입니다 ([triton-npu#1](https://github.com/PSAL-POSTECH/triton-npu/pull/1)). + +| 훅 | 왜 필요했나 | +|---|---| +| `tnpu.cycle` | 타일 하나만 gem5로 재려면 DMA를 지운 1-program 바이너리가 필요 | +| `dram_arg` | TOG 빌더가 DMA의 DRAM 쪽이 어느 커널 인자인지 알아야 함 | +| `tnpu.spike` | stage 6이 자체 생성 입력 대신 **호출자의 텐서**로 돌아야 함 | + +--- + +## 3. 핵심 설계 문제: 커널 하나가 무엇을 뜻하는가 + +``` +MLIR 경로 커널 = 루프 네스트 전체. TOG 가 루프에서 work-item 을 읽어냄 +Triton 커널 = 타일 하나. grid 는 커널 밖, launch 가 쥐고 있음 +``` + +이식의 본질적 어려움은 여기 하나로 모입니다. 그런데 TOGSim의 트레이스 계약(`docs/design/togsim_cpp_trace.md` §9.1/§9.3)이 이미 둘을 구분하고 있었습니다: + +- `togsim_kernel_tile(ctx, iv, n)` — work-item 하나 +- `togsim_kernel(ctx, shape_args, n)` — 병렬 영역의 열거 + +Triton 커널 본문은 전자에 대응하므로 **후자를 합성해서 씌우면** 계약을 그대로 만족합니다. + +### grid 를 outer loop 으로 되세우기 + +Triton 쪽에서 grid 는 커널이 아니라 **KernelSpec 에 붙어 있습니다.** `kernel_spec.grid_of(meta)` 가 Inductor 의 numel 과 고정한 블록 크기로부터 축별 ceil-div 를 계산해 `grid=(8,)` 같은 값을 spec 에 적고, 커널 본문은 그 중 자기 몫이 몇 번째인지를 `pidX/Y/Z` 인자로 받을 뿐입니다. + +기존 PyTorchSim 은 정반대를 기대합니다. `build_tog` 는 **역할 속성이 붙은 최상위 루프**를 TOG 의 루트로 잡습니다: + +```python +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") +roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] +``` + +즉 work-item 을 **루프에서 읽어냅니다.** 그런데 Triton 커널에는 그 루프가 없습니다 — grid 로 흩어져 있으니까요. 루프가 없으면 루트도 없고, TOG 가 비게 됩니다. + +그래서 spec 의 grid 를 다시 루프로 세웁니다. `_materialize_grid_loop` 이 하는 일입니다: + +``` +들어올 때 func @k(..., %pidX: i32, ...) <- 타일 하나. 루프 없음 + body(%pidX) + +나갈 때 scf.for %p = 0 to G { {outer_loop = true} + body(<%pidX 를 index_cast %p 로 치환>) + } +``` + +`WorkItem(parallel_args, grid)` 이 **어느 인자가 program id 인지**와 **축이 몇 개인지**를 들고 있습니다. 패스는 축마다 루프를 하나씩 중첩하고, 본문을 그 안으로 옮기고, pid 인자의 모든 사용처를 루프 유도변수로 바꿉니다. 마지막에 `outer_loop` 속성을 답니다 — **`build_tog` 가 찾는 바로 그 표식**이고, 이게 붙어야 합성한 루프가 TOG 의 루트가 됩니다. + +(축별 범위는 `WorkItem` 이 들고 있을 수도, 런타임으로 미룰 수도 있습니다. timing 경로는 후자를 씁니다 — 바로 아래.) + +결과적으로 Triton 이 grid 로 표현한 것과 PyTorchSim 이 outer loop 로 표현한 것이 같은 것을 가리키게 되고, 그 뒤 파이프라인(`build_skeleton` → `trace.so` → TOGSim)은 MLIR 경로와 한 글자도 다르지 않게 흘러갑니다. + +### 동적 shape이 여기서 나옵니다 + +`_materialize_grid_loop`은 축 **개수**만 컴파일에 박고 **범위**는 `shape_args`에서 읽습니다. + +``` +컴파일 시 축이 몇 개인지만 안다 → 루프 네스트 골격 생성 +런타임 실제 numel 로 grid 계산 → trace_shape.txt 로 전달 + TOGSim 이 build_trace_tilegraph 에서 읽어 shape_args 로 주입 +``` + +측정: `dynamic=True`로 n=1024 → grid 8, n=4096 → grid 32. 트레이스 재생성 없음. + +다차원 grid(Triton 제약상 최대 3D)도 지원합니다. 구현 중 두 번 틀렸고 둘 다 rank ≥ 2에서만 드러났습니다 — 종료자가 있는 블록 끝에 삽입하는 문제, bound를 루프 뒤에 만들어 dominance를 깨는 문제. 그래서 테스트가 생성된 C++가 아니라 **MLIR 모듈 자체를 verify**합니다. + +--- + +## 4. 측정 결과 + +### functional + +| 커널 | 원소 | max abs error | +|---|---:|---:| +| `x + y` | 1024 | 0.0 | +| `(x + y) * 2 - x` (Inductor가 단일 커널로 융합) | 1024 | 0.0 | + +### timing + +| 항목 | 값 | 확인 내용 | +|---|---:|---| +| 타일 compute (gem5) | 19–21 | 마커 사이 `numCycles` 실측. placeholder 아님 | +| TOGSim 총계 | 650 | DRAM 트래픽 8192 B = 8 work-item × 2 load × 512 B, 정확히 일치 | +| 기존 MLIR 경로 (동일 연산) | 251 | 같은 자릿수 | + +650 대 251은 모델 오류가 아닙니다. **lowering pass가 동기 DMA만 내보내기 때문**입니다 — 생성된 IR에 `togsim.transfer` 3개, `togsim.wait` **0개**. work-item 안에서 load → compute → store가 직렬화되어 TOGSim이 겹칠 것이 없습니다. 기존 경로는 `togsim.wait` → `togsim.memory_barrier` 태그 슬롯 기계를 갖추고 있습니다. + +### 도중에 찾은 버그: 인자 한 칸 밀림 + +functional 배선은 배관 작업일 줄 알았는데 첫 실행에서 **1024개 중 896개가 틀렸습니다.** `pid_x=0` 블록만 맞고 나머지 7개는 전부 0. + +``` +lowered MLIR @k(%arg0..2: memref<*xf32> in_ptr0, in_ptr1, out_ptr0 + %arg3: i32 xnumel <- 사용자 스칼라 + %arg4,5,6: i32 gridX,Y,Z + %arg7,8,9: i32 pidX,Y,Z ) + +wrapper 호출 k(1,&d_in_ptr0, 1,&d_in_ptr1, 1,&d_out_ptr0, 8,1,1, pid_x,pid_y,pid_z) + +---- i32 6개뿐 ----+ + xnumel 누락 +``` + +triton-shared는 사용자 스칼라를 자기 grid/pid 인자 **앞에** 둡니다. wrapper는 이를 `spec.extra["scalar_args"]`에서 읽는데 우리가 생성하는 spec에는 `extra`가 없었습니다. 인자가 밀려 `pidX`가 `pid_y`(grid 루프가 절대 바꾸지 않는 값)를 받았고, program 0이 8번 돈 셈이 됐습니다. + +**틀린 값이 쓰레기가 아니라 0으로 나온 점**이 고약합니다. 쓰레기값이면 즉시 눈에 띄지만 0은 그럴듯해 보입니다. timing 경로는 인자 위치를 lowered MLIR 시그니처에서 직접 읽어 애초에 정확했고, 그래서 functional을 붙이기 전까지 드러나지 않았습니다. + +--- + +## 5. 남은 일 — 모델 커버리지까지 + +지금 통과하는 것은 elementwise와 그 융합입니다. **목표는 기존 MLIR 경로가 돌리는 모델들을 Triton 경로로도 돌리는 것**이고, 새 기능을 얹기보다 이미 있는 테스트를 그대로 돌려 막히는 곳을 고쳐 나가는 일입니다. + +목표선은 저장소에 이미 있습니다. + +| 단계 | 대상 | 지금 | +|---|---|---| +| 1. op | `tests/ops/` — elementwise, reduce, gemm, conv, attention, view, sort, fusion, misc | elementwise만 | +| 2. 모델 | `tests/models/` — MLP, MobileNet, ResNet, ViT, Transformer, Llama, Mixtral, DeepSeek, MoE, Diffusion, Yolov5 | 미착수 | + +모델은 op의 조합이라 op 하나가 막히면 모델은 첫 커널에서 멈춥니다. 그래서 op 먼저이고, 모델은 난이도 순으로 MLP → MobileNet/ResNet → ViT/Transformer → Llama가 무난합니다. 판정 기준은 두 단계가 같습니다 — **값이 torch와 일치하고, 사이클이 나오고, 경로에 실제로 진입할 것.** + +### 1단계에서 막히는 지점 + +대표 op를 돌려 확인한 것입니다. **경로 진입** 열이 필요한 이유는, Inductor가 일부 연산을 자체 커널 대신 외부 구현으로 빼기 때문입니다 — 그 경우 값은 맞지만 시뮬레이터를 거치지 않습니다. + +| 케이스 | 경로 진입 | 결과 | +|---|---|---| +| `x + y`, `(x+y)*2 - x` | 예 | 값 일치 | +| `x.t() + 1` | 예 | **값 틀림** | +| `relu`, `softmax` | 예 | 중단 — 헬퍼 모듈 부재 | +| `exp`, `cat`, `sum(dim=1)` | 예 | 중단 — NPU lowering pass | +| `a @ b` | **아니오** | 외부 구현으로 처리됨 | + +### 할 일 + +1. **비연속 텐서 처리.** `x.t() + 1`이 조용히 틀린 값을 냅니다. 값이 틀리면서 아무 신호도 없는 유일한 항목이라 최우선입니다. 원인은 파악됐고 이식 쪽입니다. +2. **헬퍼 모듈 벤더링.** `relu`, `softmax`, `clamp`, `max`, `min` 등이 torch 안의 헬퍼를 참조하는데 lowering pass 쪽 환경에는 없습니다. 필요한 것만 옮기면 커버리지가 한 번에 크게 늘어납니다. +3. **NPU lowering pass 쪽 실패** (`exp`, `cat`, reduction). 담당(이정민)과 나눌 부분입니다. +4. **matmul을 경로 안으로.** 지금은 외부 구현으로 빠져 시뮬레이터를 거치지 않습니다. 이게 뚫려야 systolic array 경로를 볼 수 있습니다. +5. **DMA 겹침.** 값이 아니라 사이클 정확도 항목입니다(4절의 251 대 650). 기존 경로에 이미 있는 기계를 옮기는 일이라 모델 단계와 병행 가능합니다. + +1~4가 풀리면 op 스위트는 대체로 통과할 것으로 봅니다. + +### 2단계에서 새로 볼 것 + +op 단위에서는 드러나지 않다가 모델에서 처음 나오는 것들입니다. **아직 돌려보지 않았으므로 측정이 아니라 예상입니다.** + +- **컴파일 시간** — 커널이 수백 개가 될 때 캐시가 실제로 먹는지 +- **커널 사이 버퍼 재사용** — op 테스트는 커널 하나로 끝나 드러나지 않음 +- **실제 shape** — op 테스트는 대개 잘 나뉘는 크기를 씀 +- **f16 / bf16** — 지금 확인된 것은 f32뿐 +- **backward 커널** — training 경로는 forward와 형태가 다름 +- **메모리 사용량** — 모델 규모에서 설정값을 넘는지 + +### 회귀 방지 + +`tests/system/test_triton_codegen.py`가 현재 경계를 못박고 있습니다. reduction은 **거부되는 동안 통과**하도록 되어 있어서, 컴파일에 성공하면 테스트가 실패합니다 — 지원이 생겼거나(그럼 체크를 지우면 됨), 하드웨어가 하지 않을 연산을 시뮬레이션하고 있다는 뜻이기 때문입니다. 위 항목이 하나씩 풀릴 때마다 이 방식으로 경계를 옮겨 적으면 됩니다. + +--- + +모듈별 동작과 사용법은 [`triton_backend/README.md`](triton_backend/README.md)에, 이식 작업 본체는 [PyTorchSim#305](https://github.com/PSAL-POSTECH/PyTorchSim/pull/305)에 있습니다. + +측정 환경: torch 2.10.0+cpu / triton 3.6.0, `systolic_ws_128x128_c1_simple_noc_tpuv3.yml`, `vpu_num_lanes` 128. 기존 MLIR 경로는 `tests/ops/elementwise/test_add.py` 통과로 회귀 없음 확인. diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md new file mode 100644 index 00000000..8477ea4f --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -0,0 +1,220 @@ +# 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. For how the route +compares with the MLIR one and what the numbers are, see +[`../triton-codegen-route.md`](../triton-codegen-route.md). + +## 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 00000000..460cf925 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/__init__.py @@ -0,0 +1,37 @@ +"""The Triton codegen route: Inductor's Triton backend + the tnpu lowering passes. + +WHAT THIS REPLACES +------------------ +The production path emits MLIR by hand (`PyTorchSimFrontend/mlir/`, ~5,500 lines +of string emission that also decides tiling, vectorization, DMA and scratchpad +placement). This route keeps Inductor's OWN Triton codegen for "what to compute" +and hands the resulting Triton kernel to triton-npu for "how to map it onto the +NPU": + + Inductor -> TritonNPUScheduling.define_kernel (scheduling.py) + | triton kernel SOURCE TEXT + v + -> TritonNPUCodeCache.load (codecache.py) + | a tnpu KernelSpec file (kernel_spec.py) + v + -> triton-npu, in a subprocess (tnpu_bridge.py) + ttir -> ttshared -> tnpu passes -> RISC-V ELF + v + -> Spike (functional) / TOGSim (timing) + +The two routes are mutually exclusive and chosen at device-registration time by +`extension_config.CONFIG_TRITON_CODEGEN` (env `TORCHSIM_TRITON_CODEGEN=1`). +Default OFF -- nothing here is on the production path yet. + +STATUS: scaffolding. The seams are wired and each stage names precisely what it +still owes; see README.md for the gap list. Expect failures, not results. +""" + +from . import _triton_compat + +# 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() + +from .scheduling import TritonNPUScheduling # noqa: E402,F401 +from .wrapper_codegen import TritonNPUWrapperCodegen # noqa: E402,F401 diff --git a/PyTorchSimFrontend/triton_backend/_triton_compat.py b/PyTorchSimFrontend/triton_backend/_triton_compat.py new file mode 100644 index 00000000..9c863b15 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/_triton_compat.py @@ -0,0 +1,143 @@ +"""Let Inductor's Triton codegen run on a machine with no GPU. + +ONE SKEW LEFT, AND IT IS NOT A VERSION SKEW +------------------------------------------- +`triton_hash_with_backend()` asks the triton runtime driver for the *current +target*, which on a box with no GPU raises "0 active drivers". We never launch +through triton's runtime -- triton-npu compiles the kernel ahead of time to a +RISC-V ELF -- so the value is only a cache-key ingredient, and a deterministic +string does the job. + +WHAT USED TO BE HERE, AND WHY IT IS GONE +---------------------------------------- +On torch 2.8 there was a second, larger skew: Inductor targeted the triton ~3.3 +API while triton-npu pins 3.6 (3.6 is what pins LLVM 23, and both sides of +triton-npu's textual IR seam must be the same LLVM), so `triton_key` had to be +injected back into `triton.compiler.compiler`. + +torch 2.10 pins triton 3.6.0 itself and reaches that symbol through its own +compat layer (`torch._inductor.runtime.triton_compat`), so on 2.10 the versions +simply agree and the injection is a no-op. It is kept, guarded, so the module +still works if someone runs an older torch -- `_torch_handles_triton()` decides. +""" + +import functools +import hashlib +import importlib +import os +import sys + +_installed = False + + +def triton_src_dir(): + """Where tnpu's triton checkout lives (its editable install points here). + + Read out of tnpu's own `setup/versions.env` rather than guessed, so the two + repos cannot drift: that file is the single place the checkout layout is + pinned (HEXAGON_MLIR_ROOT). + """ + from PyTorchSimFrontend import extension_config + override = os.environ.get("TNPU_TRITON_SRC") + if override: + return override + + root = "/workspace/hexagon-mlir" + versions = os.path.join(extension_config.CONFIG_TNPU_DIR, "setup", "versions.env") + try: + with open(versions) as f: + for line in f: + if line.startswith("HEXAGON_MLIR_ROOT="): + root = line.split("=", 1)[1].strip() + break + except OSError: + pass + return os.path.join(root, "triton", "python") + + +def ensure_triton_importable(): + """`import triton` in THIS interpreter, borrowing tnpu's checkout if needed. + + Inductor's Triton codegen imports triton at codegen time (for metadata and + hashing), so the driver needs it even though it never compiles with it. + """ + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + cand = triton_src_dir() + if os.path.isdir(os.path.join(cand, "triton")): + sys.path.insert(0, cand) + try: + import triton # noqa: F401 + return True + except ModuleNotFoundError: + pass + return False + + +def _stable_backend_hash(): + try: + import triton + version = triton.__version__ + except Exception: # noqa: BLE001 + version = "unknown" + key = f"pytorchsim-tnpu-{version}" + return hashlib.sha256(key.encode("utf-8")).hexdigest().upper() + + +def _torch_handles_triton(): + """True when this torch already knows how to reach triton's key itself. + + torch 2.10 routes it through torch._inductor.runtime.triton_compat, which + understands triton 3.6. Older torch imports `triton_key` straight out of + triton.compiler.compiler, where 3.6 no longer defines it. + """ + try: + from torch._inductor.runtime.triton_compat import triton_key # noqa: F401 + return True + except Exception: # noqa: BLE001 + pass + try: + mod = importlib.import_module("triton.compiler.compiler") + except Exception: # noqa: BLE001 + return False + return hasattr(mod, "triton_key") + + +def install(): + """Idempotently apply the shims. Returns a short report for logging.""" + global _installed + notes = [] + if not ensure_triton_importable(): + raise ModuleNotFoundError( + f"the Triton codegen route needs `triton` importable in this " + f"interpreter (Inductor imports it during codegen). Not found, and " + f"no checkout at {triton_src_dir()}. Set TNPU_TRITON_SRC, or install " + f"triton into this environment.") + if _installed: + return notes + + if not _torch_handles_triton(): + # Pre-2.10 torch: `triton_key` is imported from triton.compiler.compiler + # by SEVERAL call sites (codecache.CacheBase.get_system, + # _triton.triton_hash_with_backend, ...), each with its own local import. + # Supplying the symbol on the triton side satisfies all of them at once + # instead of chasing every call site; it is a cache-key ingredient, so + # any stable string will do. On torch 2.10 this branch does not run. + mod = importlib.import_module("triton.compiler.compiler") + mod.triton_key = _stable_backend_hash + notes.append("injected triton.compiler.compiler.triton_key " + "(this torch predates the triton 3.6 compat layer)") + + # Separately: triton_hash_with_backend also asks the triton runtime driver + # for the current target, which needs a GPU. We compile ahead of time to a + # RISC-V ELF and never use triton's runtime, so short-circuit it. + import torch.utils._triton as _t + _t.triton_hash_with_backend = functools.cache(_stable_backend_hash) + notes.append("patched torch.utils._triton.triton_hash_with_backend " + "(no GPU target to query)") + + _installed = True + return notes diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py new file mode 100644 index 00000000..d62de72d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -0,0 +1,90 @@ +"""Compile cache for the Triton route -- the counterpart of extension_codecache. + +`triton_npu_compile` is what the generated wrapper calls, exactly where the MLIR +route calls `custom_async_compile.mlir(...)`. It compiles the Triton kernel via +tnpu and returns the callable the wrapper then invokes per launch. + + define_kernel -> triton_npu_compile(src, meta, kernel_name) -> launcher + call site -> launcher(arg0, arg1, ..., xnumel) + +Layout mirrors the MLIR route so the two are comparable: one directory per source +hash under the dump path, holding the generated tnpu kernel file and every tnpu +artifact (01-ttir.mlir ... 05-*.elf). +""" + +import os + +from filelock import FileLock +from torch._inductor.codecache import get_hash + +from PyTorchSimFrontend import extension_config +from . import functional, kernel_spec, timing, tnpu_bridge + +logger = extension_config.setup_logger() + +LOCK_TIMEOUT = 600 + + +def _write_path(src_code): + return os.path.join(extension_config.get_dump_path(), + "triton_" + get_hash(src_code.strip())[1:12]) + + +class TritonNPULauncher: + """What a compiled kernel name is bound to in the generated wrapper. + + Holds the compile result; each call is one launch of the whole grid. + """ + + def __init__(self, kernel_name, workdir, meta): + self.kernel_name = kernel_name + self.workdir = workdir + self.meta = meta + self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") + + def __call__(self, *args): + """One launch of the whole grid: run it on Spike, then time it. + + Spike runs first so the caller's output tensors hold real values even if + TOGSim fails -- the two halves are independent. + """ + if extension_config.pytorchsim_functional_mode: + written = functional.run(self.workdir, self.meta, args) + logger.info("[Spike] %s wrote %s", self.kernel_name, written) + else: + logger.warning( + "[Spike] %s: functional mode is off, so the output tensors keep " + "whatever they held", self.kernel_name) + + if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): + timing.emit_trace(self.workdir, self.meta) + result = timing.run_togsim(self.workdir, meta=self.meta, args=args) + logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) + return result + + +def triton_npu_compile(src_code, meta, kernel_name): + """Compile one Inductor-generated Triton kernel through tnpu. + + Called from the generated wrapper at module import time (same point as + `custom_async_compile.mlir`). Synchronous for now: the MLIR route's thread + pool buys nothing until the pipeline itself is proven. + """ + write_path = _write_path(src_code) + os.makedirs(write_path, exist_ok=True) + + lock = FileLock(os.path.join(write_path, ".compile.lock"), timeout=LOCK_TIMEOUT) + with lock: + spec_path = os.path.join(write_path, f"{kernel_name}_spec.py") + elf = os.path.join(write_path, f"05-{kernel_name}.elf") + if not os.path.isfile(elf): + # Before write_spec_file, which rejects exactly the kernels whose + # source is worth keeping. + with open(os.path.join(write_path, "kernel.py"), "w") as f: + f.write(src_code) # the unmodified Inductor source + kernel_spec.write_spec_file(src_code, meta, spec_path, + tnpu_bridge.tnpu_dir()) + timing.store_meta(write_path, meta) # lets the timing step run standalone + tnpu_bridge.run_pipeline(spec_path, write_path, to_stage="binary") + logger.info("[triton-npu] %s -> %s", kernel_name, write_path) + return TritonNPULauncher(kernel_name, write_path, meta) diff --git a/PyTorchSimFrontend/triton_backend/functional.py b/PyTorchSimFrontend/triton_backend/functional.py new file mode 100644 index 00000000..4dd15682 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/functional.py @@ -0,0 +1,132 @@ +"""The functional half of the Triton route: real tensors -> Spike -> real tensors. + +The timing half (timing.py) tells you how long the kernel takes; this tells you +whether it computed the right thing. tnpu's stage 6 already runs the ELF under +Spike, but on inputs it generates itself. Here the launch's own tensors are +written as the `.raw` files stage 6 reads, and the outputs are copied back: + + run(workdir, meta, args) args -> runtime/*.raw -> spike -> args + +The binary is shape-specialised -- the spec bakes the grid, the scalar values and +the memref extents in -- so a launch whose shapes differ from the compiled ones +is rejected rather than silently run against the wrong bounds. +""" + +import os +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +RUNTIME_DIR = "runtime" + + +class ShapeMismatch(RuntimeError): + """The launch does not match the shapes the binary was compiled for.""" + + +def _np_dtype(name): + import numpy as np + return np.dtype("bool" if name == "bool" else name) + + +def tensor_args(meta, args): + """[(arg_meta, tensor)] for the launch, paired by position. + + Inductor passes the tensors first and the numels after, in signature order, + so `meta["args"]` (tensors only) lines up with the leading arguments. + """ + import torch + + tensors = [a for a in args if isinstance(a, torch.Tensor)] + metas = meta["args"] + if len(tensors) != len(metas): + raise ShapeMismatch( + f"{meta['kernel_name']}: launch passed {len(tensors)} tensor(s), " + f"but the spec declares {len(metas)} ({[m['name'] for m in metas]})") + return list(zip(metas, tensors)) + + +def _check(meta, pairs): + for m, t in pairs: + if t.numel() != m["numel"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' has {t.numel()} " + f"element(s), but the binary was compiled for {m['numel']}. " + f"tnpu bakes the extents, the grid and the scalar values into " + f"the kernel, so a dynamic-shape graph reuses an ELF that does " + f"not fit. The timing path does handle this (it takes the grid " + f"at run time); set pytorchsim_functional_mode: False to study " + f"cycles alone, or keep shapes static to check values.") + if str(t.dtype).removeprefix("torch.") != m["dtype"]: + raise ShapeMismatch( + f"{meta['kernel_name']}: '{m['name']}' is {t.dtype}, but the " + f"binary was compiled for {m['dtype']}") + + +def 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"): + t.detach().to("cpu").contiguous().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") + t.copy_(torch.from_numpy(flat).view_as(t).to(t.dtype)) + written.append(m["name"]) + return written + + +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 RuntimeError( + f"[Spike] {meta['kernel_name']} failed:\n" + + (proc.stdout + proc.stderr)[-2000:]) + + return read_outputs(workdir, meta, args) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py new file mode 100644 index 00000000..e5a34400 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -0,0 +1,433 @@ +"""Inductor kernel -> tnpu KernelSpec. + +Two jobs, both of which exist because Inductor's Triton output is written for a +GPU launcher and tnpu's is written for a static, ahead-of-time pipeline: + +1. `collect_meta` -- pull everything tnpu needs out of the Inductor kernel while + we still have `V.graph`: argument names/roles/dtypes/sizes, the constexprs, + and the numels the grid is computed from. This runs at codegen time; by the + time the compile callable fires, `V.graph` is gone. + +2. `write_spec_file` -- turn the Triton source + that metadata into a kernel file + tnpu can load (`tnpu.spec.load_spec`). + +WHY THE SOURCE HAS TO BE REWRITTEN +---------------------------------- +Inductor emits, above the kernel: + + from torch._inductor.runtime import triton_heuristics + @triton_heuristics.pointwise(size_hints={'x': 1024}, ..., inductor_meta=...) + @triton.jit + def triton_npu_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): + +Neither line can survive into tnpu: + + * the tnpu triton venv has NO torch (deliberately -- tnpu/spec.py), so the + `torch._inductor.runtime` import fails on sight; + * `triton_heuristics.pointwise` is the AUTOTUNER. It picks XBLOCK at runtime + and derives `grid = cdiv(xnumel, XBLOCK)` from it. tnpu needs both to be + constants: the block size becomes a `tl.constexpr` in the ttir signature and + the grid is executed as a sequential loop by the generated C wrapper. + +So the decorator is stripped and XBLOCK is pinned as a constexpr instead. That is +not a workaround -- fixing the config at codegen time is what makes the kernel +statically describable, which is the whole premise of this route. +""" + +import math +import os +import re + +from torch._inductor.virtualized import V + +#: Triton signature token -> (torch dtype name, bytes). Only the dtypes +#: tnpu/wrapper.py can round-trip through .raw files. +_DTYPE = { + "*fp32": "float32", "*fp16": "float16", "*bf16": "bfloat16", + "*i64": "int64", "*i32": "int32", "*i8": "int8", "*i1": "bool", + "fp32": "float32", "i32": "int32", "i64": "int64", +} + + +#: Triton scalar token -> C type, for the wrapper's kernel declaration. +_C_TYPE = {"i32": "int32_t", "i64": "int64_t", "fp32": "float"} + + +class SpecIncomplete(RuntimeError): + """Metadata tnpu requires that this kernel did not provide. + + Raised with the missing field named, rather than writing a spec that fails + deeper in the pipeline where the cause is unrecoverable. + """ + + +# --------------------------------------------------------------------------- +# 1. codegen-time metadata capture +# --------------------------------------------------------------------------- +def _buffer_numel(name): + """Element count of an Inductor buffer, or None if it cannot be resolved.""" + try: + buf = V.graph.get_buffer(name) + if buf is None: + return None + size = buf.get_layout().size + n = 1 + for s in size: + n *= int(V.graph.sizevars.size_hint(s)) + return n + except Exception: # noqa: BLE001 - best effort; caller reports it as missing + return 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(): + out[arg] = ("out", 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 + args.append({ + "name": name, + "role": role, + "buffer": buf, + "dtype": _DTYPE.get(signature.get(name, ""), None), + "numel": _buffer_numel(buf) if buf else None, + }) + + # The numels Inductor appends to the call. They live in `kernel.numels`, + # keyed by iteration-space PREFIX ('x', 'y', 'r0', ...), not as xnumel/rnumel + # attributes (SIMDKernel.__init__ builds them from the tiling). These are + # what the grid is computed from. + numels = {} + for prefix, val in (getattr(kernel, "numels", None) or {}).items(): + try: + numels[f"{prefix}numel"] = int(V.graph.sizevars.size_hint(val)) + except Exception: # noqa: BLE001 - dynamic shape; reported by _grid + numels[f"{prefix}numel"] = None + + return { + "kernel_name": kernel_name, + "signature": {str(k): str(v) for k, v in signature.items()}, + "constants": {str(k): v for k, v in constants.items()}, + "args": args, + "numels": numels, + "inside_reduction": bool(getattr(kernel, "inside_reduction", False)), + "fixed_config": fixed_config_for(kernel), + } + + +#: Parallel iteration prefixes, OUTERMOST first. Inductor's `x` is the +#: contiguous axis, so it is innermost; `r*` prefixes are reductions, looped +#: inside the kernel rather than spread over the grid (prefix_is_reduction). +_PARALLEL_PREFIXES = ("z", "y", "x") + + +def _block_name(prefix): + return f"{prefix.upper()}BLOCK" + + +def parallel_axes(numels): + """Grid axes this kernel uses, outermost first.""" + return [p for p in _PARALLEL_PREFIXES if f"{p}numel" in numels] + + +def fixed_config_for(kernel): + """Block sizes pinned at codegen time. + + tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a + sequential loop, so there is no autotuner to choose the blocks later and no + runtime `grid=` callable. Fixing them here is what makes the launch shape + static. + + Tile dim 0 is the one `bank_vectorize` spreads over the lanes, so the + OUTERMOST axis gets the lane count -- a per-lane depth of 1, the shape every + tnpu baseline runs. The remaining axes get 1, which leaves the tile exactly + that verified shape and lets the grid cover the rest. It is conservative + rather than fast; choosing real tile sizes is the block-size policy gap in + README, not something to guess at here. + """ + from PyTorchSimFrontend import extension_config + lanes = int(extension_config.vpu_num_lanes) + + axes = parallel_axes(getattr(kernel, "numels", None) or {}) + cfg = {_block_name(p): (lanes if i == 0 else 1) for i, p in enumerate(axes)} + if len(axes) > 1: + # Loud, because the shape is correct but pathological: an inner block of + # 1 makes every work-item move a strided column. Fine for getting a + # multi-axis kernel through the route, misleading to benchmark. + extension_config.setup_logger().warning( + "[triton-npu] %s tiles over %s; inner blocks pinned to 1, which is " + "correct but not a tiling worth measuring", + getattr(kernel, "kernel_name", "kernel"), axes) + cfg.setdefault("XBLOCK", lanes) # a kernel with no tiling info still has x + if getattr(kernel, "inside_reduction", False): + # A reduction block is NOT free to be the lane count: the reduced axis + # has to stay inside a lane (see triton-npu kernels/reduce.py). Left + # unset on purpose so the reduction path fails loudly rather than + # silently picking a layout the hardware cannot execute. + 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 runtime setup Inductor emits at module scope. Meaningless here (the +#: kernel is compiled ahead of time to a RISC-V ELF) and its import is dropped +#: above, so the call would be a NameError. +_DROP_CALL_RE = re.compile(r"^\s*triton_helpers\.set_driver_to_gpu\(\)") +#: Anything else from triton_helpers is a real dependency -- maximum/minimum/ +#: promote_to_tensor and friends, which reductions and clamps use constantly. +_HELPER_USE_RE = re.compile(r"\btriton_helpers\.(\w+)") + + +def strip_for_tnpu(src): + """Remove everything the torch-free tnpu venv cannot import. + + Drops torch/inductor imports and the `@triton_heuristics.*(...)` decorator + (keeping `@triton.jit`), then re-adds the two imports the kernel body needs. + + Raises SpecIncomplete if the kernel still calls into `triton_helpers`: that + module lives in torch, so it has to be vendored into the tnpu venv before + such a kernel can compile. Failing here names the missing helper; letting it + through fails as a bare NameError inside tnpu's stage-1 worker instead. + """ + lines = src.splitlines() + out, i = [], 0 + while i < len(lines): + line = lines[i] + if _HEURISTIC_RE.match(line.strip()) or _HEURISTIC_RE.match(line): + # skip the whole decorator call, up to (not including) @triton.jit + while i < len(lines) and lines[i].strip() != "@triton.jit": + i += 1 + continue + if _DROP_IMPORT_RE.match(line) or _DROP_CALL_RE.match(line): + i += 1 + continue + out.append(line) + i += 1 + body = "\n".join(out) + + used = sorted(set(_HELPER_USE_RE.findall(body))) + if used: + raise SpecIncomplete( + f"kernel uses triton_helpers.{{{','.join(used)}}}, which lives in " + f"torch and the tnpu venv has no torch. Vendor a minimal " + f"triton_helpers into the tnpu venv (or into TRITON_SRC) before this " + f"kernel can compile.") + + # The generated source already imports triton itself; only add what a + # stripped module might be missing. + prefix = "" + if "import triton.language as tl" not in body: + prefix = "import triton\nimport triton.language as tl\n\n" + # tl_math is triton's own, re-exported through triton_helpers; the dropped + # torch import took it with it. + if re.search(r"\btl_math\.", body): + prefix += "from triton.language import math as tl_math\n" + + # libdevice members are @core.extern: no triton_shared implementation, so a + # call returns None and fails obscurely in stage 1. Name it here instead. + ext = sorted(set(re.findall(r"\blibdevice\.(\w+)", body))) + if ext: + raise SpecIncomplete( + f"kernel calls libdevice.{{{','.join(ext)}}}: those are extern math " + f"intrinsics with no implementation on the triton_shared backend. " + f"They need lowering to a VPU op (or a scalar fallback) before this " + f"kernel can compile.") + return prefix + body + + +def scalar_args(meta): + """User scalar parameters, in kernel order, as [(name, c_type, value)]. + + triton-shared keeps these in the lowered signature ahead of its own six + grid/pid arguments, so the wrapper must pass them or every later argument + lands one slot early -- pidX then reads pidY and only program 0 runs. + """ + numels = meta["numels"] + out = [] + for name, token in meta["signature"].items(): + if token.startswith("*") or token == "constexpr": + continue + ctype = _C_TYPE.get(token) + if ctype is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: scalar '{name}' has type {token!r}, " + f"which has no C mapping in _C_TYPE") + if numels.get(name) is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: no value for scalar '{name}' -- " + f"collect_meta resolves these from kernel.numels") + out.append((name, ctype, int(numels[name]))) + return out + + +def grid_of(meta): + """Launch grid, from the numels and the pinned block sizes, outermost first. + + 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 = parallel_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}) +from tnpu.spec import KernelSpec, Arg # noqa: E402 + +#: The rewritten Triton source, beside this file. It must be a REAL file on +#: disk, not an exec'd string: triton's @jit reads the function back with +#: inspect.getsourcefile and rejects anything else ("@jit functions should be +#: defined in a Python file"). +TRITON_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), + {triton_module!r}) + + +def kernel(): + spec = importlib.util.spec_from_file_location( + {kernel_name!r} + "_triton", TRITON_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, {kernel_name!r}) + + +def make_inputs(torch, seed=0): + g = torch.Generator().manual_seed(seed) + out = {{}} +{make_inputs_body} + return out + + +def reference(inputs): + # The Inductor route has no per-kernel torch reference: correctness is + # checked at the graph level by the test that ran torch.compile. tnpu's + # stage 7 is therefore not meaningful here and the pipeline is driven to + # stage 6 (spike) instead. + return {{}} + + +SPEC = KernelSpec( + name={kernel_name!r}, + kernel=kernel, + signature={signature!r}, + constexprs={constexprs!r}, + args=[ +{args_body} + ], + grid={grid!r}, + reference=reference, + make_inputs=make_inputs, + extra={{"scalar_args": {scalar_decls!r}, + "scalar_values": {scalar_values!r}}}, + notes="generated from Inductor triton codegen", +) +''' + + +def write_spec_file(src_code, meta, path, tnpu_dir): + """Write a tnpu kernel file for this Inductor kernel. Returns `path`.""" + missing = [a["name"] for a in meta["args"] if not a["dtype"] or not a["numel"]] + if missing: + raise SpecIncomplete( + f"{meta['kernel_name']}: no dtype/numel for {missing} -- " + f"collect_meta could not resolve them from V.graph") + + signature = dict(meta["signature"]) + constexprs = dict(meta["constants"]) + for k, v in (meta.get("fixed_config") or {}).items(): + if k not in signature: + # Inductor already fixed this one in the kernel BODY rather than + # taking it as a parameter -- a persistent reduction does that with + # R0_BLOCK. Passing it would not match the signature, and there is + # nothing left for us to choose. + continue + if v is None: + raise SpecIncomplete( + f"{meta['kernel_name']}: block size {k} is unset " + f"(fixed_config_for leaves reduction blocks unset on purpose)") + constexprs[k] = v + signature[k] = "constexpr" + + args_body = "\n".join( + f" Arg({a['name']!r}, {a['role']!r}, {a['dtype']!r}, ({a['numel']},))," + for a in meta["args"]) + make_inputs_body = "\n".join( + f" out[{a['name']!r}] = torch.randn({a['numel']}, generator=g)" + f".to(torch.{a['dtype']})" + for a in meta["args"] if a["role"] in ("in", "inout")) or " pass" + + triton_module = f"{meta['kernel_name']}_triton.py" + with open(os.path.join(os.path.dirname(path), triton_module), "w") as f: + f.write(strip_for_tnpu(src_code)) + + scalars = scalar_args(meta) + text = SPEC_TEMPLATE.format( + kernel_name=meta["kernel_name"], + tnpu_dir=tnpu_dir, + triton_module=triton_module, + signature=signature, + constexprs=constexprs, + args_body=args_body, + make_inputs_body=make_inputs_body, + grid=grid_of(meta), + scalar_decls=[(n, c) for n, c, _ in scalars], + scalar_values={n: v for n, _, v in scalars}, + ) + with open(path, "w") as f: + f.write(text) + return path diff --git a/PyTorchSimFrontend/triton_backend/scheduling.py b/PyTorchSimFrontend/triton_backend/scheduling.py new file mode 100644 index 00000000..56e7f1f8 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/scheduling.py @@ -0,0 +1,92 @@ +"""Inductor scheduling for the Triton route. + +`TritonNPUScheduling` keeps ALL of Inductor's Triton codegen -- fusion, index +expressions, masking, reductions, the kernel source itself -- and changes only +what happens to the generated source afterwards. Upstream hands it to +`async_compile.triton(...)`, which calls `triton.compile` for a GPU; we hand it +to `triton_npu_compile(...)`, which runs the triton-npu pipeline for this NPU. + +Two overrides, and nothing else: + + define_kernel emit our compile call into the wrapper instead of upstream's. + kernel_type a TritonKernel whose call site is a plain python call, because + the name is bound to our callable rather than to a triton + launcher with a `.run(grid=..., stream=...)` interface. +""" + +from torch._inductor.codegen.common import IndentedBuffer +from torch._inductor.codegen.triton import TritonKernel, TritonScheduling +from torch._inductor.utils import Placeholder, get_fused_kernel_name +from torch._inductor.virtualized import V + +from . import kernel_spec + + +class TritonNPUKernel(TritonKernel): + """A TritonKernel launched as a plain call. + + Upstream emits `kernel.run(a, b, xnumel, grid=grid(xnumel), stream=...)`, + where `grid` is resolved at RUNTIME by triton_heuristics from the autotuned + XBLOCK. There is no autotuner and no stream here: the kernel name is bound to + the callable `triton_npu_compile` returned, so the call is `kernel(a, b, n)`. + + That is also why the block sizes must be fixed at CODEGEN time -- see + kernel_spec.fixed_config_for. A grid that is only known after autotuning + cannot be written into a tnpu KernelSpec. + """ + + # **kwargs, not a fixed signature: Inductor keeps adding parameters here + # (2.10 added `deallocate_ws`). None of them apply to this route -- there is + # no triton launcher and no workspace to release -- so they are accepted and + # ignored rather than pinning us to one torch release. + def call_kernel(self, name: str, node=None, **kwargs): + wrapper = V.graph.wrapper_code + _, call_args, _, arg_types = self.args.python_argdefs() + self.add_numel_to_call_args(name, call_args, arg_types) + # add_numel_to_call_args appends the numels as SYMPY values, which the + # triton path later renders through pexpr. ExtensionWrapperCodegen joins + # call args as plain strings (mlir_codegen_backend.py:241), so render + # them here instead of handing it a sympy Integer. + call_args = [a if isinstance(a, str) else str(a) for a in call_args] + # triton=False -> PythonWrapperCodegen emits `name(args...)`, the same + # shape the MLIR route uses (mlir_common.py:627). + wrapper.generate_kernel_call(name, call_args, triton=False) + + +class TritonNPUScheduling(TritonScheduling): + kernel_type = TritonNPUKernel + + count = 0 + + def define_kernel(self, src_code, node_schedule, kernel): + wrapper = V.graph.wrapper_code + if src_code in wrapper.src_to_kernel: + return wrapper.src_to_kernel[src_code] + + fused_name = get_fused_kernel_name(node_schedule, "original_aten") + kernel_name = "_".join( + x for x in ("triton_npu", fused_name, str(TritonNPUScheduling.count)) if x + ) + TritonNPUScheduling.count += 1 + wrapper.src_to_kernel[src_code] = kernel_name + + # Upstream substitutes these two placeholders inside define_kernel; the + # source still carries them here, and the tnpu side parses the source, so + # they have to be resolved before it leaves this function. + src_code = src_code.replace(str(Placeholder.DESCRIPTIVE_NAME), kernel_name) + src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name) + + meta = kernel_spec.collect_meta(kernel, kernel_name) + + compile_wrapper = IndentedBuffer() + compile_wrapper.writeline(f"triton_npu_compile('''{src_code}''',") + compile_wrapper.writeline(f" meta={meta!r},") + compile_wrapper.writeline(f" kernel_name={kernel_name!r})") + + origins = ", ".join( + sorted({str(o) for n in node_schedule + for o in getattr(getattr(n, "node", None), "origins", ()) or ()}) + ) + wrapper.define_kernel(kernel_name, compile_wrapper.getvalue(), + f"# origins: {origins}") + return kernel_name diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py new file mode 100644 index 00000000..be28060b --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -0,0 +1,234 @@ +"""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. + + `grid_of` orders axes OUTERMOST first (z, y, x -- x is Inductor's contiguous + one), while the program-id arguments are always laid out x, y, z. The two + are zipped downstream, so the argument list is built per axis rather than as + a range. + """ + 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.parallel_axes(meta["numels"]) + # Extents are left to run time: only the axis COUNT has to be compiled in, + # and the launch knows the real numels. One trace then serves every shape. + return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], + grid=[None] * len(axes)) + + +def write_shape(workdir, meta, args=()): + """Write the grid extents the trace producer reads as shape_args. + + `args` is the launch's positional arguments; Inductor appends the numels + after the tensors, so the trailing values are them, in `meta["numels"]` + order. Falls back to the compile-time hint when they are absent. + """ + from . import kernel_spec + + numels = dict(meta["numels"]) + # Only the PARALLEL numels ride along on the call -- a reduction axis is + # looped inside the kernel, so it is not passed and must not consume one of + # the trailing values. They arrive in kernel order, which is the dict's. + passed = [k for k in numels if not k.startswith("r")] + trailing = [a for a in args if isinstance(a, int) and not isinstance(a, bool)] + if passed and len(trailing) >= len(passed): + for key, val in zip(passed, trailing[-len(passed):]): + numels[key] = val + + axes = kernel_spec.parallel_axes(numels) + + cfg = meta.get("fixed_config") or {} + grid = [] + for p in axes: + n, block = numels.get(f"{p}numel"), cfg.get(f"{p.upper()}BLOCK") + if n is None or not block: + raise ValueError(f"no extent for grid axis '{p}': {n!r} / {block!r}") + grid.append(-(-int(n) // int(block))) # ceil-div + + path = os.path.join(workdir, SHAPE_TXT) + with open(path, "w") as f: + f.write("\n".join(str(g) for g in grid) + "\n") + logger.info("[TOGSim] grid %s -> %s", grid, SHAPE_TXT) + return grid + + +def emit_trace(workdir, meta): + """Build `trace.so` + `trace_cycles.tsv` from tnpu's post-vcix IR. + + Returns the number of compute tiles the cycle table covers. + """ + from PyTorchSimFrontend.mlir.passes import build_skeleton as bs + from PyTorchSimFrontend.mlir.passes import cycle_table as ct + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + postvcix = os.path.join(workdir, "04-custom.mlir") + if not os.path.isfile(postvcix): + raise FileNotFoundError( + f"{postvcix} not found -- tnpu must have run at least to stage 4 " + f"(the post-vcix IR is what the trace is built from)") + + # Before build_skeleton: both read the post-vcix IR, which it rewrites in place. + cycles = measure_tile_cycles(workdir, meta) + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(postvcix).read(), ctx) + bs.build_skeleton(module) + compute_types = ct._compute_types(module) + n_tiles = len(compute_types) + + if cycles: + # One numCycles per compute node; pad/truncate as the MLIR route does. + cl = list(cycles) + if len(cl) != n_tiles: + logger.warning("[Gem5] returned %d cycle(s) for %d " + "tile(s); padding with the last", len(cl), n_tiles) + cl = (cl + [cl[-1]] * n_tiles)[:n_tiles] + # Systolic-array fill; only matmul tiles use it. + lanes = int(extension_config.vpu_num_lanes) + table = ct.build_cycle_table(module, cl, x_offset=lanes, w_offset=0) + else: + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + logger.warning( + "[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d " + "tiles): gem5 sampling did not produce a measurement, so " + "compute latency is NOT modelled", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), + work_item=work_item_for(meta)) + + ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) + if cycles: + logger.info("[Gem5] tile cycles: %s", table) + return n_tiles + + +def run_togsim(workdir, meta=None, args=(), attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict. + + `meta`/`args` supply the grid: the trace producer takes its loop bounds from + shape_args, so they are written out per launch rather than compiled in. + """ + from Simulator.simulator import TOGSimulator + + so = os.path.join(workdir, TRACE_SO) + if not os.path.isfile(so): + raise FileNotFoundError(f"{so} not found -- call emit_trace first") + if meta is not None: + write_shape(workdir, meta, args) + + # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its + # DIRECTORY, and reads the file itself only on the STONNE path. + handle = os.path.join(workdir, "tile_graph.onnx") + result_path = TOGSimulator.run_standalone( + handle, attribute_path or os.path.join(workdir, "attribute"), + timeout_sec=timeout_sec) + return TOGSimulator.get_result_from_file(result_path) + + +def store_meta(workdir, meta): + """Persist codegen metadata beside the artifacts, so the timing step can run + standalone.""" + with open(os.path.join(workdir, META_JSON), "w") as f: + json.dump(meta, f, indent=2) + + +def load_meta(workdir): + with open(os.path.join(workdir, META_JSON)) as f: + return json.load(f) diff --git a/PyTorchSimFrontend/triton_backend/tnpu_bridge.py b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py new file mode 100644 index 00000000..a30593d1 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/tnpu_bridge.py @@ -0,0 +1,107 @@ +"""Run the triton-npu pipeline, out of process. + +WHY A SUBPROCESS +---------------- +tnpu's passes run on LLVM 23's MLIR python bindings; this process holds LLVM 20's +(TORCHSIM_LLVM_PATH). `mlir` ships without an `__init__.py`, so it is a NAMESPACE +package whose `__path__` is the union of every `mlir/` directory on sys.path -- +two LLVMs in one interpreter silently merge and fail later with an AttributeError +from a generated dialect module (tnpu/config.py:activate_bindings documents the +exact failure). They cannot share an interpreter, so tnpu gets its own. + +The seam between them is a FILE, which is measured to work: LLVM 23 prints the +IR, and LLVM 20's bindings parse it back without complaint (verified by feeding +tnpu's 04-custom.mlir to PyTorchSim's build_tog). That is what makes the split +viable rather than merely necessary. +""" + +import os +import re +import subprocess + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + + +class TnpuError(RuntimeError): + """A tnpu stage failed. Inductor reports only str(exc), so the stage's own + diagnostic has to travel in the message.""" + + #: How a failing stage names itself: MLIR diagnostics and exception lines. + _SIGNAL = re.compile( + r"^(?!\s|Traceback|During handling|The above)" + r"(.*\berror:\s.*|.*failed to legalize.*|" + r"[\w.]*(?:Error|Exception)\b.*|.*Assertion.*)$", re.M) + #: Frames and carets: context, not the diagnostic. + _FRAME = re.compile(r'^\s|^\s*File "|^\s*\^') + + def __init__(self, message, cmd=None, output=None): + self.cmd = cmd + self.output = output + if output: + hits = [h.strip() for h in self._SIGNAL.findall(output) + if not self._FRAME.match(h)] + if not hits: + hits = [l for l in output.strip().splitlines() + if l.strip() and not self._FRAME.match(l)] + message = message + "\n " + "\n ".join(l[:300] for l in hits[-3:]) + super().__init__(message) + + +def tnpu_dir(): + d = extension_config.CONFIG_TNPU_DIR + if not os.path.isdir(d): + raise TnpuError( + f"triton-npu checkout not found at {d}. It is a separate repository " + f"and is not vendored; clone it there or set TNPU_DIR.") + return d + + +def doctor(): + """Return (ok, output) for tnpu's own toolchain check.""" + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, os.path.join(tnpu_dir(), "run.py"), "doctor"], + capture_output=True, text=True, cwd=tnpu_dir()) + return proc.returncode == 0, proc.stdout + proc.stderr + + +def run_pipeline(spec_path, workdir, to_stage="binary", from_stage="ttir", + verbose=False, timeout=1800): + """Drive tnpu's stages over `spec_path`, writing artifacts into `workdir`. + + Stops at `to_stage`. The default is `binary` (through the RISC-V ELF): + stage 6 (spike) needs the caller's real tensors as .raw files and stage 7 + compares against a per-kernel torch reference, neither of which exists on + the Inductor route -- correctness is a graph-level property there. + + Returns the workdir on success; raises TnpuError with tnpu's own stage + report (which names the failing command and its stderr) otherwise. + """ + cmd = [extension_config.CONFIG_TNPU_PYTHON, + os.path.join(tnpu_dir(), "run.py"), spec_path, + "--from", from_stage, "--to", to_stage, "--workdir", workdir] + if verbose: + cmd.append("-v") + + env = dict(os.environ) + # tnpu deliberately does not read TORCHSIM_LLVM_PATH (it would drag the + # backend back to LLVM 20 and break the textual seam), but a stale + # PYTHONPATH pointing at LLVM 20's mlir_core would still be picked up by the + # namespace package before tnpu's own activate_bindings() runs. + env.pop("PYTHONPATH", None) + + proc = subprocess.run(cmd, capture_output=True, text=True, + cwd=tnpu_dir(), env=env, timeout=timeout) + output = proc.stdout + proc.stderr + if proc.returncode != 0: + # run.py prints a stage table; the diagnostic itself only reaches + # stage.log. + log = os.path.join(workdir, "stage.log") + if os.path.isfile(log): + with open(log, errors="replace") as fh: + output += "\n" + fh.read() + raise TnpuError(f"tnpu pipeline failed (exit {proc.returncode})", + cmd=" ".join(cmd), output=output) + logger.debug("[triton-npu] %s", output) + return workdir diff --git a/PyTorchSimFrontend/triton_backend/wrapper_codegen.py b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py new file mode 100644 index 00000000..cfc6490d --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/wrapper_codegen.py @@ -0,0 +1,21 @@ +"""Wrapper codegen for the Triton route. + +Reuses `ExtensionWrapperCodegen` wholesale -- device guards, buffer allocation, +the TOGSimulator plumbing and the SRAM plan hooks are all route-independent -- +and adds the one import the generated module needs: `triton_npu_compile`, which +is to this route what `custom_async_compile` is to the MLIR one. +""" + +from PyTorchSimFrontend.mlir.mlir_codegen_backend import ExtensionWrapperCodegen + +from . import codecache + + +class TritonNPUWrapperCodegen(ExtensionWrapperCodegen): + def write_header(self): + super().write_header() + self.header.splice( + f""" + from {codecache.__name__} import triton_npu_compile + """ + ) diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 0ef98eff..b98d3305 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -40,7 +40,18 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } } if (cyc.empty()) { cyc.assign(256, 128); ovl.assign(256, 0); } - return trace_to_tilegraph(trace_so_path.c_str(), nullptr, 0, + // Shape args: the producer's grid bounds, one per axis, when the trace was + // compiled without them baked in. Same sidecar convention as the cycle table + // -- absent means the producer carries its own constants. + std::vector shape; + { + std::ifstream sh(fs::path(trace_so_path).parent_path() / "trace_shape.txt"); + int64_t v; + while (sh >> v) shape.push_back(v); + } + return trace_to_tilegraph(trace_so_path.c_str(), + shape.empty() ? nullptr : shape.data(), + (int32_t)shape.size(), bases.data(), (int)bases.size(), cyc.data(), ovl.data(), (int)cyc.size(), partition_cores.data(), (int32_t)partition_cores.size(), diff --git a/docs/triton-route-coverage.ko.md b/docs/triton-route-coverage.ko.md new file mode 100644 index 00000000..86cb2ed8 --- /dev/null +++ b/docs/triton-route-coverage.ko.md @@ -0,0 +1,461 @@ +# Triton codegen route 커버리지 측정 보고서 + +기존 PyTorchSim 테스트 스위트를 MLIR 경로가 아니라 **Triton 경로**(Inductor의 +Triton 백엔드 + triton-npu lowering pass)로 돌린 첫 측정 결과입니다. + +| | | +|---|---| +| 측정일 | 2026-08-03 | +| 브랜치 | `feature/triton-codegen` @ `6e3bd7e` | +| tnpu 핀 | `5d84caf` | +| 환경 | torch 2.10.0, triton 3.6.0 | +| 대상 | 69개 (`tests/` 전체) | +| 소요 | `-j 10` 기준 5분 (직렬 약 50분) | + +재현: + +```bash +python scripts/ci/triton_route_sweep.py --all -j 10 \ + --markdown coverage.md --artifacts failures +``` + +아래 모든 주장은 `failures/` 아래 실제 파일로 뒷받침됩니다. 확인할 수 있도록 +경로를 함께 적었습니다. + +--- + +## 1. 결론부터 + +``` +69개 테스트 +├── 11 경로를 타고 통과 ← 이것이 커버리지 수치 +├── 5 통과하지만 경로 미사용 ← 커널을 아예 안 만듦 +└── 53 실패 + ├── 17 로컬 venv 패키지 없음 (CI 이미지에는 있음) + └── 36 실제 블로커 +``` + +**16/69가 아니라 11/69입니다.** 다섯 개(`test_matmul`, `test_bmm`, `test_topk`, +`test_moe_cpu`, `test_mlir_bindings`)는 통과하지만 Triton 커널을 하나도 만들지 +않습니다. Inductor가 `mm`/`bmm`을 커널 생성 대신 extern call로 내리기 때문에, +이 테스트들은 정작 검증 대상을 한 번도 거치지 않습니다. 스윕은 이를 별도로 +기록(JSON의 `exercised`)하고 gate에서 제외합니다. 포함시키면 커버리지가 45% +부풀려집니다. + +### 통과한 11개 + +| 테스트 | 시간 | +|---|---| +| `tests/ops/elementwise/test_add.py` | 77.5s | +| `tests/ops/fusion/test_addmm_residual.py` | 33.0s | +| `tests/ops/fusion/test_matmul_scalar.py` | 11.5s | +| `tests/ops/fusion/test_matmul_vector.py` | 17.4s | +| `tests/ops/fusion/test_prologue_fusion.py` | 41.4s | +| `tests/ops/misc/test_expert_mask.py` | 11.0s | +| `tests/ops/reduce/test_batchnorm.py` | 37.7s | +| `tests/ops/view/test_view3D_2D.py` | 36.4s | +| `tests/system/test_eager.py` | 15.0s | +| `tests/system/test_stonne.py` | 9.7s | +| `tests/system/test_triton_codegen.py` | 10.0s | + +11개 중 **4개가 fusion 테스트**입니다. 이 마이그레이션에서 Inductor의 fusion은 +공짜로 얻는 절반인데, 이미 tnpu가 받아들이는 커널을 만들어내고 있다는 뜻입니다. + +--- + +## 2. 실패 하나가 어떻게 진단되는가 — softmax 전 과정 + +`tests/ops/reduce/test_softmax.py`. 스윕이 남기는 것: + +``` +failures/tests_ops_reduce_test_softmax/ + kernel.py Inductor 가 만든 Triton 커널 원본 + error.txt 버킷, 단계, 로그 마지막 60줄 +``` + +`error.txt` 첫 줄이 담당자를 지정합니다: + +``` +test: tests/ops/reduce/test_softmax.py +bucket: triton_helpers +stage: 0 kernel generated, not accepted +``` + +그리고 `kernel.py`가 이유 전부를 보여줍니다: + +```python +@triton.jit +def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr): + xnumel = 64 + r0_numel = 128 + R0_BLOCK: tl.constexpr = 128 + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < xnumel + r0_index = tl.arange(0, R0_BLOCK)[None, :] + r0_1 = r0_index + x0 = xindex + tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) + tmp1 = tl.broadcast_to(tmp0, [XBLOCK, R0_BLOCK]) + tmp3 = tl.where(xmask, tmp1, float("-inf")) + tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <-- 블로커 1 + tmp5 = tmp0 - tmp4 + tmp6 = libdevice.exp(tmp5) # <-- 블로커 2 + tmp7 = tl.broadcast_to(tmp6, [XBLOCK, R0_BLOCK]) + tmp9 = tl.where(xmask, tmp7, 0) + tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) + tmp11 = (tmp6 / tmp10) + tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) +``` + +막힌 곳은 두 군데뿐이고, 아무것도 실행하지 않고 눈으로 확인됩니다. + +- `triton_helpers.max2` — `torch._inductor.runtime`에 있는 모듈인데, tnpu venv는 + 의도적으로 torch가 없습니다. +- `libdevice.exp` — `@core.extern` 인트린식으로 triton_shared 백엔드에 구현이 + 없습니다. + +**주목할 점은 나머지가 전부 멀쩡하다는 것입니다.** 마스크가 붙은 `tl.load`, +`tl.where`, `tl.sum`, `tl.store`, 2차원 `[XBLOCK, R0_BLOCK]` 브로드캐스트가 모두 +문제없이 통과합니다. softmax는 구조적으로 막힌 게 아니라 **함수 호출 두 개**에 +막혔습니다. + +이것이 버그 리포트 전체이고, 재실행 없이 만들어집니다. + +--- + +## 3. 커널이 어디서 멈추는가 + +각 테스트를, 그 테스트의 커널이 산출물을 남긴 가장 깊은 단계에 배치했습니다. +커널이 **도달하지 못한** 단계가 그 실패의 담당자입니다. + +| 단계 | 수 | | +|---|---|---| +| — 커널 생성 전 | 26 | codegen 이전에 torch/Inductor에서 사망 | +| 0 생성 후 거절 | 16 | `kernel_spec`이 기술을 거부 | +| 1 triton → ttir | 1 | | +| 2 ttir → tts/linalg | 1 | triton-shared | +| 4 tnpu lower (DMA, lane, spad) | 6 | | +| 5 trace producer | 3 | | + +**lowering pass는 아직 병목이 아닙니다.** 53개 실패 중 tnpu pass가 IR을 거절한 +것은 **단 2건**입니다. 나머지 34개 실제 블로커는 그보다 앞 — tnpu에 넘겨주는 +우리 쪽 포트, 또는 torch 자체 — 에서 멈춥니다. 다음 작업은 대부분 seam의 우리 +쪽에 있습니다. + +--- + +## 4. 원인별 상세 (근거 포함) + +### `spec_incomplete` — 13개 · 담당: `triton_backend/kernel_spec.py` + +**libdevice 인트린식 (5개).** `@core.extern` 멤버로 triton_shared 구현이 없어 +호출하면 `None`이 반환됩니다. + +| 테스트 | 심볼 | +|---|---| +| `ops/elementwise/test_exponent.py` | `libdevice.exp` | +| `ops/elementwise/test_pointwise.py` | `libdevice.isnan` | +| `ops/elementwise/test_transcendental.py` | `libdevice.tanh` | +| `ops/reduce/test_layernorm.py` | `libdevice.rsqrt` | +| `ops/view/test_floormod_axis_split.py` | `libdevice.rsqrt` | + +이번 세션에서 진단을 고치기 전에는 이것들이 tnpu stage-1 워커 안에서 +`NameError('libdevice is not defined')` 로만 죽었습니다. 그래서 서로 다른 여섯 +개의 lowering 버그처럼 보였습니다. 지금은 이렇게 말합니다: + +``` +SpecIncomplete: kernel calls libdevice.{exp}: those are extern math intrinsics +with no implementation on the triton_shared backend. They need lowering to a +VPU op (or a scalar fallback) before this kernel can compile. +``` + +**다축 grid (4개).** `fixed_config_for`가 가장 바깥 축만 고정하기 때문에 +`YBLOCK`이 `None`이 되고 grid를 계산할 수 없습니다. 알려진 block-size 정책 +공백입니다. + +| 테스트 | 진단 | +|---|---| +| `ops/view/test_transpose2D.py` | axis `y`: ynumel=156, YBLOCK=None | +| `ops/view/test_transpose3D.py` | axis `y`: ynumel=2728, YBLOCK=None | +| `ops/fusion/test_conv_fusion.py` | axis `y`: ynumel=192, YBLOCK=None | +| `ops/conv/test_conv_view_input.py` | axis `y`: ynumel=512, YBLOCK=None | + +**reduction block 미설정 (2개)** — `R0_BLOCK`을 의도적으로 비워둡니다: +`ops/fusion/test_bmm_reduction.py`, `ops/fusion/test_matmul_reduction.py`. + +**진짜 메타데이터 구멍 (1개)** — `ops/misc/test_widen_dtype.py`: `out_ptr0`의 +dtype/numel을 `collect_meta`가 `V.graph`에서 해결하지 못했습니다. + +### `triton_helpers` — 7개 · 담당: `triton_backend` + +| 테스트 | 헬퍼 | +|---|---| +| `ops/reduce/test_softmax.py` | `max2` | +| `ops/sort/test_sort.py` | `sort_with_index` | +| `ops/elementwise/test_activation.py` | `maximum` | +| `ops/conv/test_cnn.py` | `maximum` | +| `ops/fusion/test_matmul_activation.py` | `maximum` | +| `ops/sparsity/test_sparsity.py` | `maximum` | +| `models/test_mlp.py` | `maximum` | + +7개 중 4개가 `maximum` 하나만 필요합니다. pass 수정이 아니라 작은 파일 하나를 +vendoring 하는 작업입니다. + +### `wrapper_gap` — 6개 · 담당: `triton_backend` + +전부 동일합니다: + +``` +AttributeError: 'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak' +``` + +`ops/attention/test_gqa.py`, `test_gqa_decode.py`, +`ops/fusion/test_attention_fusion.py`, `test_transformer_fusion.py`, +`models/Mixtral8x7B/test_attention.py`, `models/test_transformer.py` + +스위트의 **attention·transformer 테스트 전부**가 미구현 메서드 하나에 막혀 +있습니다. + +### `device_op` — 3개 · 담당: `PyTorchSimDevice` + +이 경로 이전부터 있던 문제입니다. MLIR 경로는 이들을 dispatcher 도달 전에 +가로챕니다. + +| 테스트 | 오류 | +|---|---| +| `ops/conv/test_conv2d.py` | `convolution_overrideable not implemented` | +| `ops/conv/test_group_conv.py` | `convolution_overrideable not implemented` | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable not implemented` | + +### `tnpu_stage` — 2개 · 담당: triton-npu lowering pass + +진짜로 pass가 IR을 거절한 유일한 두 건입니다. 산출물에 Python부터 거절된 op까지 +사슬 전체가 남아 있습니다. + +**`ops/conv/test_pool.py`** — stage 1. + +`kernel.py` 끝에, Inductor가 reduction 뒤에 붙이는 무해해 보이는 한 줄: + +```python + tmp4 = tl.sum(tmp3, 1)[:, None].to(tl.float32) + tmp5 = 49.0 + tmp6 = (tmp4 / tmp5) + tl.debug_barrier() # <-- 이것 +``` + +`01-ttir.mlir:46`에서 이렇게 됩니다: + +```mlir +%tmp4_17 = tt.expand_dims %tmp4 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> +%tmp6_18 = arith.divf %tmp4_17, %tmp6 : tensor<128x1xf32> +ttg.barrier all # <-- GPU 다이얼렉트 op +%0 = tt.splat %in_out_ptr0 : !tt.ptr -> tensor<128x1x!tt.ptr> +``` + +그리고 `triton-shared-opt`가 파싱하지 못합니다: + +``` +01-ttir.mlir:46:5: error: Dialect `ttg' not found for custom op 'ttg.barrier' +``` + +`ttg`는 GPU 다이얼렉트입니다. reduction 뒤 pointwise 커널에서 이 타깃에는 의미 +없는 배리어인데, IR에 남아 있어서 파서가 멈춥니다. + +**`ops/reduce/test_reduce.py`** — stage 2. 평범한 `(a + b).sum(dim=1)`: + +```python +tmp0 = tl.load(in_ptr0 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp1 = tl.load(in_ptr1 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp2 = tmp0 + tmp1 +tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK]) +tmp5 = tl.where(r0_mask & xmask, tmp3, 0) +``` + +`01-ttir.mlir`까지는 살아남고, linalg 변환에서 실패합니다: + +``` +error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than + the number of loops (2) of the enclosing LinalgOp +``` + +두 산출물 모두 진단과 문제의 `.mlir`을 함께 갖고 있어, 그대로 업스트림에 넘길 +수 있습니다. + +### `togsim` / `기타` — 5개 + +| 테스트 | 단계 | 내용 | +|---|---|---| +| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed`, exit 255 | +| `ops/misc/test_masked_nondividing.py` | 4 | `[Spike] triton_npu_fused_constant_pad_nd_0 failed` | +| `ops/misc/test_indirect_access.py` | 5 | TOGSim이 `index_put`에 `inf` 사이클 반환 | +| `system/test_hetro.py` | — | `KeyError: 'vpu_num_lanes'` (hetero config에 키 없음) | +| `ops/sparsity/test_sparse_core.py` | — | `TypeError: '>' between Tensor and torch.device` (테스트 쪽 버그) | + +Spike 실패 두 건이 이번 스윕에서 가장 흥미롭습니다. **동작하는 RISC-V 바이너리 +까지 컴파일된 뒤 런타임에 실패하는 유일한 사례**입니다. `test_cat`의 +`04-custom.mlir`을 보면 lowering은 제 역할을 했습니다: + +```mlir +"togsim.transfer"(%reinterpret_cast_5, %c0, %2, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast_4, %c0, %1, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast, %9, %0, %c0, %7, %c0, %c3, %c1, %c0, %13) + {vlane_split_axis = 0, ...} +``` + +입력 DMA 2개와 출력 DMA 1개, 축 0으로 lane split — `cat`이 내야 할 정확한 +모양입니다. **lowering 구조가 아니라 실행이 잘못됩니다.** + +**이 두 건에 대한 단서.** `tnpu.spike`는 `StageError: command failed with exit +code 255`만 보고하고 spike 자신의 stderr는 살아남지 못합니다. 기록된 spike +명령을 같은 workdir에서 손으로 돌리면 exit 0이 나오는데, `write_inputs`가 매 +launch마다 `runtime/*.raw`를 새로 쓰기 때문에 손으로 돌린 실행은 **옛 입력을 +재생**하기 때문입니다. 따라서 실패하는 입력은 현재 파이프라인 밖에서 재현할 수 +없습니다. `TnpuError`에 적용한 것과 같은 방식으로 spike의 stderr를 노출시키는 +것이 이 두 건 진단의 선결 조건이고, 아직 하지 않았습니다. + +### `missing_dep` — 17개 · 경로 문제 아님 + +`transformers`(5), `torchvision`(4), `matplotlib`(4), `pytest`(2), `diffusers`, +`requests`, `sklearn`. 로컬 venv에만 없는 것으로, CI 이미지에서는 실제로 +돌아갑니다. 스윕이 CI에 있어야 하는 이유이기도 합니다. + +--- + +## 5. 인프라 + +### 러너 + +`scripts/ci/triton_route_sweep.py`. `TORCHSIM_TRITON_CODEGEN`은 device 등록 +시점에 한 번 읽히므로(`PyTorchSimDevice/torch_openreg/__init__.py:30`), +**테스트 파일은 하나도 고칠 필요가 없었습니다.** 69개 전부가 이미 이 경로의 +테스트였고, 없던 것은 러너뿐이었습니다. + +세 가지 산출물: + +1. **Gate** — `scripts/ci/triton_route_passing.txt`에 현재 통과 목록. 하나라도 + 깨지면 CI 실패. 커버리지는 이 파일을 재생성해서만 늘어나므로 + (`--update-allowlist`) 조용히 줄어들 수 없습니다. +2. **Report** — 담당 레이어와 파이프라인 단계로 분류. +3. **Artifacts** — 실패 테스트당 디렉토리 하나 (2절 참고). + +더 깊이 간 실패는 더 많이 남깁니다. `tests_ops_conv_test_cnn/`에는 +`01-ttir.mlir 02-ttshared.mlir 03-adapted.mlir 04-custom.mlir kernel.py +stage.log error.txt`가 있습니다 — 멈춘 지점까지의 lowering 사슬 전체입니다. + +### 병렬화 + +테스트는 각자 독립 서브프로세스이고 자기 덤프 디렉토리, Inductor 캐시 +(`TORCHINDUCTOR_CACHE_DIR`가 `TORCHSIM_DUMP_PATH`를 따라감), TOGSim FIFO(pid +기준)를 갖습니다. 따라서 `-j`에 조율이 필요 없습니다. 프로세스가 아니라 +스레드입니다 — `run_one`은 서브프로세스를 기다리기만 합니다. +**69개 기준 약 50분 → `-j 10`에서 5분.** + +### CI + +`.github/workflows/triton_npu.yml`의 `triton-route-suite` 잡: + +- **Allowlisted tests** — gate 역할. +- **Full sweep** — `continue-on-error`. `coverage.md`를 step summary에 쓰고 + `triton-route-coverage`(results.json + failures/)를 업로드. + +잡은 PSAL Slurm 러너 팜(`PSAL-POSTECH/slurm-ghr`)에서 돕니다. `runs-on`에 +`slurm` 라벨이 있어야 하고, 이미지 빌드와 스윕은 `big`(16c/64G/2h), 나머지는 +small 버킷입니다. `docker/setup-buildx-action`은 추가하면 안 됩니다 — 러너가 +자체 빌더를 등록해 둡니다. + +--- + +## 6. 측정하면서 고친 진단 3가지 + +이 셋을 고치기 전에는 보고 인프라를 만들 수 없었습니다. 각각이 증거를 파괴하고 +있었기 때문입니다. + +**`kernel.py`가 그것을 거절하는 검사 뒤에 저장되고 있었습니다.** +`write_spec_file`은 정확히 보존할 가치가 있는 커널(`triton_helpers`, +`SpecIncomplete`)에서 예외를 던지는데, 소스 저장보다 **먼저** 실행됐습니다. +결국 흥미로운 소스일수록 버려지고 있었습니다. 순서를 뒤집었고, 이제 거절된 16개 +커널 전부의 덤프가 남습니다 — 2절의 softmax 예시가 그중 하나입니다. + +**tnpu가 "exit 1"만 보고했습니다.** `run.py`는 stage 표를 stdout에 찍고 진짜 +진단은 `stage.log`에만 씁니다. 이전: + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) +``` + +이후 (`TnpuError`가 `stage.log`를 읽음): + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) + triton.compiler.errors.CompilationError: at 8:11: + NameError('tl_math is not defined') +``` + +이 한 가지 변경으로 여섯 개 실패가 **하나의 버그**로 정리됐습니다: + +**`libdevice`와 `tl_math`가 유탄을 맞고 있었습니다.** `strip_for_tnpu`가 +`from torch...`를 지우는데, Inductor는 이 두 이름을 +`torch._inductor.runtime.triton_helpers`에서 import합니다. 그런데 이들은 torch +코드가 아니라 **triton 자체 심볼의 재수출**입니다. 커널 여섯 개가 stage 1 안에서 +맨 `NameError`로 죽고 있었습니다. + +- `tl_math`는 `triton.language`에서 다시 바인딩했습니다. `test_pointwise`가 첫 + op에서 죽던 것이 op 14개를 지나 trace producer까지 갑니다. +- `libdevice`는 재바인딩이 불가능합니다(멤버가 `@core.extern`이고 triton_shared + 구현이 없어 호출하면 `None`). `triton_helpers`와 같은 방식으로 명시적으로 + 이름을 밝히도록 했습니다. + +순효과: `tnpu_stage` 8 → 2, `spec_incomplete` 7 → 13. **같은 53개가 실패하지만 +그중 6개가 이제 참을 말합니다.** + +**별건으로**, 로컬 TOGSim 빌드가 07-20자여서 `trace_shape.txt` 지원 이전이었고, +그래서 `togsim_kernel`이 `shape_args = nullptr`로 호출되어 모든 Triton 경로 +테스트가 `trace_to_tilegraph`에서 SIGSEGV로 죽었습니다. 재빌드로 해결됐습니다 — +코드 문제가 아니고, CI는 소스에서 빌드하므로 영향이 없었습니다. 다른 사람이 낡은 +`TOGSim/build`를 갖고 있다면 알아둘 만합니다. + +--- + +## 7. 다음 작업 — 해제되는 테스트 수 기준 + +수치는 측정값이지 추정이 아닙니다. 다만 한 단계에서 풀린 테스트가 다음 단계에서 +그냥 다시 실패할 수는 있습니다. + +| # | 작업 | 해제 | 담당 | +|---|---|---|---| +| 1 | `TritonNPUWrapperCodegen.estimate_peak` 구현 | 6 | triton_backend | +| 2 | torch 없는 `triton_helpers`를 tnpu venv에 vendoring | 7 | triton_backend | +| 3 | `libdevice` 인트린식(`exp`, `tanh`, `rsqrt`, `isnan`)을 VPU op로 lowering | 5 | tnpu 또는 triton_backend | +| 4 | `fixed_config_for`에 다축 block 정책 | 4 | triton_backend | +| 5 | `ttg.barrier` + `linalg.index` rank 오류를 업스트림에 전달 | 2 | tnpu | +| 6 | spike stderr 노출 후 `cat` / `constant_pad_nd` 진단 | 2 | triton_backend → 조사 | + +**1번이 압도적으로 쌉니다** — 메서드 하나로 6개, attention/transformer 계열 +전체가 열립니다. + +**2번과 3번은 함께 해야 softmax가 열립니다.** 2절에서 봤듯 softmax는 둘 다 +필요하고, 하나만 고치면 다른 하나에서 계속 실패합니다. + +**3번은 착수 전 결정이 필요합니다**: tnpu pass에서 lowering할 것인가, +`strip_for_tnpu`에서 triton 레벨 polyfill로 대체할 것인가. 전자가 옳고 후자가 +싸며 측정을 더 빨리 풀어줍니다. + +**6번은 선결 조건이 있습니다.** 스위트에서 유일하게 "답이 틀리는" 실패이고 진짜 +lowering 버그일 가능성이 가장 높지만, spike의 stderr가 서브프로세스를 넘어오기 +전에는 진단할 수 없습니다 — 6절에서 `TnpuError`에 이미 적용한 것과 같은 +수정입니다. + +--- + +## 8. 이 측정이 말해주지 않는 것 + +- `missing_dep` 17개는 로컬 venv 사정입니다. CI 이미지에서는 실제로 돌기 때문에 + 버킷이 이동할 것입니다 — 대부분 transformer·CNN 모델이므로 아마 `wrapper_gap` + 과 `triton_helpers` 쪽으로 갑니다. +- 어떤 버킷을 풀면 그 테스트들은 **다음 실패**로 이동하는 것이지, 반드시 통과로 + 가는 것이 아닙니다. +- 이 수치는 6절의 수정이 이미 적용된 상태에서 측정한 것이라, 그 이전 실행과 + 직접 비교할 수 없습니다. diff --git a/docs/triton-route-coverage.md b/docs/triton-route-coverage.md new file mode 100644 index 00000000..51ce81d3 --- /dev/null +++ b/docs/triton-route-coverage.md @@ -0,0 +1,463 @@ +# Triton codegen route — test suite coverage + +First measurement of PyTorchSim's existing test suite running through the Triton +codegen route (Inductor's Triton backend + the triton-npu lowering passes) +instead of the MLIR route. Korean version: +[`triton-route-coverage.ko.md`](triton-route-coverage.ko.md). + +| | | +|---|---| +| Date | 2026-08-03 | +| Branch | `feature/triton-codegen` @ `8e17519` | +| tnpu pin | `5d84caf` | +| torch | 2.10.0, triton 3.6.0 | +| Tests | 69 (everything under `tests/`) | +| Runtime | 5 min at `-j 10` (~50 min serial) | + +Reproduce: + +```bash +python scripts/ci/triton_route_sweep.py --all -j 10 \ + --markdown coverage.md --artifacts failures +``` + +Every claim below is backed by a file in `failures/`. Paths are given so each +one can be checked. + +--- + +## 1. Headline + +``` +69 tests +├── 11 pass THROUGH the route ← this is the coverage number +├── 5 pass without using the route ← no kernel emitted at all +└── 53 fail + ├── 17 missing test deps (local venv only; present in the CI image) + └── 36 real blockers +``` + +**11/69, not 16/69.** Five tests pass while emitting no Triton kernel at all: +`test_matmul`, `test_bmm`, `test_topk`, `test_moe_cpu`, `test_mlir_bindings`. +Inductor sends `mm`/`bmm` to an extern call rather than generating a kernel, so +those tests never exercise the thing under test. The sweep records that +separately (`exercised` in the JSON) and keeps them out of the gate — counting +them would overstate coverage by 45%. + +### What passes + +| Test | Time | +|---|---| +| `tests/ops/elementwise/test_add.py` | 77.5s | +| `tests/ops/fusion/test_addmm_residual.py` | 33.0s | +| `tests/ops/fusion/test_matmul_scalar.py` | 11.5s | +| `tests/ops/fusion/test_matmul_vector.py` | 17.4s | +| `tests/ops/fusion/test_prologue_fusion.py` | 41.4s | +| `tests/ops/misc/test_expert_mask.py` | 11.0s | +| `tests/ops/reduce/test_batchnorm.py` | 37.7s | +| `tests/ops/view/test_view3D_2D.py` | 36.4s | +| `tests/system/test_eager.py` | 15.0s | +| `tests/system/test_stonne.py` | 9.7s | +| `tests/system/test_triton_codegen.py` | 10.0s | + +Four of the eleven are fusion tests. Inductor's fusion is the half of this +migration we get for free, and it is already producing kernels tnpu accepts. + +--- + +## 2. Worked example — how one failure is diagnosed + +`tests/ops/reduce/test_softmax.py`. The sweep leaves this behind: + +``` +failures/tests_ops_reduce_test_softmax/ + kernel.py the Inductor Triton kernel, unmodified + error.txt bucket, stage, last 60 lines +``` + +`error.txt` opens with the routing header: + +``` +test: tests/ops/reduce/test_softmax.py +bucket: triton_helpers +stage: 0 kernel generated, not accepted +``` + +And `kernel.py` is the whole reason, in 28 lines: + +```python +@triton.jit +def triton_npu_fused__softmax_0(in_ptr0, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr): + xnumel = 64 + r0_numel = 128 + R0_BLOCK: tl.constexpr = 128 + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < xnumel + r0_index = tl.arange(0, R0_BLOCK)[None, :] + r0_1 = r0_index + x0 = xindex + tmp0 = tl.load(in_ptr0 + (r0_1 + 128*x0), xmask, other=0.0) + tmp1 = tl.broadcast_to(tmp0, [XBLOCK, R0_BLOCK]) + tmp3 = tl.where(xmask, tmp1, float("-inf")) + tmp4 = triton_helpers.max2(tmp3, 1)[:, None].to(tl.float32) # <-- blocker 1 + tmp5 = tmp0 - tmp4 + tmp6 = libdevice.exp(tmp5) # <-- blocker 2 + tmp7 = tl.broadcast_to(tmp6, [XBLOCK, R0_BLOCK]) + tmp9 = tl.where(xmask, tmp7, 0) + tmp10 = tl.sum(tmp9, 1)[:, None].to(tl.float32) + tmp11 = (tmp6 / tmp10) + tl.store(out_ptr2 + (r0_1 + 128*x0), tmp11, xmask) +``` + +Two blockers, visible without running anything: + +- `triton_helpers.max2` — lives in `torch._inductor.runtime`, and the tnpu venv + deliberately has no torch. +- `libdevice.exp` — an `@core.extern` intrinsic with no triton_shared + implementation. + +Note what *is* fine: `tl.load` with a mask, `tl.where`, `tl.sum`, `tl.store`, +the 2-D `[XBLOCK, R0_BLOCK]` broadcast. Softmax is not blocked on anything +structural. It is blocked on two function calls. + +This is the whole bug report, and it needed no rerun to produce. + +--- + +## 3. Where kernels stop + +Each test is placed at the furthest stage any of its kernels produced an +artifact for. The stage a kernel *fails to reach* owns the failure. + +| Stage | Count | | +|---|---|---| +| — no kernel generated | 26 | died in torch/Inductor before codegen | +| 0 generated, rejected | 16 | `kernel_spec` refused to describe it | +| 1 triton → ttir | 1 | | +| 2 ttir → tts/linalg | 1 | triton-shared | +| 4 tnpu lower (DMA, lanes, spad) | 6 | | +| 5 trace producer | 3 | | + +**The lowering passes are not the bottleneck yet.** Only 2 of 53 failures are a +tnpu pass rejecting IR. The other 34 real blockers stop earlier — in the port +that feeds tnpu, or in torch itself. The next round of work is mostly on our +side of the seam. + +--- + +## 4. Failures by cause, with evidence + +### `spec_incomplete` — 13 · owner: `triton_backend/kernel_spec.py` + +**libdevice intrinsics (5).** `@core.extern` members with no triton_shared +implementation; a call returns `None`. + +| Test | Symbol | +|---|---| +| `ops/elementwise/test_exponent.py` | `libdevice.exp` | +| `ops/elementwise/test_pointwise.py` | `libdevice.isnan` | +| `ops/elementwise/test_transcendental.py` | `libdevice.tanh` | +| `ops/reduce/test_layernorm.py` | `libdevice.rsqrt` | +| `ops/view/test_floormod_axis_split.py` | `libdevice.rsqrt` | + +Without the diagnostic added this session, these failed as a bare +`NameError('libdevice is not defined')` inside tnpu's stage-1 worker — which is +what made them look like six separate lowering bugs. They now say: + +``` +SpecIncomplete: kernel calls libdevice.{exp}: those are extern math intrinsics +with no implementation on the triton_shared backend. They need lowering to a +VPU op (or a scalar fallback) before this kernel can compile. +``` + +**Multi-axis grid (4).** `fixed_config_for` pins only the outermost axis, so +`YBLOCK` is `None` and the grid cannot be computed. This is the known +block-size policy gap. + +| Test | Diagnostic | +|---|---| +| `ops/view/test_transpose2D.py` | axis `y`: ynumel=156, YBLOCK=None | +| `ops/view/test_transpose3D.py` | axis `y`: ynumel=2728, YBLOCK=None | +| `ops/fusion/test_conv_fusion.py` | axis `y`: ynumel=192, YBLOCK=None | +| `ops/conv/test_conv_view_input.py` | axis `y`: ynumel=512, YBLOCK=None | + +**Reduction blocks unset (2)** — `R0_BLOCK` is left unset on purpose: +`ops/fusion/test_bmm_reduction.py`, `ops/fusion/test_matmul_reduction.py`. + +**Genuine metadata hole (1)** — `ops/misc/test_widen_dtype.py`: no dtype/numel +for `out_ptr0`; `collect_meta` could not resolve it from `V.graph`. + +### `triton_helpers` — 7 · owner: `triton_backend` + +| Test | Helper | +|---|---| +| `ops/reduce/test_softmax.py` | `max2` | +| `ops/sort/test_sort.py` | `sort_with_index` | +| `ops/elementwise/test_activation.py` | `maximum` | +| `ops/conv/test_cnn.py` | `maximum` | +| `ops/fusion/test_matmul_activation.py` | `maximum` | +| `ops/sparsity/test_sparsity.py` | `maximum` | +| `models/test_mlp.py` | `maximum` | + +Four of seven want only `maximum`. The fix is a small vendored file, not a pass +change. + +### `wrapper_gap` — 6 · owner: `triton_backend` + +Every one, identically: + +``` +AttributeError: 'TritonNPUWrapperCodegen' object has no attribute 'estimate_peak' +``` + +`ops/attention/test_gqa.py`, `test_gqa_decode.py`, +`ops/fusion/test_attention_fusion.py`, `test_transformer_fusion.py`, +`models/Mixtral8x7B/test_attention.py`, `models/test_transformer.py` + +Every attention and transformer test in the suite, blocked on one unimplemented +method. + +### `device_op` — 3 · owner: `PyTorchSimDevice` + +Predates this route — the MLIR route intercepts these before the dispatcher. + +| Test | Error | +|---|---| +| `ops/conv/test_conv2d.py` | `convolution_overrideable not implemented` | +| `ops/conv/test_group_conv.py` | `convolution_overrideable not implemented` | +| `ops/attention/test_sdpa.py` | `_scaled_dot_product_fused_attention_overrideable not implemented` | + +### `tnpu_stage` — 2 · owner: triton-npu lowering passes + +The only failures that are genuinely a pass rejecting IR — and the artifacts +show the full chain from Python to the rejected op. + +**`ops/conv/test_pool.py`** — stage 1. + +`kernel.py` ends with an innocuous line Inductor emits after a reduction: + +```python + tmp4 = tl.sum(tmp3, 1)[:, None].to(tl.float32) + tmp5 = 49.0 + tmp6 = (tmp4 / tmp5) + tl.debug_barrier() # <-- this +``` + +`01-ttir.mlir:46` is what that becomes: + +```mlir +%tmp4_17 = tt.expand_dims %tmp4 {axis = 1 : i32} : tensor<128xf32> -> tensor<128x1xf32> +%tmp6_18 = arith.divf %tmp4_17, %tmp6 : tensor<128x1xf32> +ttg.barrier all # <-- GPU dialect op +%0 = tt.splat %in_out_ptr0 : !tt.ptr -> tensor<128x1x!tt.ptr> +``` + +and `triton-shared-opt` cannot parse it: + +``` +01-ttir.mlir:46:5: error: Dialect `ttg' not found for custom op 'ttg.barrier' +``` + +`ttg` is the GPU dialect. A `tl.debug_barrier()` in a pointwise-after-reduction +kernel is meaningless on this target, but it is in the IR and the parser stops +on it. + +**`ops/reduce/test_reduce.py`** — stage 2. A plain `(a + b).sum(dim=1)`: + +```python +tmp0 = tl.load(in_ptr0 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp1 = tl.load(in_ptr1 + (r0_1 + 47*x0), r0_mask & xmask, other=0.0) +tmp2 = tmp0 + tmp1 +tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK]) +tmp5 = tl.where(r0_mask & xmask, tmp3, 0) +``` + +survives `01-ttir.mlir`, then fails converting to linalg: + +``` +error: "-":101:11: 'linalg.index' op expected dim (2) to be lower than + the number of loops (2) of the enclosing LinalgOp +``` + +Both artifacts carry the diagnostic and the offending `.mlir`, so they can go +upstream as-is. + +### `togsim` / `other` — 5 + +| Test | Stage | Detail | +|---|---|---| +| `ops/view/test_cat.py` | 4 | `[Spike] triton_npu_fused_cat_0 failed`, exit 255 | +| `ops/misc/test_masked_nondividing.py` | 4 | `[Spike] triton_npu_fused_constant_pad_nd_0 failed` | +| `ops/misc/test_indirect_access.py` | 5 | TOGSim returned `inf` cycles for `index_put` | +| `system/test_hetro.py` | — | `KeyError: 'vpu_num_lanes'` (hetero config lacks the key) | +| `ops/sparsity/test_sparse_core.py` | — | `TypeError: '>' between Tensor and torch.device` (test-side bug) | + +The two Spike failures are the most interesting in the sweep: the only cases +that compile all the way to a working RISC-V binary and then fail at run time. +`test_cat`'s `04-custom.mlir` shows the lowering did its job — + +```mlir +"togsim.transfer"(%reinterpret_cast_5, %c0, %2, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast_4, %c0, %1, %c0, %7, %c0, %c2, %c1, %6) + {vlane_split_axis = 0, ...} +"togsim.transfer"(%reinterpret_cast, %9, %0, %c0, %7, %c0, %c3, %c1, %c0, %13) + {vlane_split_axis = 0, ...} +``` + +— two input DMAs and one output DMA, lane-split on axis 0, exactly the shape +`cat` should produce. It is the execution that goes wrong, not the lowering +structure. + +**Caveat on these two.** `tnpu.spike` reports only `StageError: command failed +with exit code 255`; spike's own stderr does not survive. Running the recorded +spike command by hand on the same workdir exits 0, because `write_inputs` +rewrites `runtime/*.raw` per launch and a by-hand run replays stale inputs. So +the failing input is not currently reproducible outside the pipeline. Surfacing +spike's stderr the way `TnpuError` now surfaces tnpu's is the prerequisite for +diagnosing these, and is not yet done. + +### `missing_dep` — 17 · not a route problem + +`transformers` (5), `torchvision` (4), `matplotlib` (4), `pytest` (2), +`diffusers`, `requests`, `sklearn`. Local venv only — these run for real in the +CI image, which is why the sweep belongs in CI. + +--- + +## 5. Infrastructure + +### The runner + +`scripts/ci/triton_route_sweep.py`. `TORCHSIM_TRITON_CODEGEN` is read once at +device registration (`PyTorchSimDevice/torch_openreg/__init__.py:30`), so **no +test file needed to change** — all 69 were already tests of this route. Only a +runner was missing. + +Three outputs: + +1. **Gate** — `scripts/ci/triton_route_passing.txt` lists what passes today; CI + fails if any regresses. Coverage grows by regenerating the file + (`--update-allowlist`), so it cannot silently shrink. +2. **Report** — bucketed by owning layer and by pipeline stage. +3. **Artifacts** — one directory per failing test (section 2). + +A deeper failure keeps more. `tests_ops_conv_test_cnn/` holds +`01-ttir.mlir 02-ttshared.mlir 03-adapted.mlir 04-custom.mlir kernel.py +stage.log error.txt` — the complete lowering chain up to the point it stopped. + +### Parallelism + +Tests are independent subprocesses with their own dump dir, Inductor cache +(`TORCHINDUCTOR_CACHE_DIR` follows `TORCHSIM_DUMP_PATH`) and TOGSim FIFO (keyed +by pid), so `-j` needs no coordination. Threads, not processes: `run_one` only +waits on a subprocess. **69 tests: ~50 min → 5 min at `-j 10`.** + +### CI + +`.github/workflows/triton_npu.yml`, job `triton-route-suite`: + +- **Allowlisted tests** — gates. +- **Full sweep** — `continue-on-error`; writes `coverage.md` into the step + summary and uploads `triton-route-coverage` (results.json + failures/). + +Jobs run on the PSAL Slurm runner farm (`PSAL-POSTECH/slurm-ghr`): `runs-on` +carries the `slurm` label, image builds and the sweep on `big` (16c/64G/2h), +the rest on small. Do not add `docker/setup-buildx-action` — the runner +registers its own builder. + +--- + +## 6. Three diagnostics fixed while measuring + +The reporting infrastructure could not be built until these were fixed, because +each was destroying the evidence. + +**`kernel.py` was written after the check that rejects it.** `write_spec_file` +raises for exactly the kernels worth keeping (`triton_helpers`, +`SpecIncomplete`) and ran *before* the source was saved — so the interesting +sources were the ones being thrown away. Reordered; the dump now exists for all +16 rejected kernels, including the softmax example in section 2. + +**tnpu reported "exit 1" and nothing else.** `run.py` prints a stage table to +stdout and the real diagnostic only to `stage.log`. Before: + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) +``` + +After (`TnpuError` now reads `stage.log`): + +``` +torch._inductor.exc.InductorError: TnpuError: tnpu pipeline failed (exit 1) + triton.compiler.errors.CompilationError: at 8:11: + NameError('tl_math is not defined') +``` + +That single change resolved six failures into one bug: + +**`libdevice` and `tl_math` were collateral damage.** `strip_for_tnpu` drops +`from torch...`, and Inductor imports both names from +`torch._inductor.runtime.triton_helpers` — but they are re-exports of *triton's +own* symbols, not torch code. Six kernels died as a bare `NameError` inside +stage 1. + +- `tl_math` is now rebound from `triton.language`. `test_pointwise` gets + through fourteen ops and as far as the trace producer instead of failing on + the first. +- `libdevice` cannot be rebound (its members are `@core.extern` with no + triton_shared implementation, so a call returns `None`) and is now named + explicitly, the way `triton_helpers` is. + +Net effect: `tnpu_stage` 8 → 2, `spec_incomplete` 7 → 13. The same 53 tests +fail; six of them now say something true. + +**Separately:** the local TOGSim build was from 07-20 and predated +`trace_shape.txt` support, so `togsim_kernel` was called with `shape_args = +nullptr` and every Triton-route test died with SIGSEGV in `trace_to_tilegraph`. +A rebuild fixed it — not a code problem, and CI builds from source so it was +never affected. Worth knowing if anyone else has a stale `TOGSim/build`. + +--- + +## 7. Next, ranked by tests unblocked per fix + +Counts are measured, not estimated — though a test unblocked at one stage may +simply fail at the next. + +| # | Fix | Unblocks | Owner | +|---|---|---|---| +| 1 | Implement `TritonNPUWrapperCodegen.estimate_peak` | 6 | triton_backend | +| 2 | Vendor a torch-free `triton_helpers` into the tnpu venv | 7 | triton_backend | +| 3 | Lower `libdevice` intrinsics (`exp`, `tanh`, `rsqrt`, `isnan`) to VPU ops | 5 | tnpu or triton_backend | +| 4 | Multi-axis block policy in `fixed_config_for` | 4 | triton_backend | +| 5 | Hand `ttg.barrier` + `linalg.index` rank error upstream | 2 | tnpu | +| 6 | Surface spike's stderr, then diagnose `cat` / `constant_pad_nd` | 2 | triton_backend, then investigate | + +**1 is the cheapest by a wide margin** — one method, six tests, and it opens the +entire attention/transformer family. + +**2 and 3 unblock softmax together.** Section 2 shows softmax needs both; either +alone leaves it failing on the other. + +**3 needs a decision before work starts:** lower in a tnpu pass, or substitute a +triton-level polyfill in `strip_for_tnpu`. The former is correct; the latter is +cheap and would unblock measurement sooner. + +**6 has a prerequisite.** These are the only wrong-answer failures in the suite +and the most likely to be a real lowering bug, but they cannot be diagnosed +until spike's stderr survives the subprocess — the same fix already applied to +`TnpuError` in section 6. + +--- + +## 8. Caveats + +- The 17 `missing_dep` failures are local-venv artifacts. In the CI image those + tests run for real and the buckets will shift — probably toward `wrapper_gap` + and `triton_helpers`, since most are transformer and CNN models. +- Unblocking a bucket moves its tests to the *next* failure, not necessarily to + passing. +- These numbers were taken with the section-6 fixes already applied, so they are + not comparable to a run from before them. diff --git a/scripts/ci/tnpu_base_pin.sh b/scripts/ci/tnpu_base_pin.sh new file mode 100755 index 00000000..59663836 --- /dev/null +++ b/scripts/ci/tnpu_base_pin.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +# Deterministic short pin for tagging torchsim_tnpu_base images. +# Mirrors thirdparty_base_pin.sh, over the tnpu manifest + its Dockerfile, so the +# ~1.8 GiB toolchain layer is rebuilt only when one of those two actually moves. +set -euo pipefail +ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +cd "$ROOT" +{ cat thirdparty/triton-npu.json; cat Dockerfile.tnpu; } | sha256sum | awk '{print substr($1,1,12)}' diff --git a/scripts/ci/triton_route_passing.txt b/scripts/ci/triton_route_passing.txt new file mode 100644 index 00000000..213a8d58 --- /dev/null +++ b/scripts/ci/triton_route_passing.txt @@ -0,0 +1,15 @@ +# 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. +tests/ops/elementwise/test_add.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/misc/test_expert_mask.py +tests/ops/reduce/test_batchnorm.py +tests/ops/view/test_view3D_2D.py +tests/system/test_eager.py +tests/system/test_stonne.py +tests/system/test_triton_codegen.py diff --git a/scripts/ci/triton_route_sweep.py b/scripts/ci/triton_route_sweep.py new file mode 100755 index 00000000..df8660e3 --- /dev/null +++ b/scripts/ci/triton_route_sweep.py @@ -0,0 +1,338 @@ +#!/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|linalg\.reduce|no reduction path"), + ("dynamic_shape", r"ShapeMismatch|dynamic shape|size_hint returned None"), + ("matmul_timing", r"vcix\.iv|sf\.vc\.|no compute node"), + ("togsim", r"TOGSim|trace\.so|SIGSEGV|Signals\.SIG|'vpu_num_lanes'"), + ("wrong_values", r"VALUES WRONG|allclose|Test Failed"), + ("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 "" + + +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) + env = dict(os.environ, TORCHSIM_TRITON_CODEGEN="1", TORCHSIM_DUMP_PATH=dump) + + 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) + 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. + "exercised": workdir is not None, + "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", + "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("--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 = 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" + extra = ("" if r["exercised"] else " (route not exercised)") 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/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py new file mode 100644 index 00000000..2af839bd --- /dev/null +++ b/tests/system/test_triton_codegen.py @@ -0,0 +1,177 @@ +"""Drive the Triton codegen route as far as it currently goes. + +This route is WIP (see PyTorchSimFrontend/triton_backend/README.md). The test is +written to report WHERE it stops rather than to assert success: the value right +now is a reproducible statement of the next gap, not a pass/fail gate. Register +it in .github/workflows/pytorchsim_test.yml only once the route runs end to end. + + TORCHSIM_TRITON_CODEGEN=1 python tests/system/test_triton_codegen.py +""" +import os +import sys +import traceback + +# Must be set before torch_openreg registers the Inductor backend for `npu`. +os.environ.setdefault("TORCHSIM_TRITON_CODEGEN", "1") + +import torch # noqa: E402 + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +N = 1024 + + +def build(): + def fn(x, y): + return x + y + + x = torch.randn(N) + y = torch.randn(N) + return fn, x, y + + +def check_multi_axis_grid(): + """A 2-D grid must nest one loop per axis and hand both indices to iv[]. + + Guards the multi-axis path, which the add kernel does not reach: Inductor + only uses y/z when x would overflow, so a 1-D grid exercises just the first + iteration of the nest. + """ + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + src = """ + module { + func.func @k(%arg0: memref<*xf32>, %arg1: i32, %arg2: i32) { + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : i32 + %a = arith.muli %arg1, %c8 : i32 + %b = arith.addi %a, %arg2 : i32 + %o = arith.index_cast %b : i32 to index + "togsim.dma"(%o, %c0) {arg_id = 0 : i32, base = "arg0", dims = [128], + dir = 0 : i32, elem_bits = 32 : i32, is_async = false, read_bufs = [], + strides = [1], tag_id = 0 : i32, write_bufs = [0]} : (index, index) -> () + return + } + } + """ + problems = [] + # Verify the IR the pass itself produces: a bound created after an outer loop + # would not dominate an inner loop's use of it, which only shows at rank >= 2 + # and which the emitc lowering happens to paper over. + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + l2e._materialize_grid_loop( + l2e._find_kernel(module), + l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3]), ctx) + try: + module.operation.verify() + except Exception as e: # noqa: BLE001 + problems.append(f"materialized IR does not verify: {e}") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + emitc = l2e.lower_to_emitc( + module, work_item=l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3])) + cpp = l2e.emitc_to_cpp(emitc, include_dir=l2e._default_include_dir()) + + entry = cpp.split("togsim_kernel(EmitCtx*")[-1] + if entry.count("for (") != 2: + problems.append(f"expected 2 nested loops, found {entry.count('for (')}") + if "togsim_dispatch" not in entry: + problems.append("no togsim_dispatch call") + if ", 2);" not in entry: + problems.append("dispatch does not pass 2 indices") + for p in problems: + print(f" multi-axis grid: {p}") + return not problems + + +def check_reduction_is_refused(): + """A reduction must fail LOUDLY, not compile into wrong numbers. + + tnpu has no lane-aware reduction: the scratchpad is lane-banked, so the + reduced axis has to live inside a lane, and triton-shared hands over a + linalg.reduce (plus a linalg.transpose) that no pass lowers that way. Until + one does, reaching the launcher would mean simulating a kernel whose compute + is not what the hardware would do. + + Passing this check means the attempt still stops. When the lane path lands, + this is the test to delete. + """ + x = torch.randn(128, 64) + try: + torch.compile(lambda t: t.sum(dim=1))(x.to("npu:0")) + except Exception as e: # noqa: BLE001 - any diagnosed stop is the point + first = (str(e).strip().splitlines() or [type(e).__name__])[0] + print(f" reduction stops at: {type(e).__name__}: {first[:74]}") + return True + print(" reduction COMPILED -- if the lane-aware path landed, drop this " + "check; otherwise the numbers it produces are wrong") + return False + + +def main(): + from PyTorchSimFrontend import extension_config + from PyTorchSimFrontend.triton_backend import tnpu_bridge + + print(f"multi-axis grid = " + f"{'ok' if check_multi_axis_grid() else 'FAILED'}") + print(f"reduction refused = " + f"{'ok' if check_reduction_is_refused() else 'FAILED'}") + print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") + print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") + ok, _out = tnpu_bridge.doctor() + print(f"tnpu doctor = {'ok' if ok else 'FAILED (see run.py doctor)'}") + print() + + fn, x, y = build() + expected = fn(x, y) + + opt = torch.compile(fn, backend="inductor") + try: + got = opt(x.to("npu:0"), y.to("npu:0")) + except Exception as e: # noqa: BLE001 - the point is to report the stop + print(f"STOPPED AT: {type(e).__name__}") + print() + traceback.print_exc() + print() + print("The stage reached is what this test measures; see the traceback " + "above and README.md's gap list.") + return 1 + + ok = torch.allclose(got.cpu(), expected, rtol=1e-4, atol=1e-4) + if not ok: + bad = (~torch.isclose(got.cpu(), expected, rtol=1e-4, atol=1e-4)) + print(f"VALUES WRONG: {int(bad.sum())}/{expected.numel()} elements") + print(f" got {got.cpu()[:4].tolist()}") + print(f" expected {expected[:4].tolist()}") + return 1 + print(f"values ok ({expected.numel()} elements through Spike)") + + import glob + + from PyTorchSimFrontend.triton_backend import timing + + dirs = glob.glob(os.path.join(extension_config.get_dump_path(), "triton_*")) + if not dirs: + print("no kernel directory was produced") + return 1 + workdir = max(dirs, key=os.path.getmtime) + for name in (timing.TRACE_SO, timing.CYCLE_TSV): + path = os.path.join(workdir, name) + if not os.path.isfile(path): + print(f"missing {name} in {workdir}") + return 1 + print(f" {name:18s} {os.path.getsize(path)} bytes") + print(f"\ntiming path OK ({workdir})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/thirdparty/triton-npu.json b/thirdparty/triton-npu.json new file mode 100644 index 00000000..4bac8dfb --- /dev/null +++ b/thirdparty/triton-npu.json @@ -0,0 +1,8 @@ +{ + "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. The repository is private: the workflow needs a PAT in secrets.TNPU_TOKEN, since the default Actions token is scoped to this repository.", + "triton_npu": { + "repository": "PSAL-POSTECH/triton-npu", + "ref": "5d84cafdd5292f77548a84a6c44a4d2d15c4dd4f", + "release_tag": "toolchain-llvm23" + } +}