carray 2.0.0 → 3.0.0

data/ext/carray_bincount.c

@@ -0,0 +1,255 @@
+/* ---------------------------------------------------------------------------
+
+  carray_bincount.c — dedicated bincount kernels (count-only + weighted)
+
+  Tight per-element scatter into a pre-sized 1-D output buffer.  The
+  inner loop reads labels in their native integer dtype (no cast to
+  int64), skips per-iteration bounds checking (caller validates label
+  range), and emits the output dtype directly.
+
+  Output dtype:
+    - count-only:   UInt32 if length < 2^32, else UInt64.
+    - weighted:     weights.data_type.
+
+  Mask: a masked label is skipped; for the weighted kernel a masked
+  weight is also skipped (its label contributes 0).
+
+  Caller contract (= lib/carray/methods/bincount.rb's CArray#bincount):
+    - self is integer dtype.
+    - length is the output size, pre-sized to max(length, label_max+1)
+      with label_min >= 0 already verified.
+
+  Public Ruby surface is CArray#bincount(weights:, length:) in
+  lib/carray/methods/bincount.rb; this file provides the private
+  kernels __bincount_count__(length) and
+  __bincount_weighted__(weights, length).
+
+  Design: devel/PROPOSAL_BINCOUNT_DEDICATED_KERNEL.md.
+
+--------------------------------------------------------------------------- */
+
+#include "carray.h"
+
+/* Tight inner loop: 2 mask-aware variants × 8 label dtypes × 2 output
+   dtypes = 32 specializations.  Generated via macro expansion.
+
+   Layout:
+     COUNT_KERNEL(LABEL_T, OUT_T)
+       no-mask path: for i; out[label[i]]++
+       label-mask path: for i; if (mlabel[i]) continue; out[label[i]]++
+*/
+
+#define COUNT_KERNEL(LABEL_T, OUT_T) do { \
+  const LABEL_T *lp = (const LABEL_T *) cl->ptr; \
+  OUT_T         *op = (OUT_T *)         co->ptr; \
+  if ( mlabel ) { \
+    for (i = 0; i < n; i++) { \
+      if ( mlabel[i] ) continue; \
+      op[(size_t) lp[i]]++; \
+    } \
+  } \
+  else { \
+    for (i = 0; i < n; i++) { \
+      op[(size_t) lp[i]]++; \
+    } \
+  } \
+} while (0)
+
+#define COUNT_DISPATCH_LABEL(OUT_T) do { \
+  switch ( cl->data_type ) { \
+  case CA_INT8:    COUNT_KERNEL(int8_t,   OUT_T); break; \
+  case CA_INT16:   COUNT_KERNEL(int16_t,  OUT_T); break; \
+  case CA_INT32:   COUNT_KERNEL(int32_t,  OUT_T); break; \
+  case CA_INT64:   COUNT_KERNEL(int64_t,  OUT_T); break; \
+  case CA_UINT8:   COUNT_KERNEL(uint8_t,  OUT_T); break; \
+  case CA_UINT16:  COUNT_KERNEL(uint16_t, OUT_T); break; \
+  case CA_UINT32:  COUNT_KERNEL(uint32_t, OUT_T); break; \
+  case CA_UINT64:  COUNT_KERNEL(uint64_t, OUT_T); break; \
+  default: \
+    rb_raise(rb_eCADataTypeError, \
+             "bincount: integer label array required (got %d)", \
+             cl->data_type); \
+  } \
+} while (0)
+
+/* __bincount_count__(length) -- count occurrences of each label in
+   self.  Allocates a zero-filled UInt32 (or UInt64 if length >= 2^32)
+   output of size `length`, then runs the 8-way label dispatch. */
+static VALUE
+rb_ca_bincount_count_kernel (VALUE self, VALUE rlength)
+{
+  CArray   *cl, *co;
+  VALUE     vout;
+  ca_size_t i, n, length;
+  ca_size_t shape_out[1];
+  boolean8_t *mlabel;
+  int out_type;
+
+  TypedData_Get_Struct(self, CArray, &carray_data_type, cl);
+  length = (ca_size_t) NUM2SIZET(rlength);
+  if ( length < 0 ) {
+    rb_raise(rb_eArgError, "bincount: length must be non-negative");
+  }
+
+  /* Output dtype: UInt32 default; UInt64 if length doesn't fit. */
+  out_type = (length > 0xFFFFFFFFLL) ? CA_UINT64 : CA_UINT32;
+  shape_out[0] = length;
+  vout = rb_carray_new(out_type, 1, shape_out, 0, NULL);
+  TypedData_Get_Struct(vout, CArray, &carray_data_type, co);
+
+  /* Zero-fill output (we own it, just allocated). */
+  memset(co->ptr, 0, (size_t) co->bytes * (size_t) co->elements);
+
+  if ( cl->elements == 0 ) {
+    return vout;
+  }
+
+  ca_attach(cl);
+  n      = cl->elements;
+  mlabel = cl->mask ? (boolean8_t *) cl->mask->ptr : NULL;
+
+  if ( out_type == CA_UINT32 ) {
+    COUNT_DISPATCH_LABEL(uint32_t);
+  }
+  else {
+    COUNT_DISPATCH_LABEL(uint64_t);
+  }
+
+  ca_detach(cl);
+
+  return vout;
+}
+
+#undef COUNT_KERNEL
+#undef COUNT_DISPATCH_LABEL
+
+/* --------------------------------------------------------------- */
+
+/* Weighted variant: output dtype = weights dtype.
+   Inner: out[label[i]] += weight[i].
+   Mask: skip if label[i] masked OR weight[i] masked. */
+
+#define WEIGHTED_KERNEL(LABEL_T, W_T) do { \
+  const LABEL_T *lp = (const LABEL_T *) cl->ptr; \
+  const W_T     *wp = (const W_T *)     cw->ptr; \
+  W_T           *op = (W_T *)           co->ptr; \
+  if ( mlabel && mweight ) { \
+    for (i = 0; i < n; i++) { \
+      if ( mlabel[i] || mweight[i] ) continue; \
+      op[(size_t) lp[i]] += wp[i]; \
+    } \
+  } \
+  else if ( mlabel ) { \
+    for (i = 0; i < n; i++) { \
+      if ( mlabel[i] ) continue; \
+      op[(size_t) lp[i]] += wp[i]; \
+    } \
+  } \
+  else if ( mweight ) { \
+    for (i = 0; i < n; i++) { \
+      if ( mweight[i] ) continue; \
+      op[(size_t) lp[i]] += wp[i]; \
+    } \
+  } \
+  else { \
+    for (i = 0; i < n; i++) { \
+      op[(size_t) lp[i]] += wp[i]; \
+    } \
+  } \
+} while (0)
+
+#define WEIGHTED_DISPATCH_LABEL(W_T) do { \
+  switch ( cl->data_type ) { \
+  case CA_INT8:    WEIGHTED_KERNEL(int8_t,   W_T); break; \
+  case CA_INT16:   WEIGHTED_KERNEL(int16_t,  W_T); break; \
+  case CA_INT32:   WEIGHTED_KERNEL(int32_t,  W_T); break; \
+  case CA_INT64:   WEIGHTED_KERNEL(int64_t,  W_T); break; \
+  case CA_UINT8:   WEIGHTED_KERNEL(uint8_t,  W_T); break; \
+  case CA_UINT16:  WEIGHTED_KERNEL(uint16_t, W_T); break; \
+  case CA_UINT32:  WEIGHTED_KERNEL(uint32_t, W_T); break; \
+  case CA_UINT64:  WEIGHTED_KERNEL(uint64_t, W_T); break; \
+  default: \
+    rb_raise(rb_eCADataTypeError, \
+             "bincount: integer label array required (got %d)", \
+             cl->data_type); \
+  } \
+} while (0)
+
+/* __bincount_weighted__(weights, length) -- sum `weights[i]` into
+   `out[label[i]]`.  Allocates a zero-filled output of weights.data_type
+   and size `length`, then runs the 8 label × 10 weight dispatch
+   (integer + float; the inner WEIGHTED_KERNEL macro branches over the
+   four mask combinations). */
+static VALUE
+rb_ca_bincount_weighted_kernel (VALUE self, VALUE rweights, VALUE rlength)
+{
+  CArray   *cl, *cw, *co;
+  VALUE     vout;
+  ca_size_t i, n, length;
+  ca_size_t shape_out[1];
+  boolean8_t *mlabel, *mweight;
+  int w_type;
+
+  TypedData_Get_Struct(self, CArray, &carray_data_type, cl);
+  TypedData_Get_Struct(rweights, CArray, &carray_data_type, cw);
+  length = (ca_size_t) NUM2SIZET(rlength);
+  if ( length < 0 ) {
+    rb_raise(rb_eArgError, "bincount: length must be non-negative");
+  }
+  if ( cw->elements != cl->elements ) {
+    rb_raise(rb_eArgError,
+             "bincount: weights length (%lld) doesn't match labels length (%lld)",
+             (long long) cw->elements, (long long) cl->elements);
+  }
+
+  w_type = cw->data_type;
+  shape_out[0] = length;
+  vout = rb_carray_new(w_type, 1, shape_out, 0, NULL);
+  TypedData_Get_Struct(vout, CArray, &carray_data_type, co);
+
+  memset(co->ptr, 0, (size_t) co->bytes * (size_t) co->elements);
+
+  if ( cl->elements == 0 ) {
+    return vout;
+  }
+
+  ca_attach_n(2, cl, cw);
+  n       = cl->elements;
+  mlabel  = cl->mask ? (boolean8_t *) cl->mask->ptr : NULL;
+  mweight = cw->mask ? (boolean8_t *) cw->mask->ptr : NULL;
+
+  switch ( w_type ) {
+  case CA_FLOAT64: WEIGHTED_DISPATCH_LABEL(double);   break;
+  case CA_FLOAT32: WEIGHTED_DISPATCH_LABEL(float);    break;
+  case CA_INT64:   WEIGHTED_DISPATCH_LABEL(int64_t);  break;
+  case CA_INT32:   WEIGHTED_DISPATCH_LABEL(int32_t);  break;
+  case CA_INT16:   WEIGHTED_DISPATCH_LABEL(int16_t);  break;
+  case CA_INT8:    WEIGHTED_DISPATCH_LABEL(int8_t);   break;
+  case CA_UINT64:  WEIGHTED_DISPATCH_LABEL(uint64_t); break;
+  case CA_UINT32:  WEIGHTED_DISPATCH_LABEL(uint32_t); break;
+  case CA_UINT16:  WEIGHTED_DISPATCH_LABEL(uint16_t); break;
+  case CA_UINT8:   WEIGHTED_DISPATCH_LABEL(uint8_t);  break;
+  default:
+    ca_detach_n(2, cl, cw);
+    rb_raise(rb_eCADataTypeError,
+             "bincount: weights must be numeric (got %d)", w_type);
+  }
+
+  ca_detach_n(2, cl, cw);
+
+  return vout;
+}
+
+#undef WEIGHTED_KERNEL
+#undef WEIGHTED_DISPATCH_LABEL
+
+/* --------------------------------------------------------------- */
+
+void
+Init_carray_bincount (void)
+{
+  rb_define_private_method(rb_cCArray, "__bincount_count__",
+                           rb_ca_bincount_count_kernel, 1);
+  rb_define_private_method(rb_cCArray, "__bincount_weighted__",
+                           rb_ca_bincount_weighted_kernel, 2);
+}

data/ext/carray_broadcast.c

@@ -0,0 +1,283 @@
+/* ---------------------------------------------------------------------------
+
+  Implicit broadcasting helpers and CArray#broadcast_to.
+
+    ca_broadcast_view (src, ndim, target_dim)
+        Wrap `src` in a CAStride view whose shape is `target_dim`,
+        treating size-1 axes of `src` as broadcast (stride 0).
+        Returns `src` unchanged when no expansion is needed.
+        Called by ca_broadcast_pair below and by the binop dispatcher
+        in ext/ca_obj_binop.c.
+
+    ca_broadcast_pair (&self, &other)
+        Two-sided expansion for binary ops (case A only: same ndim,
+        size-1 axes broadcast pairwise).  Leaves both operands
+        unchanged when shapes are already equal or when they are
+        incompatible (the caller's existing elements-mismatch raise
+        path then handles the error).  Called by carray_cast.c's
+        coercion path and by the binop dispatcher.
+
+    rb_ca_broadcast_to -- backs the public CArray#broadcast_to method
+        (right-to-left axis pairing; see the docstring at the function).
+
+  Case B (cross-ndim dim-prepending) is not handled by the implicit
+  helpers; users with cross-ndim operands keep using the explicit :*
+  form or #broadcast_to (which does accept cross-ndim with target axes
+  pinned to size 1).  See PROPOSAL_BROADCASTING_AND_UNBOUND.md.
+
+---------------------------------------------------------------------------- */
+
+#include "carray.h"
+
+/* Build a CAStride view of `src` with shape `target_dim`.  For each
+   axis where src.dim[i] == target_dim[i], inherit the row-major byte
+   stride; for src.dim[i] == 1 and target_dim[i] > 1, use stride 0.
+   Otherwise, raise.  If no axis needs expansion, return src as-is. */
+VALUE
+ca_broadcast_view (VALUE src, int8_t ndim, ca_size_t *target_dim)
+{
+  CArray *cs;
+  ca_size_t src_strides[CA_RANK_MAX];
+  ca_size_t new_strides[CA_RANK_MAX];
+  ca_size_t s;
+  int needs_broadcast = 0;
+  int i;
+  CAStride *view;
+  volatile VALUE obj;
+
+  TypedData_Get_Struct(src, CArray, &carray_data_type, cs);
+
+  if (cs->ndim != ndim) {
+    rb_raise(rb_eRuntimeError,
+             "broadcast: ndim mismatch (%d vs %d)",
+             (int) cs->ndim, (int) ndim);
+  }
+
+  /* Row-major byte strides for src treated as contiguous.  When src is
+     itself a CAStride, the compose-fold path at attach time collapses
+     the chain to root, so passing parent-relative row-major strides
+     is correct. */
+  s = cs->bytes;
+  for (i = ndim - 1; i >= 0; i--) {
+    src_strides[i] = s;
+    s *= cs->dim[i];
+  }
+
+  for (i = 0; i < ndim; i++) {
+    if (cs->dim[i] == target_dim[i]) {
+      new_strides[i] = src_strides[i];
+    }
+    else if (cs->dim[i] == 1) {
+      new_strides[i] = 0;
+      needs_broadcast = 1;
+    }
+    else {
+      rb_raise(rb_eRuntimeError,
+               "broadcast: cannot broadcast axis %d (%lld vs %lld)",
+               i, (long long) cs->dim[i], (long long) target_dim[i]);
+    }
+  }
+
+  if (!needs_broadcast) return src;
+
+  view = ca_stride_new(CA_OBJ_STRIDE, cs,
+                       cs->data_type, cs->bytes,
+                       ndim, target_dim, new_strides, 0);
+  ca_set_flag(view, CA_FLAG_READ_ONLY);
+  if (view->mask) {
+    ca_set_flag(view->mask, CA_FLAG_READ_ONLY);
+  }
+  obj = ca_wrap_struct(view);
+  rb_ca_set_parent(obj, src);
+  return obj;
+}
+
+/* Two-sided broadcast for binary ops (case A only).  When both
+   operands are non-scalar with equal ndim and at least one axis pair
+   is (1, N) / (N, 1) / (1, 1), wrap each in a CAStride with stride-0
+   axes so the downstream iterator sees matched shapes.  If shapes are
+   pairwise equal, no-op.  If any axis pair is incompatible (neither
+   side is 1 nor equal), no-op (caller's existing element-count check
+   raises). */
+void
+ca_broadcast_pair (volatile VALUE *self, volatile VALUE *other)
+{
+  CArray *ca, *cb;
+  ca_size_t target_dim[CA_RANK_MAX];
+  int can_broadcast = 1;
+  int needs_broadcast = 0;
+  int i;
+
+  TypedData_Get_Struct(*self,  CArray, &carray_data_type, ca);
+  TypedData_Get_Struct(*other, CArray, &carray_data_type, cb);
+
+  if (ca_is_scalar(ca) || ca_is_scalar(cb)) return;
+  if (ca->ndim != cb->ndim) return;
+  if (ca->ndim == 0) return;
+
+  for (i = 0; i < ca->ndim; i++) {
+    if (ca->dim[i] == cb->dim[i]) {
+      target_dim[i] = ca->dim[i];
+    }
+    else if (ca->dim[i] == 1) {
+      target_dim[i] = cb->dim[i];
+      needs_broadcast = 1;
+    }
+    else if (cb->dim[i] == 1) {
+      target_dim[i] = ca->dim[i];
+      needs_broadcast = 1;
+    }
+    else {
+      can_broadcast = 0;
+      break;
+    }
+  }
+
+  if (!can_broadcast || !needs_broadcast) return;
+
+  *self  = ca_broadcast_view(*self,  ca->ndim, target_dim);
+  *other = ca_broadcast_view(*other, ca->ndim, target_dim);
+}
+
+/* CArray#broadcast_to(*newdim) -- pairs axes right-to-left:
+     - source axis == target axis  -> data axis (inherit stride)
+     - target axis == 1            -> size-1 collapse, source consumed
+                                      only when it is also 1
+     - source axis == 1            -> broadcast (stride 0, source
+                                      consumed)
+     - source exhausted, target 1  -> insert size-1 axis
+     - otherwise                   -> RuntimeError
+
+   Intentionally conservative on cross-ndim: prepended axes are only
+   inserted as size-1.  Target > 1 on a prepended axis raises, in
+   keeping with PROPOSAL_BROADCASTING_AND_UNBOUND.md's strict-ndim
+   policy.
+
+   Result is tagged CA_OBJ_REPEAT (CARepeat) rather than plain CAStride
+   so introspection (`view.class`) surfaces the broadcast intent. */
+
+static VALUE
+rb_ca_broadcast_to (int argc, VALUE *argv, VALUE self)
+{
+  CArray *cs;
+  ca_size_t target_dim[CA_RANK_MAX];
+  ca_size_t new_strides[CA_RANK_MAX];
+  ca_size_t src_strides[CA_RANK_MAX];
+  ca_size_t s;
+  int8_t target_ndim;
+  int needs_view = 0;
+  int t_idx, s_idx;
+  int i;
+  CAStride *view;
+  volatile VALUE obj;
+
+  TypedData_Get_Struct(self, CArray, &carray_data_type, cs);
+
+  if (argc < 0 || argc > CA_RANK_MAX) {
+    rb_raise(rb_eArgError, "broadcast_to: invalid number of dims (%d)", argc);
+  }
+  target_ndim = (int8_t) argc;
+
+  if (target_ndim < cs->ndim) {
+    rb_raise(rb_eRuntimeError,
+             "broadcast_to: target ndim %d smaller than source ndim %d",
+             (int) target_ndim, (int) cs->ndim);
+  }
+
+  for (i = 0; i < target_ndim; i++) {
+    target_dim[i] = NUM2SIZE(argv[i]);
+    if (target_dim[i] < 0) {
+      rb_raise(rb_eArgError,
+               "broadcast_to: negative dim at axis %d", i);
+    }
+  }
+
+  /* CScalar source: build an all-stride-0 CARepeat view of any
+     target shape.  Generalises the cross-ndim rule, which is strict
+     for non-scalar sources (each source axis must pair with a target
+     axis); for a scalar the natural extension is "any target shape,
+     stride 0 everywhere". */
+  if (ca_is_scalar(cs)) {
+    for (i = 0; i < target_ndim; i++) {
+      new_strides[i] = 0;
+    }
+    view = ca_stride_new(CA_OBJ_REPEAT, cs,
+                         cs->data_type, cs->bytes,
+                         target_ndim, target_dim, new_strides, 0);
+    ca_set_flag(view, CA_FLAG_READ_ONLY);
+    if (view->mask) {
+      ca_set_flag(view->mask, CA_FLAG_READ_ONLY);
+    }
+    obj = ca_wrap_struct(view);
+    rb_ca_set_parent(obj, self);
+    return obj;
+  }
+
+  /* Row-major byte strides for source (interpreted as contiguous;
+     compose-fold at attach time handles non-contiguous parents). */
+  s = cs->bytes;
+  for (i = cs->ndim - 1; i >= 0; i--) {
+    src_strides[i] = s;
+    s *= cs->dim[i];
+  }
+
+  /* Pair right-to-left.  s_idx tracks the next source axis to be
+     consumed; -1 means source is exhausted. */
+  t_idx = target_ndim - 1;
+  s_idx = cs->ndim - 1;
+
+  while (t_idx >= 0) {
+    ca_size_t dd = target_dim[t_idx];
+    ca_size_t sd = (s_idx >= 0) ? cs->dim[s_idx] : -1;
+
+    if (sd == dd) {
+      new_strides[t_idx] = src_strides[s_idx];
+      s_idx--;
+    }
+    else if (dd == 1) {
+      new_strides[t_idx] = 0;
+      if (sd == 1) s_idx--;
+      needs_view = 1;
+    }
+    else if (sd == 1) {
+      new_strides[t_idx] = 0;
+      s_idx--;
+      needs_view = 1;
+    }
+    else {
+      /* sd == -1 (source exhausted) with dd > 1, or sd > 1 with
+         dd > 1 and not equal: cannot broadcast. */
+      rb_raise(rb_eRuntimeError,
+               "broadcast_to: cannot broadcast axis %d "
+               "(source %s, target %lld)",
+               t_idx,
+               (sd < 0 ? "(exhausted)" : "non-1"),
+               (long long) dd);
+    }
+    t_idx--;
+  }
+
+  if (s_idx >= 0) {
+    rb_raise(rb_eRuntimeError,
+             "broadcast_to: %d source axes left unmatched", s_idx + 1);
+  }
+
+  if (!needs_view) return self;
+
+  view = ca_stride_new(CA_OBJ_REPEAT, cs,
+                       cs->data_type, cs->bytes,
+                       target_ndim, target_dim, new_strides, 0);
+  ca_set_flag(view, CA_FLAG_READ_ONLY);
+  if (view->mask) {
+    ca_set_flag(view->mask, CA_FLAG_READ_ONLY);
+  }
+  obj = ca_wrap_struct(view);
+  rb_ca_set_parent(obj, self);
+  return obj;
+}
+
+void
+Init_carray_broadcast (void)
+{
+  rb_define_method(rb_cCArray, "broadcast_to", rb_ca_broadcast_to, -1);
+}