cumo 0.8.0 → 0.9.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 695bd66c93bd2511f905d74cf242a7d2d7d0ac4aea8f5915fb722e013e658e9a
-  data.tar.gz: fa51db7ef5f01745ded6d71834bde62fa238fb92cba9bb358374293a18b7577b
+  metadata.gz: 0f27ebdaa195fa967c144d6178ad3f9fbeca551162f8edacac205d5edd5d0fb6
+  data.tar.gz: 3d87faa55ab7c4c0e52bee3afdb0116d9d463fa778c408a46e200bff0cabdf81
 SHA512:
-  metadata.gz: fbad56229183188dc36b2341105a68a5ca670de3a86ac9d50887544bebf4200e80c0f571783aefe24657c609aa86bc3f6af0da8d5943b56c0c91031ea7c22244
-  data.tar.gz: a5a3e9cfa2e99873b3e7f555fec3e397ed8c9814fbd0255ccc2e0cd52ac9d337fd654741a00412f1b0eb46a777ed346fa46ae98c7e4bf45ce0d7eb8de964fa75
+  metadata.gz: 4b498a760e59725c15aa3e651ecbada01e52517833fd3132a860dbfe459e354f4afba47207cc33f2c94963f58f19a8a1d7e52ce722e51065876d637c8ddfc94b
+  data.tar.gz: bbda42b2ef7ce7ea882982a015bec06ff058ad42a2b72e4c8953e75facfcb64657c8da4b69d63666cf4c06d3c265dc7e0acf54c431450f181b39cbeebed9560e

data/CHANGELOG.md

@@ -1,3 +1,37 @@
+# 0.9.0 (2026/09/19)
+
+Breaking changes:
+
+* An integer array divided by a zero Ruby number raises `ZeroDivisionError` before the kernel runs rather than after it, so `a.inplace / 0` leaves `a` alone where it wrote over it first. The divisor was on the host all along, so the flag the device raised and the `cudaDeviceSynchronize` that read it back are gone: `Cumo::Int32.new(1, 768)` costs 2.34us for `a / 2` where it cost 9.27, and 2.20 for `a % 2` where it cost 7.88. `2 / a` and an array by an array keep the flag, the divisor really being on the device there (PR #466)
+
+Fixes:
+
+* Fix `mulsum` reducing everything where the receiver is a view that marked its axes with `:sum`, `:reduce` or `:+` and the operands have different dtypes. The marking lives on the array the method is asked of, and `mulsum` answers another class by putting a different array in the receiver's place, so `Cumo::Int8.new(4, 6)[:sum, true].mulsum(a_float)` answered a scalar where the same call with matching dtypes answers a row of six. numo answers the same wrong thing (PR #480)
+* Fix an elementwise operation whose inplace receiver shares memory with an operand that walks it differently, `a.inplace + a.transpose` answering something new on every run: the element a thread reads is one another thread is writing, and nothing orders the two on a device. An operand sitting exactly where the output does is left alone, which is the common `a.inplace + a` (PR #473)
+* Fix `to_binary` handing out whatever lay past the end of a `Cumo::Bit` array whose size is not a multiple of eight. `Cumo::Bit.new(24).fill(1)[0..3].to_binary` answered `[255]` where four bits belong to it, and a view copied first answered the memory pool's contents, so the same expression could answer differently from one run to the next. `marshal_dump` goes through here too (PR #471)
+* Fix an inplace view built from an index array reducing from the wrong rows. The copy that makes such a view contiguous carried `CUMO_NDF_INPLACE`, so for an inplace receiver it answered self and the reduction walked the index-backed view it meant to have copied. `median` and `max_index` answered their own wrong rows the same way, and inplace survives an operation, so `(v.inplace * 2).sum(axis: 2)` reached it too (PR #461)
+* Hold `layer_norm`, `rms_norm` and `softmax` to the same checks. `softmax` measured the result by its element count where the other two compare the shape, so an `allocate` that reshaped either side to the same count left every check true: reshaping the result answered the right values under the wrong shape, and reshaping the receiver normalized along the wrong axis (PR #459)
+
+Changes:
+
+* Add `rms_norm`, which normalizes a row in one kernel rather than the six it is written out of. This is `layer_norm` without the centring, the layer Llama and the models after it normalize with, so it takes no beta. On SFloat `[1, 768]` it goes from 14.2 to 2.9us and `[4096, 768]` from 122.5 to 38.5. No shape is slower (PR #458)
+* Add `silu` to `Cumo::NMath`, which takes one kernel rather than the five it is written out of: 768 elements go from 13.4 to 3.6us. It follows `torch.nn.functional.silu`, answering within one unit in the last place of the true value at the nine double points measured, and matching it at the edges (PR #460)
+* Add `quantize_symmetric`, which quantizes a row to `Cumo::Int8` in one kernel rather than the six it is written out of and hands the scale back. A launch costs about two microseconds whatever it is handed, so a short row pays for those rather than for its bytes (PR #476)
+* Speed up `mulsum`. Operands of different dtypes are read where they already are and converted one element at a time rather than cast into an array of the wider type first, so `SFloat#mulsum(Int8)` at `[4096, 4096]` goes from 384 to 53us. A broadcast operand no longer falls off the fast path, `1000x12x64` against `1000x12x1` over axis 0 going from 13.5 to 4.9us, and the thread layout takes the shorter step of the two operands rather than the first one's, so `DFloat[64,1].mulsum(DFloat[64,1024], axis: 1)` goes from 16.26 to 5.68 whichever order it is written in (PR #479, PR #477, PR #475)
+* Fix every allocation slowing down as the memory pool fragments, the arena having been rebuilt whenever an allocation emptied a bin. GPT-2 124M generating 256 tokens goes from 11.9 to 3.2 seconds in fp16, where the pool ends up with 4035 free chunks against fp32's 71 (PR #474)
+* Range check an NArray index on the device rather than copying the indices to the host and back. `Cumo::SFloat.new(50257, 768)`, `a[idx, true]`: 4096 indices go from 58.2 to 22.2us and 64 from 20.8 to 13.8. Nothing reads managed memory from the host on that path any more, which is what breaks under Ractor (sonots/cumo#180) (PR #470)
+* Copy a transposed view through the shared-memory tile the typed store has used since 0.7.0. `Cumo::SFloat.new(4096, 1024)`: `a.transpose.copy` goes from 338 to 97us, level with `a.transpose.dup` at 102, and `a.transpose.max_index(axis: 1)` from 362 to 107. Every reduction that needs a contiguous operand goes through this copy (PR #468)
+* Reduce a view built from an index array without staging it in a buffer first. Copying a `[6250, 64]` view of a `[100000, 64]` array goes from 13.9 to 7.2us and reducing it from 25.1 to 18.3 (PR #462)
+* Let `store_binary` reach a view it could only read back from. `to_binary` reads any view, so `v.store_binary(v.to_binary)` raised for the same `v` that had just produced the string. A view the bytes cannot move into directly now takes them through an array of its own class and shape (PR #472)
+* Answer `respond_to?` for the methods `Cumo::NMath` dispatches. Every one of them arrives through `method_missing`, so `respond_to?` answered false for `sqrt` and `exp` alike, which code that checks for a feature before calling it reads as absence (PR #463)
+* Take one Ruby method dispatch out of every ndloop call, the argument list always being a plain Array where `rb_obj_dup` took the general path: about 800 fewer retired instructions per operation, 3.4% of the 23,200 an operation takes (PR #465)
+* Tell which array a type check rejected and what it wanted. It said "invalid NArray type (class)", naming neither, so two checks side by side read exactly alike; it now says "gamma must be Cumo::SFloat, not Cumo::HFloat" (PR #478)
+* Say in the README that `reshape` answers a copy where slicing answers a view, priced across four shapes: on a GPU it is an allocation and a copy kernel every call, 58.1us and 12MB for a `4096x768` SFloat against 0.22us and nothing for `reshape!` (PR #457)
+* Show the changelog, source and issue links on the gem page, and require multi-factor authentication to push the gem (PR #456)
+* Share one body between the half and bfloat16 macros so a fix reaches both, 399 duplicated lines where it would have reached one. The compiled code is unchanged: the device SASS is identical (PR #464)
+* Keep the index-array test in one place so copy and store cannot drift apart, 172 generated copies of it becoming one (PR #467)
+* Drop the `Cumo::NMath::DISPATCH` rows no argument can reach. A bare Ruby numeric takes an NArray row, so the `Integer`, `Float` and `Complex` ones were never read: 944 calls over 16 kinds of argument reached them zero times (PR #469)
+
 # 0.8.0 (2026/09/16)
 
 Breaking changes:

data/README.md

@@ -178,6 +178,44 @@ x += p_dir * alpha     # and consumed there, without crossing the bus
 
 Read the value back once the loop is done, or every k iterations if it has to test something.
 
+### Reshape Copies, Reshape! Does Not
+
+`reshape` answers a copy of the whole array, never a view.
+Slicing answers a view, so the two look alike and are not:
+
+```ruby
+a = Cumo::SFloat.new(4, 6).seq
+a[0..1, true][0, 0] = 77.0   # a view, so a changes
+a.reshape(2, 12)[0, 0] = 99.0  # a copy, so a does not
+```
+
+This is what Numo does too, and numpy is where the expectation comes from: there `reshape` answers a view whenever the strides allow one.
+On a GPU the difference is an allocation and a copy kernel, paid every call.
+`reshape!` changes the receiver in place and costs neither.
+RTX 5070 Ti Laptop, `Cumo::SFloat`, 200 calls a measurement:
+
+```
+shape        reshape     reshape!    the copy allocates
+1x768        5.1 us      0.04 us          3 KB
+1024x768    39.1 us      0.18 us          3 MB
+4096x768    58.1 us      0.22 us         12 MB
+8192x768   142.7 us      0.43 us         24 MB
+```
+
+`reshape!` is host-side bookkeeping, so it stays under a microsecond whatever the array weighs.
+The `reshape` column is the copy, and it grows with the bytes.
+
+The catch is that `reshape!` changes the array everything else is holding.
+It fits a temporary the calling expression owns, and not an argument, an ivar, or anything a cache still points at:
+
+```ruby
+x = a * b            # a temporary nothing else holds
+x.reshape!(t, n, h)  # free
+```
+
+Where the array is not yours to change, the copy is the price of the shape.
+Reach for `reshape!` when a profile says the copies are worth removing, not by default.
+
 ### Ruby Floats In NMath Promote To Double
 
 `Cumo::NMath` picks the module it dispatches to from every argument it is given, and a Ruby `Float` counts as a `DFloat` there.
@@ -224,14 +262,18 @@ 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.
+`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.
+Written out of the operators they take nine launches, six, five and six, 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
+z = x.rms_norm(gamma, eps: 1e-5)           # x / sqrt(mean(x * x) + eps) * gamma
 probs = scores.softmax                     # exp(x - max) / sum(exp(x - max))
+xq, scale = x.quantize_symmetric           # scale = max(|x|) / 127, xq = round(x / scale)
 ```
 
+`quantize_symmetric` answers a `Cumo::Int8` shaped like self and the scale of every row, which is self's shape without its last axis. `xq * scale[false, :new]` is what the row stood for. A row of zeros has no scale to divide by and answers zero for both, and a row holding an infinity or a NaN answers that in its scale and zeros in the row, since neither is a value 8 bits could carry. The scale comes back in the class the reduction accumulates in, which is `Cumo::SFloat` for `Cumo::HFloat` and `Cumo::BFloat` and self's own otherwise.
+
 On an RTX 5070 Ti Laptop, against the same arithmetic spelled with operators, in microseconds:
 
 ```
@@ -248,23 +290,58 @@ shape              fused  written    fused  written    fused  written
 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
+
+rms_norm                SFloat            DFloat            HFloat
+shape              fused  written    fused  written    fused  written
+1 x 768              2.9     14.2      5.8     14.1      3.1     15.1
+256 x 768            4.6     16.7      6.9     20.2      5.2     15.5
+4096 x 768          38.5    122.5    130.8    367.0     19.2     80.8
+1 x 1000000         12.5     35.8     23.7    102.2     12.0     33.2
 ```
 
-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.
+```
+quantize_symmetric      SFloat            DFloat            HFloat
+shape              fused  written    fused  written    fused  written
+1 x 768              3.7     20.2      6.0     15.7      2.3     11.9
+256 x 768            4.7     15.0     20.2     21.3      4.7     13.0
+4096 x 768          24.4    105.4    289.5    333.3     23.8     80.3
+1 x 1000000         18.8     40.8     97.1     81.6     15.7     37.5
+```
+
+The tables were taken in separate sessions, so read each row against the row beside it and not across the tables.
+
+`quantize_symmetric` in `Cumo::DFloat` is worth less than the others and loses outright on a million elements, whatever shape they are in. The division it does per element is what costs: on this card one row of a million takes 44.0 us to divide in double against 10.8 in single, where the reduction over the same row takes 25.3 and 15.4. Six kernels give that division a kernel of its own to fill the device with, and one kernel leaves it behind the reduction.
+
+All three 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.
+All three 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.
+`layer_norm` divides the variance by the row length and not by one less, which is how the layer is defined and unlike `var`; `rms_norm` divides its mean square the same way.
+All three 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.
 
+`rms_norm` is `layer_norm` without the centring, and takes no `beta`, which is the layer Llama and the models after it normalize with.
+The two agree on a row whose mean is already zero and differ on every other row, so they are not interchangeable:
+
+```ruby
+x = Cumo::SFloat[[1.0, 2.0, 3.0, 4.0]]
+ones = Cumo::SFloat.ones(4)
+zeros = Cumo::SFloat.zeros(4)
+x.rms_norm(ones)               #=> [[0.365, 0.730, 1.095, 1.461]]
+x.layer_norm(ones, zeros)      #=> [[-1.342, -0.447, 0.447, 1.342]]
+```
+
+All three accumulate in single precision, so a half row that squares past the 65504 half holds is still answered the way double answers it.
+bfloat16 carries a float's exponent rather than a half's, so a bfloat16 row that squares past what a float holds overflows the accumulator too: `Cumo::BFloat[[-3e20, 3e20, 3e20, -3e20]].rms_norm(ones)` answers zeros where double answers ones.
+
 `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.
+`layer_norm` and `softmax` above, with `gelu_tanh` from the section below, are where a GPT-2 style 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.
+A Llama style one spends them on `rms_norm` and `silu` instead, which is why they are here.
 
 ### Two Spellings Of gelu
 
@@ -284,6 +361,34 @@ 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`.
 
+### SiLU
+
+`Cumo::NMath.silu` is `x * sigmoid(x)`, the activation Llama and the models after it use where GPT-2 uses `gelu_tanh`.
+It is also called Swish.
+
+```ruby
+Cumo::NMath.silu(x)   # x / (1 + exp(-x)), in one kernel rather than five
+```
+
+Written out of the operators it costs five launches, and one kernel runs 1.1x to 4.3x faster on an RTX 5070 Ti Laptop, in microseconds:
+
+```
+silu                    SFloat            DFloat            HFloat
+elements           fused  written    fused  written    fused  written
+768                  3.6     13.4      3.6      9.5      3.4     14.3
+786432               7.0     25.9     90.5    102.4      5.2     19.2
+16777216           353.2   1312.9   1984.7   3258.8    199.6    851.8
+```
+
+The three rows are three regimes rather than one curve: 768 elements pay for the launches, 786432 fit in L2 and move faster than this card reads from memory, and 16777216 are what it costs from DRAM.
+
+It follows `torch.nn.functional.silu`, answering within one unit in the last place of the true value at the nine double points measured, as torch does.
+At the edges it matches torch exactly: `-Float::INFINITY` answers `NaN`, since that is what infinity times zero is, and a large negative `x` answers a signed zero.
+
+Where that zero starts depends on the type, and not on `exp` alone.
+Single and bfloat16 reach it at -89, where `exp(-x)` passes what a float holds, and double never does.
+Half reaches it at -21, because the round back to half gets there first.
+
 ### Half Precision
 
 `Cumo::HFloat`, also reachable as `Cumo::Float16`, holds IEEE binary16: one sign bit, five of exponent and ten of mantissa.
@@ -430,7 +535,7 @@ The two do not contain each other, so an expression mixing them promotes to `Cum
 
 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`.
+`layer_norm`, `rms_norm`, `softmax`, `silu` 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:

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 gelu gelu_tanh].each do |f|
+   sinh cosh tanh asinh acosh atanh cbrt erf erfc sinc gelu gelu_tanh silu].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

@@ -16,6 +16,12 @@ Gem::Specification.new do |spec|
   spec.homepage      = "https://github.com/sonots/cumo"
   spec.license       = "MIT"
 
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+  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"
+
   spec.required_ruby_version = ">= 3.0.0"
 
   spec.files = `git ls-files -z`.split("\x0").reject do |f|

data/ext/cumo/cuda/memory_pool_impl.cpp

@@ -108,7 +108,16 @@ bool SingleDeviceMemoryPool::RemoveFromFreeList(size_t size, std::shared_ptr<Chu
     }
     assert(arena.size() > arena_index);
     FreeList& free_list = arena[arena_index];
-    return EraseFromFreeList(free_list, chunk);
+    if (!EraseFromFreeList(free_list, chunk)) {
+        return false;
+    }
+    // A bin left empty here would be walked by every later search, and the
+    // merge path below Free empties one whenever it takes a neighbour out.
+    if (free_list.empty()) {
+        arena.erase(arena.begin() + arena_index);
+        arena_index_map.erase(arena_index_map.begin() + arena_index);
+    }
+    return true;
 }
 
 intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
@@ -132,6 +141,7 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
 
         // find best-fit, or a smallest larger allocation
         Arena& arena = GetArena(stream_ptr);
+        ArenaIndexMap& arena_index_map = GetArenaIndexMap(stream_ptr);
         size_t arena_index = GetArenaIndex(size, stream_ptr);
         size_t arena_length = arena.size();
         for (size_t i = arena_index; i < arena_length; ++i) {
@@ -141,10 +151,15 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
             }
             chunk = PopFromFreeList(free_list);
             if (free_list.empty()) {
-                // An emptied bin stays in the arena otherwise, and every later
-                // search walks it. Dropping it here invalidates arena and
-                // free_list, so nothing below this loop may touch them.
-                CompactIndex(stream_ptr, false);
+                // Drop just this bin. Rebuilding the whole arena instead
+                // copied every other bin's free list, and a free list holds
+                // shared pointers, so a fragmented pool touched the reference
+                // count of every chunk it held on every allocation.
+                // This leaves free_list, the elements of arena_index_map and
+                // every iterator from i on dangling, so nothing below may
+                // touch them.
+                arena.erase(arena.begin() + i);
+                arena_index_map.erase(arena_index_map.begin() + i);
             }
             break;
         }
@@ -252,7 +267,7 @@ bool SingleDeviceMemoryPool::Free(intptr_t ptr) {
     return true;
 }
 
-void SingleDeviceMemoryPool::CompactIndex(cudaStream_t stream_ptr, bool free) {
+void SingleDeviceMemoryPool::CompactIndex(cudaStream_t stream_ptr) {
     // need lock ouside this function
     if (!HasArena(stream_ptr)) return;
 
@@ -266,22 +281,17 @@ void SingleDeviceMemoryPool::CompactIndex(cudaStream_t stream_ptr, bool free) {
         if (free_list.empty()) {
             continue;
         }
-        if (free) {
-            FreeList keep_list;
-            for (auto chunk : free_list) {
-                if (chunk->prev() != nullptr || chunk->next() != nullptr) {
-                    keep_list.emplace_back(chunk);
-                }
-            }
-            if (keep_list.size() == 0) {
-                continue;
+        FreeList keep_list;
+        for (auto chunk : free_list) {
+            if (chunk->prev() != nullptr || chunk->next() != nullptr) {
+                keep_list.emplace_back(chunk);
             }
-            new_arena_index_map.emplace_back(arena_index_map[arena_index]);
-            new_arena.emplace_back(keep_list);
-        } else {
-            new_arena_index_map.emplace_back(arena_index_map[arena_index]);
-            new_arena.emplace_back(free_list);
         }
+        if (keep_list.size() == 0) {
+            continue;
+        }
+        new_arena_index_map.emplace_back(arena_index_map[arena_index]);
+        new_arena.emplace_back(keep_list);
     }
     if (new_arena.empty()) {
         index_.erase(stream_ptr);
@@ -299,7 +309,7 @@ void SingleDeviceMemoryPool::FreeAllBlocks() {
     std::vector<cudaStream_t> keys(free_.size());
     transform(free_.begin(), free_.end(), keys.begin(), [](auto pair) { return pair.first; });
     for (cudaStream_t stream_ptr : keys) {
-        CompactIndex(stream_ptr, true);
+        CompactIndex(stream_ptr);
     }
 }
 
@@ -307,7 +317,7 @@ void SingleDeviceMemoryPool::FreeAllBlocks() {
 void SingleDeviceMemoryPool::FreeAllBlocks(cudaStream_t stream_ptr) {
     std::lock_guard<std::recursive_mutex> lock{mutex_};
 
-    CompactIndex(stream_ptr, true);
+    CompactIndex(stream_ptr);
 }
 
 size_t SingleDeviceMemoryPool::GetNumFreeBlocks() {

data/ext/cumo/cuda/memory_pool_impl.hpp

@@ -284,7 +284,9 @@ public:
     //         be found in the free list as the chunk is allocated.)
     bool RemoveFromFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr);
 
-    void CompactIndex(cudaStream_t stream_ptr, bool free);
+    // Drops every chunk that was never split from a larger one, and every
+    // bin left empty by that.
+    void CompactIndex(cudaStream_t stream_ptr);
 };
 
 // Memory pool for all GPU devices on the host.

data/ext/cumo/cuda/memory_pool_impl_test.cpp

@@ -337,10 +337,10 @@ public:
         assert(pool_->GetFreeBytes() == kRoundSize * 8);
     }
 
-    // RemoveFromFreeList leaves an emptied bin in the arena -- it erases the
-    // chunk and returns, where Malloc drops one through CompactIndex. The
-    // assertions below spell that out, so they describe the arena as it is
-    // rather than wait on a compaction this path is not getting.
+    // RemoveFromFreeList drops a bin its erase leaves empty, the same as
+    // Malloc. The merge path below Free empties one whenever it takes a
+    // neighbour out, so a bin left behind here would be walked by every later
+    // search.
     void TestRemoveFromFreeList() {
         Arena& arena = pool_->GetArena(stream_ptr_);
         ArenaIndexMap& arena_index_map = pool_->GetArenaIndexMap(stream_ptr_);
@@ -372,36 +372,30 @@ public:
         assert(arena_index_map[1] == 3);
         assert(arena_index_map[2] == 4);
 
-        // remove two from two
+        // remove two from two, which empties that bin
         pool_->RemoveFromFreeList(chunk2->size(), chunk2, stream_ptr_);
-        assert(arena.size() == 3);
+        assert(arena.size() == 2);
         assert(arena[0].size() == 1);
-        assert(arena[1].size() == 0);
-        assert(arena[2].size() == 1);
-        assert(arena_index_map.size() == 3);
+        assert(arena[1].size() == 1);
+        assert(arena_index_map.size() == 2);
         assert(arena_index_map[0] == 2);
-        assert(arena_index_map[1] == 3);
-        assert(arena_index_map[2] == 4);
+        assert(arena_index_map[1] == 4);
 
         pool_->RemoveFromFreeList(chunk3->size(), chunk3, stream_ptr_);
-        assert(arena.size() == 3);
+        assert(arena.size() == 1);
         assert(arena[0].size() == 1);
-        assert(arena[1].size() == 0);
-        assert(arena[2].size() == 0);
-        assert(arena_index_map.size() == 3);
+        assert(arena_index_map.size() == 1);
         assert(arena_index_map[0] == 2);
-        assert(arena_index_map[1] == 3);
-        assert(arena_index_map[2] == 4);
 
         pool_->RemoveFromFreeList(chunk4->size(), chunk4, stream_ptr_);
-        assert(arena.size() == 3);
-        assert(arena[0].size() == 0);
-        assert(arena[1].size() == 0);
-        assert(arena[2].size() == 0);
-        assert(arena_index_map.size() == 3);
-        assert(arena_index_map[0] == 2);
-        assert(arena_index_map[1] == 3);
-        assert(arena_index_map[2] == 4);
+        assert(arena.size() == 0);
+        assert(arena_index_map.size() == 0);
+
+        // A chunk that is not in the free list leaves the arena alone.
+        pool_->AppendToFreeList(chunk1->size(), chunk1, stream_ptr_);
+        assert(!pool_->RemoveFromFreeList(chunk2->size(), chunk2, stream_ptr_));
+        assert(arena.size() == 1);
+        assert(arena[0].size() == 1);
     }
 
     void TestMalloc() {

data/ext/cumo/include/cumo/check.h

@@ -11,10 +11,18 @@ extern "C" {
 #endif
 #endif
 
-#define CUMO_CHECK_NARRAY_TYPE(x,t)                            \
-    if (rb_obj_class(x)!=(t)) {                                \
-        rb_raise(rb_eTypeError,"invalid NArray type (class)"); \
+// Two of these often sit side by side, so the message has to say which array
+// was turned down as well as what it is.
+static inline void
+cumo_na_check_narray_type(VALUE x, VALUE type, const char *name)
+{
+    if (rb_obj_class(x) != type) {
+        rb_raise(rb_eTypeError, "%s must be %" PRIsVALUE ", not %" PRIsVALUE,
+                 name, type, rb_obj_class(x));
     }
+}
+
+#define CUMO_CHECK_NARRAY_TYPE(x,t,name) cumo_na_check_narray_type((x),(t),(name))
 
 #define CUMO_CHECK_SIZE_EQ(sz1,sz2)                  \
     if ((sz1) != (sz2)) {                            \

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

@@ -38,7 +38,7 @@ cumo_cuda_cudnn_check_output(VALUE out, VALUE type, size_t ndim, size_t *shape)
 {
     cumo_narray_t *na;
 
-    CUMO_CHECK_NARRAY_TYPE(out, type);
+    CUMO_CHECK_NARRAY_TYPE(out, type, "output");
     CumoGetNArray(out, na);
     CUMO_CHECK_DIM_EQ((size_t)(na->ndim), ndim);
     for (size_t idim = 0; idim < ndim; ++idim) {
@@ -75,7 +75,7 @@ cumo_cuda_cudnn_check_input(VALUE in, VALUE type, size_t ndim, size_t *shape)
 {
     cumo_narray_t *na;
 
-    CUMO_CHECK_NARRAY_TYPE(in, type);
+    CUMO_CHECK_NARRAY_TYPE(in, type, "input");
     CumoGetNArray(in, na);
     CUMO_CHECK_DIM_EQ((size_t)(na->ndim), ndim);
     for (size_t idim = 0; idim < ndim; ++idim) {
@@ -88,10 +88,7 @@ cumo_cuda_cudnn_check_input(VALUE in, VALUE type, size_t ndim, size_t *shape)
 static inline void
 cumo_cuda_cudnn_check_param_type(VALUE param, VALUE type, const char* name)
 {
-    if (rb_obj_class(param) != type) {
-        rb_raise(rb_eTypeError, "%s must be %s, not %s",
-                 name, rb_class2name(type), rb_obj_classname(param));
-    }
+    CUMO_CHECK_NARRAY_TYPE(param, type, name);
 }
 
 void

data/ext/cumo/include/cumo/indexer.h

@@ -176,6 +176,22 @@ cumo_na_make_iarray_stridx(cumo_na_loop_args_t* arg)
     return iarray;
 }
 
+// Whether any argument of the loop is backed by an index array. With
+// INDEX_LOOP an operand can arrive carrying one, and cumo_na_iarray_t holds
+// byte steps only: an indexed dimension has a step of zero there, so a loop
+// that answers yes has to address its operands through the stridx form above.
+static int
+cumo_na_loop_has_index(cumo_na_loop_t *const lp)
+{
+    int j, i;
+    for (j = 0; j < lp->narg; ++j) {
+        for (i = 0; i < lp->args[j].ndim; ++i) {
+            if (lp->args[j].iter[i].idx) return 1;
+        }
+    }
+    return 0;
+}
+
 static cumo_na_bit_iarray_stridx_t
 cumo_na_make_bit_iarray_stridx(cumo_na_loop_args_t* arg)
 {

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

@@ -54,6 +54,8 @@ void cumo_na_release_lock(VALUE); // currently do nothing
 // used in reduce methods
 VALUE cumo_na_reduce_dimension(int argc, VALUE *argv, int naryc, VALUE *naryv,
                             cumo_ndfunc_t *ndf, cumo_na_iter_func_t nan_iter);
+VALUE cumo_na_get_reduce(VALUE v);
+void cumo_na_set_reduce(VALUE v, VALUE reduce);
 
 VALUE cumo_na_reduce_options(VALUE axes, VALUE *opts, int naryc, VALUE *naryv,
                           cumo_ndfunc_t *ndf);