DictionaryMatcher: exhaustive best-atom match in one CUDA kernel

CHANGELOG.md

@@ -9,6 +9,12 @@
   k-th value along each parameter, then the `2k + 1` values around the winner
   along each; `refine=False` stops after the coarse pass. A 3 % (T1, T2) grid
   of about 37 000 atoms at `coarse=3` scores about 4 200 + 49 atoms per voxel.
+- **Exhaustive matching in one kernel on the card.** Without groups or a
+  lattice, a `DictionaryMatcher` with `top_k=1`, real float32 atoms and up to
+  16 contrasts matches float32 or complex64 signals on a card in one launch
+  (`_match.cu`): every signal against every atom, keeping only the best, so the
+  voxels-by-atoms scores are never written. Indices, scores and scales are the
+  tensor path's. Anything else is matched in tensors as before.
 
 ### Fixed
 

CMakeLists.txt

@@ -120,7 +120,7 @@ if(BLOCHSIM_CUDA)
         _pooled_kernel _pooled_adjoint_kernel
         _regress_kernel _regress_vjp_kernel
     )
-    set(_blochsim_gpu_sources src/blochsim/_gpu.cu)
+    set(_blochsim_gpu_sources src/blochsim/_gpu.cu src/blochsim/_match.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)

src/blochsim/_gpu.cu

@@ -10,6 +10,7 @@
 
 #include "_launch.hpp"
 #include "_layout.hpp"
+#include "_match.hpp"
 
 #include <cuda_runtime.h>
 
@@ -136,7 +137,41 @@ PyObject* pooled_layout_floats(PyObject*, PyObject* args) {
     return PyLong_FromLongLong(blochsim_layout::pooled_adjoint_floats(problems, event_count, n, m, width, dual != 0));
 }
 
+PyObject* match_best(PyObject*, PyObject* args) {
+    unsigned long long signals = 0, atoms = 0, index = 0, score = 0, stream = 0;
+    long long voxels = 0, count = 0;
+    int contrasts = 0, is_complex = 0, device = 0;
+    if (!PyArg_ParseTuple(args, "KKLLipKKiK", &signals, &atoms, &voxels, &count, &contrasts, &is_complex, &index,
+                          &score, &device, &stream)) {
+        return nullptr;
+    }
+    if (contrasts < 1 || contrasts > blochsim_match::MAX_CONTRASTS) {
+        Py_RETURN_FALSE;  // the caller matches in tensors
+    }
+    int previous = 0;
+    cudaError_t status = cudaGetDevice(&previous);
+    if (status == cudaSuccess && previous != device) {
+        status = cudaSetDevice(device);
+    }
+    if (status != cudaSuccess) {
+        return cuda_error(status, "selecting the device");
+    }
+    const int queued = blochsim_match::launch(
+        reinterpret_cast<const float*>(signals), reinterpret_cast<const float*>(atoms), voxels, count, contrasts,
+        is_complex != 0, reinterpret_cast<long long*>(index), reinterpret_cast<float*>(score),
+        reinterpret_cast<void*>(stream));
+    if (previous != device) {
+        cudaSetDevice(previous);
+    }
+    if (queued != cudaSuccess) {
+        return cuda_error(static_cast<cudaError_t>(queued), "matching");
+    }
+    Py_RETURN_TRUE;
+}
+
 PyMethodDef METHODS[] = {
+    {"match_best", match_best, METH_VARARGS,
+     "Queue the exhaustive best-atom match of normalized signals; False where it is not built for them."},
     {"kernels", blochsim_launch::kernel_table, METH_NOARGS,
      "Each kernel's parameter names and kinds."},
     {"launch", launch, METH_VARARGS,

src/blochsim/_match.cu

@@ -0,0 +1,154 @@
+// Exhaustive dictionary matching in one pass on the card.
+//
+// Every signal is scored against every atom, |<s, d>| for normalized real
+// atoms, and only the best atom and its score are kept: the voxels-by-atoms
+// score matrix a matching in tensors writes, reads back and reduces is never
+// stored. A thread holds VOXELS signals in registers; a block stages a tile of
+// atoms in shared memory, where every thread of the block reads the same atom
+// at once (a broadcast). So the work is the C multiply-adds per pair and
+// little else, and a whole volume is matched in a launch per chunk instead of
+// a launch per group and chunk.
+//
+// The signals are real or complex (interleaved); a complex signal scores
+// hypot(<Re s, d>, <Im s, d>), as correlate() in estimators/_grouped.py does.
+#include "_match.hpp"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace {
+
+constexpr int THREADS = 256;
+constexpr int VOXELS = 4;  // signals per thread
+constexpr int TILE = 256;  // atoms staged per pass
+
+template <int C, bool COMPLEX>
+__global__ void __launch_bounds__(THREADS) best_atom(const float* __restrict__ signals,
+                                                     const float* __restrict__ atoms, long long voxels,
+                                                     long long count, long long* __restrict__ index,
+                                                     float* __restrict__ score) {
+    constexpr int PARTS = COMPLEX ? 2 : 1;
+    __shared__ float tile[TILE * C];
+
+    const long long first = (static_cast<long long>(blockIdx.x) * THREADS * VOXELS) + threadIdx.x;
+    float s[VOXELS][PARTS][C];
+    float best[VOXELS];
+    long long where[VOXELS];
+#pragma unroll
+    for (int v = 0; v < VOXELS; ++v) {
+        const long long n = first + static_cast<long long>(v) * THREADS;
+        const bool inside = n < voxels;
+#pragma unroll
+        for (int c = 0; c < C; ++c) {
+#pragma unroll
+            for (int p = 0; p < PARTS; ++p) {
+                s[v][p][c] = inside ? signals[(n * C + c) * PARTS + p] : 0.0f;
+            }
+        }
+        best[v] = -1.0f;
+        where[v] = 0;
+    }
+
+    for (long long start = 0; start < count; start += TILE) {
+        const int here = static_cast<int>(count - start < TILE ? count - start : TILE);
+        __syncthreads();
+        for (int k = threadIdx.x; k < here * C; k += THREADS) {
+            tile[k] = atoms[start * C + k];
+        }
+        __syncthreads();
+        for (int a = 0; a < here; ++a) {
+            const float* d = tile + a * C;
+#pragma unroll
+            for (int v = 0; v < VOXELS; ++v) {
+                float value;
+                if (COMPLEX) {
+                    float re = 0.0f, im = 0.0f;
+#pragma unroll
+                    for (int c = 0; c < C; ++c) {
+                        re = fmaf(s[v][0][c], d[c], re);
+                        im = fmaf(s[v][PARTS - 1][c], d[c], im);
+                    }
+                    value = hypotf(re, im);
+                } else {
+                    float dot = 0.0f;
+#pragma unroll
+                    for (int c = 0; c < C; ++c) {
+                        dot = fmaf(s[v][0][c], d[c], dot);
+                    }
+                    value = fabsf(dot);
+                }
+                // Strictly greater: of equal scores the first atom stays.
+                if (value > best[v]) {
+                    best[v] = value;
+                    where[v] = start + a;
+                }
+            }
+        }
+    }
+
+#pragma unroll
+    for (int v = 0; v < VOXELS; ++v) {
+        const long long n = first + static_cast<long long>(v) * THREADS;
+        if (n < voxels) {
+            index[n] = where[v];
+            score[n] = best[v];
+        }
+    }
+}
+
+template <int C>
+cudaError_t launch_for(const float* signals, const float* atoms, long long voxels, long long count, bool is_complex,
+                       long long* index, float* score, cudaStream_t stream) {
+    const long long per_block = static_cast<long long>(THREADS) * VOXELS;
+    const long long blocks = (voxels + per_block - 1) / per_block;
+    if (blocks > 2147483647LL) {
+        return cudaErrorInvalidValue;
+    }
+    if (is_complex) {
+        best_atom<C, true><<<static_cast<unsigned>(blocks), THREADS, 0, stream>>>(signals, atoms, voxels, count,
+                                                                                  index, score);
+    } else {
+        best_atom<C, false><<<static_cast<unsigned>(blocks), THREADS, 0, stream>>>(signals, atoms, voxels, count,
+                                                                                   index, score);
+    }
+    return cudaGetLastError();
+}
+
+}  // namespace
+
+namespace blochsim_match {
+
+int launch(const float* signals, const float* atoms, long long voxels, long long count, int contrasts,
+           bool is_complex, long long* index, float* score, void* stream) {
+    if (voxels <= 0) {
+        return cudaSuccess;
+    }
+    const cudaStream_t queue = reinterpret_cast<cudaStream_t>(stream);
+    switch (contrasts) {
+#define BLOCHSIM_MATCH_CASE(C) \
+    case C:                    \
+        return launch_for<C>(signals, atoms, voxels, count, is_complex, index, score, queue);
+        BLOCHSIM_MATCH_CASE(1)
+        BLOCHSIM_MATCH_CASE(2)
+        BLOCHSIM_MATCH_CASE(3)
+        BLOCHSIM_MATCH_CASE(4)
+        BLOCHSIM_MATCH_CASE(5)
+        BLOCHSIM_MATCH_CASE(6)
+        BLOCHSIM_MATCH_CASE(7)
+        BLOCHSIM_MATCH_CASE(8)
+        BLOCHSIM_MATCH_CASE(9)
+        BLOCHSIM_MATCH_CASE(10)
+        BLOCHSIM_MATCH_CASE(11)
+        BLOCHSIM_MATCH_CASE(12)
+        BLOCHSIM_MATCH_CASE(13)
+        BLOCHSIM_MATCH_CASE(14)
+        BLOCHSIM_MATCH_CASE(15)
+        BLOCHSIM_MATCH_CASE(16)
+#undef BLOCHSIM_MATCH_CASE
+        default:
+            return -1;  // more contrasts than the kernel is built for: the caller matches in tensors
+    }
+}
+
+}  // namespace blochsim_match

src/blochsim/_match.hpp

@@ -0,0 +1,17 @@
+// Exhaustive best-atom matching on the card (_match.cu).
+#pragma once
+
+namespace blochsim_match {
+
+// The largest contrast count the kernel is built for.
+constexpr int MAX_CONTRASTS = 16;
+
+// For each of ``voxels`` normalized signals ``(voxels, contrasts)`` (float, or
+// interleaved complex when ``is_complex``), the index of the best of ``count``
+// normalized real atoms ``(count, contrasts)`` by |<s, d>|, and that score.
+// Queued on ``stream``. Returns a cudaError_t, or -1 when ``contrasts`` is
+// above MAX_CONTRASTS.
+int launch(const float* signals, const float* atoms, long long voxels, long long count, int contrasts,
+           bool is_complex, long long* index, float* score, void* stream);
+
+}  // namespace blochsim_match

src/blochsim/estimators/_dictionary.py

@@ -405,6 +405,18 @@ def _match_here(
                 *self._scaled(signals, indices),
             )
 
+        fused = _best_on_card(
+            normalized_signals, self.normalized_dictionary, self.top_k
+        )
+        if fused is not None:
+            indices, scores = fused
+            return (
+                indices,
+                scores,
+                self._density(signal_norm, scores, indices),
+                *self._scaled(signals, indices),
+            )
+
         score_chunks = []
         index_chunks = []
         for query in normalized_signals.split(self.query_chunk_size):
@@ -499,6 +511,62 @@ def _shaped(
 # %% private module subroutines
 
 
+def _best_on_card(
+    signals: torch.Tensor, dictionary: torch.Tensor, top_k: int
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """The exhaustive best match in one kernel on the card, or ``None`` where it does not apply.
+
+    The kernel (``_match.cu``) scores every signal against every atom and keeps
+    only the best, so the voxels-by-atoms scores are never stored: one launch
+    per chunk instead of a matrix product, an absolute value and a reduction
+    per chunk of atoms, each through memory. It takes one candidate, real
+    single-precision atoms and float32 or complex64 signals on a card where
+    the GPU kernels are installed; anything else is matched in tensors.
+    """
+    if (
+        top_k != 1
+        or signals.device.type != "cuda"
+        or dictionary.device != signals.device
+        or dictionary.dtype != torch.float32
+        or signals.dtype not in (torch.float32, torch.complex64)
+        or signals.ndim != 2
+        or dictionary.ndim != 2
+        or signals.shape[0] == 0
+    ):
+        return None
+    from .. import _gpu_launch
+
+    if not _gpu_launch.available():
+        return None
+    module = _gpu_launch._module("cuda")
+    if not hasattr(module, "match_best"):
+        return None
+    is_complex = signals.is_complex()
+    flat = (torch.view_as_real(signals) if is_complex else signals).contiguous()
+    atoms = dictionary.contiguous()
+    voxels, contrasts = int(signals.shape[0]), int(signals.shape[1])
+    indices = torch.empty((voxels, 1), dtype=torch.int64, device=signals.device)
+    scores = torch.empty((voxels, 1), dtype=torch.float32, device=signals.device)
+    stream = torch.cuda.current_stream(signals.device)
+    queued = module.match_best(
+        flat.data_ptr(),
+        atoms.data_ptr(),
+        voxels,
+        int(atoms.shape[0]),
+        contrasts,
+        is_complex,
+        indices.data_ptr(),
+        scores.data_ptr(),
+        signals.device.index
+        if signals.device.index is not None
+        else torch.cuda.current_device(),
+        stream.cuda_stream,
+    )
+    if not queued:
+        return None
+    return indices, scores
+
+
 def _prepare_parameters(
     parameters: torch.Tensor | None,
     dictionary: torch.Tensor,

tests/estimators/test_dictionary.py

@@ -186,3 +186,45 @@ def test_a_coarse_search_follows_the_matcher_to_another_device() -> None:
     assert moved.lattice.table.device == moved.dictionary.device
     found = moved.match(dictionary[:5].to(moved.dictionary.device))
     torch.testing.assert_close(found.indices[:, 0].cpu(), torch.arange(5))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.complex64])
+@pytest.mark.parametrize("contrasts", [3, 7, 16, 17])
+def test_the_card_kernel_finds_the_tensor_match(
+    dtype: torch.dtype, contrasts: int
+) -> None:
+    from blochsim.estimators import _dictionary
+
+    g = torch.Generator().manual_seed(contrasts)
+    atoms = torch.randn(1500, contrasts, generator=g).abs()
+    signals = atoms[torch.randint(0, 1500, (5000,), generator=g)]
+    signals = signals + 0.05 * torch.randn(signals.shape, generator=g)
+    if dtype.is_complex:
+        signals = signals * torch.polar(
+            torch.ones(5000, 1), torch.rand(5000, 1, generator=g) * 6.28
+        )
+    signals = signals.to(dtype).cuda()
+    matcher = DictionaryMatcher(dictionary=atoms.cuda(), query_chunk_size=1024)
+
+    fused = matcher.match(signals)
+    called = []
+    original = _dictionary._best_on_card
+    _dictionary._best_on_card = lambda *args: called.append(1)  # the tensor path
+    try:
+        tensors = matcher.match(signals)
+    finally:
+        _dictionary._best_on_card = original
+    assert called
+    if contrasts <= 16:
+        assert (
+            original(
+                signals / signals.norm(dim=-1, keepdim=True),
+                matcher.normalized_dictionary,
+                1,
+            )
+            is not None
+        )
+    assert torch.equal(fused.indices, tensors.indices)
+    torch.testing.assert_close(fused.scores, tensors.scores, rtol=1e-5, atol=1e-6)
+    torch.testing.assert_close(fused.scales, tensors.scales, rtol=1e-5, atol=1e-6)