cumo 0.8.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) {
@@ -635,7 +856,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>
+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;
@@ -648,9 +869,11 @@ 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_check_launch_holding(partial, held);
 
@@ -727,7 +950,7 @@ void cumo_reduce_split(cumo_na_reduction_arg_t arg, ReductionImpl&& impl, char*
 
 // 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>
+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;
@@ -736,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);
@@ -748,16 +971,18 @@ 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_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>
+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));
 
@@ -768,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);
@@ -777,12 +1002,12 @@ 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), held);
+        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, held);
+    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));
 }

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

@@ -23,9 +23,9 @@
 
 namespace cumo_detail {
 
-template <typename TypeIn, typename Impl, typename Apply>
+template <typename TypeIn, typename TypeOut, typename Stats, 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)
+        const TypeIn* x, TypeOut* y, Stats* stats_out, uint64_t rows, uint64_t cols, Impl impl, Apply apply)
 {
     typedef decltype(impl.Identity(0)) Accum;
     static_assert(alignof(Accum) <= 8,
@@ -38,7 +38,7 @@ __global__ void row_reduce_apply_kernel(
 
     for (uint64_t row = blockIdx.x; row < rows; row += gridDim.x) {
         const TypeIn* xr = x + row * cols;
-        TypeIn* yr = y + row * cols;
+        TypeOut* yr = y + row * cols;
         Accum accum = impl.Identity(0);
 
         for (uint64_t i = tid; i < cols; i += blockDim.x) {
@@ -50,6 +50,12 @@ __global__ void row_reduce_apply_kernel(
         reduce_in_block(accum, sdata, tid, 1, blockDim.x, true, impl);
         auto stats = impl.MapOut(sdata[0]);
 
+        // What the row was reduced to, for a caller that wants it back. One
+        // thread writes it, and the pass below reads only registers, so no
+        // barrier is owed between the two.
+        if (stats_out != NULL && tid == 0) {
+            stats_out[row] = stats;
+        }
         for (uint64_t i = tid; i < cols; i += blockDim.x) {
             yr[i] = apply(xr[i], i, stats);
         }
@@ -62,13 +68,13 @@ __global__ void row_reduce_apply_kernel(
 
 // 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>
+template <typename TypeIn, typename TypeOut, typename Stats, typename Apply>
 __global__ void row_apply_kernel(
-        const TypeIn* x, TypeIn* y, const Stats* stats, uint64_t cols, Apply apply)
+        const TypeIn* x, TypeOut* 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;
+    TypeOut* yr = y + row * cols;
     Stats st = stats[row];
 
     for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < cols;
@@ -82,9 +88,14 @@ __global__ void row_apply_kernel(
 // 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)
+//
+// pstats, when it is not NULL, takes what each row reduced to. Its elements are
+// the accumulator type, decltype(impl.MapOut(...)), which is not the element
+// type for a class whose accumulator is wider, so the buffer behind it has to
+// be sized in those. Nothing is written there for a row of no length.
+template <typename TypeIn, typename TypeOut, typename Impl, typename Apply>
+void cumo_row_reduce_apply_out(
+        char* px, char* py, char* pstats, uint64_t rows, uint64_t cols, Impl impl, Apply apply)
 {
     typedef decltype(impl.Identity(0)) Accum;
     typedef decltype(impl.MapOut(impl.Identity(0))) Stats;
@@ -111,7 +122,11 @@ void cumo_row_reduce_apply(
     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));
+        // The reduction writes the row totals wherever it is pointed, so a
+        // caller that wants them back is handed the buffer rather than a copy.
+        Stats* stats = pstats != NULL
+            ? (Stats*)pstats
+            : (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.
@@ -133,11 +148,19 @@ void cumo_row_reduce_apply(
         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);
+        // held is what the failure path frees, so it may only ever name scratch.
+        // Handing it a buffer a live Ruby array owns would return that to the
+        // pool and leave the array's own free to come back to it.
+        cumo_reduce_split<TypeIn, Stats, Impl>(arg, Impl(impl),
+                                               pstats != NULL ? NULL : (char*)stats);
+        cumo_detail::row_apply_kernel<TypeIn, TypeOut, Stats, Apply><<<apply_grid, apply_block>>>(
+                (const TypeIn*)px, (TypeOut*)py, stats, cols, apply);
+        if (pstats != NULL) {
+            cumo_cuda_runtime_check_kernel_launch();
+        } else {
+            cumo_check_launch_holding(stats);
+            cumo_cuda_runtime_free((char*)stats);
+        }
         return;
     }
 
@@ -162,9 +185,18 @@ void cumo_row_reduce_apply(
     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_detail::row_reduce_apply_kernel<TypeIn, TypeOut, Stats, Impl, Apply><<<grid_dim, block_dim, shared_mem_size>>>(
+            (const TypeIn*)px, (TypeOut*)py, (Stats*)pstats, rows, cols, impl, apply);
     cumo_cuda_runtime_check_kernel_launch();
 }
 
+// The shape layer_norm, rms_norm and softmax take: one array in, one of the
+// same type out, and nothing kept from the reduction.
+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)
+{
+    cumo_row_reduce_apply_out<TypeIn, TypeIn, Impl, Apply>(px, py, NULL, rows, cols, impl, apply);
+}
+
 #endif // CUMO_ROW_KERNEL_H