cumo 0.7.0 → 0.8.0

data/CHANGELOG.md

@@ -1,3 +1,32 @@
+# 0.8.0 (2026/09/16)
+
+Breaking changes:
+
+* `sort` and `sort_index` with `nan: true` put every NaN last. The host comparator they used answered -1 for a NaN against anything, so the partition carried elements other than the NaN to the wrong side: of 4156 cases over lengths 5 to 129, 2121 came back with the numbers out of order. `nan:` no longer changes the answer, and 1M DFloat costs 0.405 ms rather than 73.0 (PR #447)
+* `conv` and `conv_transpose` refuse a bias that is not one value per output channel. A bias of length 1 was broadcast over every channel and answered as if it had been right, and a bias of two dimensions passed whenever its size fitted (PR #443)
+* `var` and `stddev` answer NaN for a row holding an infinity, where the answer turned on where the infinity sat and how long the row was (PR #430)
+
+Fixes:
+
+* Fix every `dot` after `Cumo::CUDA::Runtime.cudaSetDevice` raising `CUBLAS_STATUS_EXECUTION_FAILED` for the life of the process. Cumo took a context of its own when it loaded, and `cudaSetDevice` binds the device's primary one, so the event inside the cuBLAS handle was left in a context nothing ran in (PR #451)
+* Fix the device memory a rejected kernel launch left in the pool, `rb_raise` leaving by longjmp past the free below it. A sort lost 24MB and a convolution given a bias of the wrong class 5MB, every time; reductions, `softmax`, `layer_norm`, `Cumo::Bit#count`, `where`, `poly` and a store from a Ruby Array lost theirs the same way (PR #439, PR #438, PR #437, PR #436, PR #434, PR #433, PR #432, PR #431)
+* Fix two cuDNN descriptor leaks, one where `batch_norm` raises after building its descriptors and one where a convolution finds no algorithm inside `CUMO_CUDNN_MAX_WORKSPACE_SIZE` (PR #441, PR #440)
+* Fix a `std::bad_alloc` from the convolution algorithm cache killing the process, `cudnn_impl.cpp` being one `extern "C"` block with no handler in it (PR #444)
+* Fix the nine cuDNN methods and `layer_norm` reading the wrong memory when `dup` answers a different class, a different shape or a non-contiguous view; `conv` took a reversed view and read past it (PR #425)
+* Fix the spelling of "propagate" and "position" in the `nan:` documentation on `sort`, `sort_index`, `median`, `argmin`, `argmax`, `min_index` and `max_index` (PR #446)
+
+Changes:
+
+* Add `Cumo::BFloat`, the bfloat16 dtype, aliased `Cumo::BFloat16`. It keeps a float's exponent and gives up mantissa bits for it, so it holds the magnitudes a `Cumo::HFloat` overflows on, which is how published weights are written. Everything the dtype has is generated: the four operations, the comparisons, the reductions, `Cumo::NMath`, `sort`, `rand`, `dot` through cuBLAS and the cuDNN methods, both with a float compute type. `BFloat` and `HFloat` promote to `SFloat` when they meet, neither format containing the other. The accelerated paths want compute capability 8.0 (PR #455, PR #453, PR #452, PR #450)
+* Add `layer_norm`, which normalizes a row in one kernel rather than the nine it is written out of, and `softmax`, which takes one rather than five. On SFloat `[1, 768]` the first goes from 19.9 to 2.9us and the second from 11.0 to 2.2, and `[4096, 768]` from 198.3 to 31.4 and 109.6 to 31.6. No shape is slower (PR #428, PR #426, PR #424, PR #423)
+* Add `gelu` and `gelu_tanh` to `Cumo::NMath`. Two of them because there are two: `gelu` is the error function form PyTorch answers by default, and `gelu_tanh` the approximation GPT-2 and the models after it were trained with. They differ by 4.1e-4 at most, too little to tell apart by eye and far too much to swap under trained weights (PR #427)
+* Take one reciprocal instead of two divides in the moments accumulator, a double divide being what this class of card charges most for; `var` goes from 172.7 to 106.5us in DFloat and `layer_norm` from 179.6 to 138.0 (PR #429)
+* Store a Ruby Array into a strided or indexed view through one buffer held by the loop rather than one a row, a row having cost 2.53 times what the same payload written contiguously does and now costing 1.04 (PR #445)
+* Stop `Cumo::RObject#poly` stopping the whole device once per element. It runs with `CUMO_NO_LOOP`, so ndloop calls its iterator once an element, and there is nothing left there to wait for (PR #435)
+* Say once what `conv` and `conv_transpose` had each grown a copy of, the run struct and the bias block, and build the bias descriptor from a shape rather than lending the narray one (PR #442)
+* Declare the gem's license as MIT. RubyGems reads the gemspec, which said BSD-3-Clause while the `LICENSE.txt` inside the same gem has carried the MIT text from the start (PR #448)
+* Skip the CUDA build matrix for a README, license or benchmark edit no build step opens, and give ccache a ceiling one build fits in (PR #454, PR #449)
+
 # 0.7.0 (2026/09/13)
 
 Breaking changes:

data/README.md

@@ -222,6 +222,68 @@ Cumo::SFloat::Math.atan2(a, 2.0)   #=> Cumo::SFloat
 
 The 0-dimensional form has no effect under Numo, where `[]` returns a Ruby Float.
 
+### Fused Operations
+
+`layer_norm` and `softmax` normalize along the last axis in one kernel each.
+Written out of the operators they take nine launches and five, and a launch costs about two microseconds whatever it is handed, so a short row pays for the launches rather than for its bytes.
+
+```ruby
+y = x.layer_norm(gamma, beta, eps: 1e-5)   # (x - mean) / sqrt(var + eps) * gamma + beta
+probs = scores.softmax                     # exp(x - max) / sum(exp(x - max))
+```
+
+On an RTX 5070 Ti Laptop, against the same arithmetic spelled with operators, in microseconds:
+
+```
+layer_norm              SFloat            DFloat            HFloat
+shape              fused  written    fused  written    fused  written
+1 x 768              3.9     30.1      9.7     30.3      3.9     29.8
+256 x 768            6.6     22.5     22.0     35.4      7.8     22.4
+4096 x 768          32.7    177.4    300.7    610.5     28.2    142.9
+1 x 1000000         19.0     59.4    125.4    175.4     22.3     68.5
+
+softmax                 SFloat            DFloat            HFloat
+shape              fused  written    fused  written    fused  written
+1 x 768              4.3     19.8      7.4     26.5      6.3     34.4
+256 x 768            8.8     13.9     23.4     35.1      6.8     18.5
+4096 x 768          49.2    122.7    318.3    549.0     28.2    109.8
+1 x 1000000         27.8     44.0    108.6    148.5     32.1     53.8
+```
+
+Both pay off in every precision, by the most where the row is short enough that the launches were all it was doing, and by the least in double, where the reduction itself costs more than the launches ever did.
+The memory clock on this card steps between 9001 and 11001 MHz under a benchmark this short, and the absolute figures move with it.
+The ratios hold across the steps.
+
+Both methods normalize along the last axis only.
+Other axes are reachable through `transpose`.
+`layer_norm` divides the variance by the row length and not by one less, which is how the layer is defined and unlike `var`.
+Both answer a new array and ignore `inplace!`, and a non-contiguous receiver is copied once rather than walked.
+`gamma` and `beta` must be one-dimensional, as long as the last axis, and of the same class as the receiver.
+
+`softmax` subtracts the row maximum before exponentiating, so a row masked with `-Float::INFINITY` answers zeros rather than `NaN`.
+A row that is entirely `-Float::INFINITY` answers `NaN`, as the written-out form does.
+
+Those two and `gelu_tanh` from the section below are where a transformer spends most of its launches.
+Decoding one token of GPT-2 124M written against Numo's API takes 640 of them, and calling these three instead removes 332: 200 for the layer norms, 84 for the activations and 48 for the softmaxes.
+
+### Two Spellings Of gelu
+
+`Cumo::NMath` answers both definitions of the GELU activation, under names of
+their own rather than a keyword, since a unary NMath function takes none.
+
+| method | formula | matches |
+| --- | --- | --- |
+| `gelu` | `0.5 * x * (1 + erf(x / sqrt(2)))` | PyTorch's default, `approximate='none'` |
+| `gelu_tanh` | `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3)))` | PyTorch's `approximate='tanh'`, and what GPT-2 was trained with |
+
+They differ by at most 4.1e-4, which is small enough to mistake one for the
+other and far too large to swap under a set of trained weights, so pick the one
+the weights were trained with.
+
+Both follow PyTorch at the edges rather than the limit: a large negative `x`
+runs out of significant digits and answers zero, and `-Float::INFINITY` answers
+`NaN`.
+
 ### Half Precision
 
 `Cumo::HFloat`, also reachable as `Cumo::Float16`, holds IEEE binary16: one sign bit, five of exponent and ten of mantissa.
@@ -262,6 +324,7 @@ That distinction is what makes half usable in a transformer's layer normalizatio
 Squaring overflows at 256, since 256 squared is already past the top, and a residual stream with one outlier feature reaches thousands.
 Writing the normalization out as `((x - mean) ** 2).mean` builds those squares as a half array and the answer is `Infinity`, while `x.var(axis: 1)` keeps them in its single-precision accumulator and answers 10208.0 against a single-precision 10209.
 `var` saturates only if the variance itself is out of range, which is a far higher bar than any one deviation being over 256, but it is still a bar.
+An input over 65504 is a different matter: it becomes an infinite element on the way in, and `var` and `stddev` answer `NaN` for a row holding one.
 The quantity to check against 65504 is the largest variance a layer produces, not the largest activation in it.
 A single outlier of size d among n values contributes only d squared over n to the variance, so a row 768 wide divides it by 768: an activation of 3000 squares to nine million but raises the variance of its row by about twelve thousand.
 An activation that looks safe therefore says nothing about whether `var` overflows, in either direction.
@@ -342,6 +405,83 @@ Keeping the running statistics in single precision is what the arithmetic wants
 Gradients are the other place half runs out of room.
 Values below `Cumo::HFloat::MIN` of 6.1e-05 fall into the subnormals and then to zero, which is what loss scaling in a training loop exists to prevent.
 
+### Brain Float
+
+`Cumo::BFloat`, also reachable as `Cumo::BFloat16`, holds bfloat16: one sign bit, eight of exponent and seven of mantissa.
+It is the other sixteen-bit float, and it differs from `Cumo::HFloat` in where those bits went.
+A bfloat16 has the exponent of a `Cumo::SFloat` and three fewer mantissa bits than a binary16, so it reaches everything a single-precision value reaches and resolves less of it:
+
+```ruby
+Cumo::BFloat[1.0e38]   #=> 9.969209968386869e+37
+Cumo::HFloat[1.0e38]   #=> Infinity
+
+Cumo::BFloat[257.0]    #=> 256.0
+Cumo::HFloat[257.0]    #=> 257.0
+```
+
+Integers are exact only to 256, against 2048 for a binary16, and the largest finite value is 3.3895314e+38 rather than 65504.
+That trade is why weights are published in it: a tensor that came out of training keeps its magnitudes, and the bits it loses are ones a trained weight does not carry.
+
+The two do not contain each other, so an expression mixing them promotes to `Cumo::SFloat`, which does:
+
+```ruby
+(Cumo::BFloat[1.0] + Cumo::HFloat[1.0]).class   #=> Cumo::SFloat
+```
+
+Everything else promotes as `Cumo::SFloat` does, so an integer array or a Ruby Float mixed in stays bfloat16 while anything wider takes over.
+
+`layer_norm`, `softmax` and both spellings of `gelu` take it, and so do the reductions, `sort`, `median`, `cumsum`, `rand` and `dot`.
+The cuDNN methods take it too: `conv`, `conv_transpose`, `conv_grad_w`, `max_pool`, `avg_pool` and the three batch norm entries.
+
+Reductions accumulate in single precision and round once at the end, exactly as they do for half, so a sum passes 256 without stopping there:
+
+```ruby
+Cumo::BFloat.new(40_000).fill(1.0).sum   #=> 39936.0
+```
+
+That 39936 is 40000 rounded to the nearest bfloat16, whose values are 256 apart up there; a bfloat16 accumulator would have answered 256.
+`sum`, `mean`, `var`, `stddev`, `rms`, `mulsum`, `dot`, `gemm` and `cumsum` all widen this way, and `prod` and `cumprod` do not, for the same reasons given under Half Precision.
+
+The wider exponent moves where the saturation described there bites.
+A variance is far from 3.4e+38, so `var` does not saturate on any input a network produces, and squaring is safe to 1.8e+19 rather than to 256.
+What replaces it is the mantissa: a bfloat16 resolves about three decimal digits, so a normalization computes its statistics well enough while the values it writes back carry that resolution and no more.
+
+#### What bfloat16 is faster at
+
+`gemm` reaches the tensor cores and lands on the same throughput as half. Square matrices on an RTX 5070 Ti Laptop, median of five runs each, one process per case:
+
+```
+              BFloat              HFloat              SFloat
+1024x1024     0.044 ms  48.9 TF   0.044 ms  48.5 TF   0.144 ms  14.9 TF
+2048x2048     0.298 ms  57.6 TF   0.303 ms  56.7 TF   1.404 ms  12.2 TF
+4096x4096     2.839 ms  48.4 TF   2.898 ms  47.4 TF   9.540 ms  14.4 TF
+```
+
+The odd-column penalty is the same as half's and for the same reason, a row starting on a two-byte boundary.
+1024x1024 times 1024x1024, with one dimension made odd at a time:
+
+```
+             all even   M odd      K odd      N odd
+BFloat        48.9 TF   48.7 TF    23.5 TF    23.0 TF
+HFloat        48.5 TF      -       24.0 TF    23.1 TF
+```
+
+So the choice between the two sixteen-bit types is about range and resolution, not speed.
+Take bfloat16 where the magnitudes came from somewhere else and binary16 where three more mantissa bits are worth having.
+
+#### The accelerated paths want Ampere
+
+cuBLAS supports `CUDA_R_16BF`, which is what `gemm` and `dot` reach, from compute capability 8.0.
+Storing, casting and elementwise arithmetic have no such floor, because every operation is computed in single precision and rounded back, and they build and run wherever cumo does.
+A `dot` on a pre-Ampere card is the one to expect trouble from.
+**This is not measured here**: the only GPU these numbers came from is a Blackwell one, and the requirement is read from cuBLAS's documentation rather than reproduced.
+
+cuDNN reaches bfloat16 as `CUDNN_DATA_BFLOAT16`, and its own bfloat16 kernels want Ampere for the same reason cuBLAS does.
+A convolution is given `CUDNN_DATA_FLOAT` to accumulate in, so it passes 256 the way a reduction does.
+The batch norm parameters are `Cumo::SFloat`, the same as they are for `Cumo::HFloat`.
+The workspace ceiling matters here as much as it does for half: left at the default 8MB, a bfloat16 convolution is no faster than a single-precision one.
+See [Raise the cuDNN workspace ceiling](#raise-the-cudnn-workspace-ceiling).
+
 ### Select a GPU device ID
 
 Set the `CUDA_VISIBLE_DEVICES=id` environment variable, or

data/bench/cumo_probe.rb

@@ -312,7 +312,7 @@ op(:scalar, 'scalar_left') { |c| c.scalar - c.sc.inplace }
 
 # NMath
 %w[sqrt exp exp2 exp10 expm1 log log2 log10 log1p sin cos tan asin acos atan
-   sinh cosh tanh asinh acosh atanh cbrt erf erfc sinc].each do |f|
+   sinh cosh tanh asinh acosh atanh cbrt erf erfc sinc gelu gelu_tanh].each do |f|
   op(:math, f, :float) { |c| XM::NMath.send(f, c.usc.inplace) }
 end
 op(:math, 'atan2', :float) { |c| XM::NMath.atan2(c.sc.inplace, c.b) }

data/cumo.gemspec

@@ -14,7 +14,7 @@ Gem::Specification.new do |spec|
   spec.summary       = %q{Cumo is CUDA aware numerical library whose interface is highly compatible with Ruby Numo}
   spec.description   = %q{Cumo is CUDA aware numerical library whose interface is highly compatible with Ruby Numo.}
   spec.homepage      = "https://github.com/sonots/cumo"
-  spec.license       = "BSD-3-Clause"
+  spec.license       = "MIT"
 
   spec.required_ruby_version = ">= 3.0.0"
 

data/ext/cumo/cuda/cudnn.c

@@ -53,6 +53,24 @@ BAD:
             env, cudnn_max_workspace_size);
 }
 
+VALUE
+cumo_cuda_cudnn_release_conv_held(VALUE held)
+{
+    cumo_cuda_cudnn_conv_held_t *h = (cumo_cuda_cudnn_conv_held_t*)held;
+
+    if (h->x_desc) { cudnnDestroyTensorDescriptor(h->x_desc); h->x_desc = 0; }
+    if (h->y_desc) { cudnnDestroyTensorDescriptor(h->y_desc); h->y_desc = 0; }
+    if (h->b_desc) { cudnnDestroyTensorDescriptor(h->b_desc); h->b_desc = 0; }
+    if (h->w_desc) { cudnnDestroyFilterDescriptor(h->w_desc); h->w_desc = 0; }
+    if (h->conv_desc) { cudnnDestroyConvolutionDescriptor(h->conv_desc); h->conv_desc = 0; }
+    // cuDNN is still reading the workspace, so the wait is the caller's to
+    // report. A free that raises from here would replace the exception being
+    // carried out, so it goes back the way an unwinding path has to.
+    cumo_cuda_runtime_return_scratch(h->workspace, h->workspace != NULL, &h->wait_status);
+    h->workspace = NULL;
+    return Qnil;
+}
+
 // Lazily initialize cudnn handle, and cache it
 cudnnHandle_t
 cumo_cuda_cudnn_handle()

data/ext/cumo/cuda/cudnn_impl.cpp

@@ -67,17 +67,11 @@ cumo_cuda_cudnn_GetConvTransposeOutDim(
 }
 
 cudnnStatus_t
-cumo_cuda_cudnn_CreateTensorDescriptor(
+cumo_cuda_cudnn_CreateTensorDescriptorFromShape(
         cudnnTensorDescriptor_t *desc,
-        VALUE a, cudnnDataType_t cudnn_dtype) {
+        int ndim, size_t *shape, cudnnDataType_t cudnn_dtype) {
     cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
-    cumo_narray_t *na;
-    CumoGetNArray(a, na);
-    int ndim = (int)(na->ndim);
-    size_t *shape = na->shape;
 
-    // only the shape is read: the descriptor is NCHW and the caller passes the
-    // contiguous copy, so a non-contiguous a still describes what is read
     status = cudnnCreateTensorDescriptor(desc);
     if (status != CUDNN_STATUS_SUCCESS) return status;
 
@@ -120,6 +114,47 @@ cumo_cuda_cudnn_CreateTensorDescriptor(
     return status;
 }
 
+// Only the shape is read: the descriptor is NCHW and the caller passes the
+// contiguous copy, so a non-contiguous a still describes what is read.
+cudnnStatus_t
+cumo_cuda_cudnn_CreateTensorDescriptor(
+        cudnnTensorDescriptor_t *desc,
+        VALUE a, cudnnDataType_t cudnn_dtype) {
+    cumo_narray_t *na;
+    CumoGetNArray(a, na);
+    return cumo_cuda_cudnn_CreateTensorDescriptorFromShape(desc, (int)(na->ndim), na->shape, cudnn_dtype);
+}
+
+// The bias is one value a channel, and cuDNN broadcasts it from a descriptor
+// shaped like the output. Shared because conv and conv_transpose add it alike.
+cudnnStatus_t
+cumo_cuda_cudnn_AddBias(cumo_cuda_cudnn_conv_run_t *r, cudnnDataType_t dtype, const void *one) {
+    size_t bias_shape[CUMO_NA_MAX_DIMENSION];
+    cumo_narray_t *nb;
+    cudnnStatus_t status;
+
+    if (!RTEST(r->b_cont)) return CUDNN_STATUS_SUCCESS;
+
+    CumoGetNArray(r->b_cont, nb);
+    bias_shape[0] = 1;
+    bias_shape[1] = nb->size;
+    for (size_t i = 0; i < r->ndim; ++i) {
+        bias_shape[i + 2] = 1;
+    }
+    status = cumo_cuda_cudnn_CreateTensorDescriptorFromShape(
+            &r->held.b_desc, (int)(r->ndim + 2), bias_shape, dtype);
+    if (status != CUDNN_STATUS_SUCCESS) return status;
+
+    return cudnnAddTensor(
+            r->handle,
+            one,
+            r->held.b_desc,
+            (void*)r->b_cont_ptr,
+            one,
+            r->held.y_desc,
+            (void*)r->y_ptr);
+}
+
 cudnnStatus_t
 cumo_cuda_cudnn_CreateFilterDescriptor(
         cudnnFilterDescriptor_t *desc,
@@ -322,6 +357,17 @@ struct AlgoCacheEntry {
     size_t memory;
     cudnnMathType_t math_type;
 };
+
+// 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.
+template <typename Map>
+void StoreAlgoCache(Map& map, const AlgoCacheKey& key, const typename Map::mapped_type& entry) {
+    try {
+        map[key] = entry;
+    } catch (...) {
+    }
+}
 }
 
 using FwdAlgoCacheMap = std::unordered_map<AlgoCacheKey, AlgoCacheEntry<cudnnConvolutionFwdAlgo_t>, AlgoCacheKeyHash>;
@@ -416,7 +462,7 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
                 perf_result,
                 (void*)workspace,
                 max_workspace_size);
-    cumo_cuda_runtime_free(workspace);
+    cumo_cuda_runtime_return_scratch(workspace, 1, NULL);
     if (status != CUDNN_STATUS_SUCCESS) return status;
     // A search that answers success with nothing to report would leave
     // perf_result untouched, and its algo and math type are used below.
@@ -424,7 +470,7 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
+    StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }
 
@@ -487,7 +533,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
                 perf_result,
                 (void*)workspace,
                 max_workspace_size);
-    cumo_cuda_runtime_free(workspace);
+    cumo_cuda_runtime_return_scratch(workspace, 1, NULL);
     if (status != CUDNN_STATUS_SUCCESS) return status;
     // A search that answers success with nothing to report would leave
     // perf_result untouched, and its algo and math type are used below.
@@ -495,7 +541,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
+    StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }
 
@@ -558,7 +604,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
                 perf_result,
                 (void*)workspace,
                 max_workspace_size);
-    cumo_cuda_runtime_free(workspace);
+    cumo_cuda_runtime_return_scratch(workspace, 1, NULL);
     if (status != CUDNN_STATUS_SUCCESS) return status;
     // A search that answers success with nothing to report would leave
     // perf_result untouched, and its algo and math type are used below.
@@ -566,7 +612,7 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
     if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
+    StoreAlgoCache(algo_cache_map, key, {perf_result->algo, perf_result->memory, perf_result->mathType});
     return status;
 }
 

data/ext/cumo/cuda/driver.c

@@ -449,15 +449,17 @@ Init_cumo_cuda_driver()
 
     // A driver API call needs a current context, and the runtime API only
     // creates its primary one once an array operation happens, so this covers
-    // the gap in between. Losing it is not fatal -- everything but a driver
-    // call made before any array operation still works -- so a device that
-    // refuses a context is left for that call to report. cuDeviceGet leaves
-    // cuDevice untouched when it fails, hence the guard.
+    // the gap in between. It has to be the primary context and not one of our
+    // own: cudaSetDevice binds the primary one, so anything built in a
+    // different context stops working the moment a caller sets the device,
+    // including the events a cuBLAS handle keeps inside itself.
+    // Losing it is not fatal -- everything but a driver call made before any
+    // array operation still works -- so a device that refuses a context is
+    // left for that call to report. cuDeviceGet leaves cuDevice untouched
+    // when it fails, hence the guard.
     if (cuDeviceGet(&cuDevice, 0) == CUDA_SUCCESS) {
-#if defined(CUDA_VERSION) && CUDA_VERSION >= 13000
-        cuCtxCreate(&context, NULL, 0, cuDevice);
-#else
-        cuCtxCreate(&context, 0, cuDevice);
-#endif
+        if (cuDevicePrimaryCtxRetain(&context, cuDevice) == CUDA_SUCCESS) {
+            cuCtxSetCurrent(context);
+        }
     }
 }

data/ext/cumo/cuda/runtime.c

@@ -2,6 +2,8 @@
 #include <assert.h>
 #include <cuda_runtime.h>
 #include "cumo/cuda/runtime.h"
+#include "cumo/cuda/memory_pool.h"
+#include "cumo/template_kernel.h"
 
 VALUE cumo_cuda_eRuntimeError;
 VALUE cumo_cuda_mRuntime;
@@ -20,6 +22,60 @@ cumo_cuda_runtime_check_kernel_launch(void)
     check_status(cudaGetLastError());
 }
 
+// A pointer given twice would be freed twice, and the second free would reach
+// past the pool to a block it still owns, so each one goes back once.
+static void
+release_held(char *p[5])
+{
+    int i, j;
+    for (i = 0; i < 5; ++i) {
+        if (p[i] == NULL) { continue; }
+        for (j = 0; j < i; ++j) {
+            if (p[j] == p[i]) { break; }
+        }
+        if (j == i) { cumo_cuda_runtime_free_no_raise(p[i]); }
+    }
+}
+
+// What a launcher holding scratch has to call instead of the plain check, for a
+// status it already took out of the slot: a library call in the window takes it
+// first otherwise, and whoever reads it first is the only one who sees it.
+// Raising is a longjmp, so the frees written below the check never run, and
+// unwinding frees through the form that cannot raise so one failure cannot hide
+// another.
+void
+cumo_cuda_runtime_check_taken_status_holding(int status, char *p0, char *p1, char *p2, char *p3, char *p4)
+{
+    if ((cudaError_t)status != cudaSuccess) {
+        char *held[5];
+        held[0] = p0; held[1] = p1; held[2] = p2; held[3] = p3; held[4] = p4;
+        release_held(held);
+    }
+    check_status((cudaError_t)status);
+}
+
+// The slot is read once here rather than peeked and read again: the frees below
+// can put a status of their own in it, and the launch is the one to report.
+void
+cumo_cuda_runtime_check_kernel_launch_holding(char *p0, char *p1, char *p2, char *p3, char *p4)
+{
+    cumo_cuda_runtime_check_taken_status_holding((int)cudaGetLastError(), p0, p1, p2, p3, p4);
+}
+
+// How scratch goes back from a handler that may be unwinding: a free that
+// raises there replaces the exception being carried, and leaves every free
+// below it unrun. wait_for_stream asks for the wait a caller needs when the
+// buffer is still being read, since the pool hands a freed chunk straight out.
+void
+cumo_cuda_runtime_return_scratch(char *ptr, int wait_for_stream, cudaError_t *status)
+{
+    if (wait_for_stream) {
+        cudaError_t wait = cudaStreamSynchronize(0);
+        if (status != NULL && *status == cudaSuccess) { *status = wait; }
+    }
+    if (ptr != NULL) { cumo_cuda_runtime_free_no_raise(ptr); }
+}
+
 int*
 cumo_cuda_runtime_error_flag_ptr(void)
 {

data/ext/cumo/cumo.c

@@ -21,6 +21,7 @@ void Init_cumo_uint16();
 void Init_cumo_uint32();
 void Init_cumo_uint64();
 void Init_cumo_hfloat();
+void Init_cumo_bfloat();
 void Init_cumo_sfloat();
 void Init_cumo_scomplex();
 void Init_cumo_dfloat();
@@ -31,6 +32,7 @@ void Init_cumo_dfloat_upcast();
 void Init_cumo_scomplex_upcast();
 void Init_cumo_sfloat_upcast();
 void Init_cumo_hfloat_upcast();
+void Init_cumo_bfloat_upcast();
 void Init_cumo_int64_upcast();
 void Init_cumo_uint64_upcast();
 void Init_cumo_int32_upcast();
@@ -167,6 +169,7 @@ Init_cumo()
     Init_cumo_scomplex();
     Init_cumo_sfloat();
     Init_cumo_hfloat();
+    Init_cumo_bfloat();
 
     Init_cumo_int64();
     Init_cumo_uint64();
@@ -186,6 +189,7 @@ Init_cumo()
     Init_cumo_scomplex_upcast();
     Init_cumo_sfloat_upcast();
     Init_cumo_hfloat_upcast();
+    Init_cumo_bfloat_upcast();
     Init_cumo_int64_upcast();
     Init_cumo_uint64_upcast();
     Init_cumo_int32_upcast();

data/ext/cumo/extconf.rb

@@ -65,6 +65,7 @@ narray/types/uint16
 narray/types/uint32
 narray/types/uint64
 narray/types/hfloat
+narray/types/bfloat
 narray/types/sfloat
 narray/types/dfloat
 narray/types/scomplex
@@ -80,6 +81,7 @@ narray/types/uint16_kernel
 narray/types/uint32_kernel
 narray/types/uint64_kernel
 narray/types/hfloat_kernel
+narray/types/bfloat_kernel
 narray/types/sfloat_kernel
 narray/types/dfloat_kernel
 narray/types/scomplex_kernel

data/ext/cumo/include/cumo/bit_reduce_kernel.h

@@ -442,7 +442,7 @@ static inline void cumo_bit_count_reduce(cumo_na_bit_reduction_arg_t arg, int in
     cumo_bit_detail::bit_count_partial_kernel<<<grid_size, block_size, shared_mem_size>>>(
         arg, p.ad, p.wa, invert, partial, p.n_split, p.chunk,
         (int)p.split_out_block_size, (int)p.split_reduce_block_size, p.unit_total_size);
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial);
 
     cumo_na_reduction_arg_t combine;
     combine.in.ptr = reinterpret_cast<char*>(partial);
@@ -452,7 +452,7 @@ static inline void cumo_bit_count_reduce(cumo_na_bit_reduction_arg_t arg, int in
     combine.in_indexer.total_size = partial_total_size;
     combine.out = arg.out;
     combine.out_indexer = arg.out_indexer;
-    cumo_reduce<uint64_t, uint64_t>(combine, cumo_bit_detail::BitCountImpl{});
+    cumo_reduce<uint64_t, uint64_t>(combine, cumo_bit_detail::BitCountImpl{}, reinterpret_cast<char*>(partial));
 
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
@@ -483,12 +483,12 @@ static inline void cumo_bit_stat_reduce(cumo_na_bit_reduction_arg_t arg, int sta
     cumo_bit_detail::bit_count_partial_kernel<<<grid_size, block_size, shared_mem_size>>>(
         arg, p.ad, p.wa, 0, partial, p.n_split, p.chunk,
         (int)p.split_out_block_size, (int)p.split_reduce_block_size, p.unit_total_size);
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial);
 
     int64_t combine_grid = (p.out_total_size + block_size - 1) / block_size;
     if (combine_grid > cumo_detail::max_grid_size) combine_grid = cumo_detail::max_grid_size;
     cumo_bit_detail::bit_stat_combine_kernel<<<combine_grid, block_size>>>(arg, p.ad, stat, partial, p.n_split);
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial);
 
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
@@ -519,12 +519,12 @@ static inline void cumo_bit_pred_reduce(cumo_na_bit_pred_reduction_arg_t arg, in
     cumo_bit_detail::bit_count_partial_kernel<<<grid_size, block_size, shared_mem_size>>>(
         arg, p.ad, p.wa, 0, partial, p.n_split, p.chunk,
         (int)p.split_out_block_size, (int)p.split_reduce_block_size, p.unit_total_size);
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial);
 
     int64_t combine_grid = (p.out_total_size + block_size - 1) / block_size;
     if (combine_grid > cumo_detail::max_grid_size) combine_grid = cumo_detail::max_grid_size;
     cumo_bit_detail::bit_pred_combine_kernel<<<combine_grid, block_size>>>(arg, p.ad, all, partial, p.n_split);
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial);
 
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }

data/ext/cumo/include/cumo/cuda/cudnn.h

@@ -97,9 +97,42 @@ cumo_cuda_cudnn_check_param_type(VALUE param, VALUE type, const char* name)
 void
 cumo_cuda_cudnn_check_status(cudnnStatus_t status);
 
+// Raises, so the methods that clean up through a label take it before they
+// hold a descriptor: a raise runs no label.
 cudnnHandle_t
 cumo_cuda_cudnn_handle();
 
+// What a convolution holds while it works. A raise runs no error label, so the
+// work goes under rb_ensure and the handler gives these back instead.
+typedef struct {
+    cudnnTensorDescriptor_t      x_desc;
+    cudnnTensorDescriptor_t      y_desc;
+    cudnnTensorDescriptor_t      b_desc;
+    cudnnFilterDescriptor_t      w_desc;
+    cudnnConvolutionDescriptor_t conv_desc;
+    char                        *workspace;
+    cudaError_t                  wait_status;
+} cumo_cuda_cudnn_conv_held_t;
+
+// Takes the held as its argument so it can be handed to rb_ensure as it is.
+VALUE
+cumo_cuda_cudnn_release_conv_held(VALUE held);
+
+// What conv and conv_transpose carry into the run the ensure above guards.
+typedef struct {
+    cumo_cuda_cudnn_conv_held_t held;
+    cudnnHandle_t handle;
+    VALUE   x_cont, w_cont, y, b_cont;
+    char   *x_cont_ptr, *w_cont_ptr, *y_ptr, *b_cont_ptr;
+    size_t  ndim;
+    int    *int_stride, *int_pad;
+    cudnnStatus_t status;
+} cumo_cuda_cudnn_conv_run_t;
+
+// Adds nothing when the run carries no bias, so callers need no guard.
+cudnnStatus_t
+cumo_cuda_cudnn_AddBias(cumo_cuda_cudnn_conv_run_t *r, cudnnDataType_t dtype, const void *one);
+
 // VALUE is Ruby Array
 static inline void
 cumo_cuda_cudnn_get_int_ary(int* int_ary, VALUE ary, size_t ndim, int default_value)
@@ -178,6 +211,13 @@ cumo_cuda_cudnn_CreateTensorDescriptor(
         VALUE a,
         cudnnDataType_t cudnn_dtype);
 
+cudnnStatus_t
+cumo_cuda_cudnn_CreateTensorDescriptorFromShape(
+        cudnnTensorDescriptor_t *desc,
+        int ndim,
+        size_t *shape,
+        cudnnDataType_t cudnn_dtype);
+
 cudnnStatus_t
 cumo_cuda_cudnn_CreateFilterDescriptor(
         cudnnFilterDescriptor_t *desc,

data/ext/cumo/include/cumo/cuda/runtime.h

@@ -28,6 +28,8 @@ cumo_cuda_runtime_check_status(cudaError_t status)
     }
 }
 
+void cumo_cuda_runtime_return_scratch(char *ptr, int wait_for_stream, cudaError_t *status);
+
 static inline void
 cumo_cuda_runtime_device_synchronize(void)
 {

data/ext/cumo/include/cumo/intern.h

@@ -66,6 +66,15 @@ VALUE cumo_na_ndloop4(cumo_ndfunc_t *nf, void *ptr, VALUE args);
 
 VALUE cumo_na_ndloop_cast_narray_to_rarray(cumo_ndfunc_t *nf, VALUE nary, VALUE fmt);
 VALUE cumo_na_ndloop_store_rarray(cumo_ndfunc_t *nf, VALUE nary, VALUE rary);
+
+// One device buffer for a whole walk, not one per row. A row that gives its
+// own back has to wait for the kernel still reading it, and that wait costs
+// more than everything else the row does. cumo_na_ndloop_store_rarray owns
+// one and hands it over as the loop's opt_ptr.
+// Opaque: an iterator asks for a buffer and reads nothing else.
+typedef struct cumo_na_stage_pool cumo_na_stage_pool_t;
+
+char *cumo_na_stage_pool_get(cumo_na_stage_pool_t *pool, size_t bytes);
 VALUE cumo_na_ndloop_store_rarray2(cumo_ndfunc_t *nf, VALUE nary, VALUE rary, VALUE opt);
 VALUE cumo_na_ndloop_inspect(VALUE nary, cumo_na_text_func_t func, VALUE opt);
 VALUE cumo_na_ndloop_with_index(cumo_ndfunc_t *nf, int argc, ...);