cumo 0.7.0 → 0.9.0

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

@@ -69,6 +69,11 @@ typedef struct {
     bool out_flat;
     bool out2_flat;
     bool out_inner;         // the out axis, not the reduce axis, runs along memory
+    // How many consecutive indices share one address, for a range that is flat
+    // once a trailing run of broadcast axes is taken off it. Zero when the
+    // range is not of that shape, and never 1: a divisor of 1 is in_out_flat.
+    int64_t in_out_div;
+    int64_t in_reduce_div;
     ssize_t in_out_step;    // bytes, or bits for a Bit input
     ssize_t in_reduce_step; // bytes, or bits for a Bit input
     ssize_t out_step;       // bytes
@@ -93,8 +98,58 @@ static inline bool axes_are_flat(const TIarray& iarray, const cumo_na_indexer_t&
     return true;
 }
 
+// axes_are_flat for a range that ends in broadcast axes. A step of 0 does not
+// move the address, so the offset of the i-th element is (i / div) * step,
+// where div is how many indices share one address. Answers a div of 1 where
+// there is no such axis, which is what axes_are_flat already describes.
+template <typename TIarray>
+static inline bool axes_are_flat_bcast(const TIarray& iarray, const cumo_na_indexer_t& indexer, int begin, int end, ssize_t* step, int64_t* div) {
+    int64_t d = 1;
+    while (end > begin && iarray.step[end - 1] == 0) {
+        d *= static_cast<int64_t>(indexer.shape[end - 1]);
+        --end;
+    }
+    if (!axes_are_flat(iarray, indexer, begin, end, step)) {
+        return false;
+    }
+    *div = d;
+    return true;
+}
+
+// Fills in the divisor form for a range axes_are_flat turned down, so that the
+// kernel spends one division on it rather than walking every dimension.
+template <typename TIarray>
+static inline int64_t reduce_addr_div(bool flat, const TIarray& iarray, const cumo_na_indexer_t& indexer, int begin, int end, ssize_t* step) {
+    ssize_t bcast_step;
+    int64_t div;
+
+    // *step stays as the caller left it unless there is a divisor to go with
+    // it, so that the two are only ever read together.
+    if (flat || !axes_are_flat_bcast(iarray, indexer, begin, end, &bcast_step, &div) || div <= 1) {
+        return 0;
+    }
+    *step = bcast_step;
+    return div;
+}
+
+// The innermost step of each of the two axis groups. out_inner is decided from
+// these, and a zip decides again from both operands', so they are handed back
+// rather than kept in cumo_reduce_addr_t: that one is a kernel parameter, and
+// the note above says what a wider one costs.
+typedef struct {
+    ssize_t out;
+    ssize_t reduce;
+} cumo_inner_steps_t;
+
+// The out axis runs along memory when its step is real and shorter than the
+// reduce axis's, which is what decides how a block shares its threads.
+static inline bool inner_steps_say_out(const cumo_inner_steps_t& inner) {
+    return inner.out != 0 &&
+        (inner.reduce == 0 || step_magnitude(inner.out) < step_magnitude(inner.reduce));
+}
+
 template <typename TArg>
-static inline cumo_reduce_addr_t make_reduce_addr(const TArg& arg, int64_t reduce_total_size) {
+static inline cumo_reduce_addr_t make_reduce_addr(const TArg& arg, int64_t reduce_total_size, cumo_inner_steps_t* inner = 0) {
     cumo_reduce_addr_t ad;
     int in_ndim = arg.in_indexer.ndim;
     ssize_t whole_step;
@@ -108,6 +163,8 @@ static inline cumo_reduce_addr_t make_reduce_addr(const TArg& arg, int64_t reduc
         ad.in_reduce_flat = true;
         ad.in_reduce_step = whole_step;
         ad.in_out_step = whole_step * reduce_total_size;
+        ad.in_out_div = 0;
+        ad.in_reduce_div = 0;
     } else {
         int split = in_ndim;
         int64_t acc = 1;
@@ -119,33 +176,74 @@ static inline cumo_reduce_addr_t make_reduce_addr(const TArg& arg, int64_t reduc
             ad.split = split;
             ad.in_reduce_flat = axes_are_flat(arg.in, arg.in_indexer, split, in_ndim, &ad.in_reduce_step);
             ad.in_out_flat = axes_are_flat(arg.in, arg.in_indexer, 0, split, &ad.in_out_step);
+            ad.in_reduce_div = reduce_addr_div(ad.in_reduce_flat, arg.in, arg.in_indexer, split, in_ndim, &ad.in_reduce_step);
+            ad.in_out_div = reduce_addr_div(ad.in_out_flat, arg.in, arg.in_indexer, 0, split, &ad.in_out_step);
         } else {
             ad.split = -1;
             ad.in_reduce_flat = false;
             ad.in_out_flat = false;
             ad.in_reduce_step = 0;
             ad.in_out_step = 0;
+            ad.in_out_div = 0;
+            ad.in_reduce_div = 0;
+        }
+        // axes_are_flat leaves the step alone where it answers false, and the
+        // kernels work one out unconditionally, so give them a zero to read.
+        if (!ad.in_reduce_flat && ad.in_reduce_div == 0) {
+            ad.in_reduce_step = 0;
+        }
+        if (!ad.in_out_flat && ad.in_out_div == 0) {
+            ad.in_out_step = 0;
         }
     }
 
+    ad.out_step = 0;
     ad.out_flat = axes_are_flat(arg.out, arg.out_indexer, 0, arg.out_indexer.ndim, &ad.out_step);
     ad.out2_flat = true;
     ad.out2_step = 0;
 
-    ssize_t out_inner_step, reduce_inner_step;
+    cumo_inner_steps_t steps;
     if (ad.in_out_flat && ad.in_reduce_flat) {
-        out_inner_step = ad.in_out_step;
-        reduce_inner_step = ad.in_reduce_step;
+        steps.out = ad.in_out_step;
+        steps.reduce = ad.in_reduce_step;
     } else {
-        out_inner_step = ad.split > 0 ? arg.in.step[ad.split - 1] : 0;
-        reduce_inner_step = (ad.split >= 0 && ad.split < in_ndim) ? arg.in.step[in_ndim - 1] : 0;
+        steps.out = ad.split > 0 ? arg.in.step[ad.split - 1] : 0;
+        steps.reduce = (ad.split >= 0 && ad.split < in_ndim) ? arg.in.step[in_ndim - 1] : 0;
+    }
+    ad.out_inner = inner_steps_say_out(steps);
+    if (inner != 0) {
+        *inner = steps;
     }
-    ad.out_inner = out_inner_step != 0 &&
-        (reduce_inner_step == 0 || step_magnitude(out_inner_step) < step_magnitude(reduce_inner_step));
 
     return ad;
 }
 
+// The shorter of two steps, counting a zero as no step at all: an operand that
+// does not move along an axis reads one address for the whole of it, so it has
+// no say in which axis runs along memory.
+static inline ssize_t shorter_step(ssize_t a, ssize_t b) {
+    if (a == 0) return b;
+    if (b == 0) return a;
+    return step_magnitude(a) < step_magnitude(b) ? a : b;
+}
+
+// A zip reduction reads both operands through one thread layout, so the layout
+// has to answer for both. Taking the shorter step of the two along each axis
+// gives the same answer whichever operand the caller wrote first, which a
+// decision read off one of them does not.
+template <typename TArg>
+static inline void make_zip_reduce_addrs(const TArg& arg, const TArg& arg2, int64_t reduce_total_size,
+        cumo_reduce_addr_t* ad, cumo_reduce_addr_t* ad2) {
+    cumo_inner_steps_t inner, inner2, both;
+
+    *ad = make_reduce_addr(arg, reduce_total_size, &inner);
+    *ad2 = make_reduce_addr(arg2, reduce_total_size, &inner2);
+
+    both.out = shorter_step(inner.out, inner2.out);
+    both.reduce = shorter_step(inner.reduce, inner2.reduce);
+    ad->out_inner = ad2->out_inner = inner_steps_say_out(both);
+}
+
 static inline void set_reduce_addr_out2(cumo_reduce_addr_t* ad, const cumo_na_reduction_arg_t& arg, const cumo_na_iarray_t& out2) {
     ad->out2_flat = axes_are_flat(out2, arg.out_indexer, 0, arg.out_indexer.ndim, &ad->out2_step);
 }
@@ -219,6 +317,7 @@ __device__ static __forceinline__ void axes_offset_pair(const TIarray& a, const
 template <bool FLAT>
 __device__ static __forceinline__ ssize_t reduce_in_out_offset(const cumo_na_iarray_t& in, const cumo_na_indexer_t& in_indexer, const cumo_reduce_addr_t& ad, int64_t i_out) {
     if (FLAT || ad.in_out_flat) return i_out * ad.in_out_step;
+    if (!FLAT && ad.in_out_div > 0) return (i_out / ad.in_out_div) * ad.in_out_step;
     if (ad.split < 0) return 0;
     return axes_offset(in, in_indexer, 0, ad.split, i_out);
 }
@@ -226,7 +325,8 @@ __device__ static __forceinline__ ssize_t reduce_in_out_offset(const cumo_na_iar
 // reduce_in_out_offset for the two operands of a zip reduction at once.
 template <bool FLAT>
 __device__ static __forceinline__ void reduce_in_out_offset_pair(const cumo_na_iarray_t& in, const cumo_na_iarray_t& in2, const cumo_na_indexer_t& in_indexer, const cumo_reduce_addr_t& ad, const cumo_reduce_addr_t& ad2, int64_t i_out, ssize_t* off, ssize_t* off2) {
-    if (!FLAT && !ad.in_out_flat && !ad2.in_out_flat && ad.split >= 0 && ad.split == ad2.split) {
+    if (!FLAT && !ad.in_out_flat && !ad2.in_out_flat && ad.in_out_div == 0 && ad2.in_out_div == 0 &&
+        ad.split >= 0 && ad.split == ad2.split) {
         axes_offset_pair(in, in2, in_indexer, 0, ad.split, i_out, off, off2);
         return;
     }
@@ -237,6 +337,7 @@ __device__ static __forceinline__ void reduce_in_out_offset_pair(const cumo_na_i
 template <bool FLAT>
 __device__ static __forceinline__ ssize_t reduce_in_offset(const cumo_na_iarray_t& in, const cumo_na_indexer_t& in_indexer, const cumo_reduce_addr_t& ad, ssize_t in_out_off, int64_t i_reduce, int64_t i_in) {
     if (FLAT || ad.in_reduce_flat) return in_out_off + i_reduce * ad.in_reduce_step;
+    if (!FLAT && ad.in_reduce_div > 0) return in_out_off + (i_reduce / ad.in_reduce_div) * ad.in_reduce_step;
     if (ad.split < 0) return axes_offset(in, in_indexer, 0, in_indexer.ndim, i_in);
     return in_out_off + axes_offset(in, in_indexer, ad.split, in_indexer.ndim, i_reduce);
 }
@@ -340,7 +441,7 @@ __device__ static __forceinline__ auto reduce_axis(const cumo_na_iarray_t& in, c
 // once, which is what mulsum wants: the product it accumulates never exists as
 // an array. The pair comes out of one broadcast, so arg.in_indexer addresses
 // both and only the steps differ, which is what in2 and ad2 carry.
-template <bool FLAT, typename TypeIn, typename ReductionImpl>
+template <bool FLAT, typename TypeIn, typename TypeIn2, typename ReductionImpl>
 __device__ static __forceinline__ auto reduce_axis_zip(const cumo_na_reduction_arg_t& arg, const cumo_na_iarray_t& in2,
         const cumo_reduce_addr_t& ad, const cumo_reduce_addr_t& ad2, ReductionImpl& impl,
         ssize_t in_out_off, ssize_t in_out_off2, int64_t i_in, int64_t begin, int64_t end,
@@ -356,7 +457,7 @@ __device__ static __forceinline__ auto reduce_axis_zip(const cumo_na_reduction_a
     TypeReduce accum = impl.Identity(0);
 
     for (; i_reduce < end; i_reduce += reduce_block_size, i_in += reduce_block_size) {
-        impl.Reduce(impl.MapIn(*reinterpret_cast<TypeIn*>(p), *reinterpret_cast<TypeIn*>(q), i_reduce), accum);
+        impl.Reduce(impl.MapIn(*reinterpret_cast<TypeIn*>(p), *reinterpret_cast<TypeIn2*>(q), i_reduce), accum);
         p = (FLAT || ad.in_reduce_flat)
             ? p + advance
             : arg.in.ptr + reduce_in_offset<FLAT>(arg.in, arg.in_indexer, ad, in_out_off, i_reduce + reduce_block_size, i_in + reduce_block_size);
@@ -402,7 +503,7 @@ __global__ static void reduction_kernel(CUMO_GRID_CONSTANT cumo_na_reduction_arg
 
 // Variant of reduction_kernel reading two inputs, for mulsum. See
 // reduce_axis_zip above.
-template <bool FLAT, typename TypeIn, typename TypeOut, typename ReductionImpl>
+template <bool FLAT, typename TypeIn, typename TypeIn2, typename TypeOut, typename ReductionImpl>
 __global__ static void reduction_zip_kernel(CUMO_GRID_CONSTANT cumo_na_reduction_arg_t arg, CUMO_GRID_CONSTANT cumo_na_iarray_t in2, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad2, int out_block_size, int reduce_block_size, ReductionImpl impl) {
     using TypeReduce = decltype(impl.Identity(0));
 
@@ -423,7 +524,7 @@ __global__ static void reduction_zip_kernel(CUMO_GRID_CONSTANT cumo_na_reduction
         reduce_in_out_offset_pair<FLAT>(arg.in, in2, arg.in_indexer, ad, ad2, i_out, &in_out_off, &in_out_off2);
         int64_t i_in = i_out * reduce_total_size + reduce_offset;
 
-        TypeReduce accum = reduce_axis_zip<FLAT,TypeIn>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, i_in, 0, reduce_total_size, reduce_offset, reduce_block_size);
+        TypeReduce accum = reduce_axis_zip<FLAT,TypeIn,TypeIn2>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, i_in, 0, reduce_total_size, reduce_offset, reduce_block_size);
 
         accum = reduce_in_block(accum, sdata, tid, out_block_size, reduce_block_size, !ad.out_inner, impl);
         if (reduce_offset == 0) {
@@ -433,6 +534,89 @@ __global__ static void reduction_zip_kernel(CUMO_GRID_CONSTANT cumo_na_reduction
     }
 }
 
+// The reduce axis of a zip whose operands both address without the indexer.
+// Splitting this out of reduce_axis_zip is what buys the speed: the general
+// path is gone from the instantiation, so the indexer, which carries shape[]
+// for CUMO_NA_MAX_DIMENSION, never has to be live.
+template <typename TypeIn, typename TypeIn2, typename ReductionImpl>
+__device__ static __forceinline__ auto reduce_axis_zip_nodim(const cumo_na_reduction_arg_t& arg, const cumo_na_iarray_t& in2,
+        const cumo_reduce_addr_t& ad, const cumo_reduce_addr_t& ad2, ReductionImpl& impl,
+        ssize_t in_out_off, ssize_t in_out_off2, int64_t begin, int64_t end,
+        int64_t reduce_offset, int64_t reduce_block_size) -> decltype(impl.Identity(0)) {
+    using TypeReduce = decltype(impl.Identity(0));
+
+    int64_t i_reduce = begin + reduce_offset;
+    char* p = arg.in.ptr + in_out_off + i_reduce * ad.in_reduce_step;
+    ssize_t advance = ad.in_reduce_step * reduce_block_size;
+
+    TypeReduce accum = impl.Identity(0);
+
+    // The two loops keep the divisor out of the body. One loop with a divisor
+    // of 1 for the flat case costs a 64-bit division on every element, which
+    // is more than the addressing it saves.
+    if (ad2.in_reduce_flat) {
+        char* q = in2.ptr + in_out_off2 + i_reduce * ad2.in_reduce_step;
+        ssize_t advance2 = ad2.in_reduce_step * reduce_block_size;
+
+        for (; i_reduce < end; i_reduce += reduce_block_size, p += advance, q += advance2) {
+            impl.Reduce(impl.MapIn(*reinterpret_cast<TypeIn*>(p), *reinterpret_cast<TypeIn2*>(q), i_reduce), accum);
+        }
+    } else {
+        int64_t div2 = ad2.in_reduce_div;
+
+        for (; i_reduce < end; i_reduce += reduce_block_size, p += advance) {
+            char* q = in2.ptr + in_out_off2 + (i_reduce / div2) * ad2.in_reduce_step;
+            impl.Reduce(impl.MapIn(*reinterpret_cast<TypeIn*>(p), *reinterpret_cast<TypeIn2*>(q), i_reduce), accum);
+        }
+    }
+    return accum;
+}
+
+// reduction_zip_kernel for the same case.
+template <typename TypeIn, typename TypeIn2, typename TypeOut, typename ReductionImpl>
+__global__ static void reduction_zip_nodim_kernel(CUMO_GRID_CONSTANT cumo_na_reduction_arg_t arg, CUMO_GRID_CONSTANT cumo_na_iarray_t in2, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad2, int out_block_size, int reduce_block_size, ReductionImpl impl) {
+    using TypeReduce = decltype(impl.Identity(0));
+
+    extern __shared__ __align__(8) char sdata_raw[];
+    TypeReduce* sdata = reinterpret_cast<TypeReduce*>(sdata_raw);
+    unsigned int tid = threadIdx.x;
+
+    int64_t out_total_size = arg.out_indexer.total_size;
+    int64_t reduce_total_size = arg.in_indexer.total_size / out_total_size;
+    int64_t out_div2 = ad2.in_out_flat ? 1 : ad2.in_out_div;
+
+    int64_t reduce_offset, out_offset;
+    reduce_thread_split(ad, tid, out_block_size, reduce_block_size, &reduce_offset, &out_offset);
+    int64_t out_base = blockIdx.x * out_block_size;
+    int64_t out_stride = gridDim.x * out_block_size;
+
+    for (int64_t i_out = out_base + out_offset; i_out < out_total_size; i_out += out_stride) {
+        ssize_t in_out_off = i_out * ad.in_out_step;
+        ssize_t in_out_off2 = (i_out / out_div2) * ad2.in_out_step;
+
+        TypeReduce accum = reduce_axis_zip_nodim<TypeIn,TypeIn2>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, 0, reduce_total_size, reduce_offset, reduce_block_size);
+
+        accum = reduce_in_block(accum, sdata, tid, out_block_size, reduce_block_size, !ad.out_inner, impl);
+        if (reduce_offset == 0) {
+            TypeOut* out_ptr = reinterpret_cast<TypeOut*>(arg.out.ptr + i_out * ad.out_step);
+            *out_ptr = impl.MapOut(accum);
+        }
+    }
+}
+
+// Whether both operands of a zip address without the indexer: one is flat, the
+// other flat or flat once its broadcast axes are divided out.
+//
+// Only the second operand may carry a divisor. The first is the receiver, and
+// a broadcast one reaches this as the argument: a.mulsum(b) with b the smaller
+// shape. Writing it the other way round leaves the general path, which answers
+// the same and takes the time the kernel below saves.
+static inline bool zip_axes_need_no_dim(const cumo_reduce_addr_t& ad, const cumo_reduce_addr_t& ad2) {
+    return ad.in_out_flat && ad.in_reduce_flat && ad.out_flat &&
+        (ad2.in_out_flat || ad2.in_out_div > 0) &&
+        (ad2.in_reduce_flat || ad2.in_reduce_div > 0);
+}
+
 // Variant of reduction_kernel for arg-reductions (argmax/argmin), which report
 // the index along the reduction axis rather than the index of an element.
 template <bool FLAT, typename TypeIn, typename TypeOut, typename ReductionImpl>
@@ -539,8 +723,45 @@ __global__ static void reduction_partial_kernel(CUMO_GRID_CONSTANT cumo_na_reduc
     }
 }
 
+// reduction_zip_partial_kernel for operands that address without the indexer.
+// A split reduction is what a small output over a long reduce axis takes, so
+// leaving this one on the general path would miss the shape that gains most.
+template <typename TypeIn, typename TypeIn2, typename TypeReduce, typename ReductionImpl>
+__global__ static void reduction_zip_nodim_partial_kernel(CUMO_GRID_CONSTANT cumo_na_reduction_arg_t arg, CUMO_GRID_CONSTANT cumo_na_iarray_t in2, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad2, 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);
+    unsigned int tid = threadIdx.x;
+
+    int64_t out_total_size = arg.out_indexer.total_size;
+    int64_t reduce_total_size = arg.in_indexer.total_size / out_total_size;
+    int64_t partial_total_size = out_total_size * n_split;
+    int64_t out_div2 = ad2.in_out_flat ? 1 : ad2.in_out_div;
+
+    int64_t reduce_offset, out_offset;
+    reduce_thread_split(ad, tid, out_block_size, reduce_block_size, &reduce_offset, &out_offset);
+    int64_t out_base = blockIdx.x * out_block_size;
+    int64_t out_stride = gridDim.x * out_block_size;
+
+    for (int64_t i = out_base + out_offset; i < partial_total_size; i += out_stride) {
+        int64_t i_out = i % out_total_size;
+        int64_t i_split = i / out_total_size;
+        int64_t begin = i_split * chunk;
+        int64_t end = begin + chunk;
+        if (end > reduce_total_size) end = reduce_total_size;
+        ssize_t in_out_off = i_out * ad.in_out_step;
+        ssize_t in_out_off2 = (i_out / out_div2) * ad2.in_out_step;
+
+        TypeReduce accum = reduce_axis_zip_nodim<TypeIn,TypeIn2>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, begin, end, reduce_offset, reduce_block_size);
+
+        accum = reduce_in_block(accum, sdata, tid, out_block_size, reduce_block_size, !ad.out_inner, impl);
+        if (reduce_offset == 0) {
+            partial[i_out * n_split + i_split] = accum;
+        }
+    }
+}
+
 // First pass of a split zip reduction. See reduction_partial_kernel above.
-template <bool FLAT, typename TypeIn, typename TypeReduce, typename ReductionImpl>
+template <bool FLAT, typename TypeIn, typename TypeIn2, typename TypeReduce, typename ReductionImpl>
 __global__ static void reduction_zip_partial_kernel(CUMO_GRID_CONSTANT cumo_na_reduction_arg_t arg, CUMO_GRID_CONSTANT cumo_na_iarray_t in2, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad, CUMO_GRID_CONSTANT cumo_reduce_addr_t ad2, 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);
@@ -565,7 +786,7 @@ __global__ static void reduction_zip_partial_kernel(CUMO_GRID_CONSTANT cumo_na_r
         reduce_in_out_offset_pair<FLAT>(arg.in, in2, arg.in_indexer, ad, ad2, i_out, &in_out_off, &in_out_off2);
         int64_t i_in = i_out * reduce_total_size + begin + reduce_offset;
 
-        TypeReduce accum = reduce_axis_zip<FLAT,TypeIn>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, i_in, begin, end, reduce_offset, reduce_block_size);
+        TypeReduce accum = reduce_axis_zip<FLAT,TypeIn,TypeIn2>(arg, in2, ad, ad2, impl, in_out_off, in_out_off2, i_in, begin, end, reduce_offset, reduce_block_size);
 
         accum = reduce_in_block(accum, sdata, tid, out_block_size, reduce_block_size, !ad.out_inner, impl);
         if (reduce_offset == 0) {
@@ -602,7 +823,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 +839,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);
@@ -635,8 +856,8 @@ 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) {
+template <typename TypeIn, typename TypeIn2, 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, 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));
@@ -648,11 +869,13 @@ TypeReduce* reduce_zip_partial_pass(cumo_na_reduction_arg_t arg, cumo_na_iarray_
     int64_t shared_mem_size = sizeof(TypeReduce) * max_block_size;
 
     if (zip_axes_are_flat(ad, ad2)) {
-        reduction_zip_partial_kernel<true,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);
+        reduction_zip_partial_kernel<true,TypeIn,TypeIn2,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);
+    } else if (zip_axes_need_no_dim(ad, ad2)) {
+        reduction_zip_nodim_partial_kernel<TypeIn,TypeIn2,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);
     } 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);
+        reduction_zip_partial_kernel<false,TypeIn,TypeIn2,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 +887,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 +913,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 +922,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 +938,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) {
+template <typename TypeIn, typename TypeIn2, typename TypeOut, typename ReductionImpl>
+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;
     }
@@ -732,8 +959,8 @@ void cumo_reduce_zip(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Reductio
     int64_t reduce_total_size = arg.in_indexer.total_size / arg.out_indexer.total_size;
     cumo_na_reduction_arg_t arg2 = arg;
     arg2.in = in2;
-    cumo_detail::cumo_reduce_addr_t ad = cumo_detail::make_reduce_addr(arg, reduce_total_size);
-    cumo_detail::cumo_reduce_addr_t ad2 = cumo_detail::make_reduce_addr(arg2, reduce_total_size);
+    cumo_detail::cumo_reduce_addr_t ad, ad2;
+    cumo_detail::make_zip_reduce_addrs(arg, arg2, reduce_total_size, &ad, &ad2);
 
     int64_t out_block_size, reduce_block_size;
     cumo_detail::reduce_block_split(ad, reduce_total_size, &out_block_size, &reduce_block_size);
@@ -744,17 +971,19 @@ void cumo_reduce_zip(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Reductio
     int64_t shared_mem_size = sizeof(decltype(impl.Identity(0))) * block_size;
 
     if (cumo_detail::zip_axes_are_flat(ad, ad2) && ad.out_flat) {
-        cumo_detail::reduction_zip_kernel<true,TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, in2, ad, ad2, out_block_size, reduce_block_size, impl);
+        cumo_detail::reduction_zip_kernel<true,TypeIn,TypeIn2,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, in2, ad, ad2, out_block_size, reduce_block_size, impl);
+    } else if (cumo_detail::zip_axes_need_no_dim(ad, ad2)) {
+        cumo_detail::reduction_zip_nodim_kernel<TypeIn,TypeIn2,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, in2, ad, ad2, out_block_size, reduce_block_size, impl);
     } 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_detail::reduction_zip_kernel<false,TypeIn,TypeIn2,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) {
+template <typename TypeIn, typename TypeIn2, typename TypeOut, typename ReductionImpl>
+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) {
@@ -764,8 +993,8 @@ void cumo_reduce_zip_split(cumo_na_reduction_arg_t arg, cumo_na_iarray_t in2, Re
     int64_t reduce_total_size = arg.in_indexer.total_size / arg.out_indexer.total_size;
     cumo_na_reduction_arg_t arg2 = arg;
     arg2.in = in2;
-    cumo_detail::cumo_reduce_addr_t ad = cumo_detail::make_reduce_addr(arg, reduce_total_size);
-    cumo_detail::cumo_reduce_addr_t ad2 = cumo_detail::make_reduce_addr(arg2, reduce_total_size);
+    cumo_detail::cumo_reduce_addr_t ad, ad2;
+    cumo_detail::make_zip_reduce_addrs(arg, arg2, reduce_total_size, &ad, &ad2);
 
     int64_t out_block_size, reduce_block_size;
     cumo_detail::reduce_block_split(ad, reduce_total_size, &out_block_size, &reduce_block_size);
@@ -773,13 +1002,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, TypeIn2, 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, TypeIn2, 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 +1016,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 +1038,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 +1059,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 +1093,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 +1116,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));
 }