cumo 0.12.0 → 0.13.0

data/ext/cumo/cuda/nvrtc.c

@@ -2,6 +2,10 @@
 #include <ruby/thread.h>
 #include <assert.h>
 #include <limits.h>
+#include <cuda.h>
+#include <dlfcn.h>
+#include <pthread.h>
+#include <stdio.h>
 #include <nvrtc.h>
 #include "cumo/cuda/nvrtc.h"
 #include "cumo/cuda/handle.h"
@@ -13,11 +17,70 @@ VALUE cumo_cuda_mNVRTC;
 
 static cumo_cuda_handle_set_t programs;
 
+#define NVRTC_FUNCTIONS(X) \
+    X(nvrtcGetErrorString) X(nvrtcVersion) X(nvrtcCreateProgram) X(nvrtcDestroyProgram) \
+    X(nvrtcCompileProgram) X(nvrtcGetPTXSize) X(nvrtcGetPTX) X(nvrtcAddNameExpression) \
+    X(nvrtcGetLoweredName) X(nvrtcGetProgramLogSize) X(nvrtcGetProgramLog)
+
+static struct {
+#define NVRTC_POINTER(f) __typeof__(&f) f;
+    NVRTC_FUNCTIONS(NVRTC_POINTER)
+#undef NVRTC_POINTER
+} nvrtc_functions;
+static pthread_once_t nvrtc_once = PTHREAD_ONCE_INIT;
+static char nvrtc_error[256];
+
+static const char*
+dl_error(void)
+{
+    const char *e = dlerror();
+    return e ? e : "unknown error";
+}
+
+static void
+load_nvrtc(void)
+{
+    char soname[32];
+    void *lib;
+
+    if (CUDA_VERSION >= 12000) {
+        snprintf(soname, sizeof(soname), "libnvrtc.so.%d", CUDA_VERSION / 1000);
+    } else if (CUDA_VERSION >= 11020) {
+        snprintf(soname, sizeof(soname), "libnvrtc.so.11.2");
+    } else {
+        snprintf(soname, sizeof(soname), "libnvrtc.so.%d.%d", CUDA_VERSION / 1000, CUDA_VERSION % 1000 / 10);
+    }
+    lib = dlopen(soname, RTLD_NOW | RTLD_LOCAL);
+    if (lib == NULL) {
+        snprintf(nvrtc_error, sizeof(nvrtc_error), "cannot load %s: %s", soname, dl_error());
+        return;
+    }
+#define NVRTC_RESOLVE(f)                                                                      \
+    nvrtc_functions.f = (__typeof__(&f))dlsym(lib, #f);                                        \
+    if (nvrtc_functions.f == NULL) {                                                           \
+        snprintf(nvrtc_error, sizeof(nvrtc_error), "%s has no %s: %s", soname, #f, dl_error()); \
+        dlclose(lib);                                                                          \
+        return;                                                                                \
+    }
+    NVRTC_FUNCTIONS(NVRTC_RESOLVE)
+#undef NVRTC_RESOLVE
+}
+
+static const __typeof__(nvrtc_functions)*
+nvrtc(void)
+{
+    pthread_once(&nvrtc_once, load_nvrtc);
+    if (nvrtc_error[0] != '\0') {
+        rb_raise(cumo_cuda_eNVRTCError, "%s", nvrtc_error);
+    }
+    return &nvrtc_functions;
+}
+
 static void
 check_status(nvrtcResult status)
 {
     if (status != 0) {
-        rb_raise(cumo_cuda_eNVRTCError, "%s (error=%d)", nvrtcGetErrorString(status), status);
+        rb_raise(cumo_cuda_eNVRTCError, "%s (error=%d)", nvrtc()->nvrtcGetErrorString(status), status);
     }
 }
 
@@ -28,7 +91,7 @@ rb_nvrtcVersion(VALUE self)
     nvrtcResult status;
     VALUE major, minor;
 
-    status = nvrtcVersion(&_major, &_minor);
+    status = nvrtc()->nvrtcVersion(&_major, &_minor);
 
     check_status(status);
     major = INT2NUM(_major);
@@ -105,7 +168,7 @@ nvrtcCreateProgram_without_gvl_cb(void *param)
 {
     struct nvrtcCreateProgramParam *p = param;
     nvrtcResult status;
-    status = nvrtcCreateProgram(p->prog, p->src, p->name, p->numHeaders, p->headers, p->includeNames);
+    status = nvrtc_functions.nvrtcCreateProgram(p->prog, p->src, p->name, p->numHeaders, p->headers, p->includeNames);
     return (void *)status;
 }
 
@@ -126,6 +189,7 @@ rb_nvrtcCreateProgram(
     const char** _includeNames = NULL;
     VALUE strings;
 
+    nvrtc();
     // nvrtcCreateProgram reads numHeaders entries out of includeNames as well,
     // so a shorter array would be read past its end.
     if (ary_len(includeNames, "include_names") != _numHeaders) {
@@ -168,7 +232,7 @@ nvrtcDestroyProgram_without_gvl_cb(void *param)
 {
     struct nvrtcDestroyProgramParam *p = param;
     nvrtcResult status;
-    status = nvrtcDestroyProgram(p->prog);
+    status = nvrtc_functions.nvrtcDestroyProgram(p->prog);
     return (void *)status;
 }
 
@@ -179,6 +243,7 @@ rb_nvrtcDestroyProgram(VALUE self, VALUE prog)
     nvrtcProgram _prog = (nvrtcProgram)cumo_cuda_handle_take(&programs, prog, "nvrtcProgram");
 
     struct nvrtcDestroyProgramParam param = {&_prog};
+    nvrtc();
     status = (nvrtcResult)rb_thread_call_without_gvl(nvrtcDestroyProgram_without_gvl_cb, &param, NULL, NULL);
 
     check_status(status);
@@ -196,7 +261,7 @@ nvrtcCompileProgram_without_gvl_cb(void *param)
 {
     struct nvrtcCompileProgramParam *p = param;
     nvrtcResult status;
-    status = nvrtcCompileProgram(p->prog, p->numOptions, p->options);
+    status = nvrtc_functions.nvrtcCompileProgram(p->prog, p->numOptions, p->options);
     return (void *)status;
 }
 
@@ -209,6 +274,7 @@ rb_nvrtcCompileProgram(VALUE self, VALUE prog, VALUE options)
     const char** _options = NULL;
     VALUE strings;
 
+    nvrtc();
     strings = rb_ary_new_capa(_numOptions);
     push_strings(strings, options, _numOptions);
 
@@ -237,12 +303,12 @@ rb_nvrtcGetPTX(VALUE self, VALUE prog)
     char *_ptx;
     VALUE ptx;
 
-    status = nvrtcGetPTXSize(_prog, &_ptxSizeRet);
+    status = nvrtc()->nvrtcGetPTXSize(_prog, &_ptxSizeRet);
     check_status(status);
 
     ptx = rb_str_new(NULL, _ptxSizeRet);
     _ptx = RSTRING_PTR(ptx);
-    status = nvrtcGetPTX(_prog, _ptx);
+    status = nvrtc()->nvrtcGetPTX(_prog, _ptx);
     check_status(status);
 
     return ptx;
@@ -262,7 +328,7 @@ rb_nvrtcAddNameExpression(VALUE self, VALUE prog, VALUE name_expression)
     nvrtcProgram _prog;
     StringValueCStr(name_expression);
     _prog = (nvrtcProgram)cumo_cuda_handle_get(&programs, prog, "nvrtcProgram");
-    check_status(nvrtcAddNameExpression(_prog, RSTRING_PTR(name_expression)));
+    check_status(nvrtc()->nvrtcAddNameExpression(_prog, RSTRING_PTR(name_expression)));
     RB_GC_GUARD(name_expression);
     return Qnil;
 }
@@ -283,7 +349,7 @@ rb_nvrtcGetLoweredName(VALUE self, VALUE prog, VALUE name_expression)
     const char *lowered = NULL;
     StringValueCStr(name_expression);
     _prog = (nvrtcProgram)cumo_cuda_handle_get(&programs, prog, "nvrtcProgram");
-    check_status(nvrtcGetLoweredName(_prog, RSTRING_PTR(name_expression), &lowered));
+    check_status(nvrtc()->nvrtcGetLoweredName(_prog, RSTRING_PTR(name_expression), &lowered));
     RB_GC_GUARD(name_expression);
     return rb_str_new_cstr(lowered);
 }
@@ -297,12 +363,12 @@ rb_nvrtcGetProgramLog(VALUE self, VALUE prog)
     char *_log;
     VALUE log;
 
-    status = nvrtcGetProgramLogSize(_prog, &_logSizeRet);
+    status = nvrtc()->nvrtcGetProgramLogSize(_prog, &_logSizeRet);
     check_status(status);
 
     log = rb_str_new(NULL, _logSizeRet);
     _log = RSTRING_PTR(log);
-    status = nvrtcGetProgramLog(_prog, _log);
+    status = nvrtc()->nvrtcGetProgramLog(_prog, _log);
     check_status(status);
 
     return log;

data/ext/cumo/cuda/runtime.c

@@ -11,8 +11,8 @@
 
 VALUE cumo_cuda_eRuntimeError;
 VALUE cumo_cuda_mRuntime;
-uint64_t cumo_cuda_sync_epoch = 0;
-uint64_t cumo_cuda_launch_epoch = 0;
+size_t cumo_cuda_sync_epoch = 0;
+size_t cumo_cuda_launch_epoch = 0;
 
 static ID id_current_stream;
 
@@ -63,7 +63,7 @@ cumo_cuda_stream_set(cudaStream_t stream)
 void
 cumo_cuda_runtime_check_kernel_launch(void)
 {
-    cumo_cuda_launch_epoch++;
+    cumo_atomic_size_fetch_add(&cumo_cuda_launch_epoch, 1);
     check_status(cudaGetLastError());
 }
 
@@ -215,7 +215,7 @@ cumo_cuda_runtime_check_taken_status_holding(int status, char *p0, char *p1, cha
 void
 cumo_cuda_runtime_check_kernel_launch_holding(char *p0, char *p1, char *p2, char *p3, char *p4)
 {
-    cumo_cuda_launch_epoch++;
+    cumo_atomic_size_fetch_add(&cumo_cuda_launch_epoch, 1);
     cumo_cuda_runtime_check_taken_status_holding((int)cudaGetLastError(), p0, p1, p2, p3, p4);
 }
 
@@ -866,6 +866,10 @@ thread_table_free(void *ptr)
         char *entry = t->entries + i * local->entry_size;
         spare_entry_t *spare;
         if (!entry_filled(entry, local->entry_size)) { continue; }
+        if (local->release != NULL) {
+            local->release(entry);
+            if (!entry_filled(entry, local->entry_size)) { continue; }
+        }
         spare = (spare_entry_t*)malloc(sizeof(spare_entry_t) + local->entry_size);
         if (spare == NULL) { continue; }
         spare->device = i;
@@ -883,11 +887,18 @@ cumo_cuda_thread_local_init(cumo_cuda_thread_local_t *local, size_t entry_size)
 {
     local->entry_size = entry_size;
     local->spare = NULL;
+    local->release = NULL;
     if (pthread_mutex_init(&local->lock, NULL) != 0 || pthread_key_create(&local->key, thread_table_free) != 0) {
         rb_raise(rb_eRuntimeError, "cannot set up a per-thread table");
     }
 }
 
+void
+cumo_cuda_thread_local_on_exit(cumo_cuda_thread_local_t *local, void (*release)(void *entry))
+{
+    local->release = release;
+}
+
 static int
 take_spare(cumo_cuda_thread_local_t *local, size_t device, char *entry)
 {

data/ext/cumo/depend.erb

@@ -40,7 +40,7 @@ list_type_rb = Dir.glob("#{__dir__}/narray/gen/def/*.rb")
 list_type_rb.each do |type_rb|
   type_name = File.basename(type_rb, ".rb")
   next if ENV['DTYPE'] and !type_name.downcase.include?(ENV['DTYPE'].downcase)
-  KERNEL_PARTS.each do |part|
+  kernel_parts_of(type_name).each do |part|
     list_type_cu << type_cu = "narray/types/" + kernel_part_basename(type_name, part) + ".cu"
 %>
 <%=type_cu%>: <%=type_rb%> $(CU_DEPENDS)

data/ext/cumo/extconf.rb

@@ -72,7 +72,7 @@ cuda/cusolver
 )
 %w[bit int8 int16 int32 int64 uint8 uint16 uint32 uint64 hfloat bfloat sfloat dfloat scomplex dcomplex robject].each do |type_name|
   srcs << "narray/types/#{type_name}"
-  srcs.concat(KERNEL_PARTS.map { |part| "narray/types/" + kernel_part_basename(type_name, part) })
+  srcs.concat(kernel_parts_of(type_name).map { |part| "narray/types/" + kernel_part_basename(type_name, part) })
 end
 
 $objs = srcs.map { |src| "#{src}.o" }
@@ -138,7 +138,7 @@ dir_config('cumo', HEADER_DIRS, LIB_DIRS)
 
 have_library('cuda')
 have_library('cudart')
-have_library('nvrtc')
+have_header('nvrtc.h')
 have_library('cublas')
 # have_library('curand')
 

data/ext/cumo/include/cumo/atomic.h

@@ -0,0 +1,27 @@
+#ifndef CUMO_ATOMIC_H
+#define CUMO_ATOMIC_H
+
+#include <stddef.h>
+
+// The compiler builtins rather than ruby/atomic.h: the size_t fetch-add
+// there arrived after the oldest Ruby cumo builds with, and this header is
+// read without ruby.h as well.
+static inline size_t
+cumo_atomic_size_fetch_add(size_t *var, size_t n)
+{
+    return __atomic_fetch_add(var, n, __ATOMIC_SEQ_CST);
+}
+
+static inline size_t
+cumo_atomic_size_load(const size_t *var)
+{
+    return __atomic_load_n(var, __ATOMIC_SEQ_CST);
+}
+
+static inline void
+cumo_atomic_size_store(size_t *var, size_t value)
+{
+    __atomic_store_n(var, value, __ATOMIC_SEQ_CST);
+}
+
+#endif /* ifndef CUMO_ATOMIC_H */

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

@@ -2,6 +2,7 @@
 #define CUMO_CUDA_RUNTIME_H
 
 #include "cumo/narray.h"
+#include "cumo/atomic.h"
 #include <pthread.h>
 #include <cuda_runtime.h>
 
@@ -25,17 +26,17 @@ cudaStream_t cumo_cuda_stream_get(VALUE stream);
 // the moment the count has moved on. cumo_cuda_runtime_device_synchronize
 // advances it, and so does a host read under a stream of the caller's, and a
 // count that is behind costs a wait rather than correctness.
-extern uint64_t cumo_cuda_sync_epoch;
+extern size_t cumo_cuda_sync_epoch;
 
 // How many times device memory was written from here: every kernel launch,
 // copy into it and library call counts. A host loop reads a staged copy,
 // and reads again once the count moves.
-extern uint64_t cumo_cuda_launch_epoch;
+extern size_t cumo_cuda_launch_epoch;
 
 static inline void
 cumo_cuda_runtime_note_device_write(void)
 {
-    cumo_cuda_launch_epoch++;
+    cumo_atomic_size_fetch_add(&cumo_cuda_launch_epoch, 1);
 }
 
 // A pinned host buffer for reading device memory back: a copy into pinned
@@ -67,11 +68,27 @@ cumo_cuda_runtime_check_status(cudaError_t status)
 
 void cumo_cuda_runtime_return_scratch(char *ptr, int wait_for_stream, cudaError_t *status);
 
+// The counts move before the wait, not after it. A view that records the
+// count while another thread is already waiting then waits itself, where
+// recording the count the wait left behind would have let it skip a wait
+// that never covered its fill.
 static inline void
 cumo_cuda_runtime_device_synchronize(void)
 {
+    cumo_atomic_size_fetch_add(&cumo_cuda_sync_epoch, 1);
     cumo_cuda_runtime_check_status(cudaDeviceSynchronize());
-    cumo_cuda_sync_epoch++;
+}
+
+static inline size_t
+cumo_cuda_sync_epoch_now(void)
+{
+    return cumo_atomic_size_load(&cumo_cuda_sync_epoch);
+}
+
+static inline size_t
+cumo_cuda_launch_epoch_now(void)
+{
+    return cumo_atomic_size_load(&cumo_cuda_launch_epoch);
 }
 
 // Asking costs less than half of waiting, and there is nothing to wait for
@@ -114,9 +131,9 @@ cumo_cuda_runtime_memcpy_to_pinned(void *dst, const void *src, size_t bytes)
 {
     cudaError_t status;
     if (cumo_cuda_stream() != 0 && !cumo_cuda_runtime_streams_idle()) {
+        cumo_atomic_size_fetch_add(&cumo_cuda_sync_epoch, 1);
         status = cudaDeviceSynchronize();
         if (status != cudaSuccess) { return status; }
-        cumo_cuda_sync_epoch++;
     }
     status = cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToHost, cumo_cuda_stream());
     if (status != cudaSuccess) { return status; }
@@ -131,9 +148,13 @@ typedef struct {
     size_t entry_size;
     pthread_mutex_t lock;
     void *spare;
+    void (*release)(void *entry);
 } cumo_cuda_thread_local_t;
 
 void cumo_cuda_thread_local_init(cumo_cuda_thread_local_t *local, size_t entry_size);
+// Called on each filled entry of a thread that ends, without the GVL. An
+// entry it empties is dropped instead of being kept for the next thread.
+void cumo_cuda_thread_local_on_exit(cumo_cuda_thread_local_t *local, void (*release)(void *entry));
 void *cumo_cuda_thread_local_get(cumo_cuda_thread_local_t *local);
 
 static inline int

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

@@ -22,6 +22,7 @@ unsigned int cumo_na_element_stride(VALUE nary);
 size_t cumo_na_dtype_element_stride(VALUE klass);
 
 char *cumo_na_get_pointer(VALUE);
+int cumo_na_has_pending_scalar(VALUE);
 char *cumo_na_get_pointer_for_write(VALUE);
 char *cumo_na_get_pointer_for_read(VALUE);
 char *cumo_na_get_pointer_for_read_write(VALUE);
@@ -113,8 +114,7 @@ VALUE cumo_na_ary_composition_dtype(VALUE ary);
 
 // defined in rand, used by the generated rand and rand_norm
 u_int64_t cumo_cuda_rand_seed(void);
-u_int64_t cumo_cuda_rand_offset(void);
-void cumo_cuda_rand_set_offset(u_int64_t offset);
+u_int64_t cumo_cuda_rand_reserve_offset(u_int64_t n);
 
 #include "ruby/version.h"
 

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

@@ -253,8 +253,8 @@ typedef struct {
     // until it has been worked out. What a view reaches never changes, so it
     // is worked out once, and it stays right even after the base has moved on.
     size_t   reach_end;
-    uint64_t index_sync_epoch; // synchronizes counted when the index fills were
-                               // issued; UINT64_MAX when that is not known
+    size_t index_sync_epoch;   // synchronizes counted when the index fills were
+                               // issued; SIZE_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

data/ext/cumo/include/cumo/types/int_macro_kernel.h

@@ -43,7 +43,7 @@ __host__ __device__ static inline dtype int_reciprocal(dtype x) {
     }
 }
 
-__device__ static dtype pow_int(dtype x, int p)
+__device__ static dtype pow_int(dtype x, int64_t p)
 {
     dtype r = m_one;
     switch(p) {

data/ext/cumo/include/cumo/types/uint_macro_kernel.h

@@ -18,7 +18,7 @@ __host__ __device__ static inline dtype int_reciprocal(dtype x) {
     }
 }
 
-__device__ static dtype pow_int(dtype x, int p)
+__device__ static dtype pow_int(dtype x, uint64_t p)
 {
     dtype r = m_one;
     switch(p) {
@@ -27,7 +27,6 @@ __device__ static dtype pow_int(dtype x, int p)
     case 2: return x*x;
     case 3: return x*x*x;
     }
-    if (p<0) return 0;
     while (p) {
         if (p&1) r *= x;
         x *= x;

data/ext/cumo/include/cumo.h

@@ -10,8 +10,8 @@ extern "C" {
 #endif
 #endif
 
-#define CUMO_VERSION "0.12.0"
-#define CUMO_VERSION_CODE 1200
+#define CUMO_VERSION "0.13.0"
+#define CUMO_VERSION_CODE 1300
 
 bool cumo_compatible_mode_enabled_p();
 bool cumo_show_warning_enabled_p();

data/ext/cumo/narray/array.c

@@ -449,8 +449,8 @@ cumo_na_s_array_shape(VALUE mod, VALUE ary)
   Cumo::DFloat.new_like(obj) returns DFloat instance.
 
   @overload new_like(obj)
-  @param [Numeric,Array,Cumo::NArray] obj
-  @return [Cumo::NArray]
+    @param [Numeric,Array,Cumo::NArray] obj
+    @return [Cumo::NArray]
   @example
     Cumo::NArray.new_like([[1,2,3],[4,5,6]])
     # => Cumo::Int32#shape=[2,3](empty)
@@ -490,8 +490,8 @@ cumo_na_s_array_type(VALUE mod, VALUE ary)
 /*
   Generate NArray object. NArray datatype is automatically selected.
   @overload [](elements)
-  @param [Numeric,Array] elements
-  @return [NArray]
+    @param [Numeric,Array] elements
+    @return [NArray]
 */
 static VALUE
 cumo_na_s_bracket(VALUE klass, VALUE ary)

data/ext/cumo/narray/data.c

@@ -230,9 +230,9 @@ check_axis(int axis, int ndim)
 /*
   Interchange two axes.
   @overload  swapaxes(axis1,axis2)
-  @param [Integer] axis1
-  @param [Integer] axis2
-  @return [Cumo::NArray]  view of NArray.
+    @param [Integer] axis1
+    @param [Integer] axis2
+    @return [Cumo::NArray]  view of NArray.
   @example
     x = Cumo::Int32[[1,2,3]]
 
@@ -452,9 +452,8 @@ cumo_na_check_reshape(int argc, VALUE *argv, VALUE self, size_t *shape)
   Raise exception if self is non-contiguous.
 
   @overload  reshape!(size0,size1,...)
-  @param sizeN [Integer] new shape
-  @return [Cumo::NArray] return self.
-  @example
+    @param sizeN [Integer] new shape
+    @return [Cumo::NArray] return self.
 */
 static VALUE
 cumo_na_reshape_bang(int argc, VALUE *argv, VALUE self)
@@ -510,9 +509,8 @@ cumo_na_reshape_bang(int argc, VALUE *argv, VALUE self)
   Returns a copied NArray.
 
   @overload  reshape(size0,size1,...)
-  @param sizeN [Integer] new shape
-  @return [Cumo::NArray] a copy carrying the new shape.
-  @example
+    @param sizeN [Integer] new shape
+    @return [Cumo::NArray] a copy carrying the new shape.
 */
 static VALUE
 cumo_na_reshape(int argc, VALUE *argv, VALUE self)
@@ -673,13 +671,13 @@ void cumo_na_diagonal_stride_index_kernel_launch(size_t *idx, ssize_t s0, size_t
 /*
   Returns a diagonal view of NArray
   @overload  diagonal([offset,axes])
-  @param [Integer] offset  Diagonal offset from the main diagonal.
-    The default is 0. k>0 for diagonals above the main diagonal,
-    and k<0 for diagonals below the main diagonal.
-  @param [Array] axes  Array of axes to be used as the 2-d sub-arrays
-    from which the diagonals should be taken. Defaults to last-two
-    axes ([-2,-1]).
-  @return [Cumo::NArray]  diagonal view of NArray.
+    @param [Integer] offset  Diagonal offset from the main diagonal.
+      The default is 0. k>0 for diagonals above the main diagonal,
+      and k<0 for diagonals below the main diagonal.
+    @param [Array] axes  Array of axes to be used as the 2-d sub-arrays
+      from which the diagonals should be taken. Defaults to last-two
+      axes ([-2,-1]).
+    @return [Cumo::NArray]  diagonal view of NArray.
   @example
     a = Cumo::DFloat.new(4,5).seq
     # => Cumo::DFloat#shape=[4,5]

data/ext/cumo/narray/gen/cogen_kernel.rb

@@ -52,13 +52,35 @@ TEMPLATE_KERNEL_PARTS = {
   "bit_count" => %w[reduce],
   "bit_reduce" => %w[reduce],
   "bit_stat" => %w[reduce],
+  "cond_binary" => %w[compare],
+  "cond_unary" => %w[compare],
+  "rand" => %w[misc],
+  "rand_norm" => %w[misc],
+  "poly" => %w[misc],
+  "quantize_symmetric" => %w[misc],
+  "unary_ret2" => %w[misc],
 }.freeze
 
+MATH_TEMPLATES = %w[unary_s binary_s ternary_s frexp].freeze
+COMMON_MATH = %w[sqrt log log2 log10 exp exp2 sin cos tan tanh erf gelu gelu_tanh silu sigmoid softplus].freeze
+
 abort "unknown kernel part: #{$kernel_part}" unless KERNEL_PARTS.include?($kernel_part)
+undeclared_math = COMMON_MATH - File.read(File.join(thisdir, "spec.rb")).scan(/^\s*math "(\w+)"/).flatten
+abort "COMMON_MATH names no math function in spec.rb: #{undeclared_math.join(', ')}" unless undeclared_math.empty?
 
 class ErbPP
   def kernel_parts
-    TEMPLATE_KERNEL_PARTS.fetch(get(:erb_base), %w[main])
+    erb_base = get(:erb_base)
+    return TEMPLATE_KERNEL_PARTS[erb_base] if TEMPLATE_KERNEL_PARTS.key?(erb_base)
+    return [COMMON_MATH.include?(@opts[:name]) ? "math" : "math_rare"] if MATH_TEMPLATES.include?(erb_base)
+    return @children.flat_map(&:kernel_parts).uniq unless @children.empty?
+    return [@opts[:type_name] == get(:class_name) ? "main" : "cast"] if %w[store_from store_bit].include?(erb_base)
+
+    %w[main]
+  end
+
+  def kernel_children
+    @children.select { |c| c.kernel_parts.include?($kernel_part) }
   end
 
   def kernel_part?(nan: false, from: nil)
@@ -87,7 +109,7 @@ end
 
 class DefModule
   def method_code
-    @children.select { |c| c.kernel_parts.include?($kernel_part) }.map { |c| c.result }.join("\n")
+    kernel_children.map { |c| c.result }.join("\n")
   end
 end
 

data/ext/cumo/narray/gen/erbpp2.rb

@@ -266,7 +266,13 @@ class DefMethod < ErbPP
 
   def init_def
     return if n_arg == :nodef
-    s = (singleton) ? "_singleton" : ""
+    s = if singleton
+          "_singleton"
+        elsif @opts[:private]
+          "_private"
+        else
+          ""
+        end
     "rb_define#{s}_method(#{define_method_args});"
   end
 

data/ext/cumo/narray/gen/kernel_parts.rb

@@ -1,6 +1,16 @@
 # frozen_string_literal: true
 
-KERNEL_PARTS = %w[main reduce reduce_nan reduce_from].freeze
+KERNEL_PARTS = %w[main reduce reduce_nan reduce_from math math_rare compare cast misc].freeze
+MATH_TYPES = %w[hfloat bfloat sfloat dfloat scomplex dcomplex].freeze
+
+def kernel_parts_of(type_name)
+  return KERNEL_PARTS.first(4) if type_name == "robject"
+
+  KERNEL_PARTS.reject do |part|
+    (%w[math math_rare].include?(part) && !MATH_TYPES.include?(type_name)) ||
+      (%w[compare misc].include?(part) && type_name == "bit")
+  end
+end
 
 def kernel_part_basename(type_name, part)
   part == "main" ? "#{type_name}_kernel" : "#{type_name}_kernel_#{part}"

data/ext/cumo/narray/gen/narray_def.rb

@@ -59,8 +59,8 @@ module NArrayMethod
     def_method(meth, "bit_reduce", **h)
   end
 
-  def accum(meth, dtype, result_class)
-    h = {dtype:dtype, result_class:result_class}
+  def accum(meth, dtype, result_class, **opts)
+    h = {dtype:dtype, result_class:result_class, **opts}
     def_method(meth, "accum", **h)
   end
 

data/ext/cumo/narray/gen/spec.rb

@@ -341,6 +341,16 @@ elsif is_int && !is_object
   accum "rms", "double", "cumo_cDFloat"
 end
 
+if %w[sfloat dfloat scomplex dcomplex].include?(get(:type_name))
+  %w[abs_sum abs_max abs_min].each do |name|
+    if is_complex
+      accum name, "rtype", "cRT", private: true, nan_forms: false
+    else
+      accum name, "dtype", "cT", private: true, nan_forms: false
+    end
+  end
+end
+
 if is_comparable
   accum "min", "dtype", "cT"
   accum "max", "dtype", "cT"