cumo 0.11.0 → 0.13.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d50aabaf6d47d63e97fc6c68726cac94af2ecb8b737a3441de618860d2f5ab6a
-  data.tar.gz: 9d7bfc02015586d0ef533802be1d48558abf93fddf362ce445c55b0b4713ab7b
+  metadata.gz: ad04bee56b44e50dfc5df236d806a034f1adef98415e6e78ffcdd73546047dfd
+  data.tar.gz: c202fc43967dc292dddeea61274b89be6095f8724a7cfb7ac37d299cbc56cc7b
 SHA512:
-  metadata.gz: 99f47260f8dbbb9320d2f74984262e328fd42c2dac0321300b6087712dd129d03f80d65a2089f12f624ea2814e59eb8400748c9bf08f3937b09c0bceb452f417
-  data.tar.gz: a3ccc77b320c0bfdaa64abed339cc22ce5099638f8a52e9abc5a44a1e144295f56581b43ebf99f5a431d54bbe7bed2fe590975abd59cb6b7e1a9008878c91f23
+  metadata.gz: 00b7285a170c241c968f2bba0ab4e80013f971c01964d8d092eefa4eb177ea06d68616326bf457b1e9a6a9c08ca8906b85456ba2ef91acbb13334292959b10a0
+  data.tar.gz: eba6aed6cdb916d1abee4d856a5496892ef5e33d3bb55cf4124fb01c3066a0324301baefe5d095798e21a0066af5d87c33ceca51142a00fec24936450b9f5946

data/.yardopts

@@ -0,0 +1,10 @@
+--no-private
+--exclude ext/cumo/narray/gen/
+-e ext/cumo/narray/gen/cogen_runner.rb
+lib/**/*.rb
+ext/**/*.c
+ext/**/*.cpp
+tmp/yard/types/*.c
+-
+*.md
+LICENSE.txt

data/3rd_party/mkmf-cu/lib/mkmf-cu/cli.rb

@@ -38,7 +38,7 @@ module MakeMakefileCuda
     # TODO(sonots): Make it possible to configure "nvcc" and additional arguments
     def nvcc_command
       s = MakeMakefileCuda::Nvcc.generate(argv)
-      cmd = "nvcc #{s}"
+      cmd = "nvcc #{s} -Xfatbin=-compress-all"
       if ENV['CUMO_NVCC_GENERATE_CODE']
         ENV['CUMO_NVCC_GENERATE_CODE'].split.each do |code|
           cmd << " --generate-code=#{code}"

data/CHANGELOG.md

@@ -1,3 +1,51 @@
+# 0.13.0 (2026/10/12)
+
+Breaking changes:
+
+* An integer array raised to a negative integer power raises `Cumo::NArray::ValueError`, where it answered 0 even for a base of 1 or -1. numo-narray-alt now raises here too, as NumPy does. A signed integer exponent held in an array is checked on the device, so such a `**` now waits for its kernel, as integer division by an array does. An exponent array that holds `2**31` or more is no longer narrowed to 32 bits, so `Cumo::Int64[1] ** Cumo::Int64[2**31]` answers `[1]` where it answered `[0]` (PR #614)
+
+Fixes:
+
+* Fix `inspect` printing `(empty)` for a scalar cast to an array, such as `Cumo::DFloat.cast(5.0)`, and for its views, which it did since 0.5.10. Its value and the operators on it were right. Fix such an array that is frozen, as `Ractor.make_shareable` leaves it, raising `FrozenError` when read, and then answering 0 to some operators and its value to others (PR #616)
+* Fix a crash, or a read of the wrong memory, when a 0-dimensional array is indexed with nothing but `false`, `:rest` or `:~`, such as `Cumo::DFloat.cast(5.0)[false]` (PR #615)
+* Fix the cuDNN algorithm caches and the counters that deferred host reads and index fills wait on, which two Ractors could update at once. They now take a lock or advance atomically (PR #613)
+* Fix two Ractors calling `rand` or `rand_norm` at the same time drawing the same numbers. A single thread draws what it drew before, so `srand` still reproduces a sequence (PR #612)
+* Fix wrong results, an illegal memory access, or a process stuck at full GPU load when two Ractors call cuSOLVER at once, as `Cumo::Linalg.svdvals` and `eigvals` do. Each thread keeps its own workspace and frees it when it ends (PR #611)
+* Fix `Cumo::Linalg.expm` of a triangular matrix whose elements off the diagonal are far larger than those on it answering 1 on the diagonal. The diagonal and the first off-diagonal are now reset to their exact values after each squaring, as scipy does (PR #609)
+* Fix `Cumo::Linalg.expm` raising `FloatDomainError: Infinity` for a finite matrix whose 1-norm overflows. It still refuses such a matrix, and now says that its 1-norm overflowed (PR #603)
+* Fix `Cumo::Linalg.lu` and `matrix_power(a, 0)` raising `NoMethodError` for a Bit matrix. `P` and the identity are Bit matrices (PR #602)
+* Fix `Cumo::Linalg.lu` raising `ShapeError` for an empty matrix. It answers `P`, `L` and `U` of m x m, m x 0 and 0 x n, the shapes a nonempty matrix gets. numo-linalg-alt raises here (PR #601)
+
+Changes:
+
+* The first elementwise call on a type loads the kernels of one family rather than every elementwise kernel of the type. The families are arithmetic, common math functions, rare math functions, comparisons, stores from other types, and random numbers with polynomials. A first SFloat add takes 28 MB of host memory where it took 86 MB (PR #610)
+* `require "cumo"` no longer loads NVRTC. It is loaded from the CUDA that cumo was built against the first time `ElementwiseKernel`, `ReductionKernel` or `RawKernel` builds a kernel, which takes about 34 MB off the RSS after `require` (PR #607, PR #606)
+* `Cumo::Linalg.norm` of ord 1, -1, Infinity and -Infinity takes the magnitudes as it reduces rather than through `abs`, which makes it about 2.4 times faster on a 4096 x 4096 DFloat (PR #605)
+* `Cumo::Linalg.norm` of ord 2, -2 and `"nuc"` over stacked matrices of up to 32 rows and columns computes them in one batched call, about 7 times faster for 1000 DFloat matrices of 8 x 8. A stacked matrix that holds NaN or Infinity answers NaN, as the matrices past 32 do (PR #604)
+* `Cumo::Linalg.expm` of a small matrix reads its three norms in one wait, which makes it about 3% faster on a 64 x 64 DFloat (PR #600)
+* The factorizations in `Cumo::Linalg` free their column-major copies once used, so a small call reuses memory the GPU has touched: `lu_solve` of a 256 x 256 SFloat is about 8% faster (PR #599)
+* `repeat` builds its index in one Array: `repeat(2)` of a million elements is 3.4 times faster, and `repeat` with an Array of counts 2.2 times. The development setup checks `lib` with rubocop-performance (PR #598, PR #597)
+* `Cumo::Linalg.eig` and `svd` give cuSOLVER pinned host memory to work in, kept for the process up to 256 MiB: `eig` of a 4096 x 4096 SFloat is about 8% faster, and `svd` of an 8192 x 2048 SFloat about 5% (PR #596)
+* `Cumo::Linalg.expm` frees its temporary matrices as it goes rather than leaving about 40 a call to the garbage collector, which makes it about 15% faster on a 1024 x 1024 SFloat of large norm (PR #595)
+* The API docs show the parameters of overloaded methods such as `cast`, `store_binary` and `layer_norm`, where YARD warned 1,419 times (PR #608)
+
+# 0.12.0 (2026/10/04)
+
+Fixes:
+
+* Fix a stream left current in a thread that ended never being destroyed. Neither `Stream#destroy` nor the finalizer could free it, so about 0.5 MiB of device memory stayed behind per such stream, and under M:N threads the next Ruby thread on the same native thread inherited it. The current stream now lives on the Ruby `Thread`, which adds about 2% to the fixed cost of a launch (PR #593)
+* Fix a device memory leak each time a thread that used cuBLAS, cuSOLVER or cuDNN ends: 64 MiB per thread for cuBLAS, 32 MiB more for cuSOLVER and about 8 MiB for cuDNN. An ending thread now leaves its handles to the next one, so device memory stays at what the most threads alive at once needed (PR #591)
+* Fix `x.inplace.cumsum` and `cumprod` answering a new array and leaving `x` as it was, which they did since 0.10.0 (PR #582)
+
+Changes:
+
+* Add `Cumo::Linalg`, the functions of `Numo::Linalg` from numo-linalg-alt 0.10.1 on cuSOLVER and cuBLAS, with the same names, arguments, return values and errors: `blas_char`, `dot`, `matmul`, `matrix_power`, `solve`, `inv`, `det`, `slogdet`, `lu`, `lu_fact`, `lu_solve`, `lu_inv`, `cholesky`, `cho_fact`, `cho_solve`, `cho_inv`, `ldl`, `qr`, `eig`, `eigvals`, `eigh`, `eigvalsh`, `svd`, `svdvals`, `pinv`, `lstsq`, `matrix_rank`, `orth`, `null_space`, `norm`, `cond` and `expm`. Load it with `require "cumo/linalg"`. Cumo builds it when the toolkit has cuSOLVER's 64-bit API, and `Cumo::CUDA::Cusolver.available?` says whether it did. `eig` and `eigvals` need `cusolverDnXgeev`, which CUDA 11 does not have, and `ldl` of a Hermitian complex matrix needs `cusolverDnXhetrf`, which CUDA 11 and 12 do not have. Where it differs from numo-linalg-alt on purpose, as in `expm`, `norm` and `cond`, the README says how (PR #594, PR #592, PR #589, PR #588, PR #587, PR #586, PR #584, PR #581, PR #580, PR #577, PR #576, PR #575, PR #574, PR #573, PR #572, PR #571, PR #570, PR #569, PR #568, PR #567)
+* `require "cumo"` sets `CUDA_MODULE_LOADING` to `LAZY` unless it is set already, as PyTorch does. Drivers before 535 otherwise load every kernel module when CUDA starts, which adds about 2.2 GB of RSS. Child processes inherit the setting (PR #590)
+* `triu`, `tril`, `triu!` and `tril!` clear the triangle of a contiguous array with an Integer `k` in one kernel rather than through an index array: `triu` of a 16 x 16 DFloat takes 0.006 ms where it took 0.49 (PR #585)
+* The first kernel of a type costs less host memory. Each type builds separate modules for its reductions, their nan-aware forms and `mulsum` against another type, so DFloat adds 88 MB of RSS for elementwise work and 137 MB once it reduces, where its first kernel of any kind added 181 MB. Every fatbin image is compressed, which takes `cumo.so` from 120 MB to 112 MB (PR #583)
+* The API docs show the NArray methods defined in C, such as `shape` and `NArray.zeros`, and keep the docstrings of the typed-class methods. `.yardopts` keeps the build-time code generators out of the docs, and the gemspec names rubydoc.info as the documentation (PR #579, PR #578)
+* The gem no longer ships the test suite, `bench/` and the files for working on cumo (PR #566)
+
 # 0.11.0 (2026/09/27)
 
 Breaking changes:
@@ -486,7 +534,7 @@ Fixes:
 * Remove unnecessary debug code
 * Fix capability list
 * Build only with supported capabilities to reduce compilation time
-* Fix SEGV when calling {mean, var, stddev, rms} on a single-element array (#154)
+* Fix SEGV when calling `mean`, `var`, `stddev` or `rms` on a single-element array (#154)
 * Suppress warning message for deprecated declarations
 * Fix variable typo in complex log2 and log10 functions (#152)
 
@@ -528,7 +576,7 @@ Fixes:
 * Backport: fix macro: STORE_BIT STORE_BIT_STEP: requires mask to leave the lowest bit
 * Backport: fix NArray::Bit#any?,all?: empty array should return false
 * Backport: fix NArray::Bit#count_true/false: empty array should return zero
-* Backport: bug in NArray::Bit; fix bit operation in tmpl_bit/{store_bit,unary,binary}.c
+* Backport: bug in NArray::Bit; fix bit operation in `tmpl_bit/{store_bit,unary,binary}.c`
 * Fix typo
 * Backport 135: Make all empty arrays equal
 * Backport: minor fixes in na_get_result_dimension(), check_index_count()

data/README.md

@@ -10,6 +10,7 @@ Cumo (pronounced "koomo") is a CUDA-aware, GPU-optimized numerical library that
 * NVIDIA GPU Compute Capability 3.5 (Kepler) or later
 * CUDA 11.0 or later
 * cuDNN 8.0 or later (optional, for the cuDNN features)
+* cuSOLVER with its 64-bit API, which comes with the CUDA Toolkit (optional, for `Cumo::Linalg`)
 
 ## Preparation
 
@@ -282,6 +283,41 @@ q.dot(kt)
 All of this is about the two-dimensional case, where cuBLAS is given one matrix.
 A batched multiplication takes another path through the same flag.
 
+### Linear Algebra
+
+`require "cumo/linalg"` brings in `Cumo::Linalg`, the functions of `Numo::Linalg` from [numo-linalg-alt](https://github.com/yoshoku/numo-linalg-alt) on the GPU.
+They keep numo-linalg-alt's names, arguments, return values and errors, and run on cuSOLVER and cuBLAS.
+
+```ruby
+require "cumo/linalg"
+
+a = Cumo::DFloat.new(3, 3).rand + Cumo::DFloat.eye(3) * 3
+b = Cumo::DFloat.new(3).rand
+x = Cumo::Linalg.solve(a, b)
+w, v = Cumo::Linalg.eigh(a + a.transpose)
+```
+
+The functions are `blas_char`, `dot`, `matmul`, `matrix_power`, `solve`, `inv`, `det`, `slogdet`, `lu`, `lu_fact`, `lu_solve`, `lu_inv`, `cholesky`, `cho_fact`, `cho_solve`, `cho_inv`, `ldl`, `qr`, `eig`, `eigvals`, `eigh`, `eigvalsh`, `svd`, `svdvals`, `pinv`, `lstsq`, `matrix_rank`, `orth`, `null_space`, `norm`, `cond` and `expm`.
+They take `SFloat`, `DFloat`, `SComplex` and `DComplex`.
+An integer or `Bit` matrix is computed in `DFloat`, as numo-linalg-alt does, and `HFloat` and `BFloat` raise `TypeError`.
+
+Cumo looks for cuSOLVER when it is built.
+Without it, the functions that need it raise `NotImplementedError`, and `Cumo::CUDA::Cusolver.available?` answers whether it is there.
+`eig` and `eigvals` also need `cusolverDnXgeev`, which CUDA 11 does not have.
+`ldl` of a Hermitian complex matrix needs `cusolverDnXhetrf`, which CUDA 11 and 12 do not have.
+
+Cumo differs from numo-linalg-alt in a few places on purpose:
+
+* Where numo-linalg-alt answers wrong, Cumo answers as numpy does: a wide matrix in `orth` and `lstsq`, a zero matrix in `pinv`, a matrix of full rank in `null_space`, an integer matrix in `qr` and `svd`, and empty matrices.
+* `expm` scales and squares the way scipy does, after Al-Mohy and Higham, which stays accurate where the norm of the matrix is large. Passing the order of the Pade approximant, as in `expm(a, 3)`, keeps numo-linalg-alt's algorithm.
+* `norm` takes the 2-norm and the nuclear norm of matrices stacked along other axes, which numo-linalg-alt refuses.
+* `cond` checks `ord` before it inverts the matrix, so an invalid `ord` raises `ArgumentError` even for a singular matrix, where numo-linalg-alt raises `LapackError`.
+* `lu` answers for an empty matrix and for a `Bit` matrix, where numo-linalg-alt raises. P keeps the shape and the class it has for any other matrix: an m x m identity for an m x 0 matrix, where scipy answers a 0 x 0 one, and a `Bit` for a `Bit` matrix. `matrix_power(a, 0)` of a `Bit` matrix likewise answers a `Bit` identity.
+
+A function waits for the current stream when it answers a Ruby number, or when it reads back whether cuSOLVER could do its work.
+`solve`, `inv` and `det` read back whether the matrix was singular, `cho_fact` whether it was positive definite, and `eigh` and `svd` whether they converged.
+`matmul`, `dot` of two matrices, `matrix_power` to a positive power, `cholesky`, `cho_solve`, `qr` and `norm` along an axis wait for nothing, so a run of them is queued without the host waiting.
+
 ### Fused Operations
 
 `layer_norm`, `rms_norm` and `softmax` normalize along the last axis in one kernel each, and `quantize_symmetric` takes it to 8-bit integers in one more.
@@ -962,6 +998,7 @@ export CUMO_ALLOW_TF32=1
 Over the convolutions of a ResNet-18 forward pass at batch 16, with the ceiling raised to 256MB, the pass gets faster and the worst layer moves from 1.4e-05 to 2.4e-04 against a double precision reference.
 
 The same flag puts `SFloat` and `SComplex` `gemm` on the tensor cores as TF32, and `dot` where it goes through `gemm`.
+`Cumo::Linalg.matrix_power` and `Cumo::Linalg.expm` multiply through `gemm`, so the flag reaches them too.
 Off, the answer is the one cuBLAS gives at single precision, bit for bit.
 On, a large `SFloat` `gemm` runs faster, and its answer is then about 3e-04 from a double precision reference, where it was 4e-07.
 The double types are not affected either way.

data/cumo.gemspec

@@ -18,6 +18,7 @@ Gem::Specification.new do |spec|
 
   spec.metadata["homepage_uri"] = spec.homepage
   spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["documentation_uri"] = "https://rubydoc.info/gems/cumo/"
   spec.metadata["changelog_uri"] = "https://github.com/sonots/cumo/blob/master/CHANGELOG.md"
   spec.metadata["bug_tracker_uri"] = "https://github.com/sonots/cumo/issues"
   spec.metadata["rubygems_mfa_required"] = "true"
@@ -25,9 +26,8 @@ Gem::Specification.new do |spec|
   spec.required_ruby_version = ">= 3.0.0"
 
   spec.files = `git ls-files -z`.split("\x0").reject do |f|
-    f.match(%r{^(.github|test|spec|features|docker|examples)/})
+    f.start_with?(*%w[.github/ test/ examples/ bin/ bench/ docker- Dockerfile Gemfile .gitignore .rubocop])
   end
-  spec.test_files    = `git ls-files -- {test,spec,features}/*`.split("\n")
   spec.bindir        = "exe"
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ["lib"]

data/ext/cumo/cuda/cublas.c

@@ -43,31 +43,22 @@ cumo_cuda_cublas_check_status(cublasStatus_t status)
     }
 }
 
-// Lazily initialize cublas handle, and cache it
+static cumo_cuda_thread_local_t handles;
+
+// One handle per thread: the stream is set on the handle before each call,
+// and another thread with another current stream must not reset it in
+// between. A thread that ends leaves its handle to the next one.
 cublasHandle_t
 cumo_cuda_cublas_handle()
 {
-    // One table per thread: the stream is set on the handle before each call,
-    // and another thread with another current stream must not reset it in
-    // between. A handle is never destroyed.
-    static __thread cublasHandle_t *handles = 0;
-    int device;
-    if (handles == 0) {
-        int i;
-        int device_count = cumo_cuda_runtime_get_device_count();
-        handles = ALLOC_N(cublasHandle_t, device_count);
-        for (i = 0; i < device_count; ++i) {
-            handles[i] = 0;
-        }
-    }
-    device = cumo_cuda_runtime_get_device();
-    if (handles[device] == 0) {
+    cublasHandle_t *handle = (cublasHandle_t*)cumo_cuda_thread_local_get(&handles);
+    if (*handle == 0) {
         // A discarded status leaves the handle NULL, and cuBLAS reports that as
         // CUBLAS_STATUS_NOT_INITIALIZED at the next call instead of the reason.
-        cumo_cuda_cublas_check_status(cublasCreate(&handles[device]));
+        cumo_cuda_cublas_check_status(cublasCreate(handle));
     }
-    cumo_cuda_cublas_check_status(cublasSetStream(handles[device], cumo_cuda_stream()));
-    return handles[device];
+    cumo_cuda_cublas_check_status(cublasSetStream(*handle, cumo_cuda_stream()));
+    return *handle;
 }
 
 VALUE
@@ -88,6 +79,8 @@ Init_cumo_cuda_cublas(void)
     VALUE mCumo = rb_define_module("Cumo");
     VALUE mCUDA = rb_define_module_under(mCumo, "CUDA");
 
+    cumo_cuda_thread_local_init(&handles, sizeof(cublasHandle_t));
+
     /*
       Document-module: Cumo::Cublas
     */

data/ext/cumo/cuda/cudnn.c

@@ -72,30 +72,21 @@ cumo_cuda_cudnn_release_conv_held(VALUE held)
     return Qnil;
 }
 
-// Lazily initialize cudnn handle, and cache it
+static cumo_cuda_thread_local_t handles;
+
+// One handle per thread, as in cublas.c: the stream set on the handle must
+// not be reset by another thread before the call.
 cudnnHandle_t
 cumo_cuda_cudnn_handle()
 {
-    // One table per thread, as in cublas.c: the stream set on the handle must
-    // not be reset by another thread before the call. A handle is never destroyed.
-    static __thread cudnnHandle_t *handles = 0;
-    int device;
-    if (handles == 0) {
-        int i;
-        int device_count = cumo_cuda_runtime_get_device_count();
-        handles = ALLOC_N(cudnnHandle_t, device_count);
-        for (i = 0; i < device_count; ++i) {
-            handles[i] = 0;
-        }
-    }
-    device = cumo_cuda_runtime_get_device();
-    if (handles[device] == 0) {
+    cudnnHandle_t *handle = (cudnnHandle_t*)cumo_cuda_thread_local_get(&handles);
+    if (*handle == 0) {
         // A discarded status leaves the handle NULL, and cuDNN reports that as
         // CUDNN_STATUS_NOT_INITIALIZED at the next call instead of the reason.
-        cumo_cuda_cudnn_check_status(cudnnCreate(&handles[device]));
+        cumo_cuda_cudnn_check_status(cudnnCreate(handle));
     }
-    cumo_cuda_cudnn_check_status(cudnnSetStream(handles[device], cumo_cuda_stream()));
-    return handles[device];
+    cumo_cuda_cudnn_check_status(cudnnSetStream(*handle, cumo_cuda_stream()));
+    return *handle;
 }
 
 #endif // CUDNN_FOUND
@@ -144,6 +135,7 @@ Init_cumo_cuda_cudnn(void)
 
     rb_define_singleton_method(mCUDNN, "available?", rb_cudnn_available_p, 0);
 #ifdef CUDNN_FOUND
+    cumo_cuda_thread_local_init(&handles, sizeof(cudnnHandle_t));
     init_max_workspace_size();
     rb_define_singleton_method(mCUDNN, "max_workspace_size", rb_cudnn_max_workspace_size, 0);
     rb_define_const(mCUDNN, "CUDNN_POOLING_MAX", INT2NUM(CUDNN_POOLING_MAX));

data/ext/cumo/cuda/cudnn_impl.cpp

@@ -11,6 +11,7 @@
 #include "cumo/cuda/runtime.h"
 #include "cumo/cuda/memory_pool.h"
 
+#include <mutex>
 #include <unordered_map>
 
 #if defined(__cplusplus)
@@ -387,13 +388,32 @@ struct AlgoCacheEntry {
 // Nothing may reach the C caller: it has no handler, so an escaping exception
 // is std::terminate. The cache only saves a later search, so one that cannot
 // grow leaves the answer this call already has.
+static std::mutex algo_cache_mutex_;
+
 template <typename Map>
 void StoreAlgoCache(Map& map, const AlgoCacheKey& key, const typename Map::mapped_type& entry) {
+    std::lock_guard<std::mutex> lock{algo_cache_mutex_};
     try {
         map[key] = entry;
     } catch (...) {
     }
 }
+
+template <typename Map, typename Perf>
+bool LookupAlgoCache(Map& map, const AlgoCacheKey& key, Perf* perf_result) {
+    std::lock_guard<std::mutex> lock{algo_cache_mutex_};
+    auto it = map.find(key);
+    if (it == map.end()) {
+        return false;
+    }
+    auto entry = it->second;
+    // clear the fields the search would have filled but the cache drops
+    *perf_result = {};
+    perf_result->algo = entry.algo;
+    perf_result->memory = entry.memory;
+    perf_result->mathType = entry.math_type;
+    return true;
+}
 }
 
 using FwdAlgoCacheMap = std::unordered_map<AlgoCacheKey, AlgoCacheEntry<cudnnConvolutionFwdAlgo_t>, AlgoCacheKeyHash>;
@@ -456,15 +476,7 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
                                 cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = fwd_algo_cache_map_;
-    // TODO: thread-safe
-    auto it = algo_cache_map.find(key);
-    if (it != algo_cache_map.end()) {
-        auto entry = it->second;
-        // clear the fields the search would have filled but the cache drops
-        *perf_result = {};
-        perf_result->algo = entry.algo;
-        perf_result->memory = entry.memory;
-        perf_result->mathType = entry.math_type;
+    if (LookupAlgoCache(algo_cache_map, key, perf_result)) {
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -495,7 +507,6 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
     if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
-    // TODO: thread-safe
     StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }
@@ -527,15 +538,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
                                 cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = bwd_data_algo_cache_map_;
-    // TODO: thread-safe
-    auto it = algo_cache_map.find(key);
-    if (it != algo_cache_map.end()) {
-        auto entry = it->second;
-        // clear the fields the search would have filled but the cache drops
-        *perf_result = {};
-        perf_result->algo = entry.algo;
-        perf_result->memory = entry.memory;
-        perf_result->mathType = entry.math_type;
+    if (LookupAlgoCache(algo_cache_map, key, perf_result)) {
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -566,7 +569,6 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
     if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
-    // TODO: thread-safe
     StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }
@@ -598,15 +600,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
                                 cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = bwd_filter_algo_cache_map_;
-    // TODO: thread-safe
-    auto it = algo_cache_map.find(key);
-    if (it != algo_cache_map.end()) {
-        auto entry = it->second;
-        // clear the fields the search would have filled but the cache drops
-        *perf_result = {};
-        perf_result->algo = entry.algo;
-        perf_result->memory = entry.memory;
-        perf_result->mathType = entry.math_type;
+    if (LookupAlgoCache(algo_cache_map, key, perf_result)) {
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -637,7 +631,6 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
     if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
-    // TODO: thread-safe
     StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }