From 5a200926d2fdcd3c0cf4e077242f58a251a9344d Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 4 Nov 2026 13:07:27 +0000 Subject: [PATCH] DictionaryMatcher: exhaustive best-atom match in one CUDA kernel Without groups or a lattice, top_k=1, real float32 atoms and up to 16 contrasts, the match of float32 or complex64 signals on a card runs as one kernel (_match.cu, bound in _gpu as match_best): each thread holds four signals in registers, atoms are staged in shared memory and read as broadcasts, and only the best atom and its score are kept. The voxels-by-atoms score matrix is never stored, and there is one launch per call instead of a product, an absolute value and a reduction per chunk of atoms (or, with groups, a few launches and a sync per group). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QLjhTxFMFDEzathbEm --- CHANGELOG.md | 6 + CMakeLists.txt | 2 +- src/blochsim/_gpu.cu | 35 ++++++ src/blochsim/_match.cu | 154 +++++++++++++++++++++++++ src/blochsim/_match.hpp | 17 +++ src/blochsim/estimators/_dictionary.py | 68 +++++++++++ tests/estimators/test_dictionary.py | 42 +++++++ 7 files changed, 323 insertions(+), 1 deletion(-) create mode 100644 src/blochsim/_match.cu create mode 100644 src/blochsim/_match.hpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 849e5fcc..ed296ef8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 0124ed93..1dd989aa 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/src/blochsim/_gpu.cu b/src/blochsim/_gpu.cu index 016359cd..6fa0d29b 100644 --- a/src/blochsim/_gpu.cu +++ b/src/blochsim/_gpu.cu @@ -10,6 +10,7 @@ #include "_launch.hpp" #include "_layout.hpp" +#include "_match.hpp" #include @@ -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(signals), reinterpret_cast(atoms), voxels, count, contrasts, + is_complex != 0, reinterpret_cast(index), reinterpret_cast(score), + reinterpret_cast(stream)); + if (previous != device) { + cudaSetDevice(previous); + } + if (queued != cudaSuccess) { + return cuda_error(static_cast(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, diff --git a/src/blochsim/_match.cu b/src/blochsim/_match.cu new file mode 100644 index 00000000..4ff4bca2 --- /dev/null +++ b/src/blochsim/_match.cu @@ -0,0 +1,154 @@ +// Exhaustive dictionary matching in one pass on the card. +// +// Every signal is scored against every atom, || 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(, ), as correlate() in estimators/_grouped.py does. +#include "_match.hpp" + +#include + +#include + +namespace { + +constexpr int THREADS = 256; +constexpr int VOXELS = 4; // signals per thread +constexpr int TILE = 256; // atoms staged per pass + +template +__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(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(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(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(v) * THREADS; + if (n < voxels) { + index[n] = where[v]; + score[n] = best[v]; + } + } +} + +template +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(THREADS) * VOXELS; + const long long blocks = (voxels + per_block - 1) / per_block; + if (blocks > 2147483647LL) { + return cudaErrorInvalidValue; + } + if (is_complex) { + best_atom<<(blocks), THREADS, 0, stream>>>(signals, atoms, voxels, count, + index, score); + } else { + best_atom<<(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(stream); + switch (contrasts) { +#define BLOCHSIM_MATCH_CASE(C) \ + case C: \ + return launch_for(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 diff --git a/src/blochsim/_match.hpp b/src/blochsim/_match.hpp new file mode 100644 index 00000000..f67d73d0 --- /dev/null +++ b/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 ||, 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 diff --git a/src/blochsim/estimators/_dictionary.py b/src/blochsim/estimators/_dictionary.py index 2007e7c6..5b88b46d 100644 --- a/src/blochsim/estimators/_dictionary.py +++ b/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, diff --git a/tests/estimators/test_dictionary.py b/tests/estimators/test_dictionary.py index a15dfcdc..005c8a23 100644 --- a/tests/estimators/test_dictionary.py +++ b/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)