cumo 0.6.0 → 0.7.0

data/ext/cumo/narray/gen/tmpl/batch_norm_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
-%>
 
 // gx, ggamma, gbeta = x.batch_norm_backward(gamma, gy, mean:, inv_std:, eps:, axis:)
 static VALUE
@@ -20,8 +8,8 @@ static VALUE
     cudnnDataType_t cudnn_dtype = <%= cudnn_dtype %>;
     cudnnStatus_t status = 0;
     cudnnHandle_t handle = 0;
-    dtype coef_one = 1;
-    dtype coef_zero = 0;
+    <%=cudnn_scalar_t%> coef_one = 1;
+    <%=cudnn_scalar_t%> coef_zero = 0;
 
     VALUE x=self, gamma, gy, mean, inv_std, eps, axis, gx, ggamma, gbeta;
     VALUE kw_hash = Qnil;
@@ -107,10 +95,10 @@ static VALUE
     }
 
     CUMO_CHECK_NARRAY_TYPE(x, cT);
-    CUMO_CHECK_NARRAY_TYPE(gamma, cT);
+    cumo_cuda_cudnn_check_param_type(gamma, <%=cudnn_param_class%>, "gamma");
     CUMO_CHECK_NARRAY_TYPE(gy, cT);
-    if (mean != Qnil) CUMO_CHECK_NARRAY_TYPE(mean, cT);
-    if (inv_std != Qnil) CUMO_CHECK_NARRAY_TYPE(inv_std, cT);
+    if (mean != Qnil) cumo_cuda_cudnn_check_param_type(mean, <%=cudnn_param_class%>, "mean");
+    if (inv_std != Qnil) cumo_cuda_cudnn_check_param_type(inv_std, <%=cudnn_param_class%>, "inv_std");
 
     x_cont = cumo_na_as_contiguous_array(x);
     gamma_cont = cumo_na_as_contiguous_array(gamma);
@@ -133,15 +121,15 @@ static VALUE
     }
     gx_ptr = cumo_na_get_offset_pointer_for_write(gx);
     if (ggamma == Qnil) {
-        ggamma = cumo_na_new(cT, gamma_ndim, gamma_shape);
+        ggamma = cumo_na_new(<%=cudnn_param_class%>, gamma_ndim, gamma_shape);
     } else {
-        cumo_cuda_cudnn_check_output(ggamma, cT, gamma_ndim, gamma_shape);
+        cumo_cuda_cudnn_check_output(ggamma, <%=cudnn_param_class%>, gamma_ndim, gamma_shape);
     }
     ggamma_ptr = cumo_na_get_offset_pointer_for_write(ggamma);
     if (gbeta == Qnil) {
-        gbeta = cumo_na_new(cT, gamma_ndim, gamma_shape);
+        gbeta = cumo_na_new(<%=cudnn_param_class%>, gamma_ndim, gamma_shape);
     } else {
-        cumo_cuda_cudnn_check_output(gbeta, cT, gamma_ndim, gamma_shape);
+        cumo_cuda_cudnn_check_output(gbeta, <%=cudnn_param_class%>, gamma_ndim, gamma_shape);
     }
     gbeta_ptr = cumo_na_get_offset_pointer_for_write(gbeta);
 
@@ -151,9 +139,9 @@ static VALUE
     mode = cumo_cuda_cudnn_GetBatchNormMode(axis_ndim, int_axis);
     status = cumo_cuda_cudnn_CreateBNTensorDescriptor(&bn_desc, x_desc, mode);
     if (status != CUDNN_STATUS_SUCCESS) goto BATCH_NORM_BACKWARD_ERROR;
-    // The derived descriptor carries x's own type -- cuDNN widens only for
-    // half, which these templates are never generated for -- and the checks
-    // above hold every parameter to x's class, so none of them needs a cast.
+    // cuDNN derives the parameter descriptor as float for a half x and as x's
+    // own type otherwise, and the checks above hold every parameter to exactly
+    // that class, so none of them needs a cast.
 
     handle = cumo_cuda_cudnn_handle();
 

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

@@ -48,6 +48,64 @@ __global__ void <%="cumo_#{c_iter}_kernel_dim#{idim}"%>(cumo_na_iarray_t a1, cum
 //<% end %>
 <% end %>
 
+// One operand of a 2-d elementwise op can be a transposed view, which leaves
+// either its read or the write strided whichever way the loop walks. Stage that
+// operand through a shared memory tile so all three arrays are touched in rows.
+// t_side says which operand is the transposed one.
+__global__ void <%="cumo_#{c_iter}_transpose_kernel"%>(char *p1, char *p2, char *p3, int* divzero, int t_side, uint64_t rows, uint64_t cols)
+{
+    CUMO_TRANSPOSE_TILE_DECL(dtype, tile);
+    char *tp = (t_side == 1) ? p1 : p2;
+    char *sp = (t_side == 1) ? p2 : p1;
+
+    CUMO_TRANSPOSE_TILE_LOOP(tile, rows, cols,
+        *(dtype*)(tp + cumo_tile_src * sizeof(dtype)),
+        dtype o = *(dtype*)(sp + cumo_tile_dst * sizeof(dtype));
+        dtype lhs = (t_side == 1) ? cumo_tile_val : o;
+        dtype rhs = (t_side == 1) ? o : cumo_tile_val;
+        cumo_check_intdivzero(rhs);
+        *(dtype*)(p3 + cumo_tile_dst * sizeof(dtype)) = m_<%=name%>(lhs, rhs);
+    );
+}
+
+// Launches the tiled kernel and answers 1 when it did. It takes over when one
+// operand is a transposed plain array and the other two are row-major ones, and
+// only once both sides reach a tile: a shorter one leaves most of each warp idle
+// while the plain loop already reads or writes it in one go.
+static int
+<%="cumo_#{c_iter}_launch_transpose"%>(cumo_na_iarray_t* a1, cumo_na_iarray_t* a2, cumo_na_iarray_t* a3, cumo_na_indexer_t* indexer, int* divzero)
+{
+    ssize_t esz = (ssize_t)sizeof(dtype);
+    uint64_t rows, cols;
+    ssize_t row, col;
+    int rows_a1, rows_a2, t_side;
+
+    if (!CUMO_TRANSPOSE_TILE_FITS(indexer)) {
+        return 0;
+    }
+    rows = indexer->shape[0];
+    cols = indexer->shape[1];
+    row = esz * (ssize_t)cols;
+    col = esz * (ssize_t)rows;
+    if (a3->step[0] != row || a3->step[1] != esz) {
+        return 0;
+    }
+    rows_a1 = (a1->step[0] == row && a1->step[1] == esz);
+    rows_a2 = (a2->step[0] == row && a2->step[1] == esz);
+    if (!rows_a2 && rows_a1 && a2->step[0] == esz && a2->step[1] == col) {
+        t_side = 2;
+    } else if (!rows_a1 && rows_a2 && a1->step[0] == esz && a1->step[1] == col) {
+        t_side = 1;
+    } else {
+        return 0;
+    }
+
+    CUMO_TRANSPOSE_LAUNCH(<%="cumo_#{c_iter}_transpose_kernel"%>, rows, cols,
+                          a1->ptr, a2->ptr, a3->ptr, divzero, t_side);
+    cumo_cuda_runtime_check_kernel_launch();
+    return 1;
+}
+
 //<% if has_scalar %>
 static void <%="cumo_#{c_iter}_kernel_dispatch"%>(cumo_na_iarray_t* a1, cumo_na_iarray_t* a2, cumo_na_iarray_t* a3, cumo_na_indexer_t* indexer, int* divzero, dtype sv, int use_scalar)
 {
@@ -69,6 +127,10 @@ static void <%="cumo_#{c_iter}_kernel_dispatch"%>(cumo_na_iarray_t* a1, cumo_na_
 void <%="cumo_#{c_iter}_kernel_launch"%>(cumo_na_iarray_t* a1, cumo_na_iarray_t* a2, cumo_na_iarray_t* a3, cumo_na_indexer_t* indexer, int* divzero)
 {
     dtype sv;
+
+    if (<%="cumo_#{c_iter}_launch_transpose"%>(a1,a2,a3,indexer,divzero)) {
+        return;
+    }
     memset(&sv, 0, sizeof(dtype));
     <%="cumo_#{c_iter}_kernel_dispatch"%>(a1,a2,a3,indexer,divzero,sv,0);
 }
@@ -82,8 +144,13 @@ void <%="cumo_#{c_iter}_s_kernel_launch"%>(cumo_na_iarray_t* a1, dtype sv, cumo_
 //<% else %>
 void <%="cumo_#{c_iter}_kernel_launch"%>(cumo_na_iarray_t* a1, cumo_na_iarray_t* a2, cumo_na_iarray_t* a3, cumo_na_indexer_t* indexer, int* divzero)
 {
-    size_t grid_dim = cumo_get_grid_dim(indexer->total_size);
-    size_t block_dim = cumo_get_block_dim(indexer->total_size);
+    size_t grid_dim, block_dim;
+
+    if (<%="cumo_#{c_iter}_launch_transpose"%>(a1,a2,a3,indexer,divzero)) {
+        return;
+    }
+    grid_dim = cumo_get_grid_dim(indexer->total_size);
+    block_dim = cumo_get_block_dim(indexer->total_size);
     switch (indexer->ndim) {
     <% (0..opt_indexer_ndim).each do |idim| %>
     case <%=idim%>:

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

@@ -24,7 +24,7 @@ static void
         CUMO_INIT_PTR(lp, 3, p4, s4);
 
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         for (i=n; i--;) {
             CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
             CUMO_GET_DATA_STRIDE(p2,s2,dtype,min);
@@ -69,7 +69,7 @@ static void
         CUMO_INIT_PTR(lp, 2, p3, s3);
 
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>_min", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         for (i=n; i--;) {
             CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
             CUMO_GET_DATA_STRIDE(p2,s2,dtype,min);
@@ -105,7 +105,7 @@ static void
         CUMO_INIT_PTR(lp, 2, p3, s3);
 
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>_max", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         for (i=n; i--;) {
             CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
             CUMO_GET_DATA_STRIDE(p2,s2,dtype,max);

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

@@ -20,7 +20,7 @@ static void
         CUMO_INIT_PTR_IDX(lp, 0, p1, s1, idx1);
         CUMO_INIT_PTR_BIT(lp, 1, a2, p2, s2);
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         if (idx1) {
             for (i=0; i<n; i++) {
                 CUMO_GET_DATA_INDEX(p1,idx1,dtype,x);

data/ext/cumo/narray/gen/tmpl/conv.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
 // dilation > 1 is not supported yet
@@ -22,8 +10,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, w, b, stride, pad, y;
     VALUE kw_hash = Qnil;
@@ -45,7 +33,7 @@ static VALUE
 
     cudnnConvolutionFwdAlgoPerf_t perf_result;
     cudnnConvolutionFwdAlgo_t algo;
-    size_t max_workspace_size = CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE;
+    size_t max_workspace_size = cumo_cuda_cudnn_max_workspace_size();
     size_t workspace_size;
     char* workspace = 0;
 
@@ -113,7 +101,7 @@ static VALUE
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_ERROR;
     status = cumo_cuda_cudnn_CreateFilterDescriptor(&w_desc, w_cont, cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_ERROR;
-    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, cudnn_dtype);
+    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, <%=cudnn_compute_dtype%>, <%=cudnn_math_type%>);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_ERROR;
 
     handle = cumo_cuda_cudnn_handle();
@@ -135,10 +123,16 @@ static VALUE
             ndim,
             cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_ERROR;
+    // The descriptor asked for a math type; use the one the search settled on.
+    status = cudnnSetConvolutionMathType(conv_desc, perf_result.mathType);
+    if (status != CUDNN_STATUS_SUCCESS) goto CONV_ERROR;
+
     algo = perf_result.algo;
     workspace_size = perf_result.memory;
 
-    workspace = cumo_cuda_runtime_malloc(max_workspace_size);
+    // The search may look at algorithms needing up to max_workspace_size,
+    // but only the one it picked has to be paid for.
+    if (workspace_size > 0) workspace = cumo_cuda_runtime_malloc(workspace_size);
     status = cudnnConvolutionForward(
             handle,
             (void*)&alpha,

data/ext/cumo/narray/gen/tmpl/conv_grad_w.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
-%>
 
 static void
 cumo_cuda_cudnn_get_sizet_ary(size_t *sizet_ary, VALUE ary, size_t ndim)
@@ -31,8 +19,8 @@ static VALUE
     cudnnDataType_t cudnn_dtype = <%= cudnn_dtype %>;
     cudnnStatus_t status = 0;
     cudnnHandle_t handle = 0;
-    dtype one = 1;
-    dtype zero = 0;
+    <%=cudnn_scalar_t%> one = 1;
+    <%=cudnn_scalar_t%> zero = 0;
 
     VALUE x=self, gy, w_shape, stride, pad, gw;
     VALUE kw_hash = Qnil;
@@ -51,7 +39,7 @@ static VALUE
 
     cudnnConvolutionBwdFilterAlgoPerf_t perf_result;
     cudnnConvolutionBwdFilterAlgo_t algo;
-    size_t max_workspace_size = CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE;
+    size_t max_workspace_size = cumo_cuda_cudnn_max_workspace_size();
     size_t workspace_size;
     char* workspace = 0;
 
@@ -124,7 +112,7 @@ static VALUE
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_GRAD_W_ERROR;
     status = cumo_cuda_cudnn_CreateFilterDescriptor(&gw_desc, gw, cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_GRAD_W_ERROR;
-    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, cudnn_dtype);
+    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, <%=cudnn_compute_dtype%>, <%=cudnn_math_type%>);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_GRAD_W_ERROR;
 
     handle = cumo_cuda_cudnn_handle();
@@ -146,10 +134,16 @@ static VALUE
             ndim,
             cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_GRAD_W_ERROR;
+    // The descriptor asked for a math type; use the one the search settled on.
+    status = cudnnSetConvolutionMathType(conv_desc, perf_result.mathType);
+    if (status != CUDNN_STATUS_SUCCESS) goto CONV_GRAD_W_ERROR;
+
     algo = perf_result.algo;
     workspace_size = perf_result.memory;
 
-    workspace = cumo_cuda_runtime_malloc(max_workspace_size);
+    // The search may look at algorithms needing up to max_workspace_size,
+    // but only the one it picked has to be paid for.
+    if (workspace_size > 0) workspace = cumo_cuda_runtime_malloc(workspace_size);
     status = cudnnConvolutionBackwardFilter(
             handle,
             (void*)&one,

data/ext/cumo/narray/gen/tmpl/conv_transpose.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
-%>
 
 // VALUE is Ruby Array
 static void
@@ -47,8 +35,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, w, b, stride, pad, out_size, y;
     VALUE kw_hash = Qnil;
@@ -70,7 +58,7 @@ static VALUE
 
     cudnnConvolutionBwdDataAlgoPerf_t perf_result;
     cudnnConvolutionBwdDataAlgo_t algo;
-    size_t max_workspace_size = CUMO_CUDA_CUDNN_DEFAULT_MAX_WORKSPACE_SIZE;
+    size_t max_workspace_size = cumo_cuda_cudnn_max_workspace_size();
     size_t workspace_size;
     char* workspace = 0;
 
@@ -140,7 +128,7 @@ static VALUE
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_TRANSPOSE_ERROR;
     status = cumo_cuda_cudnn_CreateFilterDescriptor(&w_desc, w_cont, cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_TRANSPOSE_ERROR;
-    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, cudnn_dtype);
+    status = cumo_cuda_cudnn_CreateConvolutionDescriptor(&conv_desc, ndim, int_stride, int_pad, <%=cudnn_compute_dtype%>, <%=cudnn_math_type%>);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_TRANSPOSE_ERROR;
 
     handle = cumo_cuda_cudnn_handle();
@@ -162,10 +150,16 @@ static VALUE
             ndim,
             cudnn_dtype);
     if (status != CUDNN_STATUS_SUCCESS) goto CONV_TRANSPOSE_ERROR;
+    // The descriptor asked for a math type; use the one the search settled on.
+    status = cudnnSetConvolutionMathType(conv_desc, perf_result.mathType);
+    if (status != CUDNN_STATUS_SUCCESS) goto CONV_TRANSPOSE_ERROR;
+
     algo = perf_result.algo;
     workspace_size = perf_result.memory;
 
-    workspace = cumo_cuda_runtime_malloc(max_workspace_size);
+    // The search may look at algorithms needing up to max_workspace_size,
+    // but only the one it picked has to be paid for.
+    if (workspace_size > 0) workspace = cumo_cuda_runtime_malloc(workspace_size);
     status = cudnnConvolutionBackwardData(
             handle,
             (void*)&alpha,

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

@@ -26,8 +26,30 @@ static void
   <% end %>
 
     CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%><%=j%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
 
+<% if is_half && name == 'cumsum' %>
+    // The scan this stands in for carries float, and a running sum of halves
+    // stops moving once it passes 2048, so the accumulator is float here too.
+    {
+        float acc, fy;
+
+        CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
+        acc = cumo_half2float(x);
+        CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
+        for (i--; i--;) {
+            CUMO_GET_DATA_STRIDE(p1,s1,dtype,y);
+            fy = cumo_half2float(y);
+  <% if j == '_nan' %>
+            if (acc != acc) { acc = fy; } else if (fy == fy) { acc += fy; }
+  <% else %>
+            acc += fy;
+  <% end %>
+            x = cumo_float2half(acc);
+            CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
+        }
+    }
+<% else %>
     CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
     CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
     //printf("i=%lu x=%f\n",i,x);
@@ -37,6 +59,7 @@ static void
         CUMO_SET_DATA_STRIDE(p2,s2,dtype,x);
         //printf("i=%lu x=%f\n",i,x);
     }
+<% end %>
 }
 <% end %>
 

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

@@ -1,4 +1,20 @@
 <% unless type_name == 'robject' %>
+<% widen = is_half && name == 'cumsum' %>
+<% cum_t = widen ? 'float' : 'dtype' %>
+<% if widen %>
+
+// A running sum of halves stops moving once it passes 2048, so the scan carries
+// float and each element is rounded as it is written.
+struct <%="cumo_thrust_#{name}_widen"%>
+{
+    __host__ __device__ float operator()(dtype x) const { return cumo_half2float(x); }
+};
+
+struct <%="cumo_thrust_#{name}_narrow"%>
+{
+    __host__ __device__ dtype operator()(float x) const { return cumo_float2half(x); }
+};
+<% end %>
 <% (is_float ? ["","_nan"] : [""]).each do |j| %>
 
 #if defined(__cplusplus)
@@ -14,12 +30,13 @@
 // last ulp, the same way sum already does.
 struct <%="cumo_thrust_#{name}#{j}"%>
 {
-    using first_argument_type  = dtype;
-    using second_argument_type = dtype;
-    using result_type          = dtype;
-    __host__ __device__ dtype operator()(dtype x, dtype y) const { m_<%=name%><%=j%>(x,y); return x; }
+    using first_argument_type  = <%=cum_t%>;
+    using second_argument_type = <%=cum_t%>;
+    using result_type          = <%=cum_t%>;
+    __host__ __device__ <%=cum_t%> operator()(<%=cum_t%> x, <%=cum_t%> y) const { m_<%=name%><%=j%>(x,y); return x; }
 };
 
+
 // Nothing may reach the C caller: it has no handler, so an escaping exception
 // is std::terminate. The status goes back instead and the caller raises.
 template<typename Iterator1, typename Iterator2>
@@ -27,7 +44,15 @@ static cudaError_t <%="cumo_#{type_name}_#{name}#{j}_scan"%>(Iterator1 first, It
 {
     cumo_thrust_pool_allocator alloc;
     try {
+<% if widen %>
+        thrust::inclusive_scan(thrust::cuda::par(alloc),
+            thrust::make_transform_iterator(first, <%="cumo_thrust_#{name}_widen"%>()),
+            thrust::make_transform_iterator(last, <%="cumo_thrust_#{name}_widen"%>()),
+            thrust::make_transform_output_iterator(result, <%="cumo_thrust_#{name}_narrow"%>()),
+            <%="cumo_thrust_#{name}#{j}"%>());
+<% else %>
         thrust::inclusive_scan(thrust::cuda::par(alloc), first, last, result, <%="cumo_thrust_#{name}#{j}"%>());
+<% end %>
     } catch (const thrust::system_error& e) {
         return (cudaError_t)e.code().value();
     } catch (const std::bad_alloc&) {

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

@@ -10,18 +10,17 @@ static void
     CUMO_INIT_COUNTER(lp, i);
     CUMO_INIT_PTR_IDX(lp, 0, p1, s1, idx1);
 
-    CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
 
     if (idx1) {
         for (; i--;) {
-            cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+            if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
             CUMO_GET_DATA_INDEX(p1,idx1,dtype,x);
             y = m_data_to_num(x);
             rb_yield(y);
         }
     } else {
         for (; i--;) {
-            cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+            if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
             CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
             y = m_data_to_num(x);
             rb_yield(y);

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

@@ -1,16 +1,3 @@
-static inline void
-yield_each_with_index(dtype x, size_t *c, VALUE *a, int nd, int md)
-{
-    int j;
-
-    a[0] = m_data_to_num(x);
-    for (j=0; j<=nd; j++) {
-        a[j+1] = SIZET2NUM(c[j]);
-    }
-    rb_yield(rb_ary_new4(md,a));
-}
-
-
 static void
 <%=c_iter%>(cumo_na_loop_t *const lp)
 {
@@ -23,29 +10,26 @@ 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);
     c[nd] = 0;
 
-    CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
 
     if (idx1) {
         for (; i--;) {
-            cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+            if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
             CUMO_GET_DATA_INDEX(p1,idx1,dtype,x);
-            yield_each_with_index(x,c,a,nd,md);
+            cumo_na_yield_with_index(m_data_to_num(x),c,a,nd,md);
             c[nd]++;
         }
     } else {
         for (; i--;) {
-            cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+            if (cumo_cuda_runtime_sync_if_busy()) { CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>"); }
             CUMO_GET_DATA_STRIDE(p1,s1,dtype,x);
-            yield_each_with_index(x,c,a,nd,md);
+            cumo_na_yield_with_index(m_data_to_num(x),c,a,nd,md);
             c[nd]++;
         }
     }

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

@@ -17,7 +17,7 @@ static VALUE
     if (na->ndim==0) {
         ptr = cumo_na_get_pointer_for_read(self) + cumo_na_get_offset(self);
         CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         v = m_extract(ptr);
         cumo_na_release_lock(self);
         return v;

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

@@ -12,7 +12,7 @@ static dtype
     VALUE  r, klass;
 
     CUMO_SHOW_SYNCHRONIZE_WARNING_ONCE("<%=name%>", "<%=type_name%>");
-    cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+    cumo_cuda_runtime_device_synchronize();
 
     if (CumoIsNArray(obj)) {
         CumoGetNArray(obj,na);

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

@@ -19,7 +19,7 @@ static void
         CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%>", "<%=type_name%>");
         // The index array is filled by a kernel; wait for it before reading it
         // from here.
-        cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
+        cumo_cuda_runtime_device_synchronize();
         if (idx1) {
             for (; i--;) {
                 CUMO_SET_DATA_INDEX(p1,idx1,dtype,y);

data/ext/cumo/narray/gen/tmpl/fixed_batch_norm.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
-%>
 
 // y = x.fixed_batch_norm(gamma, beta, mean, var, eps:, axis:)
 static VALUE
@@ -20,8 +8,8 @@ static VALUE
     cudnnDataType_t cudnn_dtype = <%= cudnn_dtype %>;
     cudnnStatus_t status = 0;
     cudnnHandle_t handle = 0;
-    dtype coef_one = 1;
-    dtype coef_zero = 0;
+    <%=cudnn_scalar_t%> coef_one = 1;
+    <%=cudnn_scalar_t%> coef_zero = 0;
 
     VALUE x=self, gamma, beta, mean, var, eps, axis, y;
     VALUE kw_hash = Qnil;
@@ -85,10 +73,10 @@ static VALUE
     }
 
     CUMO_CHECK_NARRAY_TYPE(x, cT);
-    CUMO_CHECK_NARRAY_TYPE(gamma, cT);
-    CUMO_CHECK_NARRAY_TYPE(beta, cT);
-    CUMO_CHECK_NARRAY_TYPE(mean, cT);
-    CUMO_CHECK_NARRAY_TYPE(var, cT);
+    cumo_cuda_cudnn_check_param_type(gamma, <%=cudnn_param_class%>, "gamma");
+    cumo_cuda_cudnn_check_param_type(beta, <%=cudnn_param_class%>, "beta");
+    cumo_cuda_cudnn_check_param_type(mean, <%=cudnn_param_class%>, "mean");
+    cumo_cuda_cudnn_check_param_type(var, <%=cudnn_param_class%>, "var");
 
     x_cont = cumo_na_as_contiguous_array(x);
     gamma_cont = cumo_na_as_contiguous_array(gamma);
@@ -115,9 +103,9 @@ static VALUE
     mode = cumo_cuda_cudnn_GetBatchNormMode(axis_ndim, int_axis);
     status = cumo_cuda_cudnn_CreateBNTensorDescriptor(&bn_desc, x_desc, mode);
     if (status != CUDNN_STATUS_SUCCESS) goto FIXED_BATCH_NORM_ERROR;
-    // The derived descriptor carries x's own type -- cuDNN widens only for
-    // half, which these templates are never generated for -- and the checks
-    // above hold every parameter to x's class, so none of them needs a cast.
+    // cuDNN derives the parameter descriptor as float for a half x and as x's
+    // own type otherwise, and the checks above hold every parameter to exactly
+    // that class, so none of them needs a cast.
 
     handle = cumo_cuda_cudnn_handle();