cumo 0.6.0 → 0.7.0

data/ext/cumo/narray/gen/tmpl/float_accum_kernel.cu

@@ -1,4 +1,8 @@
 <%= load_erb('real_accum').result(binding) %>
+<% acc = is_half ? 'float' : 'dtype' %>
+<% to_acc = is_half ? 'cumo_half2float' : '' %>
+<% from_acc = is_half ? 'cumo_float2half' : '' %>
+
 
 #if defined(__cplusplus)
 #if 0
@@ -14,47 +18,47 @@
 // divide 0 by 0 and leave the accumulator NaN.
 struct cumo_<%=type_name%>_moments_impl {
     struct Moments {
-        rtype n;
-        dtype mean;
-        rtype m2;
+        <%=acc%> n;
+        <%=acc%> mean;
+        <%=acc%> m2;
     };
-    __device__ Moments Identity(int64_t /*index*/) { return {0, m_zero, 0}; }
-    __device__ Moments MapIn(dtype in, int64_t /*index*/) { return {1, in, 0}; }
+    __device__ Moments Identity(int64_t /*index*/) { return {0, 0, 0}; }
+    __device__ Moments MapIn(dtype in, int64_t /*index*/) { return {1, <%=to_acc%>(in), 0}; }
     __device__ void Reduce(Moments next, Moments& accum) {
         if (next.n == 0) { return; }
         if (accum.n == 0) { accum = next; return; }
-        rtype n = accum.n + next.n;
-        dtype delta = m_sub(next.mean, accum.mean);
-        accum.mean = m_add(accum.mean, delta * (next.n / n));
+        <%=acc%> n = accum.n + next.n;
+        <%=acc%> delta = next.mean - accum.mean;
+        accum.mean = accum.mean + delta * (next.n / n);
         accum.m2 += next.m2 + delta * delta * accum.n * next.n / n;
         accum.n = n;
     }
 };
 
 struct cumo_<%=type_name%>_var_impl : cumo_<%=type_name%>_moments_impl {
-    __device__ rtype MapOut(Moments accum) { return accum.m2 / (accum.n - 1); }
+    __device__ rtype MapOut(Moments accum) { return <%=from_acc%>(accum.m2 / (accum.n - 1)); }
 };
 
 struct cumo_<%=type_name%>_stddev_impl : cumo_<%=type_name%>_moments_impl {
-    __device__ rtype MapOut(Moments accum) { return m_sqrt(accum.m2 / (accum.n - 1)); }
+    __device__ rtype MapOut(Moments accum) { return <%=from_acc%>(sqrt(accum.m2 / (accum.n - 1))); }
 };
 
 struct cumo_<%=type_name%>_mean_impl {
     // The reduce axis is the same length for every output, so the divisor is a
     // constant the launcher already knows rather than a count the tree carries.
-    dtype n;
-    __device__ dtype Identity(int64_t /*index*/) { return m_zero; }
-    __device__ dtype MapIn(dtype in, int64_t /*index*/) { return in; }
-    __device__ void Reduce(dtype next, dtype& accum) { accum = m_add(next, accum); }
-    __device__ dtype MapOut(dtype accum) { return m_div(accum, n); }
+    <%=acc%> n;
+    __device__ <%=acc%> Identity(int64_t /*index*/) { return 0; }
+    __device__ <%=acc%> MapIn(dtype in, int64_t /*index*/) { return <%=to_acc%>(in); }
+    __device__ void Reduce(<%=acc%> next, <%=acc%>& accum) { accum = next + accum; }
+    __device__ dtype MapOut(<%=acc%> accum) { return <%=from_acc%>(accum / n); }
 };
 
 struct cumo_<%=type_name%>_rms_impl {
-    rtype n;
-    __device__ rtype Identity(int64_t /*index*/) { return 0; }
-    __device__ rtype MapIn(dtype in, int64_t /*index*/) { return m_square(m_abs(in)); }
-    __device__ void Reduce(rtype next, rtype& accum) { accum += next; }
-    __device__ rtype MapOut(rtype accum) { return m_sqrt(accum / n); }
+    <%=acc%> n;
+    __device__ <%=acc%> Identity(int64_t /*index*/) { return 0; }
+    __device__ <%=acc%> MapIn(dtype in, int64_t /*index*/) { <%=acc%> x = <%=to_acc%>(in); return x * x; }
+    __device__ void Reduce(<%=acc%> next, <%=acc%>& accum) { accum += next; }
+    __device__ rtype MapOut(<%=acc%> accum) { return <%=from_acc%>(sqrt(accum / n)); }
 };
 
 // The nan-aware forms. A NaN maps to an accumulator that contributes nothing:
@@ -63,8 +67,8 @@ struct cumo_<%=type_name%>_rms_impl {
 // of the reduce axis.
 struct cumo_<%=type_name%>_moments_nan_impl : cumo_<%=type_name%>_moments_impl {
     __device__ Moments MapIn(dtype in, int64_t /*index*/) {
-        if (!not_nan(in)) { return {0, m_zero, 0}; }
-        return {1, in, 0};
+        if (!not_nan(in)) { return {0, 0, 0}; }
+        return {1, <%=to_acc%>(in), 0};
     }
 };
 
@@ -72,45 +76,46 @@ struct cumo_<%=type_name%>_var_nan_impl : cumo_<%=type_name%>_moments_nan_impl {
     // Every element being NaN leaves no count. numo divides by count-1 in
     // unsigned arithmetic there, so it answers +0 rather than the -0 that
     // dividing by -1 would give.
-    __device__ rtype MapOut(Moments accum) { return accum.n == 0 ? 0 : accum.m2 / (accum.n - 1); }
+    __device__ rtype MapOut(Moments accum) { return accum.n == 0 ? <%=from_acc%>(0) : <%=from_acc%>(accum.m2 / (accum.n - 1)); }
 };
 
 struct cumo_<%=type_name%>_stddev_nan_impl : cumo_<%=type_name%>_moments_nan_impl {
-    __device__ rtype MapOut(Moments accum) { return accum.n == 0 ? 0 : m_sqrt(accum.m2 / (accum.n - 1)); }
+    __device__ rtype MapOut(Moments accum) { return accum.n == 0 ? <%=from_acc%>(0) : <%=from_acc%>(sqrt(accum.m2 / (accum.n - 1))); }
 };
 
 struct cumo_<%=type_name%>_mean_nan_impl {
     struct SumAndCount {
-        dtype sum;
-        rtype n;
+        <%=acc%> sum;
+        <%=acc%> n;
     };
-    __device__ SumAndCount Identity(int64_t /*index*/) { return {m_zero, 0}; }
+    __device__ SumAndCount Identity(int64_t /*index*/) { return {0, 0}; }
     __device__ SumAndCount MapIn(dtype in, int64_t /*index*/) {
-        if (!not_nan(in)) { return {m_zero, 0}; }
-        return {in, 1};
+        if (!not_nan(in)) { return {0, 0}; }
+        return {<%=to_acc%>(in), 1};
     }
     __device__ void Reduce(SumAndCount next, SumAndCount& accum) {
-        accum.sum = m_add(next.sum, accum.sum);
+        accum.sum = next.sum + accum.sum;
         accum.n += next.n;
     }
-    __device__ dtype MapOut(SumAndCount accum) { return m_div(accum.sum, m_from_real(accum.n)); }
+    __device__ dtype MapOut(SumAndCount accum) { return <%=from_acc%>(accum.sum / accum.n); }
 };
 
 struct cumo_<%=type_name%>_rms_nan_impl {
     struct SumAndCount {
-        rtype sum;
-        rtype n;
+        <%=acc%> sum;
+        <%=acc%> n;
     };
     __device__ SumAndCount Identity(int64_t /*index*/) { return {0, 0}; }
     __device__ SumAndCount MapIn(dtype in, int64_t /*index*/) {
         if (!not_nan(in)) { return {0, 0}; }
-        return {m_square(m_abs(in)), 1};
+        <%=acc%> x = <%=to_acc%>(in);
+        return {x * x, 1};
     }
     __device__ void Reduce(SumAndCount next, SumAndCount& accum) {
         accum.sum += next.sum;
         accum.n += next.n;
     }
-    __device__ rtype MapOut(SumAndCount accum) { return m_sqrt(accum.sum / accum.n); }
+    __device__ rtype MapOut(SumAndCount accum) { return <%=from_acc%>(sqrt(accum.sum / accum.n)); }
 };
 
 <% if is_double_precision %>
@@ -154,7 +159,7 @@ extern "C" {
 
 void cumo_<%=type_name%>_mean_kernel_launch(cumo_na_reduction_arg_t* arg)
 {
-    dtype n = (dtype)(arg->in_indexer.total_size / arg->out_indexer.total_size);
+    <%=acc%> n = (<%=acc%>)(arg->in_indexer.total_size / arg->out_indexer.total_size);
     cumo_reduce_split<dtype, dtype, cumo_<%=type_name%>_mean_impl>(*arg, cumo_<%=type_name%>_mean_impl{n});
 }
 
@@ -170,7 +175,7 @@ void cumo_<%=type_name%>_stddev_kernel_launch(cumo_na_reduction_arg_t* arg)
 
 void cumo_<%=type_name%>_rms_kernel_launch(cumo_na_reduction_arg_t* arg)
 {
-    rtype n = (rtype)(arg->in_indexer.total_size / arg->out_indexer.total_size);
+    <%=acc%> n = (<%=acc%>)(arg->in_indexer.total_size / arg->out_indexer.total_size);
     cumo_reduce_split<dtype, rtype, cumo_<%=type_name%>_rms_impl>(*arg, cumo_<%=type_name%>_rms_impl{n});
 }
 

data/ext/cumo/narray/gen/tmpl/format.c

@@ -65,6 +65,6 @@ static VALUE
     cumo_ndfunc_t ndf = { <%=c_iter%>, CUMO_FULL_LOOP_NIP, 2, 1, ain, aout };
 
     rb_scan_args(argc, argv, "01", &fmt);
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
     return cumo_na_ndloop(&ndf, 2, self, fmt);
 }

data/ext/cumo/narray/gen/tmpl/format_to_a.c

@@ -44,6 +44,6 @@ static VALUE
     cumo_ndfunc_t ndf = { <%=c_iter%>, CUMO_FULL_LOOP_NIP, 3, 1, ain, aout };
 
     rb_scan_args(argc, argv, "01", &fmt);
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
     return cumo_na_ndloop_cast_narray_to_rarray(&ndf, self, fmt);
 }

data/ext/cumo/narray/gen/tmpl/gemm.c

@@ -12,6 +12,8 @@
     end
   cutype =
     case type_name
+    when 'hfloat'
+      'cumo_half'
     when 'sfloat'
       'float'
     when 'dfloat'
@@ -21,6 +23,12 @@
     when 'dcomplex'
       'cuDoubleComplex'
     end
+  # cublasGemmStridedBatchedEx multiplies in half and accumulates in float, and
+  # takes its scalars in the accumulator's type.
+  scalar_t = is_half ? 'float' : 'dtype'
+  num_to_scalar = is_half ? '(float)NUM2DBL' : 'm_num_to_data'
+  scalar_one = is_half ? '1.0f' : 'm_one'
+  scalar_zero = is_half ? '0.0f' : 'm_zero'
 %>
 
 #define ROW_SIZE(na) ((na)->shape[(na)->ndim-2])
@@ -102,18 +110,31 @@
         }                                               \
     }
 
+#include <limits.h>
+
 typedef struct {
-    dtype alpha, beta;
+    <%=scalar_t%> alpha, beta;
     int m, n, k;
 } gemm_args_t;
 
 typedef struct {
     int ld;
-    int stride; // in element count
+    long long int stride; // in element count
     cublasOperation_t trans;
     VALUE a;
 } gemm_layout_t;
 
+// cuBLAS takes these as int, so a shape that does not fit has to be refused
+// rather than wrapped into a smaller or negative one.
+static int
+gemm_int(size_t v, const char *name)
+{
+    if (v > INT_MAX) {
+        rb_raise(rb_eRangeError, "%s=%"SZF"u is too large for cuBLAS", name, v);
+    }
+    return (int)v;
+}
+
 static bool
 is_f_contiguous(VALUE a)
 {
@@ -192,7 +213,7 @@ static int
 gemm_batch_count(cumo_narray_t *na)
 {
     size_t matrix_size = ROW_SIZE(na) * COL_SIZE(na);
-    return matrix_size == 0 ? 0 : (int)(CUMO_NA_SIZE(na) / matrix_size);
+    return matrix_size == 0 ? 0 : gemm_int(CUMO_NA_SIZE(na) / matrix_size, "batch count");
 }
 
 static gemm_layout_t
@@ -210,27 +231,27 @@ make_gemm_layout(VALUE a)
     }
 
     if ((CUMO_NA_NDIM(na) == 2 && is_f_contiguous(a)) || is_batched_f_contiguous(a)) {
-        layout.ld = ROW_SIZE(na);
+        layout.ld = (int)ROW_SIZE(na);
         layout.trans = CUBLAS_OP_T;
         layout.a = a;
     } else {
-        layout.ld = COL_SIZE(na);
+        layout.ld = (int)COL_SIZE(na);
         layout.trans = CUBLAS_OP_N;  // transposed
         // force c-contiguous
         layout.a = is_c_contiguous(a) ? a : rb_funcall(a, rb_intern("dup"), 0);
     }
     // cuBLAS walks the operand by this stride once per batch, so an operand that
     // holds a single matrix has to be re-read rather than advanced past its end.
-    layout.stride = gemm_batch_count(na) == 1 ? 0 : ROW_SIZE(na) * COL_SIZE(na);
+    layout.stride = gemm_batch_count(na) == 1 ? 0 : (long long int)(ROW_SIZE(na) * COL_SIZE(na));
     return layout;
 }
 
 extern int cumo_na_debug_flag;  // narray.c
 
 static void
-print_gemm_args(gemm_args_t* g, gemm_layout_t* a_layout, gemm_layout_t* b_layout, int stridec, int batch_count)
+print_gemm_args(gemm_args_t* g, gemm_layout_t* a_layout, gemm_layout_t* b_layout, long long int stridec, int batch_count)
 {
-    printf("transb=%d transa=%d, n=%d, m=%d, k=%d, ldb=%d, lda=%d, ldc=n=%d, strideb=%d, stridea=%d stridec=%d batch_count=%d\n",
+    printf("transb=%d transa=%d, n=%d, m=%d, k=%d, ldb=%d, lda=%d, ldc=n=%d, strideb=%lld, stridea=%lld stridec=%lld batch_count=%d\n",
             (int)b_layout->trans,
             (int)a_layout->trans,
             (int)g->n,
@@ -239,9 +260,9 @@ print_gemm_args(gemm_args_t* g, gemm_layout_t* a_layout, gemm_layout_t* b_layout
             (int)b_layout->ld,
             (int)a_layout->ld,
             (int)g->n,
-            (int)b_layout->stride,
-            (int)a_layout->stride,
-            (int)stridec,
+            b_layout->stride,
+            a_layout->stride,
+            stridec,
             (int)batch_count);
 }
 
@@ -252,7 +273,8 @@ static void
     cublasHandle_t handle = 0;
     cublasStatus_t status = 0;
     cumo_narray_t* nc;
-    int stridec = 0;
+    size_t matrix_size;
+    long long int stridec = 0;
     int batch_count = 0;
 
     // Note that cuBLAS uses the column major matrix representation.
@@ -270,11 +292,38 @@ static void
     b_layout = make_gemm_layout(b);
 
     CumoGetNArray(c, nc);
-    stridec = ROW_SIZE(nc) * COL_SIZE(nc);
-    batch_count = CUMO_NA_SIZE(nc) / stridec;
+    matrix_size = ROW_SIZE(nc) * COL_SIZE(nc);
+    stridec = (long long int)matrix_size;
+    batch_count = gemm_int(CUMO_NA_SIZE(nc) / matrix_size, "batch count");
 
     if (cumo_na_debug_flag) print_gemm_args(g, &a_layout, &b_layout, stridec, batch_count);
     handle = cumo_cuda_cublas_handle();
+<% if is_half %>
+    status = cublasGemmStridedBatchedEx(
+            handle,
+            b_layout.trans,
+            a_layout.trans,
+            g->n,
+            g->m,
+            g->k,
+            &g->alpha,
+            (const void*)(cumo_na_get_pointer_for_read(b_layout.a) + cumo_na_get_offset(b_layout.a)),
+            CUDA_R_16F,
+            b_layout.ld,
+            b_layout.stride,
+            (const void*)(cumo_na_get_pointer_for_read(a_layout.a) + cumo_na_get_offset(a_layout.a)),
+            CUDA_R_16F,
+            a_layout.ld,
+            a_layout.stride,
+            &g->beta,
+            (void*)(cumo_na_get_pointer_for_write(c) + cumo_na_get_offset(c)),
+            CUDA_R_16F,
+            g->n,
+            stridec,
+            batch_count,
+            CUBLAS_COMPUTE_32F,
+            CUBLAS_GEMM_DEFAULT);
+<% else %>
     status = cublas<%=func_prefix%>gemmStridedBatched(
             handle,
             b_layout.trans,
@@ -294,6 +343,7 @@ static void
             g->n,
             stridec,
             batch_count);
+<% end %>
     cumo_cuda_cublas_check_status(status);
 }
 
@@ -350,9 +400,9 @@ static VALUE
     rb_scan_args(argc, argv, "11:", &b, &c, &kw_hash);
     rb_get_kwargs(kw_hash, kw_table, 0, 2, opts);
     alpha = cumo_cuda_cublas_option_value(opts[0],Qnil);
-    g.alpha = RTEST(alpha) ? m_num_to_data(alpha) : m_one;
+    g.alpha = RTEST(alpha) ? <%=num_to_scalar%>(alpha) : <%=scalar_one%>;
     beta = cumo_cuda_cublas_option_value(opts[1],Qnil);
-    g.beta = RTEST(beta) ? m_num_to_data(beta) : m_zero;
+    g.beta = RTEST(beta) ? <%=num_to_scalar%>(beta) : <%=scalar_zero%>;
 
     // b is handed to cuBLAS as a raw <%=cutype%>*, so another dtype would be
     // reinterpreted, and a narrower one read past its end.
@@ -368,13 +418,13 @@ static VALUE
     CHECK_NON_EMPTY(nb);
 
     if (ROW_SIZE(nb) != COL_SIZE(na)) {
-        rb_raise(cumo_na_eShapeError,"ROW_SIZE(b)=%d must equal to COL_SIZE(a)=%d",
-                (int)ROW_SIZE(nb), (int)COL_SIZE(na));
+        rb_raise(cumo_na_eShapeError,"ROW_SIZE(b)=%"SZF"u must equal to COL_SIZE(a)=%"SZF"u",
+                ROW_SIZE(nb), COL_SIZE(na));
     }
 
-    g.m = ROW_SIZE(na);
-    g.k = COL_SIZE(na);
-    g.n = COL_SIZE(nb);
+    g.m = gemm_int(ROW_SIZE(na), "row size of a");
+    g.k = gemm_int(COL_SIZE(na), "column size of a");
+    g.n = gemm_int(COL_SIZE(nb), "column size of b");
 
     // The batch dimensions reach cuBLAS as one stride per operand, so both
     // operands must hold the same number of matrices. An operand holding a
@@ -402,16 +452,17 @@ static VALUE
         CumoGetNArray(c, nc);
         CHECK_DIM_GE(nc, 2);
         if (ROW_SIZE(nc) != ROW_SIZE(na)) {
-            rb_raise(cumo_na_eShapeError,"ROW_SIZE(c)=%d must equal to ROW_SIZE(a)=%d",
-                    (int)ROW_SIZE(nc), (int)ROW_SIZE(na));
+            rb_raise(cumo_na_eShapeError,"ROW_SIZE(c)=%"SZF"u must equal to ROW_SIZE(a)=%"SZF"u",
+                    ROW_SIZE(nc), ROW_SIZE(na));
         }
         if (COL_SIZE(nc) != COL_SIZE(nb)) {
-            rb_raise(cumo_na_eShapeError,"COL_SIZE(c)=%d must equal to COL_SIZE(b)=%d",
-                    (int)COL_SIZE(nc), (int)COL_SIZE(nb));
+            rb_raise(cumo_na_eShapeError,"COL_SIZE(c)=%"SZF"u must equal to COL_SIZE(b)=%"SZF"u",
+                    COL_SIZE(nc), COL_SIZE(nb));
         }
-        if (gemm_batch_count(nc) != batch) {
+        int c_batch = gemm_batch_count(nc);
+        if (c_batch != batch) {
             rb_raise(cumo_na_eShapeError,"batch count of c=%d must equal to %d",
-                    gemm_batch_count(nc), batch);
+                    c_batch, batch);
         }
         // a and b are duplicated when they are not contiguous, but c is written
         // in place, and cuBLAS writes it as ldc=n stridec=m*n from its offset

data/ext/cumo/narray/gen/tmpl/init_class.c

@@ -9,13 +9,6 @@
     rb_define_const(<%=ns_var%>, "<%=x%>", <%=type_var%>);
   <% end %>
 
-    hCast = rb_hash_new();
-    rb_define_const(cT, "UPCAST", hCast);
-    rb_hash_aset(hCast, rb_cArray,   cT);
-    <% for x in upcast %>
-    <%= x %><% end %>
-    rb_obj_freeze(hCast);
-
     <% @children.each do |m| %>
     <%= m.init_def %><% end %>
     rb_define_singleton_method(cT, "[]", <%=find("cast").c_func%>, -2);

data/ext/cumo/narray/gen/tmpl/init_upcast.c

@@ -0,0 +1,6 @@
+    hCast = rb_hash_new();
+    rb_define_const(cT, "UPCAST", hCast);
+    rb_hash_aset(hCast, rb_cArray,   cT);
+    <% for x in upcast %>
+    <%= x %><% end %>
+    rb_obj_freeze(hCast);

data/ext/cumo/narray/gen/tmpl/inspect.c

@@ -17,6 +17,6 @@ static VALUE
 <%=c_func(0)%>(VALUE ary)
 {
     CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
     return cumo_na_ndloop_inspect(ary, <%=c_iter%>, Qnil);
 }

data/ext/cumo/narray/gen/tmpl/lib.c

@@ -37,7 +37,7 @@ extern VALUE cRT;
 void
 Init_<%=lib_name%>(void)
 {
-    VALUE hCast, <%=ns_var%>;
+    VALUE <%=ns_var%>;
 
     <%=ns_var%> = rb_define_module("Cumo");
 
@@ -48,3 +48,13 @@ Init_<%=lib_name%>(void)
 <%= c.init_def %>
 <% end %>
 }
+
+// cumo.c runs this after the last Init_cumo_*. Any class defined after this
+// one is still 0 during Init, and its entry would land under the key false.
+void
+Init_<%=lib_name%>_upcast(void)
+{
+    VALUE hCast;
+
+<% children.each do |c| %><%= c.upcast_def %><% end %>
+}

data/ext/cumo/narray/gen/tmpl/map_with_index.c

@@ -1,15 +1,7 @@
 static inline dtype
 yield_map_with_index(dtype x, size_t *c, VALUE *a, int nd, int md)
 {
-    int j;
-    VALUE y;
-
-    a[0] = m_data_to_num(x);
-    for (j=0; j<=nd; j++) {
-        a[j+1] = SIZET2NUM(c[j]);
-    }
-    y = rb_yield(rb_ary_new4(md,a));
-    return m_num_to_data(y);
+    return m_num_to_data(cumo_na_yield_with_index(m_data_to_num(x),c,a,nd,md));
 }
 
 static void
@@ -25,17 +17,14 @@ static void
     int nd, md;
 
     c = (size_t*)(lp->opt_ptr);
-    nd = lp->ndim;
-    if (nd > 0) {nd--;}
-    md = nd + 2;
+    cumo_na_with_index_dims(lp->ndim, &nd, &md);
     a = ALLOCA_N(VALUE,md);
 
     CUMO_INIT_COUNTER(lp, i);
     CUMO_INIT_PTR_IDX(lp, 0, p1, s1, idx1);
     CUMO_INIT_PTR_IDX(lp, 1, p2, s2, idx2);
 
-    CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
 
     c[nd] = 0;
     if (idx1) {
@@ -43,6 +32,7 @@ static void
             for (; i--;) {
                 CUMO_GET_DATA_INDEX(p1,idx1,dtype,x);
                 x = yield_map_with_index(x,c,a,nd,md);
+                if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
                 CUMO_SET_DATA_INDEX(p2,idx2,dtype,x);
                 c[nd]++;
             }
@@ -50,6 +40,7 @@ static void
             for (; i--;) {
                 CUMO_GET_DATA_INDEX(p1,idx1,dtype,x);
                 x = yield_map_with_index(x,c,a,nd,md);
+                if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
                 CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
                 c[nd]++;
             }
@@ -59,6 +50,7 @@ static void
             for (; i--;) {
                 CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
                 x = yield_map_with_index(x,c,a,nd,md);
+                if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
                 CUMO_SET_DATA_INDEX(p2,idx2,dtype,x);
                 c[nd]++;
             }
@@ -66,6 +58,7 @@ static void
             for (; i--;) {
                 CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
                 x = yield_map_with_index(x,c,a,nd,md);
+                if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
                 CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
                 c[nd]++;
             }

data/ext/cumo/narray/gen/tmpl/median.c

@@ -1,10 +1,11 @@
-void cumo_<%=type_name%>_median_kernel_launch(cumo_na_reduction_arg_t* arg, int flat);
+void cumo_<%=type_name%>_median_kernel_launch(cumo_na_reduction_arg_t* arg, int flat, int prnan);
 
 // The rows are sorted with one cub call and the middle of each is picked by a
-// kernel of its own. nan:true keeps the host path, as sort does: its comparator
-// leaves NaN unordered, so there is no sorted array to take a middle of.
+// kernel of its own. The sort keys put every NaN last, so nan:true only has to
+// notice that one is there.
+<% (is_float ? ["_ignan","_prnan"] : [""]).each do |j| %>
 static void
-<%=c_iter%>_kernel(cumo_na_loop_t *const lp)
+<%=c_iter%><%=j%>(cumo_na_loop_t *const lp)
 {
     cumo_na_reduction_arg_t arg = cumo_na_make_reduction_arg(lp, 1);
     ssize_t expect = sizeof(dtype);
@@ -19,41 +20,7 @@ static void
         expect *= (ssize_t)arg.in_indexer.shape[i];
     }
 
-    cumo_<%=type_name%>_median_kernel_launch(&arg, flat);
-}
-
-<% (is_float ? ["_ignan","_prnan"] : [""]).each do |j| %>
-static void
-<%=c_iter%><%=j%>(cumo_na_loop_t *const lp)
-{
-    size_t n;
-    char *p1, *p2;
-    dtype *buf;
-
-    CUMO_INIT_COUNTER(lp, n);
-    p1 = (lp->args[0]).ptr + (lp->args[0].iter[0]).pos;
-    p2 = (lp->args[1]).ptr + (lp->args[1].iter[0]).pos;
-    buf = (dtype*)p1;
-
-    CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%><%=j%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
-    <%=type_name%>_qsort<%=j%>(buf, n, sizeof(dtype));
-
-    <% if is_float %>
-    for (; n; n--) {
-        if (!isnan(buf[n-1])) break;
-    }
-    <% end %>
-
-    if (n==0) {
-        *(dtype*)p2 = buf[0];
-    }
-    else if (n%2==0) {
-        *(dtype*)p2 = (buf[n/2-1]+buf[n/2])/2;
-    }
-    else {
-        *(dtype*)p2 = buf[(n-1)/2];
-    }
+    cumo_<%=type_name%>_median_kernel_launch(&arg, flat, <%= j == "_prnan" ? 1 : 0 %>);
 }
 <% end %>
 
@@ -82,18 +49,13 @@ static VALUE
   <% if is_float %>
     ndf.func = <%=c_iter%>_ignan;
     reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, <%=c_iter%>_prnan);
-    if (ndf.func == <%=c_iter%>_ignan) {
-        ndf.func = <%=c_iter%>_kernel;
-        // or rather than assign: cumo_na_reduce_dimension may have set
-        // CUMO_NDF_KEEP_DIM by then, and assigning would drop it
-        ndf.flag |= CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INDEXER_LOOP;
-    }
   <% else %>
     ndf.func = <%=c_iter%>;
     reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, 0);
-    ndf.func = <%=c_iter%>_kernel;
-    ndf.flag |= CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INDEXER_LOOP;
   <% end %>
+    // or rather than assign: cumo_na_reduce_dimension may have set
+    // CUMO_NDF_KEEP_DIM by then, and assigning would drop it
+    ndf.flag |= CUMO_NDF_STRIDE_LOOP|CUMO_NDF_INDEXER_LOOP;
     v = cumo_na_ndloop(&ndf, 2, self, reduce);
     return <%=type_name%>_extract(v);
 }

data/ext/cumo/narray/gen/tmpl/minmax.c

@@ -20,7 +20,7 @@ static void
         CUMO_INIT_PTR(lp, 0, p1, s1);
 
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%><%=j%>", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         f_<%=name%><%=j%>(n,p1,s1,&xmin,&xmax);
 
         *(dtype*)(lp->args[1].ptr + lp->args[1].iter[0].pos) = xmin;

data/ext/cumo/narray/gen/tmpl/poly.c

@@ -27,7 +27,7 @@ static void
     dtype  x, y, a;
 
     CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
     n = lp->narg - 2;
     x = *(dtype*)CUMO_NDL_PTR(lp,0);
     y = *(dtype*)CUMO_NDL_PTR(lp,n);

data/ext/cumo/narray/gen/tmpl/pooling_backward.c

@@ -1,17 +1,5 @@
 #ifdef CUDNN_FOUND
 
-<%
-  cudnn_dtype =
-    case type_name
-    when 'sfloat'
-      'CUDNN_DATA_FLOAT'
-    when 'dfloat'
-      'CUDNN_DATA_DOUBLE'
-    else
-      # CUDNN_DATA_HALF
-      raise 'not supported'
-    end
-%>
 
 // cover_all=true is not supported with CUDNN
 // gy = x.pooling_backward(mode, y, kernel_size, stride: 1, pad: 0)
@@ -25,8 +13,8 @@ static VALUE
     cudnnDataType_t cudnn_dtype = <%= cudnn_dtype %>;
     cudnnStatus_t status = 0;
     cudnnHandle_t handle = 0;
-    dtype alpha = 1;
-    dtype beta = 0;
+    <%=cudnn_scalar_t%> alpha = 1;
+    <%=cudnn_scalar_t%> beta = 0;
 
     VALUE x=self, y, gy, mode, kernel_size, stride, pad, gx;
     VALUE kw_hash = Qnil;