Compile the GPU kernels ahead of time as CUDA and drop Triton

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -63,7 +63,7 @@ body:
       multiple: true
       options:
         - CPU
-        - CUDA (Triton kernels)
+        - CUDA
         - Both
     validations:
       required: true

.github/SECURITY.md

@@ -29,7 +29,7 @@ a signal model are all Python that runs in your process, so a malicious
 In scope is anything that turns *data* into execution or into memory
 corruption -- a sequence description, a pulse waveform, a phantom or a
 dictionary read from a file, or values passed to the simulator, reaching the
-C++ and Triton kernels. Those kernels index raw pointers, so an out-of-bounds
+C++ and CUDA kernels. Those kernels index raw pointers, so an out-of-bounds
 read or write reachable from ordinary arguments is a vulnerability and not
 merely a bug.
 

.github/skills/add-a-sequence/SKILL.md

@@ -63,7 +63,7 @@ evaluates. Adding an effect is one line in a model's property declaration, plus
 the term itself in **both** kernels.
 
 The shared parameter ABI is `src/blochsim/sequence/_parameters.py`, read by the
-Python dispatch, the C++ extension and the Triton kernels. A parameter added
+Python dispatch, the C++ extension and the GPU kernels. A parameter added
 there is added in all three or in none.
 
 ## What the change has to arrive with

.github/skills/build-and-test/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: build-and-test
-description: Compile the C++ kernels and run the BlochSim suite, including the Triton paths on a machine with no GPU. Use when asked to build, install, test, or reproduce a failure in blochsim.
+description: Compile the C++ kernels and run the BlochSim suite, including the GPU kernels on a machine with no GPU. Use when asked to build, install, test, or reproduce a failure in blochsim.
 ---
 
 # Build and test BlochSim
@@ -16,7 +16,8 @@ run alone, `pip install -e ".[test]"` is enough and much smaller.
 
 An editable install puts the Python sources on the path, so an edit under
 `src/blochsim` takes effect on the next import. The two C++ extensions are
-compiled artifacts and do not: re-run the install after editing any `.cpp`.
+compiled artifacts and do not: re-run the install after editing any `.cpp`,
+`.hpp` or `.cu`.
 
 **Read the exit status, not the output.** A failed compile leaves the
 previously built `.so` importable, so the suite runs green against a kernel
@@ -41,32 +42,15 @@ against closed forms — the shift, the RF rotation, relaxation, diffusion, flow
 spoiling, the two-pool and three-pool longitudinal steps — while `sequence/`,
 `model/`, `estimators/`, `recon/` and `optim/` cover the layers above.
 
-## The Triton paths, without a GPU
+## The GPU kernels, without a GPU
 
-The tests carrying the `interpreted` marker are deselected by `addopts`. They
-run a Triton kernel through Triton's CPU interpreter — no card, no compile,
-about a minute each — and are how the GPU plumbing is verified anywhere:
-
-```sh
-pytest tests/ -m interpreted
-TRITON_INTERPRET=1 python your_script.py     # the same trick, by hand
-```
-
-Triton has no wheel for macOS or Windows. On those platforms the default
-deselection is not an optimisation, it is the only thing that runs.
-
-## Before you time anything
-
-Kernel compiles dominate a cold GPU run, not the arithmetic. A suite that takes
-minutes on a card is mostly Triton compiling one specialization per feature
-combination it meets; the second run of the same suite is a different number
-entirely. Run the whole suite at natural boundaries rather than after every
-edit.
-
-The cache keys on the *source* of `_epg_triton.py`, not on what it means, so a
-formatting pass or a deleted dead local costs the same full recompile as a new
-`tl.constexpr`. If the suite suddenly takes an hour where it took minutes,
-check whether that file changed before looking for a performance regression.
+On Linux the install compiles the GPU kernels for the host as
+`blochsim._gpu_host`, and
+for the card as `blochsim._gpu` wherever CMake finds `nvcc`
+(`--config-settings=cmake.define.BLOCHSIM_CUDA=ON` insists on it).
+`tests/sequence/test_host_kernels.py` and `test_many_pools_host.py` run the
+host build against the C++ kernels, which is how the GPU kernels are verified
+on any machine; the CUDA tests skip themselves without a card.
 
 ## Style
 

.github/workflows/test.yml

@@ -48,14 +48,6 @@ jobs:
         if: runner.os == 'Linux'
         run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
-      # Triton reaches a runner as a dependency of the CUDA build of PyTorch,
-      # which the CPU index above deliberately leaves out, and it exists for
-      # Linux alone. Installing it is what lets the kernel-gate tests run; on
-      # the other two platforms they skip themselves.
-      - name: Install Triton (Linux only)
-        if: runner.os == 'Linux'
-        run: python -m pip install triton
-
       # Not editable and not from a wheel: the C++ kernels are compiled here,
       # on this runner, with this compiler, which is the half of the package a
       # pure-Python job would never exercise. CMake and Ninja arrive as
@@ -69,13 +61,12 @@ jobs:
           import importlib.util, pathlib
           spec = importlib.util.find_spec("blochsim")
           root = pathlib.Path(next(iter(spec.submodule_search_locations)))
-          for kernel in sorted(root.glob("_*_cpu*")):
+          for kernel in sorted(root.glob("_*_cpu*")) + sorted(root.glob("_gpu*")):
               print(kernel.name, kernel.stat().st_size, "bytes")
 
-      # What a runner with no card runs: the C++ kernels and every layer above
-      # them. The GPU tests skip themselves on ``torch.cuda.is_available()``,
-      # and the ones that reach for Triton's CPU interpreter carry the
-      # ``interpreted`` marker, which ``addopts`` deselects.
+      # What a runner with no card runs: the C++ kernels, the GPU kernels
+      # compiled for the host, and every layer above them. The GPU tests skip
+      # themselves on ``torch.cuda.is_available()``.
       - name: Run the tests
         run: pytest tests/ -n auto
 

.github/workflows/wheels.yml

@@ -12,6 +12,8 @@ on:
       - setup.py
       - "src/blochsim/*.cpp"
       - "src/blochsim/*.hpp"
+      - "src/blochsim/*.cu"
+      - "src/blochsim/*.cu.in"
       - scripts/check_wheel.py
       - .github/workflows/wheels.yml
   workflow_dispatch:
@@ -52,7 +54,9 @@ jobs:
   wheels:
     name: ${{ matrix.label }}
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 60
+    # The x86-64 manylinux wheel compiles the GPU kernels for every listed
+    # architecture.
+    timeout-minutes: 150
     strategy:
       fail-fast: false
       matrix:

.pre-commit-config.yaml

@@ -12,6 +12,8 @@ repos:
     hooks:
       - id: check-added-large-files
         args: [--maxkb=512]
+        # The EPG kernels for the card, in every mode.
+        exclude: ^src/blochsim/_epg_kernels\.hpp$
       - id: check-case-conflict
       - id: check-merge-conflict
       - id: check-symlinks

CHANGELOG.md

@@ -7,6 +7,15 @@
 - **The package is `blochsim`.** The distribution, the import name and the
   repository are `blochsim`; `import torchsim` becomes `import blochsim`, with
   the same modules and names beneath it.
+- **The GPU kernels are CUDA compiled ahead of time, and Triton is not used.**
+  The EPG, pooled and PERK kernels are C++ (`_epg_kernels.hpp`,
+  `_pools_kernels.hpp`, `_perk_kernels.hpp`) compiled by `nvcc` into
+  `blochsim._gpu`, which links the CUDA runtime statically; the x86-64
+  manylinux wheel carries it, and a source build compiles it wherever CMake
+  finds `nvcc`. No kernel is compiled at the first call. The same kernels are
+  compiled for the host as `blochsim._gpu_host`, which the suite holds to the
+  C++ kernels; the `interpreted` marker is gone. A launch is at most 1024
+  threads, so an EPG run of more than 1024 state orders on a card is refused.
 
 ### Added
 

CLAUDE.md

@@ -16,12 +16,14 @@ broken checkout (Windows without `core.symlinks`), not two documents.
 | `src/blochsim/sequence/` | The description an acquisition is assembled from — events, operators, builders — and the dispatch that turns one into a kernel launch. |
 | `src/blochsim/model/` | What a signal model *is*: the physics, the simulator that orders its events, and the binding that resolves a protocol's structure once and rebinds its values per call. |
 | `src/blochsim/_epg_cpu.cpp`, `_perk_cpu.cpp` | The CPU kernels. Every path the GPU has exists here too, and the two agree to float32 round-off. |
-| `src/blochsim/sequence/_epg_triton.py`, `estimators/_perk_triton.py` | The GPU kernels. |
+| `src/blochsim/_epg_kernels.hpp`, `_pools_kernels.hpp`, `_perk_kernels.hpp` | The GPU kernels, written over the tiles of `_tile.hpp`. `_kernels.hpp` is the table a launch reads: each kernel's parameters, and which of them size the block. |
+| `src/blochsim/_gpu.cu`, `_gpu_kernel.cu.in`, `_gpu_host.cpp`, `_gpu_launch.py` | The same kernels compiled ahead of time for the card (`_gpu`), and for the host (`_gpu_host`), which is how the suite checks them without one; and the launcher both answer to. |
+| `src/blochsim/sequence/_epg_gpu.py`, `_pools_gpu.py`, `estimators/_perk_gpu.py` | What a GPU launch is given: the tiling and the arguments. |
 | `src/blochsim/simulators/`, `estimators/`, `recon/`, `optim/` | The sequences that ship, and what is built on top of them. |
 | `tests/`, `examples/`, `docs/` | Mirrored by subpackage, executed by the gallery, built by Sphinx. |
 
 The shared parameter ABI — read by the Python dispatch, the C++ extension and
-the Triton kernels alike — is `src/blochsim/sequence/_parameters.py`. A
+the GPU kernels alike — is `src/blochsim/sequence/_parameters.py`. A
 parameter added there is added in all three places or in none.
 
 ## Commands
@@ -30,7 +32,6 @@ parameter added there is added in all three places or in none.
 pip install -e ".[dev]"     # the whole toolchain, and it compiles the kernels
 pytest tests/               # the suite, with coverage
 pytest tests/ -n auto       # across cores
-pytest tests/ -m interpreted   # the Triton paths, through Triton's CPU interpreter
 pre-commit install          # once per clone, or no hook runs on commit
 pre-commit run --all-files  # exactly what CI's Lint job runs
 bash scripts/build_docs.sh  # HTML into docs/build/html, examples executed
@@ -57,40 +58,19 @@ pip install -e ".[dev]" ; echo "exit: $?"
 python -c "import blochsim._epg_cpu as k; print(k.__file__)"
 ```
 
-**Kernel compiles dominate a cold GPU run**, not the arithmetic. A suite that
-takes minutes on a card is mostly Triton compiling one specialization per
-feature combination it meets. Run the whole suite at natural boundaries, not
-after every edit.
-
-**Triton's cache keys on the source of the kernel file, not on its meaning.**
-Reformatting `_epg_triton.py`, or deleting a dead local from it, invalidates
-every specialization exactly as a new `tl.constexpr` would: the next full run
-on a card goes from minutes to the better part of an hour, almost all of it one
-MLIR pass. Know that before letting a formatter touch that file, and say which
-number you are quoting.
-
-**Most of what the cache holds is never loaded.** Each specialization is
-written out as its `.ttir`, `.ttgir`, `.llir` and `.ptx` beside the `.cubin`
-and the metadata a launch actually reads, and on a kernel this size the
-intermediate IR is the bulk of the bytes. `TRITON_STORE_BINARY_ONLY=1` writes
-the binary and its metadata alone; every compilation stage still runs, so
-neither the generated code nor the set of specializations changes.
-`TRITON_DISABLE_LINE_INFO=1` drops the source locations embedded in both,
-which costs line attribution under `ncu`. Set `TRITON_CACHE_DIR` to somewhere
-you keep and the compile is paid once per edit of the kernel file rather than
-once per checkout.
-
-Neither knob stops the `.source` file, which carries a location per operation
-and is most of an entry on a kernel this size; nothing reads it at launch.
-Nothing evicts either, and the cache keys on the text of the kernel file, so
-every edit orphans every entry made before it.
-`scripts/prune_triton_cache.py` reports both and, with `--apply`, drops them:
-the IR is free to drop, the stale entries recompile if you ask for them again.
-
-**The `interpreted` marker is deselected by default.** Those tests run a Triton
-kernel through Triton's CPU interpreter — no GPU, no compile, about a minute
-each. `TRITON_INTERPRET=1` does the same thing by hand for a script. It is how
-the GPU plumbing is verified on a machine with no card.
+**The GPU kernels are compiled where a CUDA compiler is found.** CMake looks
+for `nvcc`, and `BLOCHSIM_CUDA=ON` or `OFF` (`--config-settings=cmake.define.BLOCHSIM_CUDA=ON`)
+overrides what it finds; `CMAKE_CUDA_ARCHITECTURES` picks the cards. A kernel
+compiles in a file of its own, twice -- bounded to 256 threads and to 1024 --
+so a build is minutes of `nvcc` spread over as many cores as Ninja is given.
+On Linux the same kernels are also compiled for the host (`_gpu_host`), one
+program at a time over host tensors; `tests/sequence/test_host_kernels.py` and
+`test_many_pools_host.py` hold them to the C++ kernels, which is how the GPU
+path is verified on a machine with no card.
+
+**A block is at most 1024 threads.** A tile is an element per thread, so an
+EPG launch with more than 1024 state orders, or a pooled one whose orders by
+pools pass 1024, is refused with the kernel's name rather than launched.
 
 **`--cov` is on by default** through `addopts`, so a bare `pytest` writes
 `coverage.xml`. It is ignored, not tracked.
@@ -129,7 +109,7 @@ written that way and each states its invariant in its module docstring. A test
 that only compares BlochSim to BlochSim proves the two agree, which was never
 in doubt.
 
-Whatever you change in one kernel, change in the other. The C++ and Triton
+Whatever you change in one kernel, change in the other. The C++ and GPU
 implementations are held to each other to float32 round-off, and a path that
 exists on one side and not the other is a bug in whichever side is missing it.
 
@@ -207,12 +187,14 @@ beside the list the version switcher reads.
 ## Packaging
 
 The build is `scikit-build-core` driving `CMakeLists.txt`; `setup.py` is a shim
-for tools that still shell out to it and configures nothing. Both kernels are
+for tools that still shell out to it and configures nothing. The kernels are
 plain CPython extensions against the **stable ABI** from 3.10 on: they call no
 PyTorch API and link no PyTorch library, which is why one `cp310-abi3` wheel
 per platform serves every supported interpreter and why that wheel is a couple
 of megabytes rather than the size of libtorch. Keep it that way — a `#include
-<torch/...>` in either `.cpp` ends all of that.
+<torch/...>` in any of them ends all of that. The GPU module links the CUDA
+runtime statically, so a machine needs the driver and nothing else, and the
+x86-64 manylinux wheel is the one built with it.
 
 Wheels are built by cibuildwheel and published to PyPI by trusted publishing on
 a `v*.*.*` tag. `scripts/check_wheel.py` loads each compiled kernel by path,

CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.26)
 
 project(blochsim LANGUAGES CXX)
 
-# Both kernels are plain CPython extensions against the stable ABI: they call
+# The kernels are plain CPython extensions against the stable ABI: they call
 # no PyTorch API and link no PyTorch library, which is why one wheel per
 # platform serves every supported interpreter and why that wheel is a couple of
 # megabytes rather than the size of libtorch.
@@ -56,3 +56,69 @@ endfunction()
 
 blochsim_add_kernel(_epg_cpu  src/blochsim/_epg_cpu.cpp)
 blochsim_add_kernel(_perk_cpu src/blochsim/_perk_cpu.cpp)
+
+# The GPU kernels compiled for the host, one program at a time. Nothing in the
+# package dispatches to them; they are how the suite checks the GPU kernels on
+# a machine with no card. Linux only, like the card's build.
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    set(_blochsim_host_default ON)
+else()
+    set(_blochsim_host_default OFF)
+endif()
+option(BLOCHSIM_HOST_KERNELS "Compile the GPU kernels for the host, for the tests"
+    ${_blochsim_host_default})
+if(BLOCHSIM_HOST_KERNELS)
+    blochsim_add_kernel(_gpu_host src/blochsim/_gpu_host.cpp)
+endif()
+
+# The GPU kernels compiled ahead of time for the card, wherever a CUDA
+# compiler is found unless BLOCHSIM_CUDA says otherwise.
+include(CheckLanguage)
+check_language(CUDA)
+if(CMAKE_CUDA_COMPILER AND CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    set(_blochsim_cuda_default ON)
+else()
+    set(_blochsim_cuda_default OFF)
+endif()
+option(BLOCHSIM_CUDA "Compile the GPU kernels for the card" ${_blochsim_cuda_default})
+
+if(BLOCHSIM_CUDA)
+    # Real code for each generation and PTX for the newest, which the driver
+    # compiles for a card newer than any listed.
+    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+        set(CMAKE_CUDA_ARCHITECTURES 75-real 80-real 86-real 89-real 90)
+    endif()
+    enable_language(CUDA)
+    find_package(CUDAToolkit REQUIRED)
+    set(CMAKE_CUDA_STANDARD 17)
+    set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+
+    # One file per kernel, so they compile in parallel.
+    set(_blochsim_kernels
+        _three_pool_table_jvp_kernel _three_pool_table_kernel
+        _epg_vjp_kernel _epg_vjp_jvp_kernel _epg_real_vjp_jvp_kernel
+        _epg_real_vjp_kernel _epg_real_kernel _epg_real_jvp_kernel
+        _epg_kernel _epg_jvp_kernel
+        _pooled_kernel _pooled_adjoint_kernel
+        _regress_kernel _regress_vjp_kernel
+    )
+    set(_blochsim_gpu_sources src/blochsim/_gpu.cu)
+    foreach(NAME IN LISTS _blochsim_kernels)
+        set(_source "${CMAKE_CURRENT_BINARY_DIR}/kernels/${NAME}.cu")
+        configure_file(src/blochsim/_gpu_kernel.cu.in "${_source}" @ONLY)
+        list(APPEND _blochsim_gpu_sources "${_source}")
+    endforeach()
+
+    python_add_library(_gpu MODULE
+        USE_SABI ${BLOCHSIM_ABI3_VERSION}
+        WITH_SOABI
+        ${_blochsim_gpu_sources}
+    )
+    target_include_directories(_gpu PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src/blochsim")
+    # The runtime is linked in, so a machine needs the driver and nothing else.
+    target_link_libraries(_gpu PRIVATE CUDA::cudart_static)
+    # Warning 221 is a double constant that rounds to zero in float, which the
+    # kernels write deliberately as the floor of a clamp.
+    target_compile_options(_gpu PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-diag-suppress=221>)
+    install(TARGETS _gpu DESTINATION blochsim)
+endif()