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