cumo 0.5.4 → 0.5.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0aea399ff1a36438f62f47a0a1c0350366971c34000343eb20376e8d3d448297
-  data.tar.gz: 8f8f8fab9df42cf544131dc38ac2a9ff9c2eda218a3be11eb2a887c26e86fac1
+  metadata.gz: 96a38fa0ad4be223020cef21720e30e9a32fcd2901798279293924df8bfd0f92
+  data.tar.gz: c826e49bea46001d6764c7d0f8ad124e5b6072543eb97a07990a26453dd71bf7
 SHA512:
-  metadata.gz: d57caea8974f8d7c8be8d62c0c5d52004bb414092e178459d29a5e5d3566999ce03899a5cdd82d5b3cdc0d4b690f4c1e9991cdeb40ec8a7eec5f6bd33157ec7b
-  data.tar.gz: 7cb0cd15e20c826806209ed91279694f41cc17f54aade1043ed3defd62fc9a9bdccadc0bc9345d1f802302399b915adc7fac1c2059bded270991644babce6d1a
+  metadata.gz: 1dba1ea3cb5a9c93044d222a9d03c975809c5b03d7925702aa3357a0244752d532cb9421b0632e71808944685c3ca44e7a0fc6032ca8c39d8083c279596dce61
+  data.tar.gz: 81c933a8f07f862b22a0c6be5e397000d4ea9476139203353907ba7de21a7aae418264fa8425b197d6302861a2788d262af2d2ecd2e34100900d2293706439f4

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+# 0.5.5 (2026/07/11)
+
+Fixes:
+
+* Backport: rename method: min_arg, max_arg => argmin, argmax
+* Backport: fix NArray#concatenate: avoid exception when concatenating empty arrays
+* Backport: fix boolean indexing in each axis
+
 # 0.5.4 (2026/05/30)
 
 Fixes:

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

@@ -86,6 +86,67 @@ __global__ static void reduction_kernel(cumo_na_reduction_arg_t arg, int out_blo
     }
 }
 
+// Variant of reduction_kernel for arg-reductions (argmax/argmin).
+//
+// Unlike reduction_kernel, which passes the flat 1-d index of an input element
+// (i.e. index of input elements, same as (min|max)_index), this passes the
+// index along the reduction axis (i_reduce) to the reduction impl. This matches
+// the spec of arg(min|max) which returns indices along the axis.
+template <typename TypeIn, typename TypeOut, typename ReductionImpl>
+__global__ static void reduction_arg_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(0));
+
+    extern __shared__ __align__(8) char sdata_raw[];
+    TypeReduce* sdata = reinterpret_cast<TypeReduce*>(sdata_raw);
+    unsigned int tid = threadIdx.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);
+        int64_t i_in = i_out * reduce_indexer_total_size + reduce_offset;
+
+        // Note that arg(min|max) returns the index along the reduction axis.
+        TypeReduce accum = impl.Identity(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));
+            impl.Reduce(impl.MapIn(*in_ptr, i_reduce), accum);
+        }
+
+        if (out_block_size <= max_block_size / 2) {
+            sdata[tid] = accum;
+            __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]);
+                    }
+                    __syncthreads();
+                }
+            }
+            accum = sdata[tid];
+            __syncthreads();
+        }
+        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);
+        }
+    }
+}
+
 }  // cumo_detail
 
 // TODO(sonots): Optimize indexer by squashing (or reducing) dimensions
@@ -110,4 +171,27 @@ void cumo_reduce(cumo_na_reduction_arg_t arg, ReductionImpl&& impl) {
     cumo_detail::reduction_kernel<TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, out_block_size, reduce_block_size, impl);
 }
 
+// Variant of cumo_reduce for arg-reductions (argmax/argmin), which returns
+// indices along the reduction axis. See reduction_arg_kernel above.
+template <typename TypeIn, typename TypeOut, typename ReductionImpl>
+void cumo_reduce_arg(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;
+
+    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;
+
+    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(0))) * block_size;
+
+    cumo_detail::reduction_arg_kernel<TypeIn,TypeOut,ReductionImpl><<<grid_size, block_size, shared_mem_size>>>(arg, out_block_size, reduce_block_size, impl);
+}
+
 #endif // CUMO_REDUCE_KERNEL_H

data/ext/cumo/include/cumo.h

@@ -10,7 +10,7 @@ extern "C" {
 #endif
 #endif
 
-#define CUMO_VERSION "0.5.4"
+#define CUMO_VERSION "0.5.5"
 #define CUMO_VERSION_CODE 51
 
 bool cumo_compatible_mode_enabled_p();

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

@@ -63,6 +63,10 @@ module NArrayMethod
     def_method(meth, "accum_index")
   end
 
+  def accum_arg(meth)
+    def_method(meth, "accum_arg")
+  end
+
   def cum(meth, cmacro)
     def_method(meth, "cum", cmacro:cmacro)
   end

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

@@ -333,6 +333,8 @@ if is_comparable
   accum "ptp", "dtype", "cT"
   accum_index "max_index"
   accum_index "min_index"
+  accum_arg   "argmax"
+  accum_arg   "argmin"
   def_method "minmax"
   def_module_function "maximum", "ewcomp", n_arg:2
   def_module_function "minimum", "ewcomp", n_arg:2

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

@@ -0,0 +1,120 @@
+<% (is_float ? ["","_nan"] : [""]).each do |nan| %>
+
+<%   [64,32].each do |i| %>
+<% unless type_name == 'robject' %>
+void cumo_<%=type_name%>_<%=name%>_int<%=i%>_kernel_launch(cumo_na_reduction_arg_t* arg);
+<% end %>
+
+#define idx_t int<%=i%>_t
+static void
+<%=c_iter%>_arg<%=i%><%=nan%>(cumo_na_loop_t *const lp)
+{
+    // TODO(sonots): Support nan in CUDA
+    <% if type_name == 'robject' || nan == '_nan' %>
+    {
+        size_t   n, idx;
+        char    *d_ptr, *o_ptr;
+        ssize_t  d_step;
+
+        CUMO_INIT_COUNTER(lp, n);
+        CUMO_INIT_PTR(lp, 0, d_ptr, d_step);
+        o_ptr = CUMO_NDL_PTR(lp,1);
+
+        CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("<%=name%><%=nan%>", "<%=type_name%>");
+        idx = f_<%=name[3..5]%>_index<%=nan%>(n,d_ptr,d_step);
+        *(idx_t*)o_ptr = (idx_t)idx;
+    }
+    <% else %>
+    {
+        cumo_na_reduction_arg_t arg = cumo_na_make_reduction_arg(lp);
+        cumo_<%=type_name%>_<%=name%>_int<%=i%>_kernel_launch(&arg);
+    }
+    <% end %>
+}
+#undef idx_t
+<% end;end %>
+
+/*
+  <%=name%>. Returns an index of the <%=name[3..5]%>imum value along the axis. See also `<%=name[3..5]%>_index`.
+<% if is_float %>
+  @overload <%=name%>(axis:nil, nan:false)
+  @param [TrueClass] nan  If true, apply NaN-aware algorithm (return NaN posision if exist).
+<% else %>
+  @overload <%=name%>(axis:nil)
+<% end %>
+  @param [Numeric,Array,Range] axis  Finds <%=name[3..5]%>imum values along the axis and returns indices along the axis.
+  @return [Integer,Cumo::Int] returns result indices.
+  @example
+      Cumo::NArray[3,4,1,2].<%=name%> => <% if name == 'argmin' %>2<% else %>1<% end %>
+ */
+static VALUE
+<%=c_func(-1)%>(int argc, VALUE *argv, VALUE self)
+{
+    <% if type_name == 'robject' %>
+    {
+        cumo_narray_t *na;
+        VALUE reduce;
+        cumo_ndfunc_arg_in_t ain[2] = {{Qnil,0},{cumo_sym_reduce,0}};
+        cumo_ndfunc_arg_out_t aout[1] = {{0,0,0}};
+        cumo_ndfunc_t ndf = {0, CUMO_STRIDE_LOOP_NIP|CUMO_NDF_FLAT_REDUCE|CUMO_NDF_EXTRACT, 2,1, ain,aout};
+
+        CumoGetNArray(self,na);
+        if (na->ndim==0) {
+            return INT2FIX(0);
+        }
+        if (na->size > (~(u_int32_t)0)) {
+            aout[0].type = cumo_cInt64;
+            ndf.func = <%=c_iter%>_arg64;
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, 0);
+        } else {
+            aout[0].type = cumo_cInt32;
+            ndf.func = <%=c_iter%>_arg32;
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, 0);
+        }
+
+        return cumo_na_ndloop(&ndf, 2, self, reduce);
+    }
+    <% else %>
+    {
+        cumo_narray_t *na;
+        VALUE reduce, ret;
+        cumo_ndfunc_arg_in_t ain[2] = {{Qnil,0},{cumo_sym_reduce,0}};
+        cumo_ndfunc_arg_out_t aout[1] = {{0,0,0}};
+        cumo_ndfunc_t ndf = {0, CUMO_STRIDE_LOOP_NIP|CUMO_NDF_FLAT_REDUCE|CUMO_NDF_EXTRACT|CUMO_NDF_INDEXER_LOOP, 2,1, ain,aout};
+
+        CumoGetNArray(self,na);
+        if (na->ndim==0) {
+            return INT2FIX(0);
+        }
+        if (na->size > (~(u_int32_t)0)) {
+            aout[0].type = cumo_cInt64;
+            ndf.func = <%=c_iter%>_arg64;
+            <% if is_float %>
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, <%=c_iter%>_arg64_nan);
+            <% else %>
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, 0);
+            <% end %>
+        } else {
+            aout[0].type = cumo_cInt32;
+            ndf.func = <%=c_iter%>_arg32;
+            <% if is_float %>
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, <%=c_iter%>_arg32_nan);
+            <% else %>
+            reduce = cumo_na_reduce_dimension(argc, argv, 1, &self, &ndf, 0);
+            <% end %>
+        }
+
+        if (cumo_na_has_idx_p(self)) {
+            VALUE copy = cumo_na_copy(self); // reduction does not support idx, make contiguous
+            ret = cumo_na_ndloop(&ndf, 2, copy, reduce);
+        } else {
+            ret = cumo_na_ndloop(&ndf, 2, self, reduce);
+        }
+        if (cumo_compatible_mode_enabled_p()) {
+            return rb_funcall(ret, rb_intern("extract_cpu"), 0);
+        } else {
+            return ret;
+        }
+    }
+    <% end %>
+}

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

@@ -0,0 +1,66 @@
+<% unless defined?($cumo_narray_gen_tmpl_accum_arg_kernel_included) %>
+<% $cumo_narray_gen_tmpl_accum_arg_kernel_included = 1 %>
+<% unless type_name == 'robject' %>
+
+<%   [64,32].each do |i| %>
+#define idx_t int<%=i%>_t
+
+#if defined(__cplusplus)
+#if 0
+{ /* satisfy cc-mode */
+#endif
+}  /* extern "C" { */
+#endif
+
+struct cumo_<%=type_name%>_argmin_int<%=i%>_impl {
+    struct MinAndArgMin {
+        dtype min;
+        idx_t argmin;
+    };
+    __device__ MinAndArgMin Identity(idx_t index) { return {DATA_MAX, index}; }
+    __device__ MinAndArgMin MapIn(dtype in, idx_t index) { return {in, index}; }
+    __device__ void Reduce(MinAndArgMin next, MinAndArgMin& accum) {
+        if (accum.min > next.min) {
+            accum = next;
+        }
+    }
+    __device__ idx_t MapOut(MinAndArgMin accum) { return accum.argmin; }
+};
+
+struct cumo_<%=type_name%>_argmax_int<%=i%>_impl {
+    struct MaxAndArgMax {
+        dtype max;
+        idx_t argmax;
+    };
+    __device__ MaxAndArgMax Identity(idx_t index) { return {DATA_MIN, index}; }
+    __device__ MaxAndArgMax MapIn(dtype in, idx_t index) { return {in, index}; }
+    __device__ void Reduce(MaxAndArgMax next, MaxAndArgMax& accum) {
+        if (accum.max < next.max) {
+            accum = next;
+        }
+    }
+    __device__ idx_t MapOut(MaxAndArgMax accum) { return accum.argmax; }
+};
+
+#if defined(__cplusplus)
+extern "C" {
+#if 0
+} /* satisfy cc-mode */
+#endif
+#endif
+
+void cumo_<%=type_name%>_argmin_int<%=i%>_kernel_launch(cumo_na_reduction_arg_t* arg)
+{
+    cumo_reduce_arg<dtype, idx_t, cumo_<%=type_name%>_argmin_int<%=i%>_impl>(*arg, cumo_<%=type_name%>_argmin_int<%=i%>_impl{});
+}
+
+void cumo_<%=type_name%>_argmax_int<%=i%>_kernel_launch(cumo_na_reduction_arg_t* arg)
+{
+    cumo_reduce_arg<dtype, idx_t, cumo_<%=type_name%>_argmax_int<%=i%>_impl>(*arg, cumo_<%=type_name%>_argmax_int<%=i%>_impl{});
+}
+
+#undef idx_t
+<% end %>
+
+<% end %>
+<% end %>

data/ext/cumo/narray/index.c

@@ -63,6 +63,7 @@ static ID cumo_id_dup;
 static ID cumo_id_bracket;
 static ID cumo_id_shift_left;
 static ID cumo_id_mask;
+static ID cumo_id_where;
 
 
 static void
@@ -159,6 +160,13 @@ cumo_na_parse_narray_index(VALUE a, int orig_dim, ssize_t size, cumo_na_index_ar
     if (CUMO_NA_NDIM(na) != 1) {
         rb_raise(rb_eIndexError, "should be 1-d NArray");
     }
+    if (rb_obj_class(a) == cumo_cBit) {
+        if (CUMO_NA_SIZE(na) != (size_t)size) {
+            rb_raise(rb_eIndexError, "Bit-NArray size mismatch");
+        }
+        a = rb_funcall(a, cumo_id_where, 0);
+        CumoGetNArray(a,na);
+    }
     n = CUMO_NA_SIZE(na);
     idx = cumo_na_new(cIndex,1,&n);
     cumo_na_store(idx,a);
@@ -1138,4 +1146,5 @@ Init_cumo_na_index()
     cumo_id_bracket     = rb_intern("[]");
     cumo_id_shift_left  = rb_intern("<<");
     cumo_id_mask        = rb_intern("mask");
+    cumo_id_where       = rb_intern("where");
 }

data/lib/cumo/narray/extra.rb

@@ -473,8 +473,10 @@ module Cumo
       arrays.each do |a|
         fst = lst
         lst = fst + (a.shape[axis - nd] || 1)
-        refs[axis] = fst...lst
-        result[*refs] = a
+        if lst > fst
+          refs[axis] = fst...lst
+          result[*refs] = a
+        end
       end
       result
     end
@@ -654,13 +656,17 @@ module Cumo
       result = self.class.zeros(*self_shape)
       lst = shape[axis]
       refs = [true] * ndim
-      refs[axis] = 0...lst
-      result[*refs] = self
+      if lst > 0
+        refs[axis] = 0...lst
+        result[*refs] = self
+      end
       arrays.each do |a|
         fst = lst
         lst = fst + (a.shape[axis - ndim] || 1)
-        refs[axis] = fst...lst
-        result[*refs] = a
+        if lst > fst
+          refs[axis] = fst...lst
+          result[*refs] = a
+        end
       end
       result
     end

data/test/narray_test.rb

@@ -344,11 +344,21 @@ class NArrayTest < Test::Unit::TestCase
           assert { a.min_index(axis: 0) == [0, 1, 2] }
           assert { a.max_index(axis: 1) == [2, 5] }
           assert { a.max_index(axis: 0) == [3, 4, 5] }
+          assert { a.argmin == 0 }
+          assert { a.argmax == 5 }
+          assert { a.argmin(axis: 1) == [0, 0] }
+          assert { a.argmin(axis: 0) == [0, 0, 0] }
+          assert { a.argmax(axis: 1) == [2, 2] }
+          assert { a.argmax(axis: 0) == [1, 1, 1] }
           assert { (a >= 3) == [[0, 0, 1], [1, 1, 1]] }
           assert { (a >  3) == [[0, 0, 0], [1, 1, 1]] }
           assert { (a <= 3) == [[1, 1, 1], [0, 0, 0]] }
           assert { (a <  3) == [[1, 1, 0], [0, 0, 0]] }
           assert { (a.eq 3) == [[0, 0, 1], [0, 0, 0]] }
+          assert { a[a.ne 3] == [1, 2, 5, 7, 11] }
+          assert { a[a[true, 2] < 5, true] == [[1, 2, 3]] }
+          assert { a[true, a[1, true] > 5] == [[2, 3], [7, 11]] }
+          assert { a[:*, (a[0, :*] % 2).eq(1)] == [[1, 3], [5, 11]] }
           assert { a.sort == src }
           assert { a.sort_index == [[0, 1, 2], [3, 4, 5]] }
           assert { a.percentile(0) == 1.0 }
@@ -716,6 +726,8 @@ class NArrayTest < Test::Unit::TestCase
                    [4, 1, 6]]]
         assert { a.max_index(2) == [[1, 5, 8], [9, 12, 17]] }
         assert { a.max(2) == [[8, 6, 5], [7, 9, 6]] }
+        assert { a.argmax(2) == [[1, 2, 2], [0, 0, 2]] }
+        assert { a.argmin(2) == [[2, 0, 0], [2, 2, 1]] }
 
         unless [Cumo::UInt64, Cumo::UInt32, Cumo::UInt16, Cumo::UInt8].include?(dtype)
           a = dtype[[[-6, -8, -5],
@@ -839,4 +851,40 @@ class NArrayTest < Test::Unit::TestCase
     assert { Cumo::SComplex[1].rms == 1.0 }
     assert { Cumo::DComplex[1].rms == 1.0 }
   end
+
+  test "concatenate with empty arrays" do
+    a = Cumo::DFloat[1, 2, 3]
+    empty = Cumo::DFloat[]
+    assert { Cumo::NArray.concatenate([empty, a]) == [1, 2, 3] }
+    assert { Cumo::NArray.concatenate([a, empty]) == [1, 2, 3] }
+    assert { Cumo::NArray.concatenate([empty, empty]) == [] }
+    assert { a.concatenate(empty) == [1, 2, 3] }
+    assert { empty.concatenate(a) == [1, 2, 3] }
+    assert { empty.concatenate(empty) == [] }
+  end
+
+  test "argmax/argmin" do
+    [Cumo::DFloat, Cumo::Int32, Cumo::UInt8].each do |dtype|
+      a = dtype[3, 4, 1, 2]
+      assert { a.argmax == 1 }
+      assert { a.argmin == 2 }
+
+      b = dtype[[3, 4, 1], [2, 0, 5]]
+      assert { b.argmax == 5 }
+      assert { b.argmin == 4 }
+      assert { b.argmax(axis: 1) == [1, 2] }
+      assert { b.argmax(axis: 0) == [0, 0, 1] }
+      assert { b.argmin(axis: 1) == [2, 1] }
+      assert { b.argmin(axis: 0) == [1, 1, 0] }
+      assert { b.at(b.argmax(axis: 0), 0..-1) == [3, 4, 5] }
+      assert { b.at(b.argmin(axis: 0), 0..-1) == [2, 0, 1] }
+    end
+
+    # NaN-aware (nan:true returns the NaN position; nan:false ignores NaN)
+    c = Cumo::DFloat[3.0, Float::NAN, 1.0, 5.0]
+    assert { c.argmax == 3 }
+    assert { c.argmax(nan: true) == 1 }
+    assert { c.argmin == 2 }
+    assert { c.argmin(nan: true) == 1 }
+  end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: cumo
 version: !ruby/object:Gem::Version
-  version: 0.5.4
+  version: 0.5.5
 platform: ruby
 authors:
 - Naotoshi Seo
@@ -157,6 +157,8 @@ files:
 - ext/cumo/narray/gen/narray_def.rb
 - ext/cumo/narray/gen/spec.rb
 - ext/cumo/narray/gen/tmpl/accum.c
+- ext/cumo/narray/gen/tmpl/accum_arg.c
+- ext/cumo/narray/gen/tmpl/accum_arg_kernel.cu
 - ext/cumo/narray/gen/tmpl/accum_binary.c
 - ext/cumo/narray/gen/tmpl/accum_binary_kernel.cu
 - ext/cumo/narray/gen/tmpl/accum_index.c
@@ -331,7 +333,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.10
+rubygems_version: 4.0.15
 specification_version: 4
 summary: Cumo is CUDA aware numerical library whose interface is highly compatible
   with Ruby Numo