tools.igrog compiled gridders; packed grid-normal kernel; empty_cuda_memcache

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

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
 

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

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.

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.
 

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.