Driver mempool

AGENTS.md

@@ -164,7 +164,9 @@ subpackages `src/cuda/` and `src/conveyor/` are optional and Python 3 only.
   and abandoned, in comments, commits and pull requests.
 - Stage files by name, one at a time. Never `git add -A` or `git add .`: trees here
   hold large untracked data and test output.
-- Commits end with `Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>`.
+- Commits end with a `Co-Authored-By:` line naming the model that actually wrote the change (for
+  example `Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>`), and say in the message when a
+  change was drafted by another worker.
 - Do not push or open pull requests unless asked; leave merging to a person.
 - Never modify a checkout you did not create (`bslz4_to_sparse`, `_llm`, `_dirty`
   are the owner's working copies).

src/cuda/_backend_drv.py

@@ -5,6 +5,13 @@
 Uses the device's primary context (shared with cupy, torch, jax).  A kernel runs on the calling
 thread's current stream (Stream.use()), as in cupy.  Every device-to-host copy (.get) waits for its
 stream before returning; .set queues on the stream given; zeros/asarray wait for their copy.
+
+Device memory comes from a memory pool of our own per device (stream-ordered allocator, CUDA 11.2
+drivers on), as cupy's pool: a free returns the block to the pool, without the device-wide wait of
+cuMemFree.  Blocks are taken and given back on the legacy default stream, so an array must not be
+dropped while work queued on another stream still uses it (every array the GPU classes hand out is
+complete).  free_pool() gives the unused blocks back to the driver.  Older drivers, and devices
+without pools: cuMemAlloc / cuMemFree.
 """
 
 import ctypes
@@ -64,6 +71,16 @@ def _sig(name, *args):
 )
 cuMemAlloc = _sig("cuMemAlloc_v2", ctypes.POINTER(ctypes.c_uint64), _SZ)
 cuMemFree = _sig("cuMemFree_v2", ctypes.c_uint64)
+try:  # stream-ordered allocator: drivers from CUDA 11.2
+    cuMemPoolCreate = _sig("cuMemPoolCreate", ctypes.POINTER(_P), _P)
+    cuMemPoolSetAttribute = _sig("cuMemPoolSetAttribute", _P, ctypes.c_int, _P)
+    cuMemPoolTrimTo = _sig("cuMemPoolTrimTo", _P, _SZ)
+    cuMemAllocFromPoolAsync = _sig(
+        "cuMemAllocFromPoolAsync", ctypes.POINTER(ctypes.c_uint64), _SZ, _P, _P
+    )
+    cuMemFreeAsync = _sig("cuMemFreeAsync", ctypes.c_uint64, _P)
+except AttributeError:
+    cuMemPoolCreate = None
 cuMemHostAlloc = _sig("cuMemHostAlloc", ctypes.POINTER(_P), _SZ, ctypes.c_uint)
 cuMemFreeHost = _sig("cuMemFreeHost", _P)
 cuMemcpyHtoDAsync = _sig("cuMemcpyHtoDAsync_v2", ctypes.c_uint64, _P, _SZ, _P)
@@ -79,6 +96,8 @@ def _sig(name, *args):
 CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X = 5
 CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL = 9
 CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE = 38
+CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED = 115
+CU_MEMPOOL_ATTR_RELEASE_THRESHOLD = 4
 CU_STREAM_NON_BLOCKING = 1
 CU_MEMHOSTALLOC_PORTABLE = 1
 
@@ -101,6 +120,13 @@ def _check(rc):
 _tls = threading.local()  # .stream: the current stream of this thread
 _ctx = {}  # device -> primary context
 _ctx_lock = threading.Lock()
+_pools = {}  # device -> our memory pool, or None (cuMemAlloc)
+
+
+class _CUmemPoolProps(ctypes.Structure):  # 88 bytes; the fields after location stay zero
+    _fields_ = [("allocType", ctypes.c_int), ("handleTypes", ctypes.c_int),
+                ("locationType", ctypes.c_int), ("locationId", ctypes.c_int),
+                ("win32SecurityAttributes", _P), ("reserved", ctypes.c_ubyte * 64)]
 
 
 def _context(dev):
@@ -117,6 +143,34 @@ def _context(dev):
         return _ctx[dev]
 
 
+def _pool(dev):
+    """Our memory pool on `dev` (made on first use, keeps freed blocks), or None."""
+    with _ctx_lock:
+        if dev not in _pools:
+            pool, ok = None, ctypes.c_int()
+            if cuMemPoolCreate is not None and cuDeviceGetAttribute(
+                ctypes.byref(ok), CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, dev
+            ) == 0 and ok.value:
+                props, pool = _CUmemPoolProps(), _P()
+                props.allocType, props.locationType, props.locationId = 1, 1, dev  # pinned, device
+                _check(cuMemPoolCreate(ctypes.byref(pool), ctypes.byref(props)))
+                keep = ctypes.c_uint64(2**64 - 1)  # never give blocks back at a synchronize
+                _check(cuMemPoolSetAttribute(pool, CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
+                                             ctypes.byref(keep)))
+            _pools[dev] = pool
+        return _pools[dev]
+
+
+def free_pool():
+    """Give the blocks no array uses back to the driver, on every device (as cupy's
+    free_all_blocks)."""
+    for dev, pool in list(_pools.items()):
+        if pool is not None:
+            with device(dev):
+                _check(cuCtxSynchronize())
+                _check(cuMemPoolTrimTo(pool, 0))
+
+
 def device_id():
     """The current device; with no current context, device 0's is made current (as the CUDA
     runtime does)."""
@@ -228,17 +282,24 @@ def __del__(self):
 
 
 class _Alloc(object):
-    """One cuMemAlloc, freed with its last view."""
+    """One block of device memory (from our pool, else cuMemAlloc), freed with its last view."""
 
     def __init__(self, nbytes):
         self.dev = device_id()
         self.ctx, p = _context(self.dev), ctypes.c_uint64()
-        _check(cuMemAlloc(ctypes.byref(p), max(int(nbytes), 1)))
+        self.pool = _pool(self.dev)
+        if self.pool is None:
+            _check(cuMemAlloc(ctypes.byref(p), max(int(nbytes), 1)))
+        else:
+            _check(cuMemAllocFromPoolAsync(ctypes.byref(p), max(int(nbytes), 1), self.pool, None))
         self.ptr = p.value
 
     def __del__(self):
         if getattr(self, "ptr", None) and _in_context is not None:  # None: interpreter shutdown
-            _in_context(self.ctx, cuMemFree, self.ptr)
+            if self.pool is None:
+                _in_context(self.ctx, cuMemFree, self.ptr)
+            else:
+                _in_context(self.ctx, cuMemFreeAsync, self.ptr, None)
 
 
 class _Borrowed(object):

src/cuda/bslz4_kernels.json

@@ -6,18 +6,22 @@
    "bslz4_split.cu"
   ],
   "nvrtc": {
+   "100": "12.9",
+   "103": "12.9",
+   "120": "12.9",
+   "121": "12.9",
    "20": "8.0",
    "35": "11.8",
-   "50": "12.6",
-   "52": "12.6",
-   "60": "12.6",
-   "61": "12.6",
-   "70": "12.6",
-   "75": "12.6",
-   "80": "12.6",
-   "86": "12.6",
-   "89": "12.6",
-   "90": "12.6"
+   "50": "12.9",
+   "52": "12.9",
+   "60": "12.9",
+   "61": "12.9",
+   "70": "12.9",
+   "75": "12.9",
+   "80": "12.9",
+   "86": "12.9",
+   "89": "12.9",
+   "90": "12.9"
   },
   "source_sha256": "043f94d0ff6c62357206e7c6b85923388afe590d930b72b9b2486b917e155bcc"
  },

test/test_cuda.py

@@ -472,6 +472,35 @@ def test_driver_backend_arrays_are_shared_without_a_copy():
         a.__dlpack__(dl_device=(2, 99))
 
 
+def test_driver_backend_reuses_freed_memory():
+    """Issue 43: a freed array's block goes back to the pool (no cuMemFree, which waits for the
+    device) and the next array of its size takes it; without a pool, cuMemAlloc."""
+    import gc
+
+    from bslz4_to_sparse.cuda import _backend_drv as drv
+
+    dev = drv.device_id()
+    host = np.arange(1 << 20, dtype=np.uint32)
+    pool = drv._pool(dev)
+    if pool is None:
+        pytest.skip("this driver or device has no memory pools")
+    a = drv.asarray(host)
+    p = a.ptr
+    del a
+    gc.collect()
+    b = drv.asarray(host[::-1])
+    assert b.ptr == p and (b.get() == host[::-1]).all()
+    del b
+    drv.free_pool()
+    drv._pools[dev] = None  # as on a driver before CUDA 11.2
+    try:
+        c = drv.asarray(host)
+        assert c._alloc.pool is None and (c.get() == host).all()
+        del c
+    finally:
+        drv._pools[dev] = pool
+
+
 def _device_of(a):
     ptr = a.data.ptr if hasattr(a, "data") and hasattr(a.data, "ptr") else getattr(a, "ptr", None)
     return None if not ptr else cp.cuda.runtime.pointerGetAttributes(ptr).device

tools/build_cuda_fatbin.py

@@ -5,16 +5,17 @@
     python3 tools/build_cuda_fatbin.py --nvrtc .../libnvrtc.so.12 --fatbinary .../bin/fatbinary \
         --out DIR
 
-writes DIR/bslz4_kernels.fatbin (a cubin for every architecture in --arch, and PTX for the last,
-which the driver compiles for newer GPUs) and DIR/bslz4_kernels.json (kernel name as _gpu.py asks
-for it, e.g. "count_above<unsigned short>", -> mangled symbol).  Every template instantiation
+writes DIR/bslz4_kernels.fatbin (a cubin for every architecture in --arch, and PTX for --ptx,
+which the driver compiles for GPUs with no cubin: newer ones, and Thor) and
+DIR/bslz4_kernels.json (kernel name as _gpu.py asks for it, e.g. "count_above<unsigned short>",
+-> mangled symbol).  Every template instantiation
 _gpu.py can ask for is built: each pixel dtype, accumulator and frames-per-read.
 
 The compiler is NVRTC with the options cupy gives it, so the machine code is the one the cupy
 backend runs (nvcc's front end gives different PTX for the same source, and ptxas a different
 schedule: 1.5% slower decoding measured on a GTX TITAN X).  libnvrtc: pip nvidia-cuda-nvrtc-cu12;
-fatbinary: the CUDA toolkit's (redist cuda_nvcc).  CUDA 12.x is the last toolkit for compute
-capability below 7.5 (Maxwell, Pascal, Volta).
+fatbinary: the CUDA toolkit's (redist cuda_nvcc).  CUDA 12.9 is the only toolkit for both
+Maxwell (5.x) and Blackwell (10.x, 12.x): CUDA 13 drops compute capability below 7.5.
 
 --nvrtc-for 35=.../cuda-11.8/lib64/libnvrtc.so --nvrtc-for 20=.../cuda-8.0/lib64/libnvrtc.so adds
 cubins, each from the toolkit given, for GPUs CUDA 12 dropped: Kepler (CUDA 11.8 is the last for
@@ -28,7 +29,7 @@
 is not on cvmfs: unpacked from NVIDIA's cuda_8.0.61_375.26_linux.run, md5 33e1bd98...), in a
 shell with no other module loaded (after jupyter-slurm, NVRTC fails: "failed to load builtins"):
 
-    module load cuda/12.6.2
+    module load cuda/12.9.1
     C118=/cvmfs/hpc.esrf.fr/software/packages/ubuntu24.04/x86_64/cuda/11.8.0
     C80=/data/id11/jon/testcrunchGPU/toolkits/cuda-8.0
     LD_LIBRARY_PATH=$C80/lib64:$C118/lib64:$CUDA_HOME/lib64 python3 tools/build_cuda_fatbin.py \
@@ -51,7 +52,8 @@
 FRAMES = [1, 2]  # _gpu.FRAMES_PER_READ and 1
 FILES = ["nvcomp_lz4.cuh", "shuffle_warp.cuh", "bslz4_split.cu"]  # concatenated as by _gpu.py
 PLAIN = ["unblock", "decompress", "scan_counts", "choose_level", "scan_images"]  # extern "C"
-ARCHS = "50,52,60,61,70,75,80,86,89,90"
+ARCHS = "50,52,60,61,70,75,80,86,89,90,100,103,120,121"
+PTX = "90"  # the JIT runs it on any GPU from 9.0, Thor (11.0) included
 OPTIONS = ["--std=c++17", "-ftz=true", "--device-as-default-execution-space"]  # as cupy
 
 
@@ -123,6 +125,7 @@ def main():
     ap.add_argument("--fatbinary", default="fatbinary")
     ap.add_argument("--out", required=True)
     ap.add_argument("--arch", default=ARCHS, help="compute capabilities, e.g. 50,52,61")
+    ap.add_argument("--ptx", default=PTX, help="compute capability of the PTX")
     ap.add_argument("--nvrtc-for", action="append", default=[], metavar="ARCH=LIBNVRTC",
                     help="an extra cubin from another toolkit, e.g. 35=.../libnvrtc.so.11.2")
     a = ap.parse_args()
@@ -138,7 +141,7 @@ def main():
     with tempfile.TemporaryDirectory() as tmp:
         images, table = [], None
         for c in old + archs + [None]:
-            arch = "sm_%s" % c if c else "compute_%s" % archs[-1]
+            arch = "sm_%s" % c if c else "compute_%s" % a.ptx
             data, lowered = nvs.get(c, nv).compile(src, names, arch)
             if table is not None and lowered != table:
                 raise SystemExit("lowered names differ between architectures")
@@ -147,7 +150,7 @@ def main():
             with open(path, "wb") as fh:
                 fh.write(data)
             images.append("--image3=kind=%s,sm=%s,file=%s" % ("elf" if c else "ptx",
-                                                              c or archs[-1], path))
+                                                              c or a.ptx, path))
         subprocess.check_call([a.fatbinary, "--create=" + fatbin, "-64"] + images)
     table.update({n: n for n in PLAIN})
     built = {c: "%d.%d" % nvs.get(c, nv).version for c in old + archs}  # test_cuda_fatbin.py
@@ -158,7 +161,7 @@ def main():
         fh.write("\n")
     print("wrote %s (%.1f MB, %d kernels, sm_%s + compute_%s PTX)"
           % (fatbin, os.path.getsize(fatbin) / 1e6, len(table) - 1, ",".join(old + archs),
-             archs[-1]))
+             a.ptx))
 
 
 if __name__ == "__main__":