cumo 0.5.9 → 0.5.10

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f211eb1858b22b8418a8494efb431afeca60f023de409d74e6bdb1ae1fc3753a
-  data.tar.gz: '085be7ea42f37e235945e62783c51522c206c22eb539557b965812cc71ce422d'
+  metadata.gz: bcb51638cb488e5abcb0bf885744ea1478f4daefd4945f6af8becf252bc85f29
+  data.tar.gz: 75aed5ea382f37a823992e29ef2f3fe9ee31a7b7c11350d2541e5f4fc0c0fb59
 SHA512:
-  metadata.gz: 19a6ca3c65b190e2347624232fa33302f4485e84d90df7de6972d727d76353e20a15a00bf8a7c8c96d281a97bbd7f8084016b8c2961e6781e4d94b00b17287d1
-  data.tar.gz: 2df512d0a5e2b201cd82cc91d4dc258f4d3463097e2969250ab0aa81ba591abaaa89ea0c6d1e818d8b85f32d698a6b9cc529d961a8e432be540fb7f4144f2e70
+  metadata.gz: ee03167a0966008c99598d90623757772060c09bad859af35d313dda831975461cba49c8da251a4793e6a9b933cd5ec75588038381d55e492398d0b4e2686559
+  data.tar.gz: 2e2b3e1a6f0f304be8b2375be48061ba73926973d4ec966e49d0b50f666d88d7078cec3c59ea797401bd72b91ec96a4f115bf2ed49485dc1d7ca3b6a29fb1c48

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

@@ -40,7 +40,9 @@ module MakeMakefileCuda
       s = MakeMakefileCuda::Nvcc.generate(argv)
       cmd = "nvcc #{s}"
       if ENV['CUMO_NVCC_GENERATE_CODE']
-        cmd << " --generate-code=#{ENV['CUMO_NVCC_GENERATE_CODE']}"
+        ENV['CUMO_NVCC_GENERATE_CODE'].split.each do |code|
+          cmd << " --generate-code=#{code}"
+        end
       else
         capability = nil
         if find_executable('nvidia-smi')

data/CHANGELOG.md

@@ -1,3 +1,43 @@
+# 0.5.10 (2026/08/23)
+
+Breaking changes:
+
+* An indexed assignment to a frozen view raises instead of writing through it; `store` and `fill` already raised (PR #306)
+* The cuDNN and cuBLAS entry points turn away arguments they used to read or write out of bounds: `pooling_backward` requires `y` and `gy` to carry its own dtype and the pooling output shape of `x`, `gemm` requires a contiguous `c`, `batch_norm` requires the reduced size to cover x's channels, and `conv_grad_w` raises on a mismatched `gy` rather than asserting (PR #303, PR #302, PR #300, PR #278)
+* `seq` on the unsigned types wraps a negative start and a start past 2**63 the way `fill` and `cast` do, where they collapsed to zero and clamped to INT64_MAX (PR #295, PR #294)
+* A negative power of an unsigned array answers zero instead of spinning the device until SIGKILL (PR #293)
+* A subscript that is a Cumo::RObject subclass, and an `expand_dims` past `CUMO_NA_MAX_DIMENSION`, raise instead of being taken (PR #272, PR #270)
+
+Fixes:
+
+* Fix `[]=` writing through a frozen view, the store going to a derived view whose data object is the root rather than the receiver (PR #306)
+* Fix `at()` reading the accumulator in place of the subscript on a reversed view, which addressed far in front of the buffer (PR #305)
+* Fix `cumo_cuda_cudnn_CreateBNTensorDescriptor` returning before it derived the descriptor, an inverted status check that no caller reached (PR #304)
+* Fix `batch_norm`, `fixed_batch_norm` and `batch_norm_backward` reading and writing past their parameters, the sizes being checked against the reduced shape while cuDNN reaches x's channel count (PR #303)
+* Fix `gemm` writing past its allocation when `c` is a non-contiguous inplace view, and over whatever the pool put next to it (PR #302)
+* Fix `pooling_backward` reading past `y` and `gy`, and returning what it read in `gx` (PR #300)
+* Fix `pow_int` negating INT32_MIN, which overflows and leaves the loop unbounded (PR #299)
+* Fix the device free running for a subscript that allocated no index array (PR #297)
+* Fix a NULL dereference from `at()` given a scalar subscript (PR #271)
+* Fix an out-of-bounds write through an unchecked cuDNN output array (PR #269)
+* Fix `gemm` reading past an operand that carries fewer batch dimensions (PR #268)
+
+Changes:
+
+* Pass a numeric operand to the kernel instead of casting it to a 0-dimensional array first, in the operators, `clip`, `pow`, the comparisons, the NMath functions and the coerced left-hand side; GPT-2 124M decode drops from 787 to 677 kernel launches a token (PR #296, PR #291, PR #290, PR #289, PR #288, PR #287)
+* Run the NMath binary functions through the indexer loop instead of one kernel launch per row; 11x on a transposed operand and up to 60x on the cheap kernels (PR #292)
+* Copy a reduction operand only when it carries an index array, where any non-contiguous view was copied whole (PR #286)
+* Run the nan-aware reductions, the nan-aware index reductions, `kahan_sum` and `Bit#mask` on the GPU (PR #277, PR #276, PR #275, PR #273)
+* Take a transposed batch of matrices with the cuBLAS transpose flag rather than duplicating it, which saves a temporary the size of the operand (PR #283)
+* Emit indexer accessors up to eight dimensions (PR #284)
+* Drop a free-list bin from the pool arena when `Malloc` empties it (PR #279)
+* Check the cuDNN header and version, not just the library (PR #280)
+* Document every method that returns a 0-dimensional NArray, and what keeping scalars on the device buys (PR #298, PR #282)
+* Build two GPU architectures in CI instead of every one (PR #285)
+* Add `fiddle` to the Gemfile (PR #281)
+* Remove TODO comments whose questions have been answered (PR #274)
+* Cover the Int32 array exponent in the INT_MIN power test (PR #301)
+
 # 0.5.9 (2026/08/20)
 
 Breaking changes:

data/Gemfile

@@ -5,6 +5,7 @@ source "https://rubygems.org"
 gemspec
 
 gem 'extconf_compile_commands_json'
+gem 'fiddle'
 gem 'rake-compiler'
 gem 'test-unit'
 gem 'yard'

data/README.md

@@ -9,6 +9,7 @@ Cumo (pronounced "koomo") is a CUDA-aware, GPU-optimized numerical library that
 * Ruby 3.0 or later
 * NVIDIA GPU Compute Capability 3.5 (Kepler) or later
 * CUDA 11.0 or later
+* cuDNN 8.0 or later (optional, for the cuDNN features)
 
 ## Preparation
 
@@ -97,17 +98,31 @@ a = xm::DFloat.new(3,5).seq
 
 ### Incompatibility With Numo
 
-The following methods behave incompatibly with Numo by default for performance reasons:
+Numo returns a Ruby numeric object wherever a result is 0-dimensional, while Cumo returns the 0-dimensional NArray itself.
+Cumo differs in this way to avoid synchronization and minimize CPU ⇄ GPU data transfer.
+That is not only a cost of the port; see [Keeping Scalars On The Device](#keeping-scalars-on-the-device) for what it buys.
 
-* `extract`
-* `[]`
-* `count_true`
-* `count_false`
+The methods affected are:
 
-Numo returns a Ruby numeric object for 0-dimensional NArray, while Cumo returns the 0-dimensional NArray instead of a Ruby numeric object.
-Cumo differs in this way to avoid synchronization and minimize CPU ⇄ GPU data transfer.
+* `[]` and `extract`
+* `count_true` and `count_false`
+* reductions down to a single value: `sum`, `prod`, `mean`, `stddev`, `var`, `rms`, `min`, `max`, `ptp`, `minmax`, `median`, `mulsum`, `dot`, `inner`
+* index reductions: `max_index`, `min_index`, `argmax`, `argmin`
+
+A 0-dimensional `Cumo::Bit` is truthy even when it holds 0, because Ruby treats every object but `nil` and `false` as true.
+Comparing two scalars therefore takes the wrong branch without raising anything:
+
+```ruby
+a = Cumo::SFloat[5.0]
+a[0] < 1.0                  #=> Cumo::Bit#shape=[] holding 0
+(a[0] < 1.0) ? :yes : :no   #=> :yes, where Numo gives :no
+```
+
+`assert_operator(a[0], :<, 1.0)` passes for the same reason, so a test suite written for Numo can stay green against Cumo while asserting nothing.
+Read the value back to the host before branching on it, or run under `compatible_mode`.
 
 Set the `CUMO_COMPATIBLE_MODE` environment variable to `ON` to force Numo NArray compatibility (for worse performance).
+Running a Numo test suite that way keeps its assertions meaningful.
 
 You may enable or disable `compatible_mode` as:
 
@@ -126,6 +141,43 @@ You can also use the following methods which behave like Numo's NArray methods.
 * `count_true_cpu`
 * `count_false_cpu`
 
+```ruby
+a.aref_cpu(0) < 1.0   #=> false in either mode
+Float(a.sum)          #=> 7.0 in either mode
+```
+
+They are methods on an NArray, so chaining one onto a result that `compatible_mode` has already turned into a Ruby object, as in `a.sum.extract_cpu`, raises `NoMethodError` while the mode is on.
+`Kernel#Float` and `Kernel#Integer` read either representation, and read Numo's too, so they are what code that runs against both libraries wants.
+
+### Keeping Scalars On The Device
+
+The 0-dimensional return is what lets an iterative loop stay on the GPU.
+Reading a scalar back to the host waits for everything queued behind it, so every read caps how far ahead the GPU is allowed to run.
+What a read costs is not a fixed price either: it is however much work happens to be queued when it is taken.
+
+`bench/cg_bench.rb` prices this with a conjugate gradient solve, 200 iterations over a 512x512 grid on an RTX 5070 Ti Laptop:
+
+```
+scalars        convergence test    us/iter    readbacks/iter
+Ruby Floats    every iteration       136.7        2.02
+Ruby Floats    never                 134.4        2.02
+0-dim NArray   every iteration       142.6        1.02
+0-dim NArray   every 20th             65.6        0.06
+0-dim NArray   never                  60.3        0.02
+```
+
+Written with Ruby Floats the loop reads back twice an iteration whatever the convergence test does, since `alpha` needs `pap` and `beta` needs `rs_new` as Floats.
+Thinning the test cannot get under that floor, and keeping the scalars as 0-dimensional NArrays buys nothing on its own.
+The two only pay together, and together they are worth 2.1x.
+The relative residual is identical in every row.
+
+```ruby
+alpha = rs_old / pap   # a 0-dimensional NArray, divided on the device
+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.
+
 ### Select a GPU device ID
 
 Set the `CUDA_VISIBLE_DEVICES=id` environment variable, or
@@ -201,7 +253,14 @@ ln -sf "$HOME/opt/ccache/bin/ccache" "$HOME/opt/ccache/bin/nvcc"
 bundle exec env CUMO_NVCC_GENERATE_CODE=arch=compute_60,code=sm_60 rake compile
 ```
 
+Separate the entries with a space to build for more than one architecture:
+
+```
+bundle exec env CUMO_NVCC_GENERATE_CODE="arch=compute_75,code=sm_75 arch=compute_121,code=sm_121" rake compile
+```
+
 This is useful even on development because it makes it possible to skip JIT compilation of PTX to cubin during runtime.
+Without it, and without an `nvidia-smi` to read the local compute capability from, the build covers every architecture the CUDA version supports.
 
 ### Run tests with gdb
 

data/ext/cumo/cuda/cudnn_impl.cpp

@@ -555,7 +555,7 @@ cumo_cuda_cudnn_CreateBNTensorDescriptor(
 {
     cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
     status = cudnnCreateTensorDescriptor(desc);
-    if (status == CUDNN_STATUS_SUCCESS) return status;
+    if (status != CUDNN_STATUS_SUCCESS) return status;
 
     status = cudnnDeriveBNTensorDescriptor(*desc, x_desc, mode);
     return status;

data/ext/cumo/cuda/memory_pool_impl.cpp

@@ -140,7 +140,12 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
                 continue;
             }
             chunk = PopFromFreeList(free_list);
-            // TODO(sonots): compact_index
+            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);
+            }
             break;
         }
 

data/ext/cumo/cuda/memory_pool_impl_test.cpp

@@ -137,6 +137,8 @@ public:
 
     void Run() {
         TearDown(); SetUp(); TestGetRoundedSize();
+        TearDown(); SetUp(); TestMallocCompactsEmptiedBin();
+        TearDown(); SetUp(); TestMallocCompactsLastBin();
         TearDown(); SetUp(); TestGetBinIndex();
         TearDown(); SetUp(); TestGetArenaIndexWithHugeSize();
         TearDown(); SetUp(); TestMallocOnAnotherStream();
@@ -290,6 +292,51 @@ public:
         assert(arena_index_map[2] == 4);
     }
 
+    // Malloc drops a bin it just emptied, so a later search does not walk it.
+    void TestMallocCompactsEmptiedBin() {
+        Arena& arena = pool_->GetArena(stream_ptr_);
+        ArenaIndexMap& arena_index_map = pool_->GetArenaIndexMap(stream_ptr_);
+
+        // three bins, one chunk each, none of them adjacent in memory
+        for (int k = 2; k <= 4; ++k) {
+            auto mem = std::make_shared<Memory>(kRoundSize * k);
+            auto chunk = std::make_shared<Chunk>(mem, 0, mem->size(), stream_ptr_);
+            pool_->AppendToFreeList(chunk->size(), chunk, stream_ptr_);
+        }
+        assert(arena.size() == 3);
+        assert(arena_index_map.size() == 3);
+
+        // takes the whole chunk of the middle bin, which then has nothing left
+        intptr_t ptr = pool_->Malloc(kRoundSize * 3, stream_ptr_);
+        assert(ptr != 0);
+        assert(arena.size() == 2);
+        assert(arena_index_map.size() == 2);
+        assert(arena_index_map[0] == 1);
+        assert(arena_index_map[1] == 3);
+
+        pool_->Free(ptr);
+        assert(arena.size() == 3);
+    }
+
+    // Emptying the last bin leaves no arena at all, and the pool has to carry
+    // on: the split remainder of that very allocation goes back into a new one.
+    void TestMallocCompactsLastBin() {
+        auto mem = std::make_shared<Memory>(kRoundSize * 8);
+        auto chunk = std::make_shared<Chunk>(mem, 0, mem->size(), stream_ptr_);
+        pool_->AppendToFreeList(chunk->size(), chunk, stream_ptr_);
+        assert(pool_->GetArena(stream_ptr_).size() == 1);
+
+        intptr_t ptr = pool_->Malloc(kRoundSize * 2, stream_ptr_);
+        assert(ptr != 0);
+        // the remaining 6 units are back in a free list
+        assert(pool_->GetNumFreeBlocks() == 1);
+        assert(pool_->GetFreeBytes() == kRoundSize * 6);
+
+        pool_->Free(ptr);
+        assert(pool_->GetNumFreeBlocks() == 1);
+        assert(pool_->GetFreeBytes() == kRoundSize * 8);
+    }
+
     // TODO(sonots): Fix after implementing compaction
     void TestRemoveFromFreeList() {
         Arena& arena = pool_->GetArena(stream_ptr_);

data/ext/cumo/extconf.rb

@@ -164,7 +164,36 @@ have_library('nvrtc')
 have_library('cublas')
 # have_library('cusolver')
 # have_library('curand')
-if have_library('cudnn') # TODO(sonots): cuDNN version check
+
+# cuDNN 8 is the first release that supports CUDA 11, which cumo requires.
+CUDNN_MIN_MAJOR = 8
+
+def have_cudnn?
+  unless have_header('cudnn.h')
+    message("cuDNN header not found; building without cuDNN features\n")
+    return false
+  end
+  ok = checking_for("cuDNN #{CUDNN_MIN_MAJOR} or later") do
+    try_compile(<<-SRC)
+#include <cudnn.h>
+#if CUDNN_MAJOR < #{CUDNN_MIN_MAJOR}
+#error "cuDNN is too old"
+#endif
+int main(void) { return 0; }
+    SRC
+  end
+  unless ok
+    message("cuDNN is older than #{CUDNN_MIN_MAJOR}.0; building without cuDNN features\n")
+    return false
+  end
+  unless have_library('cudnn')
+    message("cuDNN library not found; building without cuDNN features\n")
+    return false
+  end
+  true
+end
+
+if have_cudnn?
   $CFLAGS << " -DCUDNN_FOUND"
   $CXXFLAGS << " -DCUDNN_FOUND"
 end

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

@@ -44,6 +44,59 @@ extern VALUE cumo_na_eShapeError;
                  (int)(nd1), (int)(nd2));            \
     }
 
+// An output array given by the caller is written through a descriptor built
+// from another operand, or as if it were contiguous, so cuDNN never learns how
+// long it really is. It has to be checked here instead.
+static inline void
+cumo_cuda_cudnn_check_output(VALUE out, VALUE type, size_t ndim, size_t *shape)
+{
+    cumo_narray_t *na;
+
+    CUMO_CUDA_CUDNN_CHECK_NARRAY_TYPE(out, type);
+    CumoGetNArray(out, na);
+    CUMO_CUDA_CUDNN_CHECK_DIM_EQ((size_t)(na->ndim), ndim);
+    for (size_t idim = 0; idim < ndim; ++idim) {
+        CUMO_CUDA_CUDNN_CHECK_SIZE_EQ(na->shape[idim], shape[idim]);
+    }
+    if (cumo_na_check_contiguous(out) != Qtrue) {
+        rb_raise(cumo_na_eShapeError, "output NArray must be contiguous");
+    }
+}
+
+// cuDNN derives the batch-norm parameter descriptor from x with a SPATIAL mode,
+// so it reaches x's channel count in gamma, beta and the rest however the axis
+// reduced them. Normally the reduced size is the larger of the two and only
+// part of each buffer is used, but an axis that folds the channel away makes it
+// smaller, and cuDNN has no way to know.
+static inline void
+cumo_cuda_cudnn_check_reduced_size(size_t reduced_total_size, size_t x_ndim, size_t *x_shape)
+{
+    size_t channel_size = (x_ndim == 1) ? x_shape[0] : x_shape[1];
+
+    if (reduced_total_size < channel_size) {
+        rb_raise(cumo_na_eShapeError,
+                 "the axis reduces to %d, fewer than the %d channels cuDNN accesses",
+                 (int)reduced_total_size, (int)channel_size);
+    }
+}
+
+// An input array is read through a descriptor built from another operand, so
+// cuDNN reads whatever length that descriptor claims rather than the length the
+// array really has. A non-contiguous one is copied before it is read, so unlike
+// an output it does not have to be contiguous here.
+static inline void
+cumo_cuda_cudnn_check_input(VALUE in, VALUE type, size_t ndim, size_t *shape)
+{
+    cumo_narray_t *na;
+
+    CUMO_CUDA_CUDNN_CHECK_NARRAY_TYPE(in, type);
+    CumoGetNArray(in, na);
+    CUMO_CUDA_CUDNN_CHECK_DIM_EQ((size_t)(na->ndim), ndim);
+    for (size_t idim = 0; idim < ndim; ++idim) {
+        CUMO_CUDA_CUDNN_CHECK_SIZE_EQ(na->shape[idim], shape[idim]);
+    }
+}
+
 void
 cumo_cuda_cudnn_check_status(cudnnStatus_t status);
 

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

@@ -221,7 +221,7 @@ cumo_na_make_reduction_arg(cumo_na_loop_t* lp_user, int out_arg)
 
 #endif  // #ifndef __CUDACC__
 
-#define CUMO_NA_INDEXER_OPTIMIZED_NDIM 4
+#define CUMO_NA_INDEXER_OPTIMIZED_NDIM 8
 
 #ifdef __CUDACC__
 
@@ -247,6 +247,10 @@ cumo_na_indexer_set_dim##NDIM(cumo_na_indexer_t* indexer, uint64_t i) { \
     } \
 }
 
+CUMO_NA_INDEXER_SET(8)
+CUMO_NA_INDEXER_SET(7)
+CUMO_NA_INDEXER_SET(6)
+CUMO_NA_INDEXER_SET(5)
 CUMO_NA_INDEXER_SET(4)
 CUMO_NA_INDEXER_SET(3)
 CUMO_NA_INDEXER_SET(2)
@@ -280,6 +284,10 @@ cumo_na_iarray_at_dim##NDIM(cumo_na_iarray_t* iarray, cumo_na_indexer_t* indexer
     return ptr; \
 }
 
+CUMO_NA_IARRAY_AT(8)
+CUMO_NA_IARRAY_AT(7)
+CUMO_NA_IARRAY_AT(6)
+CUMO_NA_IARRAY_AT(5)
 CUMO_NA_IARRAY_AT(4)
 CUMO_NA_IARRAY_AT(3)
 CUMO_NA_IARRAY_AT(2)
@@ -313,6 +321,10 @@ cumo_na_bit_iarray_at_dim##NDIM(cumo_na_bit_iarray_t* iarray, cumo_na_indexer_t*
     return pos; \
 }
 
+CUMO_NA_BIT_IARRAY_AT(8)
+CUMO_NA_BIT_IARRAY_AT(7)
+CUMO_NA_BIT_IARRAY_AT(6)
+CUMO_NA_BIT_IARRAY_AT(5)
 CUMO_NA_BIT_IARRAY_AT(4)
 CUMO_NA_BIT_IARRAY_AT(3)
 CUMO_NA_BIT_IARRAY_AT(2)
@@ -354,6 +366,10 @@ cumo_na_bit_iarray_stridx_at_dim##NDIM(cumo_na_bit_iarray_stridx_t* iarray, cumo
     return pos; \
 }
 
+CUMO_NA_BIT_IARRAY_STRIDX_AT(8)
+CUMO_NA_BIT_IARRAY_STRIDX_AT(7)
+CUMO_NA_BIT_IARRAY_STRIDX_AT(6)
+CUMO_NA_BIT_IARRAY_STRIDX_AT(5)
 CUMO_NA_BIT_IARRAY_STRIDX_AT(4)
 CUMO_NA_BIT_IARRAY_STRIDX_AT(3)
 CUMO_NA_BIT_IARRAY_STRIDX_AT(2)
@@ -401,6 +417,10 @@ cumo_na_iarray_stridx_at_dim##NDIM(cumo_na_iarray_stridx_t* iarray, cumo_na_inde
     return ptr; \
 }
 
+CUMO_NA_IARRAY_STRIDX_AT(8)
+CUMO_NA_IARRAY_STRIDX_AT(7)
+CUMO_NA_IARRAY_STRIDX_AT(6)
+CUMO_NA_IARRAY_STRIDX_AT(5)
 CUMO_NA_IARRAY_STRIDX_AT(4)
 CUMO_NA_IARRAY_STRIDX_AT(3)
 CUMO_NA_IARRAY_STRIDX_AT(2)

data/ext/cumo/include/cumo/narray.h

@@ -471,7 +471,10 @@ cumo_na_has_idx_p(VALUE obj)
     if (CUMO_NA_TYPE(na) == CUMO_NARRAY_VIEW_T) {
         CumoGetNArrayView(obj, nv);
         for (; i < nv->base.ndim; ++i) {
-            if (nv->stridx[i].index) {
+            // A stride is tagged with its low bit set, so the raw union is
+            // non-null for every dimension a view has and reading it as a
+            // pointer answers true for all of them.
+            if (CUMO_SDX_IS_INDEX(nv->stridx[i])) {
                 return true;
             }
         }

data/ext/cumo/include/cumo/reduce_kernel.h

@@ -350,7 +350,7 @@ __global__ static void reduction_pair_kernel(cumo_na_reduction_arg_t arg, cumo_n
 // The scratch is laid out by output so the combine pass can read it flat, but
 // threads are handed out by output first, so a block still covers consecutive
 // outputs when those are the ones running along memory.
-template <bool FLAT, typename TypeIn, typename TypeReduce, typename ReductionImpl>
+template <bool FLAT, bool ARG, typename TypeIn, typename TypeReduce, typename ReductionImpl>
 __global__ static void reduction_partial_kernel(cumo_na_reduction_arg_t arg, cumo_reduce_addr_t ad, TypeReduce* partial, int64_t n_split, int64_t chunk, int out_block_size, int reduce_block_size, ReductionImpl impl) {
     extern __shared__ __align__(8) char sdata_raw[];
     TypeReduce* sdata = reinterpret_cast<TypeReduce*>(sdata_raw);
@@ -374,7 +374,7 @@ __global__ static void reduction_partial_kernel(cumo_na_reduction_arg_t arg, cum
         ssize_t in_out_off = reduce_in_out_offset<FLAT>(arg, ad, i_out);
         int64_t i_in = i_out * reduce_total_size + begin + reduce_offset;
 
-        TypeReduce accum = reduce_axis<FLAT,false,TypeIn>(arg, ad, impl, in_out_off, i_in, begin, end, reduce_offset, reduce_block_size);
+        TypeReduce accum = reduce_axis<FLAT,ARG,TypeIn>(arg, ad, impl, in_out_off, i_in, begin, end, reduce_offset, reduce_block_size);
 
         accum = reduce_in_block(accum, sdata, tid, out_block_size, impl);
         if (reduce_offset == 0) {
@@ -410,7 +410,7 @@ struct reduce_combine {
 // Allocates the scratch and runs the first pass, then points arg2's input at
 // the partials so the caller can combine them with its own second pass. The
 // caller frees the returned buffer.
-template <typename TypeIn, typename TypeReduce, typename ReductionImpl>
+template <typename TypeIn, typename TypeReduce, typename ReductionImpl, bool ARG = false>
 TypeReduce* reduce_partial_pass(cumo_na_reduction_arg_t arg, cumo_reduce_addr_t ad, int64_t n_split, int64_t reduce_total_size, cumo_na_reduction_arg_t* arg2, ReductionImpl& impl) {
     int64_t chunk = (reduce_total_size + n_split - 1) / n_split;
     int64_t partial_total_size = arg.out_indexer.total_size * n_split;
@@ -423,9 +423,9 @@ TypeReduce* reduce_partial_pass(cumo_na_reduction_arg_t arg, cumo_reduce_addr_t
     int64_t shared_mem_size = sizeof(TypeReduce) * max_block_size;
 
     if (ad.in_out_flat && ad.in_reduce_flat) {
-        reduction_partial_kernel<true,TypeIn,TypeReduce,ReductionImpl><<<grid_size, max_block_size, shared_mem_size>>>(arg, ad, partial, n_split, chunk, out_block_size, reduce_block_size, impl);
+        reduction_partial_kernel<true,ARG,TypeIn,TypeReduce,ReductionImpl><<<grid_size, max_block_size, shared_mem_size>>>(arg, ad, partial, n_split, chunk, out_block_size, reduce_block_size, impl);
     } else {
-        reduction_partial_kernel<false,TypeIn,TypeReduce,ReductionImpl><<<grid_size, max_block_size, shared_mem_size>>>(arg, ad, partial, n_split, chunk, out_block_size, reduce_block_size, impl);
+        reduction_partial_kernel<false,ARG,TypeIn,TypeReduce,ReductionImpl><<<grid_size, max_block_size, shared_mem_size>>>(arg, ad, partial, n_split, chunk, out_block_size, reduce_block_size, impl);
     }
     cumo_cuda_runtime_check_kernel_launch();
 
@@ -439,7 +439,6 @@ TypeReduce* reduce_partial_pass(cumo_na_reduction_arg_t arg, cumo_reduce_addr_t
 
 }  // cumo_detail
 
-// TODO(sonots): Optimize indexer by squashing (or reducing) dimensions
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
 void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
     if (arg.out_indexer.total_size == 0) {
@@ -581,4 +580,34 @@ void cumo_reduce_arg(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
     cumo_cuda_runtime_check_kernel_launch();
 }
 
+// cumo_reduce_split for the arg form. The partials already carry indices along
+// the reduce axis, so the second pass combines them with the plain kernel.
+// Only for an impl whose Identity ignores its index argument.
+template <typename TypeIn, typename TypeOut, typename ReductionImpl>
+void cumo_reduce_arg_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
+    using TypeReduce = decltype(impl.Identity(0));
+
+    if (arg.out_indexer.total_size == 0) {
+        return;
+    }
+
+    int64_t reduce_total_size = arg.in_indexer.total_size / arg.out_indexer.total_size;
+    cumo_detail::cumo_reduce_addr_t ad = cumo_detail::make_reduce_addr(arg, reduce_total_size);
+
+    int64_t out_block_size, reduce_block_size;
+    cumo_detail::reduce_block_split(ad, reduce_total_size, &out_block_size, &reduce_block_size);
+    int64_t out_block_num = (arg.out_indexer.total_size + out_block_size - 1) / out_block_size;
+
+    int64_t n_split = cumo_detail::reduce_split_count(reduce_total_size, out_block_num);
+    if (n_split < 2) {
+        cumo_reduce_arg<TypeIn, TypeOut, ReductionImpl>(arg, std::forward<ReductionImpl>(impl));
+        return;
+    }
+
+    cumo_na_reduction_arg_t arg2 = arg;
+    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl, true>(arg, ad, n_split, reduce_total_size, &arg2, impl);
+    cumo_reduce<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(arg2, cumo_detail::reduce_combine<ReductionImpl>{impl});
+    cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
+}
+
 #endif // CUMO_REDUCE_KERNEL_H

data/ext/cumo/include/cumo/types/complex_kernel.h

@@ -363,7 +363,7 @@ __host__ __device__ static inline dtype c_pow(dtype x, dtype y)
 }
 
 /* only internal use (called by c_pow_int) */
-__host__ __device__ static inline dtype c_pow_positive_int(dtype x, int p)
+__host__ __device__ static inline dtype c_pow_positive_int(dtype x, unsigned int p)
 {
     dtype z = c_one();
     if (p==2) {return c_square(x);}
@@ -377,13 +377,15 @@ __host__ __device__ static inline dtype c_pow_positive_int(dtype x, int p)
     return z;
 }
 
+// The magnitude is taken in unsigned, since -p overflows for INT_MIN and the
+// wrapped negative would shift down to -1 and loop for ever.
 __host__ __device__ static inline dtype c_pow_int(dtype x, int p)
 {
     if (p<0) {
-        x = c_pow_positive_int(x,-p);
+        x = c_pow_positive_int(x,-(unsigned int)p);
         return c_reciprocal(x);
     } else {
-        return c_pow_positive_int(x,p);
+        return c_pow_positive_int(x,(unsigned int)p);
     }
 }
 

data/ext/cumo/include/cumo/types/float_macro_kernel.h

@@ -123,7 +123,7 @@ extern double pow(double, double);
 #define m_frexp(x,exp) frexp(x,exp)
 
 /* only internal use (called by pow_int) */
-__host__ __device__ static inline dtype pow_positive_int(dtype x, int p)
+__host__ __device__ static inline dtype pow_positive_int(dtype x, unsigned int p)
 {
     dtype r=1;
     switch(p) {
@@ -133,7 +133,7 @@ __host__ __device__ static inline dtype pow_positive_int(dtype x, int p)
     case 3: return x*x*x;
     case 4: x=x*x; return x*x;
     }
-    if (p>64) return pow(x,p);
+    if (p>64) return pow(x,(double)p);
     while (p) {
         if (p&1) r *= x;
         x *= x;
@@ -142,10 +142,12 @@ __host__ __device__ static inline dtype pow_positive_int(dtype x, int p)
     return r;
 }
 
+// The magnitude is taken in unsigned, since -p overflows for INT_MIN and the
+// wrapped negative would shift down to -1 and loop for ever.
 __host__ __device__ static inline dtype pow_int(dtype x, int p)
 {
-    if (p<0) return 1/pow_positive_int(x, -p);
-    return pow_positive_int(x, p);
+    if (p<0) return 1/pow_positive_int(x, -(unsigned int)p);
+    return pow_positive_int(x, (unsigned int)p);
 }
 
 __host__ __device__ static inline dtype f_seq(dtype x, dtype y, double c)

data/ext/cumo/include/cumo/types/int_macro_kernel.h

@@ -6,10 +6,6 @@
 #define m_sign(x)    (((x)==0) ? 0 : (((x)>0) ? 1 : -1))
 
 __host__ __device__ static inline dtype m_abs(dtype x) {
-    // TODO(sonots): How to handle in CUDA kernel?
-    // if (x==DATA_MIN) {
-    //     rb_raise(cumo_na_eValueError, "cannot convert the minimum integer");
-    // }
     return (x<0)?-x:x;
 }
 

data/ext/cumo/include/cumo/types/robject_kernel.h

@@ -1,6 +1,3 @@
-//TODO(sonots):
-//typedef VALUE dtype;
-//typedef VALUE rtype;
 typedef void* dtype;
 typedef void* rtype;
 

data/ext/cumo/include/cumo/types/uint_macro_kernel.h

@@ -27,6 +27,7 @@ __device__ static dtype pow_int(dtype x, int p)
     case 2: return x*x;
     case 3: return x*x*x;
     }
+    if (p<0) return 0;
     while (p) {
         if (p&1) r *= x;
         x *= x;

data/ext/cumo/include/cumo.h

@@ -10,8 +10,8 @@ extern "C" {
 #endif
 #endif
 
-#define CUMO_VERSION "0.5.9"
-#define CUMO_VERSION_CODE 59
+#define CUMO_VERSION "0.5.10"
+#define CUMO_VERSION_CODE 510
 
 bool cumo_compatible_mode_enabled_p();
 bool cumo_show_warning_enabled_p();

data/ext/cumo/narray/data.c

@@ -103,7 +103,7 @@ cumo_na_copy(VALUE self)
     // Everything else goes through the indexer so that ndloop hands the whole
     // view over at once: walking the outer dimensions itself costs a
     // synchronization per step when an operand carries an index array.
-    if (rb_obj_class(self) != cumo_cRObject) {
+    if (!rb_obj_is_kind_of(self, cumo_cRObject)) {
         ndf.func = iter_copy_bytes_indexer;
         ndf.flag = CUMO_STRIDE_LOOP|CUMO_NDF_INDEXER_LOOP;
     }

data/ext/cumo/narray/gen/spec.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 def_id "cast"
+def_id "_pending_scalar", "pending_scalar"
 def_id "mulsum"
 def_id "to_a"
 if is_complex