diff --git a/CMakeLists.txt b/CMakeLists.txt index 9239fffb..1fdd54fa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -328,6 +328,7 @@ set(BARTORCH_SOURCES src/csrc/ops/iter.c src/csrc/ops/attention.c src/csrc/ops/groupnorm.c + src/csrc/ops/grog.c src/csrc/substitute/backend.c src/csrc/substitute/ref_blas.c @@ -343,7 +344,7 @@ if(OpenMP_C_FOUND AND (WIN32 OR APPLE)) endif() if(BARTORCH_CUDA) - list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu src/csrc/ops/attention.cu src/csrc/ops/groupnorm.cu) + list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu src/csrc/ops/attention.cu src/csrc/ops/groupnorm.cu src/csrc/ops/grog.cu) # The passes around a volume's transform, which cuFFT links into the # transform when a plan is made (src/csrc/ops/fft_callbacks.cu). It takes them as diff --git a/docs/api/functions.md b/docs/api/functions.md index c7e34289..dbaf9ca8 100644 --- a/docs/api/functions.md +++ b/docs/api/functions.md @@ -95,6 +95,7 @@ recorded by autograd. | {obj}`~bartorch.cuda_available` | Whether the library was built with CUDA and a device is present | | {obj}`~bartorch.set_cuda_streams` | Number of CUDA streams BART uses | | {obj}`~bartorch.use_cuda_memcache` | Whether BART caches freed device memory | +| {obj}`~bartorch.empty_cuda_memcache` | Return BART's cached device memory to the driver | ## Build information diff --git a/src/bartorch/_abi.py b/src/bartorch/_abi.py index a07e224d..fda6e1f4 100644 --- a/src/bartorch/_abi.py +++ b/src/bartorch/_abi.py @@ -356,6 +356,9 @@ class Encoding(ctypes.Structure): "bartorch_gn_built", "bartorch_gn_workspace", "bartorch_gn_silu", + "bartorch_grog_grid", + "bartorch_grog_cuda_built", + "bartorch_grog_grid_cuda", "bartorch_on_device", ) @@ -1337,6 +1340,52 @@ def bind(lib: ctypes.CDLL) -> ctypes.CDLL: ctypes.c_void_p, ctypes.c_void_p, ] + lib.bartorch_grog_grid.restype = ctypes.c_int + lib.bartorch_grog_grid.argtypes = [ + ctypes.c_int, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_float, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int64, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int64, + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_int64, + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_float), + ] + lib.bartorch_grog_cuda_built.restype = ctypes.c_int + lib.bartorch_grog_cuda_built.argtypes = [] + lib.bartorch_grog_grid_cuda.restype = ctypes.c_int + lib.bartorch_grog_grid_cuda.argtypes = [ + ctypes.c_int, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_float, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int64, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int64, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int64), + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.bartorch_on_device.restype = ctypes.c_int lib.bartorch_on_device.argtypes = [ctypes.c_void_p] return lib diff --git a/src/bartorch/_cuda.py b/src/bartorch/_cuda.py index 7328ba98..b3db5095 100644 --- a/src/bartorch/_cuda.py +++ b/src/bartorch/_cuda.py @@ -20,6 +20,7 @@ "available", "device", "device_count", + "empty_memcache", "free_memory", "ordered", "set_streams", @@ -81,6 +82,11 @@ def use_memcache(enable: bool) -> None: library().bartorch_cuda_use_memcache(int(bool(enable))) +def empty_memcache() -> None: + """Hand the device memory BART keeps for its own reuse back to the driver.""" + library().bartorch_cuda_memcache_clear_all() + + @contextmanager def ordered(device: torch.device): """Run BART on *device*, ordered against torch's current stream. diff --git a/src/bartorch/_marshal.py b/src/bartorch/_marshal.py index 3d8e57ef..b91707d2 100644 --- a/src/bartorch/_marshal.py +++ b/src/bartorch/_marshal.py @@ -137,6 +137,11 @@ def float_pointer(ptr: int): return ctypes.cast(ctypes.c_void_p(ptr), ctypes.POINTER(ctypes.c_float)) +def int64_pointer(ptr: int): + """An ``int64_t*`` at the address ``ptr``, for an entry point that takes one.""" + return ctypes.cast(ctypes.c_void_p(ptr), ctypes.POINTER(ctypes.c_int64)) + + def null_apply() -> ctypes.Array: """The NULL an apply callback the caller did not give is passed as. diff --git a/src/bartorch/_settings.py b/src/bartorch/_settings.py index 532d9322..b6fdb155 100644 --- a/src/bartorch/_settings.py +++ b/src/bartorch/_settings.py @@ -15,6 +15,7 @@ "bart_version", "build_info", "cuda_available", + "empty_cuda_memcache", "get_debug_level", "set_copy_inputs", "set_cuda_streams", @@ -70,6 +71,25 @@ def set_cuda_streams(n: int) -> None: _cuda.set_streams(n) +def empty_cuda_memcache() -> None: + """Hand the device memory BART's cache holds back to the driver. + + The cache keeps what an operator freed for its next application, which + makes an iterative solve's allocations cheap; between stages that run + elsewhere -- a solve, then a network -- it holds memory nothing uses. + Clearing it there frees that memory without slowing the solve. A no-op + without CUDA. + + Examples + -------- + >>> x = optim.CG(maxiter=3, device="cuda").inverse(b, A, x) + >>> bartorch.empty_cuda_memcache() + """ + from bartorch._lib import library + + library().bartorch_cuda_memcache_clear_all() + + def use_cuda_memcache(enable: bool) -> None: """Set whether BART keeps freed device memory for its own reuse. diff --git a/src/bartorch/tools/_igrog.py b/src/bartorch/tools/_igrog.py index 03107128..83681074 100644 --- a/src/bartorch/tools/_igrog.py +++ b/src/bartorch/tools/_igrog.py @@ -3,7 +3,7 @@ from __future__ import annotations import math -from dataclasses import dataclass +from dataclasses import dataclass, field import torch @@ -30,6 +30,7 @@ class _Gridded: sample: torch.Tensor shape: tuple noise: torch.Tensor + also: list = field(default_factory=list) def weights(self) -> torch.Tensor: """``1 / sqrt(noise)`` at unit mean square: the weights that whiten the gridded samples.""" @@ -69,46 +70,182 @@ def _lagrange4(t: torch.Tensor) -> torch.Tensor: ) -def _interpolate(cc: _Correlation, table: torch.Tensor, disp: torch.Tensor) -> torch.Tensor: - """``Q`` at displacements ``(..., d)``, as ``(..., C, C)``.""" +def _interpolate( + rows: torch.Tensor, size: int, cc: _Correlation, disp: torch.Tensor +) -> torch.Tensor: + """``Q`` at displacements ``(..., d)`` as ``(..., C C)``, from ``(nodes, C C)`` rows.""" d, U, R = disp.shape[-1], cc.upsample, cc.reach u = (disp.double() + R) * U base = torch.floor(u).long() - w = _lagrange4(u - base) - coils, size = table.shape[0], table.shape[2] - flat_table = table.reshape(coils, coils, -1) + w = _lagrange4(u - base).to(rows.real.dtype) # (..., d, 4) shape = disp.shape[:-1] - out = torch.zeros(*shape, coils, coils, dtype=torch.complex128, device=disp.device) + out = torch.zeros(*shape, rows.shape[1], dtype=rows.dtype, device=disp.device) + # Node index and weight of each tap along each axis; a tap is one node per axis. + strides = [size ** (d - 1 - a) for a in range(d)] + node = [ + (base[..., a, None] + torch.arange(-1, 3, device=disp.device)) * strides[a] + for a in range(d) + ] for taps in torch.cartesian_prod(*[torch.arange(4)] * d).reshape(-1, d).tolist(): - flat = torch.zeros(shape, dtype=torch.long, device=disp.device) - wt = torch.ones(shape, dtype=torch.float64, device=disp.device) - for a in range(d): - flat = flat * size + (base[..., a] + taps[a] - 1) + flat = node[0][..., taps[0]] + wt = w[..., 0, taps[0]] + for a in range(1, d): + flat = flat + node[a][..., taps[a]] wt = wt * w[..., a, taps[a]] - out += wt[..., None, None] * flat_table[:, :, flat].movedim((0, 1), (-2, -1)) + out.add_(rows.index_select(0, flat.reshape(-1)).view(*shape, -1) * wt[..., None]) return out def _kernels( - cc: _Correlation, table: torch.Tensor, offsets: torch.Tensor, lam: float + cc: _Correlation, rows: torch.Tensor, size: int, offsets: torch.Tensor, lam: float ) -> torch.Tensor: - """Wiener kernels ``(M, C, D C)`` of sources at ``offsets`` ``(M, D, d)`` from their targets.""" - M, D, _ = offsets.shape - C = table.shape[0] - disp = (offsets[:, :, None] - offsets[:, None]).reshape(M, D * D, -1) - Rss = ( - _interpolate(cc, table, disp) - .reshape(M, D, D, C, C) - .permute(0, 1, 3, 2, 4) - .reshape(M, D * C, D * C) + """Wiener kernels ``(M, C, D C)`` of sources at ``offsets`` ``(M, D, d)`` from their targets. + + ``Q(-δ) = Q(δ)^H`` and ``Q(0)`` is one node, so only the source pairs ``i < j`` + and the source-to-target displacements are interpolated. + """ + M, D, d = offsets.shape + C = int(round(rows.shape[1] ** 0.5)) + pairs = [(i, j) for i in range(D) for j in range(i + 1, D)] + disp = torch.stack( + [offsets[:, i] - offsets[:, j] for i, j in pairs] + [-offsets[:, k] for k in range(D)], 1 ) + q = _interpolate(rows, size, cc, disp).view(M, len(pairs) + D, C, C) + zero = _interpolate(rows, size, cc, torch.zeros(1, d, device=offsets.device)).view(C, C) + Rss = torch.empty(M, D, D, C, C, dtype=rows.dtype, device=offsets.device) + for k in range(D): + Rss[:, k, k] = zero + for p, (i, j) in enumerate(pairs): + Rss[:, i, j] = q[:, p] + Rss[:, j, i] = q[:, p].mH + Rss = Rss.permute(0, 1, 3, 2, 4).reshape(M, D * C, D * C) Rss = 0.5 * (Rss + Rss.mH) - Rts = _interpolate(cc, table, -offsets).permute(0, 2, 1, 3).reshape(M, C, D * C) + Rts = q[:, len(pairs) :].permute(0, 2, 1, 3).reshape(M, C, D * C) tr = torch.diagonal(Rss, dim1=-2, dim2=-1).real.mean(-1) Rss = Rss + (lam * tr)[:, None, None] * torch.eye(D * C, dtype=Rss.dtype, device=Rss.device) return torch.cholesky_solve(Rts.mH, torch.linalg.cholesky(Rss)).mH +def _grog_grid(): + """The library's host gridder, ``bartorch_grog_grid``.""" + from bartorch._lib import library + + return library().bartorch_grog_grid + + +_GROG_ERRORS = { + 1: "a missing argument", + 2: "a bad shape (sources x coils is at most 256)", + 3: "out of memory", + 4: "an index out of range", + 5: "a kernel system that is not positive definite", + 6: "a displacement beyond the correlation table", +} + + +def _native(cc, lam, tr, sets, shot, sample, src, outs, noise): + """Grid the targets ``(shot, sample)`` of every set in ``sets`` with the host loop. + + ``sets`` are ``(C, S, P)`` complex64 on the host; ``outs`` the ``(C, M)`` + results and ``noise`` ``(M,)``, written in place. + """ + d = tr.shape[-1] + C, D = cc.table.shape[0], src.shape[1] + table = cc.table.movedim((0, 1), (-2, -1)).contiguous() # (nodes..., C, C) + traj = tr.float().contiguous() + shot, sample = shot.contiguous(), sample.contiguous() + src = src.contiguous() + from bartorch._marshal import float_pointer, int64_pointer, pointers + + err = _grog_grid()( + d, + C, + D, + cc.reach, + cc.upsample, + float(lam), + table.data_ptr(), + traj.shape[0], + traj.shape[1], + float_pointer(traj.data_ptr()), + len(sets), + pointers([a.data_ptr() for a in sets]), + shot.numel(), + int64_pointer(shot.data_ptr()), + int64_pointer(sample.data_ptr()), + int64_pointer(src.data_ptr()), + pointers([o.data_ptr() for o in outs]), + float_pointer(noise.data_ptr()), + ) + if err: + raise RuntimeError(f"igrog: the host gridder stopped on {_GROG_ERRORS.get(err, err)}") + + +def _card_grid(cc, lam, tr, sets, shot, sample, src, card, batch): + """The card gridder over ``batch`` targets at a time; ``None`` when this build has none. + + ``sets`` are ``(C, S, P)`` complex64 on ``card``; returns the ``(C, M)`` + results of each set and ``(M,)`` noise gains, on the host. + """ + from bartorch._lib import library + from bartorch._marshal import float_pointer, int64_pointer + + lib = library() + C, D = cc.table.shape[0], src.shape[1] + if not lib.bartorch_grog_cuda_built() or C * D > 48: + return None + d = tr.shape[-1] + table = cc.table.movedim((0, 1), (-2, -1)).contiguous().to(card) + traj = tr.float().contiguous().to(card) + data = torch.tensor([a.data_ptr() for a in sets], dtype=torch.int64, device=card) + M = shot.numel() + outs = [torch.empty(a.shape[0], M, dtype=torch.complex64) for a in sets] + noise = torch.empty(M, dtype=torch.float32) + status = torch.zeros(1, dtype=torch.int32, device=card) + with torch.cuda.device(card): + stream = torch.cuda.current_stream(card).cuda_stream + for b in range(0, M, batch): + e = min(M, b + batch) + sh, sa = shot[b:e].to(card), sample[b:e].to(card) + sr = src[b:e].contiguous().to(card) + res = [torch.empty(a.shape[0], e - b, dtype=torch.complex64, device=card) for a in sets] + nb = torch.empty(e - b, dtype=torch.float32, device=card) + ptrs = torch.tensor([o.data_ptr() for o in res], dtype=torch.int64, device=card) + err = lib.bartorch_grog_grid_cuda( + d, + C, + D, + cc.reach, + cc.upsample, + float(lam), + table.data_ptr(), + traj.shape[0], + traj.shape[1], + float_pointer(traj.data_ptr()), + len(sets), + data.data_ptr(), + e - b, + int64_pointer(sh.data_ptr()), + int64_pointer(sa.data_ptr()), + int64_pointer(sr.data_ptr()), + ptrs.data_ptr(), + float_pointer(nb.data_ptr()), + status.data_ptr(), + stream, + ) + if err == -2: + return None + if err: + raise RuntimeError(f"igrog: the card gridder failed with CUDA error {err}") + for o, r in zip(outs, res): + o[:, b:e] = r.cpu() + noise[b:e] = nb.cpu() + code = int(status.item()) + if code: + raise RuntimeError(f"igrog: the card gridder stopped on {_GROG_ERRORS.get(code, code)}") + return outs, noise + + def _sources(traj: torch.Tensor, shot: torch.Tensor, sample: torch.Tensor, sources: int, spacing): """Sample indices ``(M, sources)`` of each target's sources along its readout.""" S, P, _ = traj.shape @@ -142,8 +279,10 @@ def igrog( spacing: float | None = 1.0, reach: int = 4, upsample: int = 8, - batch: int = 65536, + batch: int = 1 << 20, device=None, + double: bool = False, + also=(), ): r"""Non-Cartesian samples moved to their nearest grid point by closed-form GRAPPA (iGROG). @@ -192,11 +331,19 @@ def igrog( Displacements the correlation is tabulated to, in cells. upsample : int, default=8 Table steps per cell. - batch : int, default=65536 - Kernels solved at once. + batch : int, default=1048576 + Targets per pass on a card: what it holds besides the samples. device : torch.device, default=None - Where the kernels are solved (a card is much faster); the result is on - the host. ``None`` uses the samples' device. + Where the kernels are solved; the result is on the host. ``None`` + uses the samples' device. Each target's kernel is solved in single + precision by a compiled loop: one target per thread on a card, the + targets spread over the cores on the host. The card holds the samples + and ``batch`` targets' indices and results, nothing else. + double : bool, default=False + Solve the kernels in double precision, by torch (slow, for checks). + also : sequence of torch.Tensor, default=() + More samples shaped as ``kspace``, such as a draw of pure noise, moved + by the same kernels without solving them again. Returns ------- @@ -205,7 +352,8 @@ def igrog( ``(M,)``, the ``shot`` (flat index of ``*shots``) and ``sample`` each came from, ``noise`` ``(M,)`` (each value's noise variance per coil for unit white noise on the samples) and ``weights()``, the weights of unit - mean square that whiten the values. + mean square that whiten the values; ``also`` holds the values of the + extra samples, in their order. References ---------- @@ -249,18 +397,58 @@ def igrog( raise ValueError(f"sources reach {far:.2f} cells apart; the table reaches {lim:.2f}") out = torch.empty(y.shape[0], shot.numel(), dtype=torch.complex64) + extra = [a.reshape(a.shape[0], -1, P) for a in also] + extra_out = [torch.empty(a.shape[0], shot.numel(), dtype=torch.complex64) for a in extra] noise = torch.empty(shot.numel(), dtype=torch.float32) which = torch.zeros_like(shot) if group is None else group.reshape(-1).cpu()[shot] + + def store(idx, values, gains): + for o, r in zip([out, *extra_out], values): + o[:, idx] = r + noise[idx] = gains + + groups = [(cc, torch.nonzero(which == gi).squeeze(1)) for gi, cc in enumerate(tables)] + groups = [(cc, idx) for cc, idx in groups if idx.numel()] + if not double and card.type == "cuda": + sets = [a.to(card, torch.complex64).contiguous() for a in [y, *extra]] + got = [] + for cc, idx in groups: + got.append(_card_grid(cc, lam, tr, sets, shot[idx], sample[idx], src[idx], card, batch)) + if got[-1] is None: # no card gridder in this build, or a shape it does not take + break + del sets + if all(g is not None for g in got): + for (_, idx), g in zip(groups, got): + store(idx, *g) + return _Gridded(out, lin, shot, sample, tuple(shape), noise, extra_out) + elif not double: + sets = [a.to("cpu", torch.complex64).contiguous() for a in [y, *extra]] + for cc, idx in groups: + outs = [torch.empty(a.shape[0], idx.numel(), dtype=torch.complex64) for a in sets] + gains = torch.empty(idx.numel(), dtype=torch.float32) + _native(cc, lam, tr, sets, shot[idx], sample[idx], src[idx], outs, gains) + store(idx, outs, gains) + return _Gridded(out, lin, shot, sample, tuple(shape), noise, extra_out) + # Torch: double precision, or a card this build has no gridder for. + step = min(batch, 65536) # its (B, D C, D C) systems are what bound the batch for gi, cc in enumerate(tables): - table = cc.table.to(card) + # The table as (nodes, C C) rows, so each tap gathers contiguous rows. + dtype = torch.complex128 if double else torch.complex64 + size = cc.table.shape[2] + rows = cc.table.reshape(cc.table.shape[0] ** 2, -1).T.contiguous().to(card, dtype) idx = torch.nonzero(which == gi).squeeze(1) - for b in range(0, idx.numel(), batch): - ib = idx[b : b + batch] + for b in range(0, idx.numel(), step): + ib = idx[b : b + step] s, sb = shot[ib], src[ib] - W = _kernels(cc, table, off[ib].to(card), lam) # (B, C, D C) + W = _kernels(cc, rows, size, off[ib].to(card), lam) # (B, C, D C) v = y[:, s[:, None], sb].permute(1, 2, 0).reshape(len(s), -1) # (B, D C), source-major out[:, ib] = torch.einsum("bcj,bj->cb", W, v.to(card, W.dtype)).to( "cpu", torch.complex64 ) noise[ib] = (W.abs().square().sum((1, 2)) / W.shape[1]).float().cpu() - return _Gridded(out, lin, shot, sample, tuple(shape), noise) + for a, o in zip(extra, extra_out): + v = a[:, s[:, None], sb].permute(1, 2, 0).reshape(len(s), -1) + o[:, ib] = torch.einsum("bcj,bj->cb", W, v.to(card, W.dtype)).to( + "cpu", torch.complex64 + ) + return _Gridded(out, lin, shot, sample, tuple(shape), noise, extra_out) diff --git a/src/csrc/include/bartorch.h b/src/csrc/include/bartorch.h index 59e4a1da..df5662c9 100644 --- a/src/csrc/include/bartorch.h +++ b/src/csrc/include/bartorch.h @@ -987,6 +987,43 @@ BARTORCH_API int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t int groups, const void* weight, const void* bias, const void* scale, const void* shift, float eps, int silu, int channels_last, int dtype, void* workspace, void* stream); +/* + * iGROG gridding on the host, for bartorch.tools.igrog: each target sample + * (shot[m], sample[m]) is moved to its nearest grid point by a closed-form + * Wiener GRAPPA kernel over its sources src[m, :] (samples of the same shot), + * built in single precision from the coil cross-correlation table and solved + * by Cholesky with the ridge lam * trace / (sources * coils). + * + * table (2 reach upsample + 1)^ndim x (coils, coils) complex float, + * axes in trajectory order, Q at displacement d on the table grid + * traj (shots, samples, ndim) float, grid units + * data sets pointers to (coils, shots, samples) complex float + * out sets pointers to (coils, targets) complex float + * noise (targets) float, sum |W|^2 / coils, or NULL + * + * ndim is 2 or 3 and sources x coils at most 256. Returns 0, or 1 for a + * NULL argument, 2 a bad shape, 3 out of memory, 4 an index out of range, + * 5 a Cholesky failure, 6 a displacement beyond the table. + */ +BARTORCH_API int bartorch_grog_grid(int ndim, int64_t coils, int64_t sources, int64_t reach, int64_t upsample, + float lam, const void* table, int64_t shots, int64_t samples, const float* traj, int64_t sets, + const void* const* data, int64_t targets, const int64_t* shot, const int64_t* sample, + const int64_t* src, void* const* out, float* noise); + +/* + * The same on the card, one target per thread: every array is device memory, + * data and out are device arrays of sets device pointers, and status is one + * device int the caller zeroes, raised to the codes above (4, 5 or 6) by a + * target that fails. sources x coils is at most 48. Runs on the cudaStream_t + * passed as a void* and returns 0 or the CUDA error of the launch, -2 for a + * shape it does not take; without CUDA, cuda_built() is 0 and the grid returns -1. + */ +BARTORCH_API int bartorch_grog_cuda_built(void); +BARTORCH_API int bartorch_grog_grid_cuda(int ndim, int64_t coils, int64_t sources, int64_t reach, + int64_t upsample, float lam, const void* table, int64_t shots, int64_t samples, const float* traj, + int64_t sets, const void* data, int64_t targets, const int64_t* shot, const int64_t* sample, + const int64_t* src, const void* out, float* noise, void* status, void* stream); + /* Whether a pointer is device memory; always false without CUDA. */ BARTORCH_API int bartorch_on_device(const void* ptr); diff --git a/src/csrc/ops/grid.c b/src/csrc/ops/grid.c index 472a7ac2..fff94781 100644 --- a/src/csrc/ops/grid.c +++ b/src/csrc/ops/grid.c @@ -70,7 +70,7 @@ extern void bartorch_cb_grid_forward(struct bartorch_cb_grid* p, complex float* const complex float* src, const complex float* map); extern void bartorch_cb_grid_inverse(struct bartorch_cb_grid* p, complex float* dst, complex float* volume, const complex float* bank, const complex float* map); -extern int bartorch_cuda_contract_grid(bart_dim_t L, bart_dim_t B, int R, complex float* bank, const complex float* K); +extern int bartorch_cuda_contract_grid(bart_dim_t L, bart_dim_t B, int R, complex float* bank, const float* K, int real); /* What the kernels between the gathered spectrum and a table of samples read * of the transform's layout (kernels.cu). */ @@ -110,7 +110,10 @@ struct grid_s { /* The normal through cuFFT's callbacks: the axes it transforms, the * places of a plane the pattern keeps, and the kernel at each of them - * (NULL where it is one), on the host. */ + * (NULL where it is one), on the host. The kernel is Hermitian, so it + * is kept as its upper triangle, entry (i, j), i <= j, of place l at + * `(i + j (j + 1) / 2) kept + l`: real floats where it is real (a real + * basis), else complex floats as pairs. */ bart_flags_t flags; bart_dim_t R; bart_dim_t vol; @@ -120,7 +123,8 @@ struct grid_s { bart_dim_t words; unsigned int* mask; int* prefix; - complex float* kernel; + float* kernel; + bool kreal; complex float* mod[3]; /* The same on the card, made the first time the normal is asked for @@ -129,7 +133,7 @@ struct grid_s { struct bartorch_cb_grid* cb; unsigned int* dmask; int* dprefix; - complex float* dkernel; + float* dkernel; complex float* dmod[3]; /* Whether the normal is the closed form or the two applications. */ @@ -280,10 +284,64 @@ static complex float* grid_kernel(bart_dim_t kdims[DIMS], const bart_dim_t pat_d /* What the callbacks need, on the host: the places of a plane some frame * samples, as a bit per place and a count before each word, and the kernel at - * those places alone, `kernel[(l R + r) R + c]` -- or no kernel, where every - * kept place's is one. */ + * those places alone, as the upper triangle grid_keep_kernel keeps -- or no + * kernel, where every kept place's is one. */ static void grid_mods(struct grid_s* d); +/* Entries of a Hermitian R x R kernel's upper triangle. */ +static bart_dim_t tri_count(bart_dim_t R) +{ + return R * (R + 1) / 2; +} + +/* Where entry (r, c) of place l is kept, and whether it is the conjugate of what is kept. */ +static bart_dim_t tri_at(bart_dim_t kept, bart_dim_t l, bart_dim_t r, bart_dim_t c, bool* conjugate) +{ + bart_dim_t lo = (r < c) ? r : c; + bart_dim_t hi = (r < c) ? c : r; + *conjugate = (r > c); + return (lo + hi * (hi + 1) / 2) * kept + l; +} + +/* The upper triangle `packed`, complex, as the kernel: real floats where every + * entry is real, the complex pairs otherwise. Takes `packed` over. */ +static void grid_keep_kernel(struct grid_s* d, complex float* packed, bart_dim_t count) +{ + bool real = true; + + for (bart_dim_t i = 0; (i < count) && real; i++) + real = (0.f == cimagf(packed[i])); + + d->kreal = real; + + if (!real) { + + d->kernel = (float*)packed; + return; + } + + float* r = xmalloc((size_t)MAX(1, count) * sizeof(float)); + + for (bart_dim_t i = 0; i < count; i++) + r[i] = crealf(packed[i]); + + xfree(packed); + d->kernel = r; +} + +/* Entry (r, c) of the kernel at kept place l. */ +static complex float grid_kernel_at(const struct grid_s* d, bart_dim_t l, bart_dim_t r, bart_dim_t c) +{ + if (NULL == d->kernel) + return 1.f; + + bool cj; + bart_dim_t i = tri_at(d->kept, l, r, c, &cj); + complex float v = d->kreal ? d->kernel[i] : ((const complex float*)d->kernel)[i]; + + return cj ? conjf(v) : v; +} + static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const complex float* K, const bart_dim_t pat_dims[DIMS], const complex float* pattern) { @@ -325,7 +383,7 @@ static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const d->kept = n; d->plane = places; - complex float* kernel = xmalloc((size_t)MAX(1, n * R * R) * sizeof(complex float)); + complex float* kernel = xmalloc((size_t)MAX(1, n * tri_count(R)) * sizeof(complex float)); bool ones = (1 == R); bart_dim_t l = 0; @@ -335,12 +393,13 @@ static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const if (0 == (d->mask[p >> 5] & (1u << (p & 31)))) continue; - for (bart_dim_t r = 0; r < R; r++) - for (bart_dim_t c = 0; c < R; c++) { + for (bart_dim_t c = 0; c < R; c++) + for (bart_dim_t r = 0; r <= c; r++) { complex float v = K[p + places * (r + R * c)]; + bool cj; - kernel[(l * R + r) * R + c] = v; + kernel[tri_at(n, l, r, c, &cj)] = v; ones = ones && (1.f == v); } @@ -350,10 +409,12 @@ static void grid_compress(struct grid_s* d, const bart_dim_t kdims[DIMS], const if (ones) { xfree(kernel); - kernel = NULL; - } + d->kernel = NULL; - d->kernel = kernel; + } else { + + grid_keep_kernel(d, kernel, n * tri_count(R)); + } grid_mods(d); } @@ -403,10 +464,10 @@ static bool grid_ready(struct grid_s* d, const void* ref) if (NULL != d->kernel) { - bart_dim_t kd[1] = { d->kept * d->R * d->R }; + bart_dim_t kd[1] = { d->kept * tri_count(d->R) * (d->kreal ? 1 : 2) }; - d->dkernel = md_alloc_gpu(1, kd, CFL_SIZE); - md_copy(1, kd, d->dkernel, d->kernel, CFL_SIZE); + d->dkernel = md_alloc_gpu(1, kd, sizeof(float)); + md_copy(1, kd, d->dkernel, d->kernel, sizeof(float)); } for (int a = 0; a < 3; a++) { @@ -556,7 +617,7 @@ static void grid_fused(struct grid_s* d, complex float* dst, const complex float grid_in(d, bank + r * per, volume, hybrid, src + c * coil_step + r * coeff_step, m); if (NULL != d->dkernel) - bartorch_cuda_contract_grid(d->kept, d->batch, (int)d->R, bank, d->dkernel); + bartorch_cuda_contract_grid(d->kept, d->batch, (int)d->R, bank, d->dkernel, d->kreal); for (bart_dim_t r = 0; r < d->R; r++) grid_out(d, dst + c * coil_step + r * coeff_step, volume, hybrid, bank + r * per, m); @@ -1008,7 +1069,7 @@ static void points_host_normal(const struct grid_s* d, complex float* dst, const complex float acc = 0.f; for (bart_dim_t cc = 0; cc < R; cc++) - acc += ((NULL == d->kernel) ? 1.f : d->kernel[(l * R + r) * R + cc]) * in[cc]; + acc += grid_kernel_at(d, l, r, cc) * in[cc]; v[r * C * d->vol] = acc; } @@ -1298,6 +1359,7 @@ static struct grid_s* grid_state(const bart_dim_t cim_dims[DIMS], d->mask = NULL; d->prefix = NULL; d->kernel = NULL; + d->kreal = false; d->tried = false; d->cb = NULL; d->dmask = NULL; @@ -1618,16 +1680,17 @@ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t n += __builtin_popcount(bits); } - /* kernel[(u R + r) R + c] = sum over the entries e of place u of - * |w_e|^2 conj(B[t, r]) B[t, c], t the frame of e. */ - complex float* kernel = xmalloc((size_t)MAX(1, U * R * R) * sizeof(complex float)); + /* Entry (r, c) at place u = sum over the entries e of place u of + * |w_e|^2 conj(B[t, r]) B[t, c], t the frame of e: Hermitian, so + * the upper triangle alone. */ + complex float* kernel = xmalloc((size_t)MAX(1, U * tri_count(R)) * sizeof(complex float)); bool ones = (1 == R); #pragma omp parallel for reduction(&&:ones) for (bart_dim_t u = 0; u < U; u++) { - for (bart_dim_t r = 0; r < R; r++) { - for (bart_dim_t c = 0; c < R; c++) { + for (bart_dim_t c = 0; c < R; c++) { + for (bart_dim_t r = 0; r <= c; r++) { complex double acc = 0.; @@ -1640,7 +1703,9 @@ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t acc += a * ((NULL == bt) ? 1.f : (conjf(bt[t * R + r]) * bt[t * R + c])); } - kernel[(u * R + r) * R + c] = (complex float)acc; + bool cj; + + kernel[tri_at(U, u, r, c, &cj)] = (complex float)acc; ones = ones && (1. == acc); } } @@ -1649,10 +1714,12 @@ static struct grid_s* sampled_state(const bart_dim_t cim_dims[DIMS], bart_dim_t if (ones) { xfree(kernel); - kernel = NULL; - } + d->kernel = NULL; - d->kernel = kernel; + } else { + + grid_keep_kernel(d, kernel, U * tri_count(R)); + } grid_mods(d); } diff --git a/src/csrc/ops/grog.c b/src/csrc/ops/grog.c new file mode 100644 index 00000000..d4af7e09 --- /dev/null +++ b/src/csrc/ops/grog.c @@ -0,0 +1,283 @@ +/* + * iGROG gridding with closed-form kernels on the host, for bartorch.tools.igrog. + * Per target, in single precision: + * + * Q at the D(D-1)/2 source-pair displacements and the D source offsets + * (cubic Lagrange in the table; Q(-d) = Q(d)^H gives the other half), + * R_ss + mu I by Cholesky, W = R_ts (R_ss + mu I)^-1, target = W s, + * and the target's noise gain sum |W|^2 / C. + * + * Every target is independent, so they are spread over OpenMP threads. + * Ported from fast-encoding's grog.cpp. + */ +#include +#include +#include +#include + +#include "include/bartorch.h" + +/* Largest sources x coils the per-target scratch holds. */ +#define GROG_MAX_N 256 + +static void lagrange4(float t, float w[4]) +{ + w[0] = -t * (t - 1) * (t - 2) / 6; + w[1] = (t + 1) * (t - 1) * (t - 2) / 2; + w[2] = -(t + 1) * t * (t - 2) / 2; + w[3] = (t + 1) * t * (t - 1) / 6; +} + +/* out (C, C) = Q(disp); 0, or -1 when disp leaves the table. */ +static int table_at(const complex float* table, int ndim, int64_t C, int64_t reach, int64_t upsample, + const float* disp, complex float* out) +{ + const int64_t n = 2 * reach * upsample + 1; + int64_t base[3] = { 0, 0, 0 }; + float w[3][4] = { { 1, 0, 0, 0 }, { 1, 0, 0, 0 }, { 1, 0, 0, 0 } }; + + for (int a = 0; a < ndim; a++) { + const float u = (disp[a] + (float)reach) * (float)upsample; + const float f = floorf(u); + base[a] = (int64_t)f - 1; + if (base[a] < 0 || base[a] + 3 >= n) + return -1; + lagrange4(u - f, w[a]); + } + + const int64_t CC = C * C; + for (int64_t k = 0; k < CC; k++) + out[k] = 0.; + + const int t1 = ndim > 1 ? 4 : 1; + const int t2 = ndim > 2 ? 4 : 1; + for (int i0 = 0; i0 < 4; i0++) + for (int i1 = 0; i1 < t1; i1++) + for (int i2 = 0; i2 < t2; i2++) { + int64_t flat = base[0] + i0; + float wt = w[0][i0]; + if (ndim > 1) { + flat = flat * n + base[1] + i1; + wt *= w[1][i1]; + } + if (ndim > 2) { + flat = flat * n + base[2] + i2; + wt *= w[2][i2]; + } + const complex float* q = table + flat * CC; + for (int64_t k = 0; k < CC; k++) + out[k] += wt * q[k]; + } + return 0; +} + +/* In-place Cholesky of a Hermitian positive definite (n, n), lower triangle; 0 or -1. */ +static int cholesky(complex float* A, int64_t n) +{ + for (int64_t j = 0; j < n; j++) { + float d = crealf(A[j * n + j]); + for (int64_t k = 0; k < j; k++) + d -= crealf(A[j * n + k] * conjf(A[j * n + k])); + if (!(d > 0)) + return -1; + d = sqrtf(d); + A[j * n + j] = d; + for (int64_t i = j + 1; i < n; i++) { + complex float s = A[i * n + j]; + for (int64_t k = 0; k < j; k++) + s -= A[i * n + k] * conjf(A[j * n + k]); + A[i * n + j] = s / d; + } + } + return 0; +} + +static void cholesky_solve(const complex float* L, int64_t n, complex float* b) +{ + for (int64_t i = 0; i < n; i++) { + complex float s = b[i]; + for (int64_t k = 0; k < i; k++) + s -= L[i * n + k] * b[k]; + b[i] = s / crealf(L[i * n + i]); + } + for (int64_t i = n - 1; i >= 0; i--) { + complex float s = b[i]; + for (int64_t k = i + 1; k < n; k++) + s -= conjf(L[k * n + i]) * b[k]; + b[i] = s / crealf(L[i * n + i]); + } +} + +struct grog_scratch { + complex float* A; /* (n, n) R_ss + mu I, then its Cholesky factor */ + complex float* z; /* (n) one kernel row */ + complex float* q; /* (C, C) one interpolated Q */ + complex float* rts; /* (C, n) R_ts */ + complex float* acc; /* (sets) */ + float off[GROG_MAX_N * 3]; +}; + +struct grog_table { + const complex float* q; + int ndim; + int64_t C, D, reach, upsample; +}; + +/* R_ss + mu I factored in s->A and R_ts in s->rts for the sources at s->off; 0 or an error code. */ +static int grog_system(const struct grog_table* t, float lam, struct grog_scratch* s) +{ + const int64_t C = t->C, D = t->D, n = D * C; + + /* R_ss: block (i, k) = Q(d_i - d_k), (k, i) its adjoint */ + for (int64_t i = 0; i < D; i++) + for (int64_t k = i; k < D; k++) { + float dd[3]; + for (int a = 0; a < t->ndim; a++) + dd[a] = s->off[i * 3 + a] - s->off[k * 3 + a]; + if (table_at(t->q, t->ndim, C, t->reach, t->upsample, dd, s->q)) + return 6; + for (int64_t c = 0; c < C; c++) + for (int64_t c2 = 0; c2 < C; c2++) { + s->A[(i * C + c) * n + k * C + c2] = s->q[c * C + c2]; + s->A[(k * C + c2) * n + i * C + c] = conjf(s->q[c * C + c2]); + } + } + float tr = 0; + for (int64_t i = 0; i < n; i++) + tr += crealf(s->A[i * n + i]); + const float mu = lam * tr / (float)n; + for (int64_t i = 0; i < n; i++) + s->A[i * n + i] += mu; + if (cholesky(s->A, n)) + return 5; + + /* R_ts[c, (i, c2)] = Q_{c c2}(-d_i) */ + for (int64_t i = 0; i < D; i++) { + float dd[3]; + for (int a = 0; a < t->ndim; a++) + dd[a] = -s->off[i * 3 + a]; + if (table_at(t->q, t->ndim, C, t->reach, t->upsample, dd, s->q)) + return 6; + for (int64_t c = 0; c < C; c++) + for (int64_t c2 = 0; c2 < C; c2++) + s->rts[c * n + i * C + c2] = s->q[c * C + c2]; + } + return 0; +} + +/* + * W[c, :] = (A^-1 R_ts[c, :]^H)^H; out_c = W[c, :] x for each set's sources x + * at (shot, js) in data (coils, shots, P); returns sum |W|^2. + */ +static float grog_apply(const struct grog_table* t, struct grog_scratch* s, int64_t sets, + const void* const* data, int64_t shots, int64_t P, int64_t shot, const int64_t* js, + void* const* out, int64_t targets, int64_t m) +{ + const int64_t C = t->C, D = t->D, n = D * C; + float gain = 0; + + for (int64_t c = 0; c < C; c++) { + for (int64_t k = 0; k < n; k++) + s->z[k] = conjf(s->rts[c * n + k]); + cholesky_solve(s->A, n, s->z); + for (int64_t e = 0; e < sets; e++) + s->acc[e] = 0; + for (int64_t i = 0; i < D; i++) + for (int64_t c2 = 0; c2 < C; c2++) { + const complex float w = conjf(s->z[i * C + c2]); + const int64_t at = (c2 * shots + shot) * P + js[i]; + for (int64_t e = 0; e < sets; e++) + s->acc[e] += w * ((const complex float*)data[e])[at]; + gain += crealf(w * conjf(w)); + } + for (int64_t e = 0; e < sets; e++) + ((complex float*)out[e])[c * targets + m] = s->acc[e]; + } + return gain; +} + +int bartorch_grog_grid(int ndim, int64_t coils, int64_t sources, int64_t reach, int64_t upsample, float lam, + const void* table, int64_t shots, int64_t samples, const float* traj, int64_t sets, + const void* const* data, int64_t targets, const int64_t* shot, const int64_t* sample, + const int64_t* src, void* const* out, float* noise) +{ + const int64_t C = coils, D = sources, P = samples, n = D * C; + + if (!table || !traj || !data || !shot || !sample || !src || !out || sets < 1) + return 1; + if ((ndim != 2 && ndim != 3) || C <= 0 || D <= 0 || D > P || n > GROG_MAX_N) + return 2; + + const struct grog_table t = { table, ndim, C, D, reach, upsample }; + int bad = 0; + +#pragma omp parallel + { + struct grog_scratch s; + s.A = malloc(sizeof(complex float) * (size_t)(n * n)); + s.z = malloc(sizeof(complex float) * (size_t)n); + s.q = malloc(sizeof(complex float) * (size_t)(C * C)); + s.rts = malloc(sizeof(complex float) * (size_t)(C * n)); + s.acc = malloc(sizeof(complex float) * (size_t)sets); + + if (!s.A || !s.z || !s.q || !s.rts || !s.acc) { +#pragma omp atomic write + bad = 3; + } + +#pragma omp for schedule(dynamic, 256) + for (int64_t m = 0; m < targets; m++) { + int err; +#pragma omp atomic read + err = bad; + if (err) + continue; + + const int64_t sh = shot[m], j = sample[m]; + const int64_t* js = src + m * D; + err = (sh < 0 || sh >= shots || j < 0 || j >= P) ? 4 : 0; + for (int64_t i = 0; i < D && !err; i++) + err = (js[i] < 0 || js[i] >= P) ? 4 : 0; + + if (!err) { + const float* tj = traj + (sh * P + j) * ndim; + for (int64_t i = 0; i < D; i++) + for (int a = 0; a < ndim; a++) + s.off[i * 3 + a] = traj[(sh * P + js[i]) * ndim + a] - nearbyintf(tj[a]); + err = grog_system(&t, lam, &s); + } + if (err) { +#pragma omp atomic write + bad = err; + continue; + } + const float gain = grog_apply(&t, &s, sets, data, shots, P, sh, js, out, targets, m); + if (noise) + noise[m] = gain / (float)C; + } + free(s.A); + free(s.z); + free(s.q); + free(s.rts); + free(s.acc); + } + return bad; +} + +#ifndef USE_CUDA + +/* Without CUDA the card gridder (grog.cu) is absent. */ +int bartorch_grog_cuda_built(void) { return 0; } + +int bartorch_grog_grid_cuda(int ndim, int64_t coils, int64_t sources, int64_t reach, int64_t upsample, float lam, + const void* table, int64_t shots, int64_t samples, const float* traj, int64_t sets, const void* data, + int64_t targets, const int64_t* shot, const int64_t* sample, const int64_t* src, const void* out, + float* noise, void* status, void* stream) +{ + (void)ndim; (void)coils; (void)sources; (void)reach; (void)upsample; (void)lam; (void)table; (void)shots; + (void)samples; (void)traj; (void)sets; (void)data; (void)targets; (void)shot; (void)sample; (void)src; + (void)out; (void)noise; (void)status; (void)stream; + return -1; +} + +#endif diff --git a/src/csrc/ops/grog.cu b/src/csrc/ops/grog.cu new file mode 100644 index 00000000..b6d5889f --- /dev/null +++ b/src/csrc/ops/grog.cu @@ -0,0 +1,326 @@ +/* + * iGROG gridding on the card, for bartorch.tools.igrog: the host loop of + * grog.c with one target per warp. The warp's 32 lanes interpolate Q from + * the table into the packed lower triangle of R_ss + mu I in shared memory + * (Q(-d) = Q(d)^H, so only the blocks Q(d_r - d_s), r >= s; a lane per + * coil pair, so the table rows are read coalesced), factor it by Cholesky + * (a lane per row), and solve W a row at a time (a lane per entry), + * interpolating that row of R_ts as it goes; the row is then applied to every + * set by a warp reduction. Single precision, no workspace: the card holds + * the inputs and the results, and each warp n (n + 1) / 2 + n complex floats + * and 3 sources floats of shared memory (about 1.6 KB at 18). + * + * An earlier version kept a target per thread with the system in local + * memory; at 1.8 KB a thread that streamed through L2 and ran at about + * 1.4 M targets a second on an A40. + */ +#include +#include + +#include "include/bartorch.h" + +namespace { + +constexpr int WARPS = 4; /* per block */ +constexpr unsigned FULL = 0xffffffffu; + +struct cf { + float x, y; +}; + +__device__ __forceinline__ cf operator+(cf a, cf b) { return { a.x + b.x, a.y + b.y }; } +__device__ __forceinline__ cf operator-(cf a, cf b) { return { a.x - b.x, a.y - b.y }; } +__device__ __forceinline__ cf operator*(cf a, cf b) { return { a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x }; } +__device__ __forceinline__ cf operator*(float s, cf a) { return { s * a.x, s * a.y }; } +__device__ __forceinline__ cf conj(cf a) { return { a.x, -a.y }; } +/* a * conj(b) */ +__device__ __forceinline__ cf mulc(cf a, cf b) { return { a.x * b.x + a.y * b.y, a.y * b.x - a.x * b.y }; } + +__device__ __forceinline__ float warp_sum(float v) +{ + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(FULL, v, o); + return v; +} + +__device__ __forceinline__ cf warp_sum(cf v) +{ + return { warp_sum(v.x), warp_sum(v.y) }; +} + +struct params { + int ndim; + int C, D, reach, upsample; + float lam; + const cf* table; + int64_t shots, P; + const float* traj; + int sets; + const cf* const* data; + int64_t targets; + const int64_t* shot; + const int64_t* sample; + const int64_t* src; + cf* const* out; + float* noise; + int* status; +}; + +__device__ void lagrange4(float t, float w[4]) +{ + w[0] = -t * (t - 1) * (t - 2) / 6; + w[1] = (t + 1) * (t - 1) * (t - 2) / 2; + w[2] = -(t + 1) * t * (t - 2) / 2; + w[3] = (t + 1) * t * (t - 1) / 6; +} + +/* Where disp falls in the table: the first node and the weights along each axis; false outside it. */ +struct taps { + int64_t base[3]; + float w[3][4]; +}; + +__device__ bool locate(const params& p, const float* disp, taps& t) +{ + const int64_t n = 2 * (int64_t)p.reach * p.upsample + 1; + for (int a = 0; a < 3; a++) { + t.base[a] = 0; + t.w[a][0] = 1; + t.w[a][1] = t.w[a][2] = t.w[a][3] = 0; + } + for (int a = 0; a < p.ndim; a++) { + const float u = (disp[a] + (float)p.reach) * (float)p.upsample; + const float f = floorf(u); + t.base[a] = (int64_t)f - 1; + if (t.base[a] < 0 || t.base[a] + 3 >= n) + return false; + lagrange4(u - f, t.w[a]); + } + return true; +} + +/* Q_e(disp), entry e = c C + c2 of the (C, C) matrix at the taps t. */ +__device__ cf interpolate(const params& p, const taps& t, int e) +{ + const int64_t n = 2 * (int64_t)p.reach * p.upsample + 1; + const int64_t CC = (int64_t)p.C * p.C; + const int t1 = p.ndim > 1 ? 4 : 1, t2 = p.ndim > 2 ? 4 : 1; + cf acc = { 0.f, 0.f }; + for (int i0 = 0; i0 < 4; i0++) + for (int i1 = 0; i1 < t1; i1++) + for (int i2 = 0; i2 < t2; i2++) { + int64_t flat = t.base[0] + i0; + float wt = t.w[0][i0]; + if (p.ndim > 1) { + flat = flat * n + t.base[1] + i1; + wt *= t.w[1][i1]; + } + if (p.ndim > 2) { + flat = flat * n + t.base[2] + i2; + wt *= t.w[2][i2]; + } + acc = acc + wt * p.table[flat * CC + e]; + } + return acc; +} + +/* Packed lower triangle: entry (i, k), k <= i. */ +__device__ __forceinline__ int tri(int i, int k) { return i * (i + 1) / 2 + k; } + +/* Every lane of a warp calls it with the same values, so every branch below is uniform across the warp. */ +__device__ void fail(int lane, int* status, int code) +{ + if (lane == 0) + atomicMax(status, code); +} + +__global__ void __launch_bounds__(WARPS * 32) grog_kernel(params p) +{ + extern __shared__ float2 shared[]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int64_t m = (int64_t)blockIdx.x * WARPS + warp; + if (m >= p.targets) + return; + + const int C = p.C, D = p.D, n = C * D, nd = p.ndim; + const int per_warp = n * (n + 1) / 2 + n + (3 * D + 1) / 2; + cf* A = (cf*)(shared + (size_t)warp * per_warp); /* packed lower triangle */ + cf* z = A + n * (n + 1) / 2; /* one row of W */ + float* off = (float*)(z + n); /* (D, 3) source offsets */ + + const int64_t s = p.shot[m], j = p.sample[m]; + const int64_t* js = p.src + m * D; + bool bad = s < 0 || s >= p.shots || j < 0 || j >= p.P; + for (int i = 0; i < D && !bad; i++) + bad = js[i] < 0 || js[i] >= p.P; + if (bad) + return fail(lane, p.status, 4); + + const float* tj = p.traj + (s * p.P + j) * nd; + for (int e = lane; e < D * nd; e += 32) { + const int i = e / nd, a = e % nd; + off[i * 3 + a] = p.traj[(s * p.P + js[i]) * nd + a] - nearbyintf(tj[a]); + } + __syncwarp(); + + /* lower block (r, k), r >= k: Q(d_r - d_k); a lane per coil pair */ + for (int r = 0; r < D; r++) + for (int k = 0; k <= r; k++) { + float dd[3] = { 0, 0, 0 }; + for (int a = 0; a < nd; a++) + dd[a] = off[r * 3 + a] - off[k * 3 + a]; + taps t; + if (!locate(p, dd, t)) + return fail(lane, p.status, 6); + for (int e = lane; e < C * C; e += 32) { + const int i = r * C + e / C, l = k * C + e % C; + if (l <= i) + A[tri(i, l)] = interpolate(p, t, e); + } + } + /* every -d_k must be in the table too, checked here while the warp is uniform */ + for (int k = 0; k < D; k++) { + float dd[3] = { 0, 0, 0 }; + for (int a = 0; a < nd; a++) + dd[a] = -off[k * 3 + a]; + taps t; + if (!locate(p, dd, t)) + return fail(lane, p.status, 6); + } + __syncwarp(); + + float tr = 0; + for (int i = lane; i < n; i += 32) + tr += A[tri(i, i)].x; + const float mu = p.lam * warp_sum(tr) / (float)n; + for (int i = lane; i < n; i += 32) + A[tri(i, i)].x += mu; + __syncwarp(); + + /* Cholesky, left-looking: column c, a lane per row below it */ + for (int c = 0; c < n; c++) { + float part = 0; + for (int k = lane; k < c; k += 32) + part += A[tri(c, k)].x * A[tri(c, k)].x + A[tri(c, k)].y * A[tri(c, k)].y; + const float d2 = A[tri(c, c)].x - warp_sum(part); + if (!(d2 > 0)) + return fail(lane, p.status, 5); + const float d = sqrtf(d2), inv = 1.f / d; + for (int i = c + 1 + lane; i < n; i += 32) { + cf acc = A[tri(i, c)]; + for (int k = 0; k < c; k++) + acc = acc - mulc(A[tri(i, k)], A[tri(c, k)]); + A[tri(i, c)] = inv * acc; + } + __syncwarp(); + if (lane == 0) + A[tri(c, c)] = { d, 0.f }; + __syncwarp(); + } + + /* W[c, :] = (A^-1 R_ts[c, :]^H)^H, R_ts[c, (k, c2)] = Q_{c c2}(-d_k); out_c = W[c, :] x per set */ + float gain = 0; + for (int c = 0; c < C; c++) { + for (int e = lane; e < n; e += 32) { + const int k = e / C, c2 = e % C; + float dd[3] = { 0, 0, 0 }; + for (int a = 0; a < nd; a++) + dd[a] = -off[k * 3 + a]; + taps t; + locate(p, dd, t); /* inside, checked above */ + z[e] = conj(interpolate(p, t, c * C + c2)); + } + __syncwarp(); + /* L y = z, a column at a time: a lane per row below it */ + for (int i = 0; i < n; i++) { + const cf zi = (1.f / A[tri(i, i)].x) * z[i]; + __syncwarp(); + if (lane == 0) + z[i] = zi; + for (int r = i + 1 + lane; r < n; r += 32) + z[r] = z[r] - A[tri(r, i)] * zi; + __syncwarp(); + } + /* L^H w = y, from the bottom */ + for (int i = n - 1; i >= 0; i--) { + const cf zi = (1.f / A[tri(i, i)].x) * z[i]; + __syncwarp(); + if (lane == 0) + z[i] = zi; + for (int r = lane; r < i; r += 32) + z[r] = z[r] - conj(A[tri(i, r)]) * zi; + __syncwarp(); + } + float g = 0; + for (int e = lane; e < n; e += 32) + g += z[e].x * z[e].x + z[e].y * z[e].y; + gain += warp_sum(g); + for (int e = 0; e < p.sets; e++) { + const cf* x = p.data[e]; + cf acc = { 0.f, 0.f }; + for (int q = lane; q < n; q += 32) { + const int i = q / C, c2 = q % C; + acc = acc + mulc(x[((int64_t)c2 * p.shots + s) * p.P + js[i]], z[q]); + } + acc = warp_sum(acc); + if (lane == 0) + p.out[e][(int64_t)c * p.targets + m] = acc; + } + __syncwarp(); + } + if (p.noise && lane == 0) + p.noise[m] = gain / (float)C; +} + +} // namespace + +extern "C" { + +int bartorch_grog_cuda_built(void) +{ + return 1; +} + +int bartorch_grog_grid_cuda(int ndim, int64_t coils, int64_t sources, int64_t reach, int64_t upsample, float lam, + const void* table, int64_t shots, int64_t samples, const float* traj, int64_t sets, const void* data, + int64_t targets, const int64_t* shot, const int64_t* sample, const int64_t* src, const void* out, + float* noise, void* status, void* stream) +{ + const int64_t n = coils * sources; + if ((ndim != 2 && ndim != 3) || coils < 1 || sources < 1 || sources > samples || sets < 1) + return -2; + if (n > 48 || (targets + WARPS - 1) / WARPS > 2147483647LL) + return -2; + if (targets < 1) + return 0; + + params p; + p.ndim = ndim; + p.C = (int)coils; + p.D = (int)sources; + p.reach = (int)reach; + p.upsample = (int)upsample; + p.lam = lam; + p.table = (const cf*)table; + p.shots = shots; + p.P = samples; + p.traj = traj; + p.sets = (int)sets; + p.data = (const cf* const*)data; + p.targets = targets; + p.shot = shot; + p.sample = sample; + p.src = src; + p.out = (cf* const*)out; + p.noise = noise; + p.status = (int*)status; + + /* per warp: the packed triangle, one row of W, and the offsets, in float2 units */ + const size_t per_warp = (size_t)(n * (n + 1) / 2 + n + (3 * sources + 1) / 2); + const size_t bytes = per_warp * WARPS * sizeof(float2); + const int64_t blocks = (targets + WARPS - 1) / WARPS; + grog_kernel<<<(unsigned)blocks, WARPS * 32, bytes, (cudaStream_t)stream>>>(p); + return (int)cudaGetLastError(); +} + +} // extern "C" diff --git a/src/csrc/ops/kernels.cu b/src/csrc/ops/kernels.cu index 09d777aa..e525cd75 100644 --- a/src/csrc/ops/kernels.cu +++ b/src/csrc/ops/kernels.cu @@ -274,14 +274,17 @@ extern "C" int bartorch_cuda_contract_upper_real(bart_dim_t L, int R, _Complex f /* A coil's gathered spectrum under a Cartesian normal, laid out with the * coefficients slowest, the batch next and the kept places fastest, * multiplied at each place by that place's coefficients-by-coefficients - * kernel and written back over itself. Entry (r, c) of the kernel at place l - * is `K[(l R + r) R + c]`: the kernel varies over the transformed plane only, - * so every batch reads the same one. */ -__global__ static void kern_contract_grid(bart_dim_t L, bart_dim_t B, int R, cuFloatComplex* bank, const cuFloatComplex* K) + * kernel and written back over itself. The kernel is Hermitian and kept as + * its upper triangle: entry (i, j), i <= j, at place l is at + * `(i + j (j + 1) / 2) L + l`, real floats with `real`, complex pairs + * otherwise; the kernel varies over the transformed plane only, so every + * batch reads the same one. */ +__global__ static void kern_contract_grid(bart_dim_t L, bart_dim_t B, int R, cuFloatComplex* bank, const float* K, int real) { bart_dim_t start = threadIdx.x + (bart_dim_t)blockDim.x * blockIdx.x; bart_stride_t stride = (bart_dim_t)blockDim.x * gridDim.x; bart_dim_t n = L * B; + const cuFloatComplex* Kc = (const cuFloatComplex*)K; for (bart_dim_t i = start; i < n; i += stride) { @@ -296,20 +299,39 @@ __global__ static void kern_contract_grid(bart_dim_t L, bart_dim_t B, int R, cuF cuFloatComplex acc = make_cuFloatComplex(0.f, 0.f); - for (int c = 0; c < R; c++) - acc = cuCaddf(acc, cuCmulf(K[(l * R + r) * R + c], in[c])); + for (int c = 0; c < R; c++) { + + int lo = (r < c) ? r : c; + int hi = (r < c) ? c : r; + bart_dim_t at = (bart_dim_t)(lo + hi * (hi + 1) / 2) * L + l; + cuFloatComplex k; + + if (real) { + + k = make_cuFloatComplex(K[at], 0.f); + + } else { + + k = Kc[at]; + + if (r > c) + k = cuConjf(k); + } + + acc = cuCaddf(acc, cuCmulf(k, in[c])); + } bank[r * n + i] = acc; } } } -extern "C" int bartorch_cuda_contract_grid(bart_dim_t L, bart_dim_t B, int R, _Complex float* bank, const _Complex float* K) +extern "C" int bartorch_cuda_contract_grid(bart_dim_t L, bart_dim_t B, int R, _Complex float* bank, const float* K, int real) { if ((R < 1) || (R > CONTRACT_MAX)) return -1; - kern_contract_grid<<>>(L, B, R, (cuFloatComplex*)bank, (const cuFloatComplex*)K); + kern_contract_grid<<>>(L, B, R, (cuFloatComplex*)bank, K, real); CUDA_KERNEL_ERROR; diff --git a/tests/test_cuda.py b/tests/test_cuda.py index 7169f0c2..69dd9c3f 100644 --- a/tests/test_cuda.py +++ b/tests/test_cuda.py @@ -33,6 +33,11 @@ def test_free_memory_reports_minus_one_without_a_device(): assert bartorch._cuda.free_memory() == -1 +def test_emptying_barts_device_cache_is_safe_with_or_without_a_device(): + bartorch.empty_cuda_memcache() # a no-op without CUDA, a release with it + bartorch.empty_cuda_memcache() + + def test_a_device_tensor_is_refused_with_a_message_naming_the_remedy(): if bartorch._cuda.available() or not torch.cuda.is_available(): pytest.skip("a device is usable here, so the refusal does not apply") diff --git a/tests/test_igrog.py b/tests/test_igrog.py index 9da44d21..853a39f3 100644 --- a/tests/test_igrog.py +++ b/tests/test_igrog.py @@ -75,3 +75,72 @@ def test_groups_pick_their_own_calibration(): assert torch.allclose( one.values, two.values, rtol=1e-3, atol=1e-4 * float(one.values.abs().max()) ) + + +def test_single_precision_matches_double(): + """The host loop solves the kernels in single precision; the gridded samples barely move.""" + n = 32 + _, kspace = _object(n) + angles = torch.linspace(0, math.pi, 24, dtype=torch.float64)[:-1] + 0.1 + s = torch.arange(-10, 10.5, 0.5, dtype=torch.float64) + traj = torch.stack( + [s[None] * torch.cos(angles)[:, None], s[None] * torch.sin(angles)[:, None]], -1 + ) + y = kspace(traj).to(torch.complex64) + c = torch.arange(-8, 8, dtype=torch.float64) + ky, kx = torch.meshgrid(c, c, indexing="ij") + calib = kspace(torch.stack([kx, ky], -1)).to(torch.complex64) + double = tools.igrog(y, traj.float(), (n, n), calib, double=True) + single = tools.igrog(y, traj.float(), (n, n), calib, double=False) + assert torch.equal(double.index, single.index) + rel = (single.values - double.values).norm() / double.values.norm() + assert rel < 1e-3, rel + assert torch.allclose(single.noise, double.noise, rtol=1e-2) + + +def test_also_moves_extra_samples_by_the_same_kernels(): + n = 32 + _, kspace = _object(n) + angles = torch.tensor([0.3, 1.1, 2.0, 2.7], dtype=torch.float64) + s = torch.arange(-10, 10.5, 0.5, dtype=torch.float64) + traj = torch.stack( + [s[None] * torch.cos(angles)[:, None], s[None] * torch.sin(angles)[:, None]], -1 + ) + y = kspace(traj).to(torch.complex64) + c = torch.arange(-8, 8, dtype=torch.float64) + ky, kx = torch.meshgrid(c, c, indexing="ij") + calib = kspace(torch.stack([kx, ky], -1)).to(torch.complex64) + noise = torch.randn(y.shape, dtype=torch.complex64, generator=torch.Generator().manual_seed(1)) + both = tools.igrog(y, traj.float(), (n, n), calib, also=[noise]) + alone = tools.igrog(noise, traj.float(), (n, n), calib) + assert len(both.also) == 1 + # Same kernels (they come from the calibration and the trajectory only). + torch.testing.assert_close(both.also[0], alone.values) + + +def test_the_card_gridder_matches_the_host_loop(): + import pytest + + from bartorch._lib import library + + if not torch.cuda.is_available() or not library().bartorch_grog_cuda_built(): + pytest.skip("no card, or a build without CUDA") + n = 32 + _, kspace = _object(n) + angles = torch.tensor([0.3, 1.1, 2.0, 2.7], dtype=torch.float64) + s = torch.arange(-10, 10.5, 0.5, dtype=torch.float64) + traj = torch.stack( + [s[None] * torch.cos(angles)[:, None], s[None] * torch.sin(angles)[:, None]], -1 + ) + y = kspace(traj).to(torch.complex64) + c = torch.arange(-8, 8, dtype=torch.float64) + ky, kx = torch.meshgrid(c, c, indexing="ij") + calib = kspace(torch.stack([kx, ky], -1)).to(torch.complex64) + noise = torch.randn(y.shape, dtype=torch.complex64, generator=torch.Generator().manual_seed(1)) + host = tools.igrog(y, traj.float(), (n, n), calib, also=[noise]) + card = tools.igrog(y, traj.float(), (n, n), calib, also=[noise], device="cuda") + assert torch.equal(host.index, card.index) + # Both solve in single precision; white noise is where it shows (about 2e-4 from double). + for a, b, tol in [(host.values, card.values, 1e-4), (host.also[0], card.also[0], 1e-3)]: + assert (a - b).norm() / a.norm() < tol + torch.testing.assert_close(card.noise, host.noise, rtol=1e-3, atol=0) diff --git a/tests/test_linop_mri.py b/tests/test_linop_mri.py index 5b58cb77..06aa629e 100644 --- a/tests/test_linop_mri.py +++ b/tests/test_linop_mri.py @@ -959,6 +959,22 @@ def test_points_are_the_dense_spectrum_at_the_positions_and_the_normal_is_the_de assert (got - want).abs().max() / want.abs().max() < 1e-5 +def test_points_normal_with_a_real_basis_keeps_a_real_kernel_and_is_the_two_applications(): + """A real basis makes the kernel real symmetric, kept as R (R + 1) / 2 floats a place.""" + coils, z, y, x, frames, coeffs, shots = 3, 4, 6, 5, 7, 3, 30 + torch.manual_seed(23) + maps, image = _rand(coils, z, y, x), _rand(coeffs, z, y, x) + b = torch.randn(coeffs, frames).to(torch.complex64) + positions, weights = _points(frames, shots, (z, y, x), 24) + A = linop.CartesianSense(maps, (coeffs, z, y, x), positions=positions, weights=weights, basis=b) + slow = linop.CartesianSense( + maps, (coeffs, z, y, x), positions=positions, weights=weights, basis=b, toeplitz=False + ) + want = slow.adjoint(slow(image)) + got = A.normal(image) + assert (got - want).abs().max() / want.abs().max() < 1e-5 + + def test_points_in_2d_without_a_basis_or_weights(maps): torch.manual_seed(23) positions = torch.stack([torch.randint(0, Y, (12,)), torch.randint(0, X, (12,))], -1)