cumo 0.1.0 → 0.1.1

data/ext/cumo/include/cumo/narray_kernel.h

@@ -49,13 +49,23 @@ extern "C" {
 # endif
 #endif
 
+#ifndef SZF
 #define SZF PRI_SIZE_PREFIX // defined in ruby.h
+#endif
 
 #if   SIZEOF_LONG==8
-# define NUM2INT64(x) NUM2LONG(x)
-# define INT642NUM(x) LONG2NUM(x)
-# define NUM2UINT64(x) NUM2ULONG(x)
-# define UINT642NUM(x) ULONG2NUM(x)
+# ifndef NUM2INT64
+#  define NUM2INT64(x) NUM2LONG(x)
+# endif
+# ifndef INT642NUM
+#  define INT642NUM(x) LONG2NUM(x)
+# endif
+# ifndef NUM2UINT64
+#  define NUM2UINT64(x) NUM2ULONG(x)
+# endif
+# ifndef UINT642NUM
+#  define UINT642NUM(x) ULONG2NUM(x)
+# endif
 # ifndef PRId64
 #  define PRId64 "ld"
 # endif
@@ -63,10 +73,18 @@ extern "C" {
 #  define PRIu64 "lu"
 # endif
 #elif SIZEOF_LONG_LONG==8
-# define NUM2INT64(x) NUM2LL(x)
-# define INT642NUM(x) LL2NUM(x)
-# define NUM2UINT64(x) NUM2ULL(x)
-# define UINT642NUM(x) ULL2NUM(x)
+# ifndef NUM2INT64
+#  define NUM2INT64(x) NUM2LL(x)
+# endif
+# ifndef INT642NUM
+#  define INT642NUM(x) LL2NUM(x)
+# endif
+# ifndef NUM2UINT64
+#  define NUM2UINT64(x) NUM2ULL(x)
+# endif
+# ifndef UINT642NUM
+#  define UINT642NUM(x) ULL2NUM(x)
+# endif
 # ifndef PRId64
 #  define PRId64 "lld"
 # endif
@@ -76,10 +94,18 @@ extern "C" {
 #endif
 
 #if   SIZEOF_LONG==4
-# define NUM2INT32(x) NUM2LONG(x)
-# define INT322NUM(x) LONG2NUM(x)
-# define NUM2UINT32(x) NUM2ULONG(x)
-# define UINT322NUM(x) ULONG2NUM(x)
+# ifndef NUM2INT32
+#  define NUM2INT32(x) NUM2LONG(x)
+# endif
+# ifndef INT322NUM
+#  define INT322NUM(x) LONG2NUM(x)
+# endif
+# ifndef NUM2UINT32
+#  define NUM2UINT32(x) NUM2ULONG(x)
+# endif
+# ifndef UINT322NUM
+#  define UINT322NUM(x) ULONG2NUM(x)
+# endif
 # ifndef PRId32
 #  define PRId32 "ld"
 # endif
@@ -87,10 +113,18 @@ extern "C" {
 #  define PRIu32 "lu"
 # endif
 #elif SIZEOF_INT==4
-# define NUM2INT32(x) NUM2INT(x)
-# define INT322NUM(x) INT2NUM(x)
-# define NUM2UINT32(x) NUM2UINT(x)
-# define UINT322NUM(x) UINT2NUM(x)
+# ifndef NUM2INT32
+#  define NUM2INT32(x) NUM2INT(x)
+# endif
+# ifndef INT322NUM
+#  define INT322NUM(x) INT2NUM(x)
+# endif
+# ifndef NUM2UINT32
+#  define NUM2UINT32(x) NUM2UINT(x)
+# endif
+# ifndef UINT322NUM
+#  define UINT322NUM(x) UINT2NUM(x)
+# endif
 # ifndef PRId32
 #  define PRId32 "d"
 # endif
@@ -109,32 +143,27 @@ extern "C" {
 # define TRUE    1
 #endif
 
-typedef struct { float dat[2]; }  scomplex;
-typedef struct { double dat[2]; } dcomplex;
-typedef int fortran_integer;
-
-#define REAL(x) ((x).dat[0])
-#define IMAG(x) ((x).dat[1])
+typedef struct { float dat[2]; }  cumo_scomplex;
+typedef struct { double dat[2]; } cumo_dcomplex;
 
-extern int na_debug_flag;
+#define CUMO_REAL(x) ((x).dat[0])
+#define CUMO_IMAG(x) ((x).dat[1])
 
-#define NARRAY_DATA_T     0x1
-#define NARRAY_VIEW_T     0x2
-#define NARRAY_FILEMAP_T  0x3
+extern int cumo_na_debug_flag;
 
-//#define NA_MAX_DIMENSION (int)(sizeof(VALUE)*8-2)
-#define NA_MAX_DIMENSION 12
-#define NA_MAX_ELMSZ     65535
+#define CUMO_NARRAY_DATA_T     0x1
+#define CUMO_NARRAY_VIEW_T     0x2
+#define CUMO_NARRAY_FILEMAP_T  0x3
 
-typedef unsigned int BIT_DIGIT;
-#define BYTE_BIT_DIGIT sizeof(BIT_DIGIT)
-#define NB     (sizeof(BIT_DIGIT)*8)
-#define BALL   (~(BIT_DIGIT)0)
-#define SLB(n) (((n)==NB)?~(BIT_DIGIT)0:(~(~(BIT_DIGIT)0<<(n))))
+//#define CUMO_NA_MAX_DIMENSION (int)(sizeof(VALUE)*8-2)
+#define CUMO_NA_MAX_DIMENSION 12
+#define CUMO_NA_MAX_ELMSZ     65535
 
-#define ELEMENT_BIT_SIZE  "ELEMENT_BIT_SIZE"
-#define ELEMENT_BYTE_SIZE "ELEMENT_BYTE_SIZE"
-#define CONTIGUOUS_STRIDE "CONTIGUOUS_STRIDE"
+typedef unsigned int CUMO_BIT_DIGIT;
+#define CUMO_BYTE_BIT_DIGIT sizeof(CUMO_BIT_DIGIT)
+#define CUMO_NB     (sizeof(CUMO_BIT_DIGIT)*8)
+#define CUMO_BALL   (~(CUMO_BIT_DIGIT)0)
+#define CUMO_SLB(n) (((n)==CUMO_NB)?~(CUMO_BIT_DIGIT)0:(~(~(CUMO_BIT_DIGIT)0<<(n))))
 
 #include "cumo/indexer.h"
 #include "cumo/intern_kernel.h"

data/ext/cumo/include/cumo/ndloop.h

@@ -1,29 +1,29 @@
 #ifndef CUMO_NDLOOP_H
 #define CUMO_NDLOOP_H
 
-typedef struct NA_LOOP_ITER {
+typedef struct {
     ssize_t    pos; // - required for each dimension.
     ssize_t    step;
     size_t    *idx;
-} na_loop_iter_t;
+} cumo_na_loop_iter_t;
 
-typedef struct NA_LOOP_ARGS {
+typedef struct {
     VALUE    value;
-    ssize_t  elmsz;
+    ssize_t  elmsz;      // element size in bytes, e.g., 4 for int, 8 for double
     char    *ptr;
     //char    *buf_ptr;  //
     int      ndim;       // required for each argument.
     // ssize_t pos; - not required here.
     size_t  *shape;
-    na_loop_iter_t *iter;  // moved from na_loop_t
-} na_loop_args_t;
+    cumo_na_loop_iter_t *iter;  // moved from cumo_na_loop_t
+} cumo_na_loop_args_t;
 
 // pass this structure to user iterator
-typedef struct NA_LOOP {
+typedef struct {
     int  narg;
     int  ndim;             // n of user dimention used at user function.
     size_t *n;             // n of elements for each dim (=shape)
-    na_loop_args_t *args;  // for each arg
+    cumo_na_loop_args_t *args;  // for each arg
     VALUE  option;
     void  *opt_ptr;
     VALUE  err_type;
@@ -31,65 +31,65 @@ typedef struct NA_LOOP {
                            // 3 for sum(), 1 for sum(axis: 1), 2 for sum(axis: [1,2])
     VALUE  reduce;         // dimension indicies to reduce in reduction kernel (in bits), e.g., for an array of shape:
                            // [2,3,4], 111b for sum(), 010b for sum(axis: 1), 110b for sum(axis: [1,2])
-} na_loop_t;
+} cumo_na_loop_t;
 
 
 // ------------------ ndfunc -------------------------------------------
 
-#define NDF_HAS_LOOP            (1<<0) // x[i]
-#define NDF_STRIDE_LOOP         (1<<1) // *(x+stride*i)
-#define NDF_INDEX_LOOP          (1<<2) // *(x+idx[i])
-#define NDF_KEEP_DIM            (1<<3)
-#define NDF_INPLACE             (1<<4)
-#define NDF_ACCEPT_BYTESWAP     (1<<5)
+#define CUMO_NDF_HAS_LOOP            (1<<0) // x[i]
+#define CUMO_NDF_STRIDE_LOOP         (1<<1) // *(x+stride*i)
+#define CUMO_NDF_INDEX_LOOP          (1<<2) // *(x+idx[i])
+#define CUMO_NDF_KEEP_DIM            (1<<3)
+#define CUMO_NDF_INPLACE             (1<<4)
+#define CUMO_NDF_ACCEPT_BYTESWAP     (1<<5)
 
-#define NDF_FLAT_REDUCE         (1<<6)
-#define NDF_EXTRACT             (1<<7)
-#define NDF_CUM                 (1<<8)
+#define CUMO_NDF_FLAT_REDUCE         (1<<6)
+#define CUMO_NDF_EXTRACT             (1<<7)
+#define CUMO_NDF_CUM                 (1<<8)
 
-#define NDF_INDEXER_LOOP        (1<<9) // Cumo custom. Use cumo own indexer.
+#define CUMO_NDF_INDEXER_LOOP        (1<<9) // Cumo custom. Use cumo own indexer.
 
-#define FULL_LOOP       (NDF_HAS_LOOP|NDF_STRIDE_LOOP|NDF_INDEX_LOOP|NDF_INPLACE)
-#define FULL_LOOP_NIP   (NDF_HAS_LOOP|NDF_STRIDE_LOOP|NDF_INDEX_LOOP)
-#define STRIDE_LOOP     (NDF_HAS_LOOP|NDF_STRIDE_LOOP|NDF_INPLACE)
-#define STRIDE_LOOP_NIP (NDF_HAS_LOOP|NDF_STRIDE_LOOP)
-#define NO_LOOP         0
+#define CUMO_FULL_LOOP       (CUMO_NDF_HAS_LOOP|CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INDEX_LOOP|CUMO_NDF_INPLACE)
+#define CUMO_FULL_LOOP_NIP   (CUMO_NDF_HAS_LOOP|CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INDEX_LOOP)
+#define CUMO_STRIDE_LOOP     (CUMO_NDF_HAS_LOOP|CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INPLACE)
+#define CUMO_STRIDE_LOOP_NIP (CUMO_NDF_HAS_LOOP|CUMO_NDF_STRIDE_LOOP)
+#define CUMO_NO_LOOP         0
 
-#define OVERWRITE Qtrue // used for CASTABLE(t)
+#define CUMO_OVERWRITE Qtrue // used for CASTABLE(t)
 
-#define NDF_TEST(nf,fl)  ((nf)->flag & (fl))
-#define NDF_SET(nf,fl)  {(nf)->flag |= (fl);}
+#define CUMO_NDF_TEST(nf,fl)  ((nf)->flag & (fl))
+#define CUMO_NDF_SET(nf,fl)  {(nf)->flag |= (fl);}
 
-#define NDF_ARG_READ_ONLY   1
-#define NDF_ARG_WRITE_ONLY  2
-#define NDF_ARG_READ_WRITE  3
+#define CUMO_NDF_ARG_READ_ONLY   1
+#define CUMO_NDF_ARG_WRITE_ONLY  2
+#define CUMO_NDF_ARG_READ_WRITE  3
 
 // type of user function
-typedef void (*na_iter_func_t) _((na_loop_t *const));
-typedef VALUE (*na_text_func_t) _((char *ptr, size_t pos, VALUE opt));
-//typedef void (*) void (*loop_func)(ndfunc_t*, na_md_loop_t*))
+typedef void (*cumo_na_iter_func_t) _((cumo_na_loop_t *const));
+typedef VALUE (*cumo_na_text_func_t) _((char *ptr, size_t pos, VALUE opt));
+//typedef void (*) void (*loop_func)(cumo_ndfunc_t*, cumo_na_md_loop_t*))
 
 
-typedef struct NDF_ARG_IN {
+typedef struct {
     VALUE   type;    // argument types
     int     dim;     // # of dimension of argument handled by user function
                      // if dim==-1, reduce dimension
-} ndfunc_arg_in_t;
+} cumo_ndfunc_arg_in_t;
 
-typedef struct NDF_ARG_OUT {
+typedef struct {
     VALUE   type;    // argument types
     int     dim;     // # of dimension of argument handled by user function
     size_t *shape;
-} ndfunc_arg_out_t;
+} cumo_ndfunc_arg_out_t;
 
 // spec of user function
-typedef struct NDFUNCTION {
-    na_iter_func_t func;    // user function
+typedef struct {
+    cumo_na_iter_func_t func;    // user function
     unsigned int flag;      // what kind of loop user function supports
     int nin;                // # of arguments
     int nout;               // # of results
-    ndfunc_arg_in_t *ain;   // spec of input arguments
-    ndfunc_arg_out_t *aout; // spec of output result
-} ndfunc_t;
+    cumo_ndfunc_arg_in_t *ain;   // spec of input arguments
+    cumo_ndfunc_arg_out_t *aout; // spec of output result
+} cumo_ndfunc_t;
 
 #endif /* CUMO_NDLOOP_H */

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

@@ -7,6 +7,11 @@
 
 #include "cumo/indexer.h"
 
+namespace cumo_detail {
+
+static constexpr int64_t max_block_size = 512;
+static constexpr int64_t max_grid_size = 0x7fffffff;
+
 static inline int64_t round_up_to_power_of_2(int64_t x) {
     --x;
     x |= x >> 1;
@@ -18,109 +23,88 @@ static inline int64_t round_up_to_power_of_2(int64_t x) {
     return x + 1;
 }
 
-#define _REDUCE(offset) \
-    if (tid < offset) { \
-        impl.Reduce(sdata[(tid + offset)], sdata[tid]); \
-    }
-
-// reference: cupy reduction kernel
+// Reference: cupy reduction kernel
+// Note that reduction and out axis are inverse with cupy. Former axes are out axes, latters are reduce axes.
 
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-__global__ static void reduction_kernel(na_reduction_arg_t arg, ReductionImpl impl) {
-    na_iarray_t& in_iarray = arg.in;
-    na_iarray_t& out_iarray = arg.out;
-    na_indexer_t& in_indexer = arg.in_indexer;
-    na_indexer_t& out_indexer = arg.out_indexer;
-    na_indexer_t& reduce_indexer = arg.reduce_indexer;
+__global__ static void reduction_kernel(cumo_na_reduction_arg_t arg, int out_block_size, int reduce_block_size, ReductionImpl impl) {
+    cumo_na_iarray_t& in_iarray = arg.in;
+    cumo_na_iarray_t& out_iarray = arg.out;
+    cumo_na_indexer_t& in_indexer = arg.in_indexer;
+    cumo_na_indexer_t& out_indexer = arg.out_indexer;
 
     using TypeReduce = decltype(impl.Identity());
 
     extern __shared__ __align__(8) char sdata_raw[];
-    TypeReduce* sdata = (TypeReduce*)sdata_raw;
+    TypeReduce* sdata = reinterpret_cast<TypeReduce*>(sdata_raw);
     unsigned int tid = threadIdx.x;
-    unsigned int block_size = blockDim.x;  // number of threads
 
-    for (uint64_t i_out = blockIdx.x; i_out < out_indexer.total_size; i_out += gridDim.x) {
+    int64_t reduce_indexer_total_size = in_indexer.total_size / out_indexer.total_size;
+    int64_t reduce_offset = tid / out_block_size; // # of cols == # of elems
+
+    int64_t out_offset = tid % out_block_size; // # of rows
+    int64_t out_base = blockIdx.x * out_block_size; // # of rows
+    int64_t out_stride = gridDim.x * out_block_size; // # of rows
+
+    for (int64_t i_out = out_base + out_offset; i_out < out_indexer.total_size; i_out += out_stride) {
         cumo_na_indexer_set_dim(&out_indexer, i_out);
         TypeReduce accum = impl.Identity();
 
-        for (int8_t i_out_dim = 0; i_out_dim < out_indexer.ndim; ++i_out_dim) {
-            in_indexer.index[i_out_dim] = out_indexer.index[i_out_dim];
-        }
-        for (auto i_reduce = tid; i_reduce < reduce_indexer.total_size; i_reduce += block_size) {
-            cumo_na_indexer_set_dim(&reduce_indexer, i_reduce);
-            for (int8_t i_reduce_dim = 0; i_reduce_dim < reduce_indexer.ndim; ++i_reduce_dim) {
-                in_indexer.index[out_indexer.ndim + i_reduce_dim] = reduce_indexer.index[i_reduce_dim];
-            }
+        int64_t i_in = i_out * reduce_indexer_total_size + reduce_offset;
+        for (int64_t i_reduce = reduce_offset; i_reduce < reduce_indexer_total_size; i_reduce += reduce_block_size, i_in += reduce_block_size) {
+            cumo_na_indexer_set_dim(&in_indexer, i_in);
             TypeIn* in_ptr = reinterpret_cast<TypeIn*>(cumo_na_iarray_at_dim(&in_iarray, &in_indexer));
-            uint64_t i_in = in_ptr - reinterpret_cast<TypeIn*>(in_iarray.ptr);
-            impl.Reduce(impl.MapIn(*in_ptr, i_in), accum);
+            // Note that spec of (min|max)_index of cumo is different with arg(min|max) of cupy.
+            // Cumo returns index of input elements, CuPy returns index of reduction axis.
+            impl.Reduce(impl.MapIn(*in_ptr, in_ptr - reinterpret_cast<TypeIn*>(in_iarray.ptr)), accum);
+            //printf("threadId.x:%d blockIdx.x:%d blockDim.x:%d gridDim.x:%d accum:%d i_in:%ld i_reduce:%ld i_out:%ld in:%p(%d)\n", threadIdx.x, blockIdx.x, blockDim.x, gridDim.x, accum, i_in, i_reduce, i_out, in_ptr, *in_ptr);
         }
 
-        if (block_size >= 2) {
+        if (out_block_size <= max_block_size / 2) {
             sdata[tid] = accum;
             __syncthreads();
-
-            if (block_size > 2) {
-                if (block_size > 4) {
-                    if (block_size > 8) {
-                        if (block_size > 16) {
-                            if (block_size > 32) {
-                                if (block_size > 64) {
-                                    if (block_size > 128) {
-                                        if (block_size > 256) {
-                                            _REDUCE(256);
-                                            __syncthreads();
-                                        }
-                                        _REDUCE(128);
-                                        __syncthreads();
-                                    }
-                                    _REDUCE(64);
-                                    __syncthreads();
-                                }
-                                _REDUCE(32);
-                                __syncthreads();
-                            }
-                            _REDUCE(16);
-                            __syncthreads();
-                        }
-                        _REDUCE(8);
-                        __syncthreads();
+            // NOTE: Compiler optimizes to unroll this loop
+            for (int stride = max_block_size / 2; stride > 0; stride >>= 1) {
+                if (out_block_size <= stride) {
+                    if (tid < stride) {
+                        impl.Reduce(sdata[tid + stride], sdata[tid]);
                     }
-                    _REDUCE(4);
                     __syncthreads();
                 }
-                _REDUCE(2);
-                __syncthreads();
             }
-            _REDUCE(1);
-            accum = sdata[0];
+            accum = sdata[tid];
+            __syncthreads();
         }
-        if (tid == 0) {
+        if (reduce_offset == 0 && i_out < out_indexer.total_size) {
             TypeOut* out_ptr = reinterpret_cast<TypeOut*>(cumo_na_iarray_at_dim(&out_iarray, &out_indexer));
             *out_ptr = impl.MapOut(accum);
-            //printf("threadId.x:%d blockIdx.x:%d blockDim.x:%d gridDim.x:%d block_size:%d accum:%d out:%p(%d)\n", threadIdx.x, blockIdx.x, blockDim.x, gridDim.x, block_size, accum, out_ptr, *out_ptr);
+            //printf("threadId.x:%d blockIdx.x:%d blockDim.x:%d gridDim.x:%d accum:%d i_out:%ld out:%p(%d)\n", threadIdx.x, blockIdx.x, blockDim.x, gridDim.x, accum, i_out, out_ptr, *out_ptr);
         }
     }
 }
 
-#undef _REDUCE
-
-static constexpr size_t max_block_size = 512;
+}  // cumo_detail
 
+// TODO(sonots): Optimize indexer by squashing (or reducing) dimensions
 template <typename TypeIn, typename TypeOut, typename ReductionImpl>
-void cumo_reduce(na_reduction_arg_t arg, ReductionImpl&& impl) {
-    na_indexer_t& out_indexer = arg.out_indexer;
-    na_indexer_t& reduce_indexer = arg.reduce_indexer;
+void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
+    cumo_na_indexer_t& in_indexer = arg.in_indexer;
+    cumo_na_indexer_t& out_indexer = arg.out_indexer;
 
-    using TypeReduce = decltype(impl.Identity());
+    if (out_indexer.total_size == 0) {
+        return;
+    }
+
+    int64_t reduce_total_size_pow2 = cumo_detail::round_up_to_power_of_2(std::max(size_t{1}, in_indexer.total_size / out_indexer.total_size));
+    int64_t reduce_block_size = std::min(cumo_detail::max_block_size, reduce_total_size_pow2);
+    int64_t out_block_size = cumo_detail::max_block_size / reduce_block_size;
+    int64_t out_block_num = (out_indexer.total_size + out_block_size - 1) / out_block_size;
 
-    size_t block_size = round_up_to_power_of_2(std::max(int64_t{1}, static_cast<int64_t>(reduce_indexer.total_size)));
-    block_size = std::min(max_block_size, block_size);
-    size_t grid_size = out_indexer.total_size;
-    size_t shared_mem_size = sizeof(TypeReduce) * block_size;
+    int64_t block_size = cumo_detail::max_block_size;
+    int64_t grid_size = std::min(cumo_detail::max_grid_size, out_block_num);
+    int64_t shared_mem_size = sizeof(decltype(impl.Identity())) * block_size;
 
-    reduction_kernel<TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, impl);
+    cumo_detail::reduction_kernel<TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, out_block_size, reduce_block_size, impl);
 }
 
 #endif // CUMO_REDUCE_KERNEL_H