cumo 0.6.0 → 0.7.0

data/ext/cumo/cuda/cudnn_impl.cpp

@@ -158,7 +158,8 @@ cumo_cuda_cudnn_CreateConvolutionDescriptor(
         size_t ndim,
         int* int_stride,
         int* int_pad,
-        cudnnDataType_t cudnn_dtype) {
+        cudnnDataType_t compute_dtype,
+        cudnnMathType_t math_type) {
     cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
     int int_dilation[CUMO_NA_MAX_DIMENSION];
     for (size_t idim = 0; idim < ndim; ++idim) {
@@ -178,7 +179,7 @@ cumo_cuda_cudnn_CreateConvolutionDescriptor(
                 int_dilation[0],
                 int_dilation[1],
                 CUDNN_CROSS_CORRELATION,
-                cudnn_dtype);
+                compute_dtype);
     } else {
         status = cudnnSetConvolutionNdDescriptor(
                 *desc,
@@ -187,10 +188,14 @@ cumo_cuda_cudnn_CreateConvolutionDescriptor(
                 int_stride,
                 int_dilation,
                 CUDNN_CROSS_CORRELATION,
-                cudnn_dtype);
+                compute_dtype);
     }
+    if (status != CUDNN_STATUS_SUCCESS) return status;
 
-    return status;
+    // Tensor cores are not reached under CUDNN_DEFAULT_MATH, and asking for
+    // them where they do not apply is not an error, so every dtype names what
+    // it wants.
+    return cudnnSetConvolutionMathType(*desc, math_type);
 }
 
 cudnnStatus_t
@@ -304,9 +309,49 @@ struct AlgoCacheKeyHash {
     }
 };
 
-using FwdAlgoCacheMap = std::unordered_map<AlgoCacheKey, std::pair<cudnnConvolutionFwdAlgo_t, size_t>, AlgoCacheKeyHash>;
-using BwdDataAlgoCacheMap = std::unordered_map<AlgoCacheKey, std::pair<cudnnConvolutionBwdDataAlgo_t, size_t>, AlgoCacheKeyHash>;
-using BwdFilterAlgoCacheMap = std::unordered_map<AlgoCacheKey, std::pair<cudnnConvolutionBwdFilterAlgo_t, size_t>, AlgoCacheKeyHash>;
+// The math type belongs with the algorithm: the descriptor's own setting is
+// only a permission, and cuDNN reports which one the chosen algorithm actually
+// used. Handing that algorithm to the convolution under a different math type
+// is rejected, or answers something else.
+// A template cannot be declared with C linkage, and this file's exports are
+// wrapped in extern "C".
+extern "C++" {
+template <typename Algo>
+struct AlgoCacheEntry {
+    Algo algo;
+    size_t memory;
+    cudnnMathType_t math_type;
+};
+}
+
+using FwdAlgoCacheMap = std::unordered_map<AlgoCacheKey, AlgoCacheEntry<cudnnConvolutionFwdAlgo_t>, AlgoCacheKeyHash>;
+using BwdDataAlgoCacheMap = std::unordered_map<AlgoCacheKey, AlgoCacheEntry<cudnnConvolutionBwdDataAlgo_t>, AlgoCacheKeyHash>;
+using BwdFilterAlgoCacheMap = std::unordered_map<AlgoCacheKey, AlgoCacheEntry<cudnnConvolutionBwdFilterAlgo_t>, AlgoCacheKeyHash>;
+
+// Every search builds the same key. Leaving each of the three to do it by hand
+// is how the device id came to be set in one of them and not the others.
+static AlgoCacheKey
+MakeAlgoCacheKey(
+        cumo_narray_t* nx, cumo_narray_t* nw, cumo_narray_t* ny,
+        int* int_stride, int* int_pad, size_t ndim,
+        cudnnDataType_t cudnn_dtype, size_t max_workspace_size)
+{
+    auto key = AlgoCacheKey{};
+    cumo_cuda_runtime_check_status(cudaGetDevice(&(key.device_id)));
+    key.ndim = ndim;
+    for (size_t idim = 0; idim < ndim + 2; ++idim) {
+        key.x_shape[idim] = nx->shape[idim];
+        key.w_shape[idim] = nw->shape[idim];
+        key.y_shape[idim] = ny->shape[idim];
+    }
+    for (size_t idim = 0; idim < ndim; ++idim) {
+        key.pad[idim] = int_pad[idim];
+        key.stride[idim] = int_stride[idim];
+    }
+    key.dtype = cudnn_dtype;
+    key.max_workspace_size = max_workspace_size;
+    return key;
+}
 
 static FwdAlgoCacheMap fwd_algo_cache_map_{};
 static BwdDataAlgoCacheMap bwd_data_algo_cache_map_{};
@@ -335,28 +380,19 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
     CumoGetNArray(w, nw);
     CumoGetNArray(y, ny);
 
-    auto key = AlgoCacheKey{};
-    cumo_cuda_runtime_check_status(cudaGetDevice(&(key.device_id)));
-    key.ndim = ndim;
-    for (size_t idim = 0; idim < ndim + 2; ++idim) {
-        key.x_shape[idim] = nx->shape[idim];
-        key.w_shape[idim] = nw->shape[idim];
-        key.y_shape[idim] = ny->shape[idim];
-    }
-    for (size_t idim = 0; idim < ndim; ++idim) {
-        key.pad[idim]= int_pad[idim];
-        key.stride[idim]= int_stride[idim];
-    }
-    key.dtype = cudnn_dtype;
-    key.max_workspace_size = max_workspace_size;
+    auto key = MakeAlgoCacheKey(nx, nw, ny, int_stride, int_pad, ndim,
+                                cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = fwd_algo_cache_map_;
     // TODO: thread-safe
     auto it = algo_cache_map.find(key);
     if (it != algo_cache_map.end()) {
-        auto pair = it->second;
-        perf_result->algo = pair.first;
-        perf_result->memory = pair.second;
+        auto entry = it->second;
+        // clear the fields the search would have filled but the cache drops
+        *perf_result = {};
+        perf_result->algo = entry.algo;
+        perf_result->memory = entry.memory;
+        perf_result->mathType = entry.math_type;
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -382,10 +418,13 @@ cumo_cuda_cudnn_FindConvolutionForwardAlgorithm(
                 max_workspace_size);
     cumo_cuda_runtime_free(workspace);
     if (status != CUDNN_STATUS_SUCCESS) return status;
-    assert(returned_algo_count == 1);
+    // A search that answers success with nothing to report would leave
+    // perf_result untouched, and its algo and math type are used below.
+    if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
+    if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory};
+    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
     return status;
 }
 
@@ -412,27 +451,19 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
     CumoGetNArray(w, nw);
     CumoGetNArray(y, ny);
 
-    auto key = AlgoCacheKey{};
-    key.ndim = ndim;
-    for (size_t idim = 0; idim < ndim + 2; ++idim) {
-        key.x_shape[idim] = nx->shape[idim];
-        key.w_shape[idim] = nw->shape[idim];
-        key.y_shape[idim] = ny->shape[idim];
-    }
-    for (size_t idim = 0; idim < ndim; ++idim) {
-        key.pad[idim]= int_pad[idim];
-        key.stride[idim]= int_stride[idim];
-    }
-    key.dtype = cudnn_dtype;
-    key.max_workspace_size = max_workspace_size;
+    auto key = MakeAlgoCacheKey(nx, nw, ny, int_stride, int_pad, ndim,
+                                cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = bwd_data_algo_cache_map_;
     // TODO: thread-safe
     auto it = algo_cache_map.find(key);
     if (it != algo_cache_map.end()) {
-        auto pair = it->second;
-        perf_result->algo = pair.first;
-        perf_result->memory = pair.second;
+        auto entry = it->second;
+        // clear the fields the search would have filled but the cache drops
+        *perf_result = {};
+        perf_result->algo = entry.algo;
+        perf_result->memory = entry.memory;
+        perf_result->mathType = entry.math_type;
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -458,10 +489,13 @@ cumo_cuda_cudnn_FindConvolutionBackwardDataAlgorithm(
                 max_workspace_size);
     cumo_cuda_runtime_free(workspace);
     if (status != CUDNN_STATUS_SUCCESS) return status;
-    assert(returned_algo_count == 1);
+    // A search that answers success with nothing to report would leave
+    // perf_result untouched, and its algo and math type are used below.
+    if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
+    if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory};
+    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
     return status;
 }
 
@@ -488,27 +522,19 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
     CumoGetNArray(gy, ngy);
     CumoGetNArray(gw, ngw);
 
-    auto key = AlgoCacheKey{};
-    key.ndim = ndim;
-    for (size_t idim = 0; idim < ndim + 2; ++idim) {
-        key.x_shape[idim] = nx->shape[idim];
-        key.w_shape[idim] = ngw->shape[idim];
-        key.y_shape[idim] = ngy->shape[idim];
-    }
-    for (size_t idim = 0; idim < ndim; ++idim) {
-        key.pad[idim]= int_pad[idim];
-        key.stride[idim]= int_stride[idim];
-    }
-    key.dtype = cudnn_dtype;
-    key.max_workspace_size = max_workspace_size;
+    auto key = MakeAlgoCacheKey(nx, ngw, ngy, int_stride, int_pad, ndim,
+                                cudnn_dtype, max_workspace_size);
 
     auto& algo_cache_map = bwd_filter_algo_cache_map_;
     // TODO: thread-safe
     auto it = algo_cache_map.find(key);
     if (it != algo_cache_map.end()) {
-        auto pair = it->second;
-        perf_result->algo = pair.first;
-        perf_result->memory = pair.second;
+        auto entry = it->second;
+        // clear the fields the search would have filled but the cache drops
+        *perf_result = {};
+        perf_result->algo = entry.algo;
+        perf_result->memory = entry.memory;
+        perf_result->mathType = entry.math_type;
         return CUDNN_STATUS_SUCCESS;
     }
 
@@ -534,10 +560,13 @@ cumo_cuda_cudnn_FindConvolutionBackwardFilterAlgorithm(
                 max_workspace_size);
     cumo_cuda_runtime_free(workspace);
     if (status != CUDNN_STATUS_SUCCESS) return status;
-    assert(returned_algo_count == 1);
+    // A search that answers success with nothing to report would leave
+    // perf_result untouched, and its algo and math type are used below.
+    if (returned_algo_count < 1) return CUDNN_STATUS_NOT_SUPPORTED;
+    if (perf_result->status != CUDNN_STATUS_SUCCESS) return perf_result->status;
 
     // TODO: thread-safe
-    algo_cache_map[key] = {perf_result->algo, perf_result->memory};
+    algo_cache_map[key] = {perf_result->algo, perf_result->memory, perf_result->mathType};
     return status;
 }
 

data/ext/cumo/cuda/runtime.c

@@ -5,6 +5,7 @@
 
 VALUE cumo_cuda_eRuntimeError;
 VALUE cumo_cuda_mRuntime;
+uint64_t cumo_cuda_sync_epoch = 0;
 #define eRuntimeError cumo_cuda_eRuntimeError
 #define mRuntime cumo_cuda_mRuntime
 
@@ -42,7 +43,7 @@ cumo_cuda_runtime_error_flag_new(void)
 bool
 cumo_cuda_runtime_error_flag_get(int *flag)
 {
-    check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
     return (*flag != 0);
 }
 
@@ -189,9 +190,7 @@ rb_cudaSetDevice(VALUE self, VALUE device)
 static VALUE
 rb_cudaDeviceSynchronize(VALUE self)
 {
-    cudaError_t status;
-    status = cudaDeviceSynchronize();
-    check_status(status);
+    cumo_cuda_runtime_device_synchronize();
     return Qnil;
 }
 

data/ext/cumo/cumo.c

@@ -20,11 +20,27 @@ void Init_cumo_uint8();
 void Init_cumo_uint16();
 void Init_cumo_uint32();
 void Init_cumo_uint64();
+void Init_cumo_hfloat();
 void Init_cumo_sfloat();
 void Init_cumo_scomplex();
 void Init_cumo_dfloat();
 void Init_cumo_dcomplex();
 void Init_cumo_robject();
+void Init_cumo_dcomplex_upcast();
+void Init_cumo_dfloat_upcast();
+void Init_cumo_scomplex_upcast();
+void Init_cumo_sfloat_upcast();
+void Init_cumo_hfloat_upcast();
+void Init_cumo_int64_upcast();
+void Init_cumo_uint64_upcast();
+void Init_cumo_int32_upcast();
+void Init_cumo_uint32_upcast();
+void Init_cumo_int16_upcast();
+void Init_cumo_uint16_upcast();
+void Init_cumo_int8_upcast();
+void Init_cumo_uint8_upcast();
+void Init_cumo_bit_upcast();
+void Init_cumo_robject_upcast();
 void Init_cumo_na_math();
 void Init_cumo_na_rand();
 void Init_cumo_na_array();
@@ -150,6 +166,7 @@ Init_cumo()
     Init_cumo_dfloat();
     Init_cumo_scomplex();
     Init_cumo_sfloat();
+    Init_cumo_hfloat();
 
     Init_cumo_int64();
     Init_cumo_uint64();
@@ -163,6 +180,23 @@ Init_cumo()
     Init_cumo_bit();
     Init_cumo_robject();
 
+    // Every class exists now, so the UPCAST tables can name all of them.
+    Init_cumo_dcomplex_upcast();
+    Init_cumo_dfloat_upcast();
+    Init_cumo_scomplex_upcast();
+    Init_cumo_sfloat_upcast();
+    Init_cumo_hfloat_upcast();
+    Init_cumo_int64_upcast();
+    Init_cumo_uint64_upcast();
+    Init_cumo_int32_upcast();
+    Init_cumo_uint32_upcast();
+    Init_cumo_int16_upcast();
+    Init_cumo_uint16_upcast();
+    Init_cumo_int8_upcast();
+    Init_cumo_uint8_upcast();
+    Init_cumo_bit_upcast();
+    Init_cumo_robject_upcast();
+
     Init_cumo_na_math();
 
     Init_cumo_na_rand();

data/ext/cumo/extconf.rb

@@ -64,6 +64,7 @@ narray/types/uint8
 narray/types/uint16
 narray/types/uint32
 narray/types/uint64
+narray/types/hfloat
 narray/types/sfloat
 narray/types/dfloat
 narray/types/scomplex
@@ -78,6 +79,7 @@ narray/types/uint8_kernel
 narray/types/uint16_kernel
 narray/types/uint32_kernel
 narray/types/uint64_kernel
+narray/types/hfloat_kernel
 narray/types/sfloat_kernel
 narray/types/dfloat_kernel
 narray/types/scomplex_kernel

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

@@ -21,7 +21,14 @@ extern VALUE cumo_cuda_eCUDNNError;
 
 extern VALUE cumo_na_eShapeError;
 
-#define CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE 8 * 1024 * 1024
+#define CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE (8 * 1024 * 1024)
+
+// How much scratch cuDNN may use to pick a convolution algorithm. The fastest
+// half algorithms are the ones that need the most, so the default keeps every
+// dtype on the algorithms it has always used and CUMO_CUDNN_MAX_WORKSPACE_SIZE
+// raises the ceiling for whoever wants them.
+size_t
+cumo_cuda_cudnn_max_workspace_size();
 
 // An output array given by the caller is written through a descriptor built
 // from another operand, or as if it were contiguous, so cuDNN never learns how
@@ -76,6 +83,17 @@ cumo_cuda_cudnn_check_input(VALUE in, VALUE type, size_t ndim, size_t *shape)
     }
 }
 
+// cuDNN derives the batch norm parameter descriptor from x, and widens it to
+// float for a half x, so the parameters do not always take x's own class.
+static inline void
+cumo_cuda_cudnn_check_param_type(VALUE param, VALUE type, const char* name)
+{
+    if (rb_obj_class(param) != type) {
+        rb_raise(rb_eTypeError, "%s must be %s, not %s",
+                 name, rb_class2name(type), rb_obj_classname(param));
+    }
+}
+
 void
 cumo_cuda_cudnn_check_status(cudnnStatus_t status);
 
@@ -172,7 +190,8 @@ cumo_cuda_cudnn_CreateConvolutionDescriptor(
         size_t ndim,
         int* int_stride,
         int* int_pad,
-        cudnnDataType_t cudnn_dtype);
+        cudnnDataType_t compute_dtype,
+        cudnnMathType_t math_type);
 
 cudnnStatus_t
 cumo_cuda_cudnn_CreatePoolingDescriptor(

data/ext/cumo/include/cumo/cuda/cumo_thrust.hpp

@@ -8,6 +8,7 @@
 #include <thrust/inner_product.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/transform_output_iterator.h>
 #include <thrust/iterator/permutation_iterator.h>
 #include <thrust/reduce.h>
 #include <thrust/scan.h>

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

@@ -13,6 +13,13 @@ extern "C" {
 
 extern VALUE cumo_cuda_eRuntimeError;
 
+// How many times the whole device has been seen to settle, which is what a host
+// read of managed memory needs. One settling covers every kernel and copy issued
+// before it, so code that recorded the count when it queued work can skip a wait
+// the moment the count has moved on. Only cumo_cuda_runtime_device_synchronize
+// advances it, and a count that is behind costs a wait rather than correctness.
+extern uint64_t cumo_cuda_sync_epoch;
+
 static inline void
 cumo_cuda_runtime_check_status(cudaError_t status)
 {
@@ -21,6 +28,27 @@ cumo_cuda_runtime_check_status(cudaError_t status)
     }
 }
 
+static inline void
+cumo_cuda_runtime_device_synchronize(void)
+{
+    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_sync_epoch++;
+}
+
+// Asking costs less than half of waiting, and there is nothing to wait for
+// whenever the block stayed off the device. Neither answer advances the settle
+// count: stream 0 going quiet is not the whole device settling, and a host read
+// of managed memory needs the latter where concurrentManagedAccess is 0.
+static inline int
+cumo_cuda_runtime_sync_if_busy(void)
+{
+    if (cudaStreamQuery(0) == cudaSuccess) {
+        return 0;
+    }
+    cumo_cuda_runtime_check_status(cudaStreamSynchronize(0));
+    return 1;
+}
+
 static inline int
 cumo_cuda_runtime_get_device_count()
 {

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

@@ -33,7 +33,9 @@ char* cumo_na_get_offset_pointer_for_read_write(VALUE);
 
 void cumo_na_copy_flags(VALUE src, VALUE dst);
 
+int cumo_na_last_dim_with_elements(const cumo_narray_t *na, int start_dim);
 VALUE cumo_na_check_ladder(VALUE self, int start_dim);
+void cumo_na_set_newaxis_strides(cumo_narray_view_t *na2, const int *newaxis, int n_newaxis, int ndim, ssize_t elmsz);
 VALUE cumo_na_check_contiguous(VALUE self);
 VALUE cumo_na_as_contiguous_array(VALUE a);
 
@@ -78,8 +80,9 @@ void cumo_na_parse_enumerator_step(VALUE enum_obj, VALUE *pstep);
 
 // used in aref, aset
 int cumo_na_get_result_dimension(VALUE self, int argc, VALUE *argv, ssize_t stride, size_t *pos_idx);
+void cumo_na_index_check_stray(cumo_narray_view_t *nv);
 void cumo_na_index_mark_filled(cumo_narray_view_t *nv);
-void cumo_na_index_wait_fill(const cumo_narray_view_t *nv);
+void cumo_na_index_wait_fill(cumo_narray_view_t *nv);
 VALUE cumo_na_aref_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, size_t pos);
 VALUE cumo_na_at_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, size_t pos);
 

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

@@ -189,6 +189,7 @@ extern const rb_data_type_t cumo_na_data_type;
 extern VALUE cumo_cBit;
 extern VALUE cumo_cDFloat;
 extern VALUE cumo_cSFloat;
+extern VALUE cumo_cHFloat;
 extern VALUE cumo_cDComplex;
 extern VALUE cumo_cSComplex;
 extern VALUE cumo_cInt64;
@@ -220,7 +221,6 @@ typedef struct {
     unsigned char ndim;     // # of dimensions
     unsigned char type;
     unsigned char flag[2];  // flags
-    unsigned short elmsz;    // element size
     size_t   size;          // # of total elements
     size_t  *shape;         // # of elements for each dimension
     VALUE    reduce;
@@ -230,7 +230,6 @@ typedef struct {
 typedef struct {
     cumo_narray_t base;
     char    *ptr;
-    bool     owned;
 } cumo_narray_data_t;
 
 
@@ -255,6 +254,12 @@ typedef struct {
     size_t   reach_end;
     uint64_t index_sync_epoch; // synchronizes counted when the index fills were
                                // issued; UINT64_MAX when that is not known
+    // A view made from another one can point at the index arrays it already
+    // has rather than copy them. index_owner keeps that view alive for as long
+    // as this one borrows from it, and index_owned says which dimensions are
+    // this view's own to free.
+    VALUE    index_owner;
+    uint64_t index_owned;
 } cumo_narray_view_t;
 
 
@@ -372,7 +377,6 @@ _cumo_na_get_narray_t(VALUE obj, unsigned char cumo_na_type)
 #define CUMO_NA_DATA(na)             ((cumo_narray_data_t*)(na))
 #define CUMO_NA_VIEW(na)             ((cumo_narray_view_t*)(na))
 #define CUMO_NA_DATA_PTR(na)         (CUMO_NA_DATA(na)->ptr)
-#define CUMO_NA_DATA_OWNED(na)       (CUMO_NA_DATA(na)->owned)
 #define CUMO_NA_VIEW_DATA(na)        (CUMO_NA_VIEW(na)->data)
 #define CUMO_NA_VIEW_OFFSET(na)      (CUMO_NA_VIEW(na)->offset)
 #define CUMO_NA_VIEW_STRIDX(na)      (CUMO_NA_VIEW(na)->stridx)
@@ -507,9 +511,61 @@ typedef unsigned int CUMO_BIT_DIGIT;
 #define CUMO_BALL   (~(CUMO_BIT_DIGIT)0)
 #define CUMO_SLB(n) (((n)==CUMO_NB)?~(CUMO_BIT_DIGIT)0:(~(~(CUMO_BIT_DIGIT)0<<(n))))
 
+// Bits are packed into digits, so a buffer holding n of them is as long as the
+// digits they fill, rounded up. n has to be positive: an empty array has no
+// buffer, and the subtraction below would wrap.
+#define CUMO_BIT_BYTES(n) \
+    ((((n) - 1) / 8 / sizeof(CUMO_BIT_DIGIT) + 1) * sizeof(CUMO_BIT_DIGIT))
+
 #include "cumo/ndloop.h"
 #include "cumo/intern.h"
 
+// An index array a view holds is freed only if it was set through here, so a
+// view that puts one in stridx any other way leaks it. The bit names the slot,
+// so code that permutes stridx has to permute index_owned with it.
+static inline void
+cumo_na_index_own(cumo_narray_view_t *nv, int i, size_t *idx)
+{
+    CUMO_SDX_SET_INDEX(nv->stridx[i], idx);
+    nv->index_owned |= (uint64_t)1 << i;
+}
+
+// One dimension of nv now points at an index array that from reaches, so from
+// has to outlive nv. A lender that owns nothing itself is skipped for its own
+// lender, which keeps the common chain one link long. A lender that owns even
+// one dimension is named directly, since only it keeps that one alive.
+//
+// A view borrows every index array or none. Borrowing some while building
+// others holds the lender for arrays this view has already replaced, which
+// costs more memory than the copy it saves. One lender is recorded, so a view
+// may not borrow from two.
+static inline void
+cumo_na_index_borrow(cumo_narray_view_t *nv, VALUE from)
+{
+    cumo_narray_view_t *nv1;
+
+    if (RTEST(nv->index_owner)) {
+        return;
+    }
+    CumoGetNArrayView(from, nv1);
+    nv->index_owner = (nv1->index_owned == 0 && RTEST(nv1->index_owner))
+        ? nv1->index_owner : from;
+}
+
+// Called once a view built from another one has all its dimensions. A view that
+// owns index arrays built them here, so a fill is in flight and the first read
+// waits. One that borrows them all is as settled as the view it borrowed from.
+static inline void
+cumo_na_index_mark_derived(cumo_narray_view_t *nv, cumo_narray_view_t *nv1)
+{
+    if (nv->index_owned != 0) {
+        cumo_na_index_mark_filled(nv);
+    } else {
+        cumo_na_index_check_stray(nv);
+        nv->index_sync_epoch = nv1->index_sync_epoch;
+    }
+}
+
 // for Ractor support code
 #ifndef HAVE_RB_EXT_RACTOR_SAFE
 #   undef RUBY_TYPED_FROZEN_SHAREABLE

data/ext/cumo/include/cumo/template.h

@@ -36,17 +36,24 @@
 
 // pos is the first position in iteration order, so a step below zero walks down
 // from there and rebasing the pointer onto pos would underflow the position of
-// every element past the first word. A step of zero comes with an index, whose
-// entries only ever move forward.
+// every element past the first word. An index holds positions relative to pos
+// and one of them is below zero as soon as the view it came from walks an axis
+// backwards, so pos + idx is right in unsigned arithmetic only while pos is
+// whole. The test reads the iterator rather than the caller's copy of step,
+// which is how a template declaring it as size_t used to lose the sign.
+#define CUMO_REBASE_BIT_PTR( lp, i, ad, ps )                 \
+    if (((lp)->args[i].iter[0]).step >= 0 &&                 \
+        ((lp)->args[i].iter[0]).idx == NULL) {               \
+        ad += ps/CUMO_NB;                                    \
+        ps %= CUMO_NB;                                       \
+    }
+
 #define CUMO_INIT_PTR_BIT( lp, i, ad, ps, st )               \
     {                                                   \
         ps = ((lp)->args[i].iter[0]).pos;                       \
         st = ((lp)->args[i].iter[0]).step;                      \
         ad = (CUMO_BIT_DIGIT*)(((lp)->args[i]).ptr);            \
-        if (st >= 0) {                                       \
-            ad += ps/CUMO_NB;                                \
-            ps %= CUMO_NB;                                   \
-        }                                                    \
+        CUMO_REBASE_BIT_PTR(lp, i, ad, ps);                     \
     }
 
 #define CUMO_INIT_PTR_BIT_IDX( lp, i, ad, ps, st, id )       \
@@ -55,10 +62,7 @@
         st = ((lp)->args[i].iter[0]).step;                      \
         id = ((lp)->args[i].iter[0]).idx;                       \
         ad = (CUMO_BIT_DIGIT*)(((lp)->args[i]).ptr);            \
-        if (st >= 0) {                                       \
-            ad += ps/CUMO_NB;                                \
-            ps %= CUMO_NB;                                   \
-        }                                                    \
+        CUMO_REBASE_BIT_PTR(lp, i, ad, ps);                     \
     }
 
 // A where compaction takes three launches, so below this many elements walking
@@ -179,4 +183,25 @@ cumo_na_store_rary_fetch(VALUE ary, size_t i, VALUE *x)
     return true;
 }
 
+// A 0-dimensional walk still yields one index, so the depth floors at 0 rather
+// than going negative and writing c[-1].
+static inline void
+cumo_na_with_index_dims(int ndim, int *nd, int *md)
+{
+    *nd = (ndim > 0) ? ndim - 1 : ndim;
+    *md = *nd + 2;
+}
+
+static inline VALUE
+cumo_na_yield_with_index(VALUE x, size_t *c, VALUE *a, int nd, int md)
+{
+    int j;
+
+    a[0] = x;
+    for (j=0; j<=nd; j++) {
+        a[j+1] = SIZET2NUM(c[j]);
+    }
+    return rb_yield(rb_ary_new4(md,a));
+}
+
 #endif /* ifndef CUMO_TEMPLATE_H */