Wave 2 polish: fix shader build, remove hardcoded paths, add CI (#1)

- compile.sh: target Vulkan 1.1 (SPIR-V 1.3) so the 12 subgroup-using
  quantized kernels (matmul_q*k*, matmul_gpuq*) actually compile; the default
  `glslangValidator -V` emits SPIR-V 1.0 and failed on them. Added set -euo
  pipefail and a non-zero exit on any failure.
- vulkan_engine.cpp: drop the hardcoded /home/raz/... shader path (leaked a
  developer path and broke on every other machine). Fall back to the
  TORCH_VULKAN_SHADER_DIR env var instead.
- __init__.py: export the resolved bundled-shader dir into
  TORCH_VULKAN_SHADER_DIR so the lazily-constructed VulkanEngine resolves the
  same shaders. Aligned the module docstring with the README (.to("vulkan")
  is the supported path; torch.randn(device=) / .vulkan() are partial).
- tests: removed hardcoded /home/raz/builds/pytorch-gfx1150 sys.path inserts;
  honour TORCH_VULKAN_PYTORCH_PATH instead. Fixed the misleading "relu isn't
  implemented" fallback test (relu IS wired) and split it into a real relu
  correctness test plus a genuine CPU-fallback test using an unimplemented op.
- Removed dead imports (numpy, time, sys) and unused locals flagged by ruff.
- Added pyproject.toml (ruff config) and .github/workflows/ci.yml that
  compiles all shaders and runs ruff + py_compile. The GPU build/test suite is
  not run in CI (no GPU on hosted runners).

Co-authored-by: Peter Clemente III <peterc3@live.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 3 people

.github/workflows/ci.yml

@@ -0,0 +1,68 @@
+name: CI
+
+on:
+  push:
+    branches: [master, main]
+  pull_request:
+
+jobs:
+  shaders:
+    name: Compile SPIR-V shaders
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install glslang
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y glslang-tools
+
+      # Compile every .comp to a scratch dir and fail if any shader does not
+      # compile. This guards the regression that the default `glslangValidator -V`
+      # (SPIR-V 1.0) could not build the subgroup-using quantized kernels, which
+      # need --target-env vulkan1.1 (SPIR-V 1.3).
+      - name: Compile all shaders
+        working-directory: csrc/shaders
+        run: |
+          set -euo pipefail
+          out="$(mktemp -d)"
+          fail=0
+          for comp in *.comp; do
+            echo "Compiling $comp"
+            if ! glslangValidator --target-env vulkan1.1 -V "$comp" -o "$out/${comp%.comp}.spv"; then
+              echo "::error file=csrc/shaders/$comp::shader failed to compile"
+              fail=1
+            fi
+          done
+          [ "$fail" -eq 0 ] || { echo "One or more shaders failed to compile"; exit 1; }
+          echo "Compiled $(ls "$out"/*.spv | wc -l) shaders."
+
+  python-lint:
+    name: Python lint + syntax
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install ruff
+        run: pip install ruff
+
+      # py_compile catches syntax errors without needing the compiled _C
+      # extension (which requires a Vulkan GPU + Kompute + custom PyTorch and
+      # cannot run on a stock CI runner).
+      - name: Syntax check
+        run: |
+          python -m py_compile setup.py persistent_pipeline.py \
+            torch_vulkan/__init__.py tests/*.py
+
+      - name: Ruff
+        run: ruff check .
+
+# NOTE: The C++ extension build (CMake + Torch + Kompute) and the runtime test
+# suite require a Vulkan-capable GPU and are not run here -- GitHub-hosted
+# runners have no GPU. Build/test verification is done locally on the target
+# hardware (AMD Radeon 890M / RADV). This workflow verifies what can be checked
+# without a GPU: shader compilation and Python static checks.

README.md

@@ -78,7 +78,16 @@ attention + KV cache, kv_store.
 
 ```bash
 # Requires a Vulkan 1.2+ driver (RADV recommended) and glslangValidator for shader build.
-cd build && cmake .. && make -j$(nproc)
+
+# 1. (Re)compile the GLSL shaders to SPIR-V (pre-built .spv files are committed,
+#    so this is only needed if you edit a shader). Targets Vulkan 1.1 / SPIR-V 1.3
+#    because the quantized kernels use subgroup ops.
+bash csrc/shaders/compile.sh
+
+# 2. Build the C++ extension.
+cmake -S . -B build && cmake --build build -j
+
+# 3. Install the Python package (build_ext also invokes CMake).
 pip install -e .
 
 python -c "import torch, torch_vulkan; print(torch_vulkan.device_name())"

csrc/shaders/compile.sh

@@ -1,14 +1,30 @@
 #!/bin/bash
 # Compile all GLSL compute shaders to SPIR-V
 # Requires: glslangValidator (from Vulkan SDK or `pacman -S glslang`)
+#
+# We target Vulkan 1.1, which produces SPIR-V 1.3. Several of the quantized
+# matmul kernels (matmul_q*k*, matmul_gpuq*) use GL_KHR_shader_subgroup
+# reductions, and those subgroup ops require SPIR-V >= 1.3. The default
+# `glslangValidator -V` emits SPIR-V 1.0 and fails on them, so the explicit
+# --target-env is required to compile the full shader set.
+set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
 cd "$SCRIPT_DIR"
 
+fail=0
 for comp in *.comp; do
     spv="${comp%.comp}.spv"
     echo "Compiling $comp -> $spv"
-    glslangValidator -V "$comp" -o "$spv"
+    if ! glslangValidator --target-env vulkan1.1 -V "$comp" -o "$spv"; then
+        echo "  ERROR: failed to compile $comp" >&2
+        fail=1
+    fi
 done
 
+if [ "$fail" -ne 0 ]; then
+    echo "One or more shaders failed to compile." >&2
+    exit 1
+fi
+
 echo "Done. $(ls *.spv | wc -l) shaders compiled."

csrc/torch_vulkan.cpp

@@ -62,6 +62,12 @@ C10_REGISTER_GUARD_IMPL(PrivateUse1, VulkanGuardImpl);
 
 // Python bindings
 void set_shader_dir(const std::string& path) {
+  // Set the shader dir on the Kompute-based VulkanContext (used by all the
+  // registered ops). The raw VulkanEngine (used only by the not-yet-wired
+  // mm_raw path) is constructed lazily and reads its directory from the
+  // TORCH_VULKAN_SHADER_DIR env var, which __init__.py points at this same
+  // directory -- so we deliberately do NOT touch VulkanEngine::instance()
+  // here, to avoid eagerly creating a second Vulkan device at import time.
   VulkanContext::instance().set_shader_dir(path);
 }
 

csrc/vulkan_engine.cpp

@@ -2,6 +2,7 @@
 #include <fstream>
 #include <stdexcept>
 #include <cstring>
+#include <cstdlib>
 #include <algorithm>
 
 namespace torch_vulkan {
@@ -12,7 +13,16 @@ VulkanEngine& VulkanEngine::instance() {
 }
 
 VulkanEngine::VulkanEngine() {
-    shaderDir_ = "/home/raz/projects/torch-vulkan/csrc/shaders/";
+    // The shader directory is normally set at import time via
+    // setShaderDir() (driven by torch_vulkan._set_shader_dir in __init__.py).
+    // As a fallback for direct/standalone use, honour TORCH_VULKAN_SHADER_DIR
+    // so the path is never hardcoded to a developer's machine.
+    if (const char* env = std::getenv("TORCH_VULKAN_SHADER_DIR")) {
+        shaderDir_ = env;
+        if (!shaderDir_.empty() && shaderDir_.back() != '/') {
+            shaderDir_ += '/';
+        }
+    }
     initVulkan();
 }
 

persistent_pipeline.py

@@ -8,8 +8,6 @@
 """
 
 import torch
-import time
-import numpy as np
 
 
 class PersistentLayerPipeline:

pyproject.toml

@@ -0,0 +1,13 @@
+[tool.ruff]
+# torch-vulkan is a thin Python shim over a C++/Vulkan extension; keep linting
+# focused on real problems (pyflakes + import sorting) rather than style churn.
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+# E402 (module-level import not at top of file) is intentionally allowed: the
+# package must import `torch` first, then load the `_C` extension that registers
+# the PrivateUse1 backend, and finally alias the module into `torch.vulkan` --
+# all of which happen mid-module by design.
+ignore = ["E402"]

setup.py

@@ -8,9 +8,8 @@
 
 import os
 import subprocess
-import sys
 
-from setuptools import setup, Extension
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
 

tests/bench_layer.py

@@ -4,14 +4,19 @@
 Focuses on the hot ops: mm, add, gelu that repeat across layers.
 """
 
-import sys
 import os
+import sys
 import time
 
-sys.path.insert(0, "/home/raz/builds/pytorch-gfx1150")
+# Allow pointing at a custom-built PyTorch without hardcoding a developer's
+# path. Set TORCH_VULKAN_PYTORCH_PATH if needed.
+_custom_torch = os.environ.get("TORCH_VULKAN_PYTORCH_PATH")
+if _custom_torch:
+    sys.path.insert(0, _custom_torch)
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 
 import torch
+
 import torch_vulkan
 
 
@@ -23,12 +28,12 @@ def bench_mm_repeated(M, K, N, iters=20):
     b = b_cpu.to("vulkan")
 
     # Warmup (first call = cache miss)
-    c = torch.mm(a, b)
+    torch.mm(a, b)
 
     # Benchmark (subsequent calls should hit cache)
     t0 = time.perf_counter()
     for _ in range(iters):
-        c = torch.mm(a, b)
+        torch.mm(a, b)
     elapsed = (time.perf_counter() - t0) / iters
     print(f"  mm [{M}x{K}] @ [{K}x{N}]: {elapsed*1000:.2f} ms/call")
     return elapsed
@@ -41,11 +46,11 @@ def bench_add_repeated(N, iters=20):
     a = a_cpu.to("vulkan")
     b = b_cpu.to("vulkan")
 
-    c = torch.add(a, b)  # warmup
+    torch.add(a, b)  # warmup
 
     t0 = time.perf_counter()
     for _ in range(iters):
-        c = torch.add(a, b)
+        torch.add(a, b)
     elapsed = (time.perf_counter() - t0) / iters
     print(f"  add [{N}]: {elapsed*1000:.2f} ms/call")
     return elapsed
@@ -56,11 +61,11 @@ def bench_gelu_repeated(N, iters=20):
     a_cpu = torch.randn(N)
     a = a_cpu.to("vulkan")
 
-    c = torch.nn.functional.gelu(a)  # warmup
+    torch.nn.functional.gelu(a)  # warmup
 
     t0 = time.perf_counter()
     for _ in range(iters):
-        c = torch.nn.functional.gelu(a)
+        torch.nn.functional.gelu(a)
     elapsed = (time.perf_counter() - t0) / iters
     print(f"  gelu [{N}]: {elapsed*1000:.2f} ms/call")
     return elapsed

tests/test_algo_cache.py

@@ -4,15 +4,19 @@
 repeated dispatches with the same tensor buffers hit the cache.
 """
 
-import sys
 import os
+import sys
 import time
 
-# Use the custom-built PyTorch
-sys.path.insert(0, "/home/raz/builds/pytorch-gfx1150")
+# Allow pointing at a custom-built PyTorch (e.g. a local APU build) without
+# hardcoding a developer's path. Set TORCH_VULKAN_PYTORCH_PATH if needed.
+_custom_torch = os.environ.get("TORCH_VULKAN_PYTORCH_PATH")
+if _custom_torch:
+    sys.path.insert(0, _custom_torch)
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 
 import torch
+
 import torch_vulkan