cumo 0.7.0 → 0.8.0

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

@@ -190,6 +190,7 @@ extern VALUE cumo_cBit;
 extern VALUE cumo_cDFloat;
 extern VALUE cumo_cSFloat;
 extern VALUE cumo_cHFloat;
+extern VALUE cumo_cBFloat;
 extern VALUE cumo_cDComplex;
 extern VALUE cumo_cSComplex;
 extern VALUE cumo_cInt64;

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

@@ -602,7 +602,7 @@ struct reduce_combine {
 // 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, 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) {
+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, char* held = 0) {
     int64_t chunk = (reduce_total_size + n_split - 1) / n_split;
     int64_t partial_total_size = arg.out_indexer.total_size * n_split;
     TypeReduce* partial = reinterpret_cast<TypeReduce*>(cumo_cuda_runtime_malloc(sizeof(TypeReduce) * partial_total_size));
@@ -618,7 +618,7 @@ TypeReduce* reduce_partial_pass(cumo_na_reduction_arg_t arg, cumo_reduce_addr_t
     } else {
         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();
+    cumo_check_launch_holding(partial, held);
 
     arg2->in.ptr = reinterpret_cast<char*>(partial);
     arg2->in.step[0] = sizeof(TypeReduce);
@@ -636,7 +636,7 @@ static inline bool zip_axes_are_flat(const cumo_reduce_addr_t& ad, const cumo_re
 
 // First pass of a split zip reduction. See reduce_partial_pass above.
 template <typename TypeIn, typename TypeReduce, typename ReductionImpl>
-TypeReduce* reduce_zip_partial_pass(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, cumo_reduce_addr_t ad, cumo_reduce_addr_t ad2, int64_t n_split, int64_t reduce_total_size, cumo_na_reduction_arg_t* arg2, ReductionImpl& impl) {
+TypeReduce* reduce_zip_partial_pass(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, cumo_reduce_addr_t ad, cumo_reduce_addr_t ad2, int64_t n_split, int64_t reduce_total_size, cumo_na_reduction_arg_t* arg2, ReductionImpl& impl, char* held = 0) {
     int64_t chunk = (reduce_total_size + n_split - 1) / n_split;
     int64_t partial_total_size = arg.out_indexer.total_size * n_split;
     TypeReduce* partial = reinterpret_cast<TypeReduce*>(cumo_cuda_runtime_malloc(sizeof(TypeReduce) * partial_total_size));
@@ -652,7 +652,7 @@ TypeReduce* reduce_zip_partial_pass(cumo_na_reduction_arg_t arg, cumo_na_iarray_
     } else {
         reduction_zip_partial_kernel<false,TypeIn,TypeReduce,ReductionImpl><<<grid_size, max_block_size, shared_mem_size>>>(arg, in2, ad, ad2, partial, n_split, chunk, out_block_size, reduce_block_size, impl);
     }
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(partial, held);
 
     arg2->in.ptr = reinterpret_cast<char*>(partial);
     arg2->in.step[0] = sizeof(TypeReduce);
@@ -664,8 +664,12 @@ TypeReduce* reduce_zip_partial_pass(cumo_na_reduction_arg_t arg, cumo_na_iarray_
 
 }  // cumo_detail
 
+// held0 and held1 are every scratch buffer still outstanding, the caller's
+// included. The checks free them before they raise, since rb_raise is a longjmp
+// and nothing written after it runs. The split forms pass their own partial as
+// well as whatever they were handed.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
+void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl, char* held0 = 0, char* held1 = 0) {
     if (arg.out_indexer.total_size == 0) {
         return;
     }
@@ -686,7 +690,7 @@ void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
     } else {
         cumo_detail::reduction_kernel<false,TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, ad, out_block_size, reduce_block_size, impl);
     }
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(held0, held1);
 }
 
 // Runs the reduce axis of one output across several blocks and combines their
@@ -695,7 +699,7 @@ void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
 // case. Only for an impl whose Identity ignores its index argument — see
 // reduce_combine above.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
+void cumo_reduce_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl, char* held = 0) {
     using TypeReduce = decltype(impl.Identity(0));
 
     if (arg.out_indexer.total_size == 0) {
@@ -711,20 +715,20 @@ void cumo_reduce_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
 
     int64_t n_split = cumo_detail::reduce_split_count(reduce_total_size, out_block_num);
     if (n_split < 2) {
-        cumo_reduce<TypeIn, TypeOut, ReductionImpl>(arg, std::forward<ReductionImpl>(impl));
+        cumo_reduce<TypeIn, TypeOut, ReductionImpl>(arg, std::forward<ReductionImpl>(impl), held);
         return;
     }
 
     cumo_na_reduction_arg_t arg2 = arg;
-    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl>(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});
+    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl>(arg, ad, n_split, reduce_total_size, &arg2, impl, held);
+    cumo_reduce<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(arg2, cumo_detail::reduce_combine<ReductionImpl>{impl}, reinterpret_cast<char*>(partial), held);
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
 
 // Variant of cumo_reduce reading two inputs, for mulsum. in2 describes the same
 // shape as arg.in, since the one in_indexer addresses both.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_zip(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, ReductionImpl&& impl) {
+void cumo_reduce_zip(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, ReductionImpl&& impl, char* held0 = 0, char* held1 = 0) {
     if (arg.out_indexer.total_size == 0) {
         return;
     }
@@ -748,13 +752,13 @@ void cumo_reduce_zip(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Reductio
     } else {
         cumo_detail::reduction_zip_kernel<false,TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, in2, ad, ad2, out_block_size, reduce_block_size, impl);
     }
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(held0, held1);
 }
 
 // cumo_reduce_split for a zip reduction. The first pass reads both operands and
 // the combine pass has only accumulators left, so it is the plain one.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_zip_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, ReductionImpl&& impl) {
+void cumo_reduce_zip_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, ReductionImpl&& impl, char* held = 0) {
     using TypeReduce = decltype(impl.Identity(0));
 
     if (arg.out_indexer.total_size == 0) {
@@ -773,13 +777,13 @@ void cumo_reduce_zip_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Re
 
     int64_t n_split = cumo_detail::reduce_split_count(reduce_total_size, out_block_num);
     if (n_split < 2) {
-        cumo_reduce_zip<TypeIn, TypeOut, ReductionImpl>(arg, in2, std::forward<ReductionImpl>(impl));
+        cumo_reduce_zip<TypeIn, TypeOut, ReductionImpl>(arg, in2, std::forward<ReductionImpl>(impl), held);
         return;
     }
 
     cumo_na_reduction_arg_t combine = arg;
-    TypeReduce* partial = cumo_detail::reduce_zip_partial_pass<TypeIn, TypeReduce, ReductionImpl>(arg, in2, ad, ad2, n_split, reduce_total_size, &combine, impl);
-    cumo_reduce<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(combine, cumo_detail::reduce_combine<ReductionImpl>{impl});
+    TypeReduce* partial = cumo_detail::reduce_zip_partial_pass<TypeIn, TypeReduce, ReductionImpl>(arg, in2, ad, ad2, n_split, reduce_total_size, &combine, impl, held);
+    cumo_reduce<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(combine, cumo_detail::reduce_combine<ReductionImpl>{impl}, reinterpret_cast<char*>(partial), held);
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
 
@@ -787,7 +791,7 @@ void cumo_reduce_zip_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Re
 // See reduction_pair_kernel above. out2 has to describe the same shape as
 // arg.out, since the one out_indexer addresses both.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_pair(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2, ReductionImpl&& impl) {
+void cumo_reduce_pair(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2, ReductionImpl&& impl, char* held0 = 0, char* held1 = 0) {
     if (arg.out_indexer.total_size == 0) {
         return;
     }
@@ -809,12 +813,12 @@ void cumo_reduce_pair(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2, Reduct
     } else {
         cumo_detail::reduction_pair_kernel<false,TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, out2, ad, out_block_size, reduce_block_size, impl);
     }
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(held0, held1);
 }
 
 // cumo_reduce_split for the two-output form. Same constraint on Identity.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_pair_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2, ReductionImpl&& impl) {
+void cumo_reduce_pair_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2, ReductionImpl&& impl, char* held = 0) {
     using TypeReduce = decltype(impl.Identity(0));
 
     if (arg.out_indexer.total_size == 0) {
@@ -830,20 +834,20 @@ void cumo_reduce_pair_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t out2,
 
     int64_t n_split = cumo_detail::reduce_split_count(reduce_total_size, out_block_num);
     if (n_split < 2) {
-        cumo_reduce_pair<TypeIn, TypeOut, ReductionImpl>(arg, out2, std::forward<ReductionImpl>(impl));
+        cumo_reduce_pair<TypeIn, TypeOut, ReductionImpl>(arg, out2, std::forward<ReductionImpl>(impl), held);
         return;
     }
 
     cumo_na_reduction_arg_t arg2 = arg;
-    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl>(arg, ad, n_split, reduce_total_size, &arg2, impl);
-    cumo_reduce_pair<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(arg2, out2, cumo_detail::reduce_combine<ReductionImpl>{impl});
+    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl>(arg, ad, n_split, reduce_total_size, &arg2, impl, held);
+    cumo_reduce_pair<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(arg2, out2, cumo_detail::reduce_combine<ReductionImpl>{impl}, reinterpret_cast<char*>(partial), held);
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
 
 // Variant of cumo_reduce for arg-reductions (argmax/argmin), which returns
 // indices along the reduction axis. See reduction_arg_kernel above.
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce_arg(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
+void cumo_reduce_arg(cumo_na_reduction_arg_t arg, ReductionImpl&& impl, char* held0 = 0, char* held1 = 0) {
     if (arg.out_indexer.total_size == 0) {
         return;
     }
@@ -864,14 +868,14 @@ void cumo_reduce_arg(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
     } else {
         cumo_detail::reduction_arg_kernel<false,TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, ad, out_block_size, reduce_block_size, impl);
     }
-    cumo_cuda_runtime_check_kernel_launch();
+    cumo_check_launch_holding(held0, held1);
 }
 
 // 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) {
+void cumo_reduce_arg_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl, char* held = 0) {
     using TypeReduce = decltype(impl.Identity(0));
 
     if (arg.out_indexer.total_size == 0) {
@@ -887,13 +891,13 @@ void cumo_reduce_arg_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
 
     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));
+        cumo_reduce_arg<TypeIn, TypeOut, ReductionImpl>(arg, std::forward<ReductionImpl>(impl), held);
         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});
+    TypeReduce* partial = cumo_detail::reduce_partial_pass<TypeIn, TypeReduce, ReductionImpl, true>(arg, ad, n_split, reduce_total_size, &arg2, impl, held);
+    cumo_reduce<TypeReduce, TypeOut, cumo_detail::reduce_combine<ReductionImpl>>(arg2, cumo_detail::reduce_combine<ReductionImpl>{impl}, reinterpret_cast<char*>(partial), held);
     cumo_cuda_runtime_free(reinterpret_cast<char*>(partial));
 }
 

data/ext/cumo/include/cumo/row_kernel.h

@@ -0,0 +1,170 @@
+#ifndef CUMO_ROW_KERNEL_H
+#define CUMO_ROW_KERNEL_H
+
+#include "cumo/reduce_kernel.h"
+
+// Reducing each row of an array to a few numbers and then walking the row again
+// to write it out is the shape layer norm and softmax share. One block takes one
+// row, so what the reduction found stays in shared memory between the two passes
+// instead of going back to global memory as an array of its own: written out of
+// mean, minus, divide and the rest, a layer norm costs nine kernels and eight
+// intermediate arrays.
+//
+// Impl is a reduction impl as reduce_kernel.h means it, whose MapOut answers
+// whatever the second pass needs. Apply turns one element into one element,
+// taking its position along the row so that a per-column operand can be read.
+//
+// Two things Impl has to hold to, neither of which the compiler can check.
+// It must not read the index its Identity and MapIn are handed: the kernel
+// below always passes 0, while the split path passes the element's real offset,
+// so an impl that reads it answers differently depending on which path a shape
+// takes. And it must be default constructible, since cumo_reduce_split's
+// reduce_combine names its accumulator as decltype(Impl().Identity(0)).
+
+namespace cumo_detail {
+
+template <typename TypeIn, typename Impl, typename Apply>
+__global__ void row_reduce_apply_kernel(
+        const TypeIn* x, TypeIn* y, uint64_t rows, uint64_t cols, Impl impl, Apply apply)
+{
+    typedef decltype(impl.Identity(0)) Accum;
+    static_assert(alignof(Accum) <= 8,
+                  "the shared memory below is declared __align__(8), as every other "
+                  "dynamic one in cumo is, so a wider accumulator would land unaligned");
+
+    extern __shared__ __align__(8) char cumo_row_smem[];
+    Accum* sdata = reinterpret_cast<Accum*>(cumo_row_smem);
+    unsigned int tid = threadIdx.x;
+
+    for (uint64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+        const TypeIn* xr = x + row * cols;
+        TypeIn* yr = y + row * cols;
+        Accum accum = impl.Identity(0);
+
+        for (uint64_t i = tid; i < cols; i += blockDim.x) {
+            impl.Reduce(impl.MapIn(xr[i], 0), accum);
+        }
+        // The tree leaves the row's total in sdata[0] behind a barrier of its
+        // own, and the second pass needs it in every thread, so the broadcast is
+        // a read of that slot rather than the value handed back.
+        reduce_in_block(accum, sdata, tid, 1, blockDim.x, true, impl);
+        auto stats = impl.MapOut(sdata[0]);
+
+        for (uint64_t i = tid; i < cols; i += blockDim.x) {
+            yr[i] = apply(xr[i], i, stats);
+        }
+        // Nobody may still be in the pass above when the next row rewrites
+        // sdata. Removing this leaves every value test passing and is only
+        // caught by compute-sanitizer --tool racecheck.
+        __syncthreads();
+    }
+}
+
+// The second half of the split path. blockIdx.y names the row, so finding one
+// costs no division, and this path is only taken where rows is small.
+template <typename TypeIn, typename Stats, typename Apply>
+__global__ void row_apply_kernel(
+        const TypeIn* x, TypeIn* y, const Stats* stats, uint64_t cols, Apply apply)
+{
+    uint64_t row = blockIdx.y;
+    const TypeIn* xr = x + row * cols;
+    TypeIn* yr = y + row * cols;
+    Stats st = stats[row];
+
+    for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < cols;
+            i += (uint64_t)blockDim.x * gridDim.x) {
+        yr[i] = apply(xr[i], i, st);
+    }
+}
+
+}  // namespace cumo_detail
+
+// Reduces each row of a contiguous rows x cols array with impl and writes the
+// row back through apply. Both arrays are laid out the same way and neither may
+// be the other.
+template <typename TypeIn, typename Impl, typename Apply>
+void cumo_row_reduce_apply(
+        char* px, char* py, uint64_t rows, uint64_t cols, Impl impl, Apply apply)
+{
+    typedef decltype(impl.Identity(0)) Accum;
+    typedef decltype(impl.MapOut(impl.Identity(0))) Stats;
+    // A block that takes several rows amortizes its own launch, so the grid is
+    // held below the row count on purpose. This is a tuning number and not the
+    // ceiling for this axis, which is CUMO_MAX_GRID_DIM: a million rows of 8
+    // take 1286.0us at a block per row and 1177.5 here, and two million of 4
+    // take 2357.5 against 2094.4.
+    static const uint64_t max_row_blocks = 65536;
+    int64_t want;
+    unsigned int block_dim, grid_dim;
+    size_t shared_mem_size;
+
+    if (rows == 0 || cols == 0) {
+        return;
+    }
+
+    // A row gets one block, so a row too long for one block to walk leaves the
+    // rest of the device idle unless there are other rows to fill it: one row of
+    // a million takes 695.8us that way, against 69.6 for the nine kernels a
+    // layer norm otherwise costs. Past that the reduction goes through the split
+    // machinery, which walks the row with the whole device, and a second kernel
+    // applies what it found.
+    if (rows < (uint64_t)cumo_detail::min_grid_size &&
+            cols > (uint64_t)(cumo_detail::max_block_size * cumo_detail::min_reduce_per_thread)) {
+        cumo_na_reduction_arg_t arg;
+        Stats* stats = (Stats*)cumo_cuda_runtime_malloc(rows * sizeof(Stats));
+        // rows is below min_grid_size to be here, which is well inside the y
+        // limit, but that is a threshold from reduce_kernel.h and not a promise
+        // about this axis, so the clamp is written out rather than assumed.
+        dim3 apply_grid((unsigned int)cumo_get_grid_dim(cols),
+                        (unsigned int)(rows < CUMO_MAX_GRID_DIM_Y ? rows : CUMO_MAX_GRID_DIM_Y));
+        unsigned int apply_block = (unsigned int)cumo_get_block_dim(cols);
+
+        memset(&arg, 0, sizeof(arg));
+        arg.in.ptr = px;
+        arg.in.step[0] = (ssize_t)(cols * sizeof(TypeIn));
+        arg.in.step[1] = (ssize_t)sizeof(TypeIn);
+        arg.in_indexer.ndim = 2;
+        arg.in_indexer.total_size = rows * cols;
+        arg.in_indexer.shape[0] = rows;
+        arg.in_indexer.shape[1] = cols;
+        arg.out.ptr = (char*)stats;
+        arg.out.step[0] = (ssize_t)sizeof(Stats);
+        arg.out_indexer.ndim = 1;
+        arg.out_indexer.total_size = rows;
+        arg.out_indexer.shape[0] = rows;
+
+        cumo_reduce_split<TypeIn, Stats, Impl>(arg, Impl(impl), (char*)stats);
+        cumo_detail::row_apply_kernel<TypeIn, Stats, Apply><<<apply_grid, apply_block>>>(
+                (const TypeIn*)px, (TypeIn*)py, stats, cols, apply);
+        cumo_check_launch_holding(stats);
+        cumo_cuda_runtime_free((char*)stats);
+        return;
+    }
+
+    // How many threads to spend on a row depends on how many rows there are to
+    // go around. Past enough rows to fill the device, giving each thread a run
+    // of the row rather than an element or two keeps the tree short: 4096 rows
+    // of 768 take 84.7us at a thread per element and 31.4 at sixteen. A single
+    // row has no other block to overlap with and wants every thread it can use,
+    // where the same sixteen cost 5.6us against 2.8.
+    if (rows < (uint64_t)cumo_detail::min_grid_size) {
+        want = cumo_detail::round_up_to_power_of_2((int64_t)cols);
+    } else {
+        want = cumo_detail::round_up_to_power_of_2(
+                ((int64_t)cols + cumo_detail::min_reduce_per_thread - 1) / cumo_detail::min_reduce_per_thread);
+    }
+    // reduce_in_block walks a power-of-two tree, and at one thread it returns
+    // without writing sdata at all, which would leave the broadcast reading
+    // whatever shared memory held. A whole warp is the smallest block here.
+    if (want < cumo_detail::warp_size) want = cumo_detail::warp_size;
+    if (want > cumo_detail::max_block_size) want = cumo_detail::max_block_size;
+    block_dim = (unsigned int)want;
+    grid_dim = (unsigned int)(rows < max_row_blocks ? rows : max_row_blocks);
+    shared_mem_size = block_dim * sizeof(Accum);
+
+    cumo_detail::row_reduce_apply_kernel<TypeIn, Impl, Apply><<<grid_dim, block_dim, shared_mem_size>>>(
+            (const TypeIn*)px, (TypeIn*)py, rows, cols, impl, apply);
+    cumo_cuda_runtime_check_kernel_launch();
+}
+
+#endif // CUMO_ROW_KERNEL_H

data/ext/cumo/include/cumo/row_method.h

@@ -0,0 +1,48 @@
+#ifndef CUMO_ROW_METHOD_H
+#define CUMO_ROW_METHOD_H
+
+#include <ruby.h>
+#include "cumo/narray.h"
+#include "cumo/check.h"
+
+/* Take every pointer, then measure, then run these. Taking one runs allocate,
+ * which is Ruby and free to resize an operand, replace its buffer, or free one
+ * whose pointer was already taken, and the last of those leaves a size that
+ * still matches. Re-taking a pointer answers a different one only in that case.
+ */
+#define CUMO_ROW_CHECK_READ_BUFFER(a, ptr, what)                        \
+    do {                                                                \
+        if (cumo_na_get_offset_pointer_for_read(a) != (ptr)) {          \
+            rb_raise(rb_eRuntimeError,                                  \
+                     "the buffer of %s was replaced while it was measured", \
+                     (what));                                           \
+        }                                                               \
+    } while (0)
+
+#define CUMO_ROW_CHECK_WRITE_BUFFER(a, ptr, what)                       \
+    do {                                                                \
+        if (cumo_na_get_offset_pointer_for_write(a) != (ptr)) {         \
+            rb_raise(rb_eRuntimeError,                                  \
+                     "the buffer of %s was replaced while it was measured", \
+                     (what));                                           \
+        }                                                               \
+    } while (0)
+
+#define CUMO_ROW_CHECK_SAME_SHAPE(nout, nin)                            \
+    do {                                                                \
+        cumo_narray_t *cumo_row_out = (nout);                           \
+        cumo_narray_t *cumo_row_in = (nin);                             \
+        int cumo_row_i;                                                 \
+        if (cumo_row_out->ndim != cumo_row_in->ndim) {                  \
+            rb_raise(cumo_na_eShapeError,                               \
+                     "self or the result was reshaped while it was measured"); \
+        }                                                               \
+        for (cumo_row_i = 0; cumo_row_i < cumo_row_in->ndim; ++cumo_row_i) { \
+            if (cumo_row_out->shape[cumo_row_i] != cumo_row_in->shape[cumo_row_i]) { \
+                rb_raise(cumo_na_eShapeError,                           \
+                         "self or the result was reshaped while it was measured"); \
+            }                                                           \
+        }                                                               \
+    } while (0)
+
+#endif /* CUMO_ROW_METHOD_H */

data/ext/cumo/include/cumo/template_kernel.h

@@ -176,14 +176,40 @@ cumo_get_block_dim(size_t n)
     return (block_dim == 0) ? 1 : block_dim;
 }
 
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
 // Raises the error a kernel launch reported, if any. Defined in cuda/runtime.c
 // because raising needs ruby.h, which the .cu translation units do not include.
 void cumo_cuda_runtime_check_kernel_launch(void);
 
+// The forms to call while scratch is held, which free what they are given
+// before they raise. Pass every buffer that is still outstanding, including the
+// caller's: a raise here leaves through longjmp and nothing below it runs.
+void cumo_cuda_runtime_check_kernel_launch_holding(char *p0, char *p1, char *p2, char *p3, char *p4);
+void cumo_cuda_runtime_check_taken_status_holding(int status, char *p0, char *p1, char *p2, char *p3, char *p4);
+void cumo_cuda_runtime_return_scratch(char *ptr, int wait_for_stream, cudaError_t *status);
+
+#if defined(__cplusplus)
+static inline void cumo_check_launch_holding(void *p0 = 0, void *p1 = 0, void *p2 = 0, void *p3 = 0, void *p4 = 0)
+{
+    cumo_cuda_runtime_check_kernel_launch_holding((char*)p0, (char*)p1, (char*)p2, (char*)p3, (char*)p4);
+}
+static inline void cumo_check_status_holding(cudaError_t status, void *p0 = 0, void *p1 = 0, void *p2 = 0, void *p3 = 0, void *p4 = 0)
+{
+    cumo_cuda_runtime_check_taken_status_holding((int)status, (char*)p0, (char*)p1, (char*)p2, (char*)p3, (char*)p4);
+}
+#endif
+
 // Scratch memory for a kernel that needs somewhere to put partial results.
 // Declared here rather than including cuda/memory_pool.h, which needs ruby.h.
 char* cumo_cuda_runtime_malloc(size_t size);
 void cumo_cuda_runtime_free(char *ptr);
 
+#if defined(__cplusplus)
+}
+#endif
+
 
 #endif /* ifndef CUMO_TEMPLATE_KERNEL_H */

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

@@ -0,0 +1,82 @@
+#ifndef CUMO_BF16_DEF_H
+#define CUMO_BF16_DEF_H
+
+// A struct rather than a bare uint16, so that C code which reaches for an
+// arithmetic operator on an element fails to compile instead of computing on
+// the bit pattern. Its layout matches __nv_bfloat16, which the kernels use.
+typedef struct { unsigned short x; } cumo_bfloat;
+
+static inline float cumo_bfloat2float(cumo_bfloat b)
+{
+    union { unsigned int u; float f; } o;
+    o.u = (unsigned int)b.x << 16;
+    return o.f;
+}
+
+static inline cumo_bfloat cumo_float2bfloat(float f)
+{
+    union { unsigned int u; float f; } v;
+    cumo_bfloat b;
+
+    v.f = f;
+    // bfloat16 keeps the exponent of a float, so every finite value is the top
+    // 16 bits rounded to nearest even. A NaN carries its payload in the low
+    // bits, and truncating one leaves an all-ones exponent over a zero
+    // mantissa, which reads as an infinity.
+    if ((v.u & 0x7f800000u) == 0x7f800000u && (v.u & 0x007fffffu) != 0u) {
+        b.x = (unsigned short)(((v.u >> 16) & 0x8000u) | 0x7fc0u);
+        return b;
+    }
+    b.x = (unsigned short)((v.u + 0x7fffu + ((v.u >> 16) & 1u)) >> 16);
+    return b;
+}
+
+static inline cumo_bfloat cumo_double2bfloat(double d)
+{
+    union { unsigned long long u; double d; } v;
+    cumo_bfloat b;
+    unsigned long long u, sign, mant, round_bits, tie;
+    unsigned int shift, out;
+    int exp;
+
+    // Going through float would round twice, and the first rounding can land
+    // on a tie the second one then moves the wrong way.
+    v.d = d;
+    u = v.u;
+    sign = (u >> 48) & 0x8000ull;
+    u &= 0x7fffffffffffffffull;
+
+    if (u > 0x7ff0000000000000ull) {
+        b.x = (unsigned short)(sign | 0x7fc0u);
+        return b;
+    }
+    if (u >= 0x47f0000000000000ull) {
+        b.x = (unsigned short)(sign | 0x7f80u);
+        return b;
+    }
+    if (u < 0x3790000000000000ull) {
+        b.x = (unsigned short)sign;
+        return b;
+    }
+
+    exp = (int)(u >> 52) - 1023;
+    mant = u & 0xfffffffffffffull;
+    if (exp >= -126) {
+        out = (unsigned int)(exp + 127) << 7;
+        shift = 45;
+    } else {
+        mant |= 0x10000000000000ull;
+        shift = (unsigned int)(-81 - exp);
+        out = 0;
+    }
+    round_bits = mant & ((1ull << shift) - 1ull);
+    tie = 1ull << (shift - 1);
+    out += (unsigned int)(mant >> shift);
+    if (round_bits > tie || (round_bits == tie && ((mant >> shift) & 1ull))) {
+        out += 1;
+    }
+    b.x = (unsigned short)(sign | out);
+    return b;
+}
+
+#endif // CUMO_BF16_DEF_H

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

@@ -0,0 +1,28 @@
+#ifndef CUMO_BF16_DEF_KERNEL_H
+#define CUMO_BF16_DEF_KERNEL_H
+
+#include <cuda_bf16.h>
+
+typedef __nv_bfloat16 cumo_bfloat;
+
+// The arithmetic on __nv_bfloat16 starts at sm_80, but the conversions have no
+// such floor, so every operation is computed in float. For + - * / and sqrt
+// that answers what the native instruction would: a float carries three times
+// the mantissa bits of a bfloat16, so the second rounding cannot land on a tie
+// the first one moved.
+__host__ __device__ static inline float cumo_bfloat2float(cumo_bfloat b)
+{
+    return __bfloat162float(b);
+}
+
+__host__ __device__ static inline cumo_bfloat cumo_float2bfloat(float f)
+{
+    return __float2bfloat16(f);
+}
+
+__host__ __device__ static inline cumo_bfloat cumo_double2bfloat(double d)
+{
+    return __double2bfloat16(d);
+}
+
+#endif // CUMO_BF16_DEF_KERNEL_H