carray 3.0.1 → 3.0.3

data/ext/ca_sweep_engine.c

@@ -22,45 +22,161 @@
 #include "ca_for_buffer.h"
 #include <string.h>
 
+/* "[2, 3]" for an operand's shape. */
+static VALUE
+ca_sweep_shape_str (CArray *ca)
+{
+  volatile VALUE s = rb_str_new_cstr("[");
+  int k;
+  for (k = 0; k < ca->ndim; k++) {
+    if (k > 0) rb_str_cat_cstr(s, ", ");
+    rb_str_catf(s, "%lld", (long long) ca->dim[k]);
+  }
+  rb_str_cat_cstr(s, "]");
+  return s;
+}
+
 void
-ca_sweep_acquire (ca_sweep_state_t *st)
+ca_sweep_refuse_shapes (CArray *a, CArray *b)
 {
+  volatile VALUE sa = ca_sweep_shape_str(a);
+  volatile VALUE sb = ca_sweep_shape_str(b);
+  rb_raise(rb_eArgError,
+           "shape mismatch between operands (%s and %s); a scalar pairs "
+           "with any array, two arrays only when their shapes agree",
+           StringValueCStr(sa), StringValueCStr(sb));
+}
+
+int
+ca_sweep_same_shape (CArray *a, CArray *b)
+{
+  int k;
+  if (a->ndim != b->ndim) return 0;
+  for (k = 0; k < a->ndim; k++) {
+    if (a->dim[k] != b->dim[k]) return 0;
+  }
+  return 1;
+}
+
+/* Pair the operands: a scalar walks with stride 0, every other operand
+ * must have the shape of the first non-scalar one, which is returned (NULL
+ * when all operands are scalars).  Sets stride[] and n_kernel.  Allocates
+ * nothing, so it runs before any operand is acquired. */
+static CArray *
+ca_sweep_pair_operands (ca_sweep_state_t *st)
+{
+  CArray *donor = NULL;
   int k_op;
-  int any_input_mask = 0;
+  st->n_kernel = 1;
+  for (k_op = 0; k_op < st->n_ops; k_op++) {
+    CArray *ca = st->cx[k_op];
+    if (ca_is_scalar(ca)) {
+      st->stride[k_op] = 0;
+      continue;
+    }
+    st->stride[k_op] = ca->bytes;
+    if (!donor) {
+      donor        = ca;
+      st->n_kernel = ca->elements;
+    } else if (!ca_sweep_same_shape(donor, ca)) {
+      ca_sweep_refuse_shapes(donor, ca);
+    }
+  }
+  return donor;
+}
 
-  if ((int) strlen(st->fsync) != st->n_ops) {
-    rb_raise(rb_eRuntimeError,
-             "[BUG] invalid length of fsync arg in %s (expected %d)",
-             st->src_label ? st->src_label : "ca_sweep_acquire",
-             st->n_ops);
+/* Non-zero when INPUT k_op lives in the same root as an OUTPUT other
+ * than itself.  Read in place, or chunk by chunk after earlier chunks
+ * were written, it would show the kernel cells already written; such an
+ * input is read whole before the walk.  An input that is the output
+ * itself is read in place: each cell is read before it is written. */
+static int
+ca_sweep_input_overlaps_output (ca_sweep_state_t *st, int k_op)
+{
+  int j;
+  CArray *ca = st->cx[k_op];
+  if (st->fsync[k_op] != '0' || ca_is_scalar(ca)) return 0;
+  for (j = 0; j < st->n_ops; j++) {
+    if (st->fsync[j] == '1' && st->cx[j] != ca &&
+        ca_shares_root(ca, st->cx[j])) {
+      return 1;
+    }
   }
+  return 0;
+}
 
+/* Non-zero when any INPUT operand carries a mask.  Asks, allocates
+ * nothing. */
+static int
+ca_sweep_any_input_mask (ca_sweep_state_t *st)
+{
+  int k_op;
   for (k_op = 0; k_op < st->n_ops; k_op++) {
-    st->base[k_op]      = NULL;
+    if (st->fsync[k_op] == '0' && ca_has_mask(st->cx[k_op])) return 1;
+  }
+  return 0;
+}
+
+/* Give back everything acquire has taken so far, without syncing: detach
+ * attached operands, free owned scratch (xmalloc on the whole-buffer path,
+ * arena on the chunked one) and the mask buffers.  Called when a read
+ * raises part way through, so that the raise leaves nothing attached or
+ * allocated behind it. */
+static void
+ca_sweep_abort (ca_sweep_state_t *st, int chunked)
+{
+  int k_op;
+  for (k_op = st->n_ops - 1; k_op >= 0; k_op--) {
+    if (st->attached[k_op]) {
+      st->attached[k_op] = 0;
+      ca_detach(st->cx[k_op]);
+    } else if (st->owned_buf[k_op]) {
+      if (chunked) ca_lazy_arena_release(st->owned_buf[k_op]);
+      else         xfree(st->owned_buf[k_op]);
+    }
     st->owned_buf[k_op] = NULL;
-    st->attached[k_op]  = 0;
   }
-  st->m0       = NULL;
-  st->n_kernel = 1;
-  /* CAREFUL: do not reset st->no_mask here — caller sets it before
-   * acquire and the NO_MASK guard below consumes it. */
+  if (st->m0) {
+    if (chunked) ca_lazy_arena_release(st->m0);
+    else         xfree(st->m0);
+    st->m0 = NULL;
+  }
+  if (chunked && st->mask_scratch) {
+    ca_lazy_arena_release(st->mask_scratch);
+    st->mask_scratch = NULL;
+  }
+}
+
+typedef struct {
+  ca_sweep_state_t *st;
+  int               any_input_mask;
+  boolean8_t       *ms;          /* per-operand mask staging, whole path */
+} ca_sweep_acquire_ctx_t;
+
+/* The part of acquire that takes resources and reads operands, and so can
+ * raise holding them.  Runs under rb_protect; every buffer is recorded in
+ * the state (or the ctx) before the read that may raise. */
+static VALUE
+ca_sweep_acquire_body (VALUE arg)
+{
+  ca_sweep_acquire_ctx_t *ctx = (ca_sweep_acquire_ctx_t *) arg;
+  ca_sweep_state_t *st = ctx->st;
+  int k_op;
 
   /* Per-operand acquire:
    *   OUTPUT (fsync == '1')       -> ca_attach + base = ca->ptr
    *   INPUT alias                 -> ca_attach + base = ca->ptr
-   *   INPUT non-alias non-scalar  -> xmalloc scratch + ca_xfer_all
+   *   INPUT non-alias non-scalar,
+   *   or sharing an OUTPUT's root -> xmalloc scratch + ca_xfer_all
    * INPUT non-alias never attaches the operand itself; the scratch
    * copy carries the values into the kernel. */
   for (k_op = 0; k_op < st->n_ops; k_op++) {
     CArray *ca = st->cx[k_op];
-    if (st->fsync[k_op] == '1') {
+    if (st->fsync[k_op] == '1' ||
+        (ca_attach_is_alias(ca) && ! ca_sweep_input_overlaps_output(st, k_op))) {
       ca_attach(ca);
-      st->base[k_op]     = (char *) ca->ptr;
       st->attached[k_op] = 1;
-    } else if (ca_attach_is_alias(ca)) {
-      ca_attach(ca);
       st->base[k_op]     = (char *) ca->ptr;
-      st->attached[k_op] = 1;
     } else {
       ca_size_t bytes_total = ca->elements * ca->bytes;
       st->owned_buf[k_op] = xmalloc(bytes_total);
@@ -69,36 +185,9 @@ ca_sweep_acquire (ca_sweep_state_t *st)
     }
   }
 
-  /* compute n_kernel (= broadcast shape) and per-cell strides */
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    CArray *ca = st->cx[k_op];
-    if (ca_is_scalar(ca)) {
-      st->stride[k_op] = 0;
-    } else {
-      st->stride[k_op] = ca->bytes;
-      if (st->n_kernel == 1) {
-        st->n_kernel = ca->elements;
-      } else if (ca->elements != st->n_kernel) {
-        rb_raise(rb_eRuntimeError, "data size mismatch in operation");
-      }
-    }
-  }
-
   /* iter mask m0 = OR of INPUT operand masks, gathered via ca_xfer_all
      so no operand mask attach happens.  Stays NULL if no INPUT masks. */
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    if (st->fsync[k_op] == '0' && ca_has_mask(st->cx[k_op])) {
-      any_input_mask = 1;
-      break;
-    }
-  }
-  if (any_input_mask && st->no_mask) {
-    rb_raise(rb_eRuntimeError,
-             "%s: masked INPUT not allowed in NO_MASK form "
-             "(use the *_MASKED form to handle masked cells explicitly)",
-             st->src_label ? st->src_label : "ca_sweep_acquire");
-  }
-  if (any_input_mask) {
+  if (ctx->any_input_mask) {
     st->m0 = xmalloc(st->n_kernel);
     memset(st->m0, 0, st->n_kernel);
     for (k_op = 0; k_op < st->n_ops; k_op++) {
@@ -111,11 +200,10 @@ ca_sweep_acquire (ca_sweep_state_t *st)
         ca_xfer_all(ca->mask, &bit, CA_XFER_GET);
         if (bit) memset(st->m0, 1, st->n_kernel);
       } else {
-        boolean8_t *ms = xmalloc(st->n_kernel);
         ca_size_t j;
-        ca_xfer_all(ca->mask, ms, CA_XFER_GET);
-        for (j = 0; j < st->n_kernel; j++) st->m0[j] |= ms[j];
-        xfree(ms);
+        if (!ctx->ms) ctx->ms = xmalloc(st->n_kernel);
+        ca_xfer_all(ca->mask, ctx->ms, CA_XFER_GET);
+        for (j = 0; j < st->n_kernel; j++) st->m0[j] |= ctx->ms[j];
       }
     }
   }
@@ -126,21 +214,130 @@ ca_sweep_acquire (ca_sweep_state_t *st)
     for (k_op = 0; k_op < st->n_ops; k_op++) {
       CArray *ca = st->cx[k_op];
       if (st->fsync[k_op] != '1') continue;
-      ca_update_mask(ca);
-      if (!ca->mask) ca_create_mask(ca);
+      ca_create_mask(ca);
       memcpy(ca->mask->ptr, st->m0, st->n_kernel);
     }
   }
+  return Qnil;
 }
 
 void
-ca_sweep_release (ca_sweep_state_t *st)
+ca_sweep_acquire (ca_sweep_state_t *st)
 {
+  ca_sweep_acquire_ctx_t ctx;
   int k_op;
-  /* sync OUTPUTs (reverse order) */
+  int tag = 0;
+
+  if ((int) strlen(st->fsync) != st->n_ops) {
+    rb_raise(rb_eRuntimeError,
+             "[BUG] invalid length of fsync arg in %s (expected %d)",
+             st->src_label ? st->src_label : "ca_sweep_acquire",
+             st->n_ops);
+  }
+
+  for (k_op = 0; k_op < st->n_ops; k_op++) {
+    st->base[k_op]      = NULL;
+    st->owned_buf[k_op] = NULL;
+    st->attached[k_op]  = 0;
+  }
+  st->m0       = NULL;
+
+  /* Refusals come first, while nothing is held. */
+  ca_sweep_pair_operands(st);
+  ctx.st             = st;
+  ctx.ms             = NULL;
+  ctx.any_input_mask = ca_sweep_any_input_mask(st);
+
+  rb_protect(ca_sweep_acquire_body, (VALUE) &ctx, &tag);
+  if (ctx.ms) xfree(ctx.ms);
+  if (tag) {
+    ca_sweep_abort(st, 0);
+    rb_jump_tag(tag);
+  }
+}
+
+void
+ca_sweep_run (ca_sweep_state_t *st, VALUE (*walk)(VALUE), VALUE arg)
+{
+  int tag = 0;
+  rb_protect(walk, arg, &tag);
+  if (tag) {
+    ca_sweep_abort(st, 0);
+    rb_jump_tag(tag);
+  }
+  ca_sweep_release(st);
+}
+
+void
+ca_sweep_run_chunked (ca_sweep_state_t *st, VALUE (*walk)(VALUE), VALUE arg)
+{
+  int tag = 0;
+  rb_protect(walk, arg, &tag);
+  if (tag) {
+    /* a raise inside ca_sweep_next_chunk has already given back */
+    if (st->chunked_state != 3) {
+      ca_sweep_abort(st, 1);
+      st->chunked_state = 3;
+    }
+    rb_jump_tag(tag);
+  }
+  ca_sweep_release_chunked(st);
+}
+
+static VALUE
+ca_sweep_sync_one (VALUE arg)
+{
+  ca_sync((CArray *) arg);
+  return Qnil;
+}
+
+/* Sync the OUTPUTs (reverse order).  Writing back through a view's slots
+ * can raise (a CAObject's store hook); the raise is held so that release
+ * still detaches and frees everything, and ca_sweep_reraise propagates
+ * it afterwards.  An entity's sync does nothing that can raise, so it is
+ * called directly: the outputs this library allocates itself pay nothing. */
+static int
+ca_sweep_sync_outputs (ca_sweep_state_t *st, volatile VALUE *error)
+{
+  int k_op, first = 0;
   for (k_op = st->n_ops - 1; k_op >= 0; k_op--) {
-    if (st->fsync[k_op] == '1') ca_sync(st->cx[k_op]);
+    CArray *ca = st->cx[k_op];
+    if (st->fsync[k_op] != '1') {
+      continue;
+    }
+    if (ca_is_entity(ca)) {
+      ca_sync(ca);
+    }
+    else {
+      int tag = 0;
+      rb_protect(ca_sweep_sync_one, (VALUE) ca, &tag);
+      if (tag) {
+        if (! first) {
+          first  = tag;
+          *error = rb_errinfo();
+        }
+        rb_set_errinfo(Qnil);
+      }
+    }
   }
+  return first;
+}
+
+static void
+ca_sweep_reraise (int tag, VALUE error)
+{
+  if (tag) {
+    rb_set_errinfo(error);
+    rb_jump_tag(tag);
+  }
+}
+
+void
+ca_sweep_release (ca_sweep_state_t *st)
+{
+  int k_op;
+  volatile VALUE error = Qnil;
+  int tag = ca_sweep_sync_outputs(st, &error);
   /* detach attached / xfree owned buffers (reverse order) */
   for (k_op = st->n_ops - 1; k_op >= 0; k_op--) {
     if (st->attached[k_op]) {
@@ -154,86 +351,32 @@ ca_sweep_release (ca_sweep_state_t *st)
     xfree(st->m0);
     st->m0 = NULL;
   }
+  ca_sweep_reraise(tag, error);
 }
 
 /* ===== Chunked path implementation ===== */
 
-void
-ca_sweep_acquire_chunked (ca_sweep_state_t *st)
+static VALUE
+ca_sweep_acquire_chunked_body (VALUE arg)
 {
+  ca_sweep_acquire_ctx_t *ctx = (ca_sweep_acquire_ctx_t *) arg;
+  ca_sweep_state_t *st = ctx->st;
   int k_op;
-  int any_input_mask = 0;
-  CArray *shape_donor = NULL;
-
-  if ((int) strlen(st->fsync) != st->n_ops) {
-    rb_raise(rb_eRuntimeError,
-             "[BUG] invalid length of fsync arg in %s (expected %d)",
-             st->src_label ? st->src_label : "ca_sweep_acquire_chunked",
-             st->n_ops);
-  }
-
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    st->base[k_op]      = NULL;
-    st->base_orig[k_op] = NULL;
-    st->owned_buf[k_op] = NULL;
-    st->attached[k_op]  = 0;
-  }
-  st->m0            = NULL;
-  st->mask_scratch  = NULL;
-  st->n_kernel      = 1;
-  st->chunk_off     = 0;
-  st->chunk_n       = 0;
-  st->chunk_n_max   = 0;
-  st->inner         = 1;
-  st->chunked_state = 0;
-
-  /* compute broadcast shape (n_kernel + strides) from operand shapes.
-   * scalar operands collapse to stride 0; non-scalar operands must agree
-   * on element count.  shape_donor is the first non-scalar operand and
-   * defines the chunking inner-axis size. */
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    CArray *ca = st->cx[k_op];
-    if (ca_is_scalar(ca)) {
-      st->stride[k_op] = 0;
-    } else {
-      st->stride[k_op] = ca->bytes;
-      if (st->n_kernel == 1) {
-        st->n_kernel = ca->elements;
-        shape_donor  = ca;
-      } else if (ca->elements != st->n_kernel) {
-        rb_raise(rb_eRuntimeError, "data size mismatch in operation");
-      }
-    }
-  }
-
-  /* chunk-size policy: inner = donor's product of dims[1..]; chunk_n_max
-   * = compute_n on donor->bytes (= type-dependent 32KB target). */
-  if (shape_donor) {
-    st->inner = ca_chunk_inner_size(shape_donor);
-    st->chunk_n_max = ca_chunk_compute_n(st->n_kernel, st->inner,
-                                         shape_donor->bytes);
-  } else {
-    /* all-scalar: single 1-cell chunk */
-    st->inner       = 1;
-    st->chunk_n_max = 1;
-  }
 
   /* per-operand acquire:
    *   OUTPUT (fsync == '1') : ca_attach + base_orig = ca->ptr (legitimate)
    *   INPUT  alias          : ca_attach + base_orig = ca->ptr (zero-copy)
-   *   INPUT  non-alias non-scalar : arena chunk scratch of chunk_n_max bytes
+   *   INPUT  non-alias non-scalar,
+   *   or sharing an OUTPUT's root : arena chunk scratch of chunk_n_max bytes
    *   INPUT  scalar         : ca_attach + base_orig = ca->ptr (1 cell)
    */
   for (k_op = 0; k_op < st->n_ops; k_op++) {
     CArray *ca = st->cx[k_op];
-    if (st->fsync[k_op] == '1') {
+    if (st->fsync[k_op] == '1' || ca_is_scalar(ca) ||
+        (ca_attach_is_alias(ca) && ! ca_sweep_input_overlaps_output(st, k_op))) {
       ca_attach(ca);
-      st->base_orig[k_op] = (char *) ca->ptr;
       st->attached[k_op]  = 1;
-    } else if (ca_is_scalar(ca) || ca_attach_is_alias(ca)) {
-      ca_attach(ca);
       st->base_orig[k_op] = (char *) ca->ptr;
-      st->attached[k_op]  = 1;
     } else {
       /* non-alias non-scalar INPUT: arena scratch sized for chunk_n_max */
       ca_size_t scratch_bytes = st->chunk_n_max * ca->bytes;
@@ -251,19 +394,7 @@ ca_sweep_acquire_chunked (ca_sweep_state_t *st)
    * scaling with the operand, which is the thing this path exists to avoid:
    * at 1 byte per cell it is an eighth of an f64 operand, but an eighth of
    * unbounded is still unbounded. */
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    if (st->fsync[k_op] == '0' && ca_has_mask(st->cx[k_op])) {
-      any_input_mask = 1;
-      break;
-    }
-  }
-  if (any_input_mask && st->no_mask) {
-    rb_raise(rb_eRuntimeError,
-             "%s: masked INPUT not allowed in NO_MASK form "
-             "(use the *_MASKED form to handle masked cells explicitly)",
-             st->src_label ? st->src_label : "ca_sweep_acquire_chunked");
-  }
-  if (any_input_mask) {
+  if (ctx->any_input_mask) {
     st->m0           = (boolean8_t *) ca_lazy_arena_acquire(st->chunk_n_max);
     st->mask_scratch = (boolean8_t *) ca_lazy_arena_acquire(st->chunk_n_max);
     memset(st->m0, 0, st->chunk_n_max);
@@ -273,8 +404,7 @@ ca_sweep_acquire_chunked (ca_sweep_state_t *st)
     for (k_op = 0; k_op < st->n_ops; k_op++) {
       CArray *ca = st->cx[k_op];
       if (st->fsync[k_op] != '1') continue;
-      ca_update_mask(ca);
-      if (!ca->mask) ca_create_mask(ca);
+      ca_create_mask(ca);
     }
   }
 
@@ -283,6 +413,75 @@ ca_sweep_acquire_chunked (ca_sweep_state_t *st)
    * once more in release), never at acquire.  Author per-cell m_out writes
    * land in m0 during the chunk loop and must be captured after that loop
    * has run, not before it. */
+  return Qnil;
+}
+
+void
+ca_sweep_acquire_chunked (ca_sweep_state_t *st)
+{
+  ca_sweep_acquire_ctx_t ctx;
+  int k_op;
+  int tag = 0;
+  CArray *shape_donor = NULL;
+
+  if ((int) strlen(st->fsync) != st->n_ops) {
+    rb_raise(rb_eRuntimeError,
+             "[BUG] invalid length of fsync arg in %s (expected %d)",
+             st->src_label ? st->src_label : "ca_sweep_acquire_chunked",
+             st->n_ops);
+  }
+
+  for (k_op = 0; k_op < st->n_ops; k_op++) {
+    st->base[k_op]      = NULL;
+    st->base_orig[k_op] = NULL;
+    st->owned_buf[k_op] = NULL;
+    st->attached[k_op]  = 0;
+  }
+  st->m0            = NULL;
+  st->mask_scratch  = NULL;
+  st->chunk_off     = 0;
+  st->chunk_n       = 0;
+  st->chunk_n_max   = 0;
+  st->inner         = 1;
+  st->chunked_state = 0;
+
+  /* shape_donor is the first non-scalar operand and defines the chunking
+   * inner-axis size. */
+  shape_donor = ca_sweep_pair_operands(st);
+
+  /* chunk-size policy: inner = donor's product of dims[1..]; chunk_n_max
+   * = compute_n on donor->bytes (= type-dependent 32KB target). */
+  if (shape_donor) {
+    st->inner = ca_chunk_inner_size(shape_donor);
+    st->chunk_n_max = ca_chunk_compute_n(st->n_kernel, st->inner,
+                                         shape_donor->bytes);
+  } else {
+    /* all-scalar: single 1-cell chunk */
+    st->inner       = 1;
+    st->chunk_n_max = 1;
+  }
+
+  /* An INPUT sharing an OUTPUT's root is read whole before anything is
+     written: the walk becomes a single chunk, so its data and its mask
+     are gathered before the first flush. */
+  for (k_op = 0; k_op < st->n_ops; k_op++) {
+    if (ca_sweep_input_overlaps_output(st, k_op) && st->n_kernel > 0) {
+      st->chunk_n_max = st->n_kernel;
+      break;
+    }
+  }
+
+  /* Refusals come first, while nothing is held. */
+  ctx.st             = st;
+  ctx.ms             = NULL;
+  ctx.any_input_mask = ca_sweep_any_input_mask(st);
+
+  rb_protect(ca_sweep_acquire_chunked_body, (VALUE) &ctx, &tag);
+  if (tag) {
+    ca_sweep_abort(st, 1);
+    st->chunked_state = 3;
+    rb_jump_tag(tag);
+  }
 }
 
 /* OR one INPUT operand's mask for the current chunk into m0.  Gathered via
@@ -313,9 +512,7 @@ ca_sweep_gather_chunk_mask (ca_sweep_state_t *st, ca_size_t off, ca_size_t n)
   }
 }
 
-/* Copy the chunk just finished out to every OUTPUT operand's mask.  Runs
- * after the author's loop over that chunk, so m_out writes into m0 are
- * carried through. */
+/* Copy the chunk just finished out to every OUTPUT operand's mask. */
 static void
 ca_sweep_flush_chunk_mask (ca_sweep_state_t *st)
 {
@@ -329,18 +526,46 @@ ca_sweep_flush_chunk_mask (ca_sweep_state_t *st)
   }
 }
 
+/* Set up base[] for the chunk at chunk_off / chunk_n: gather each
+   non-alias INPUT into its arena scratch, then the chunk's mask. */
+static VALUE
+ca_sweep_next_chunk_body (VALUE arg)
+{
+  ca_sweep_state_t *st = (ca_sweep_state_t *) arg;
+  ca_size_t off = st->chunk_off;
+  ca_size_t n   = st->chunk_n;
+  int k_op;
+
+  for (k_op = 0; k_op < st->n_ops; k_op++) {
+    CArray *ca = st->cx[k_op];
+    if (st->stride[k_op] == 0) {
+      /* scalar: stride 0, base is the single-cell ptr (base_orig) */
+      st->base[k_op] = st->base_orig[k_op];
+    } else if (st->base_orig[k_op]) {
+      /* alias INPUT or OUTPUT: walk through ca->ptr by chunk_off */
+      st->base[k_op] = st->base_orig[k_op] + off * st->stride[k_op];
+    } else {
+      /* non-alias INPUT: per-chunk gather into owned_buf (arena) */
+      ca_chunked_gather(ca, off, n, st->owned_buf[k_op]);
+      st->base[k_op] = st->owned_buf[k_op];
+    }
+  }
+
+  if (st->m0) ca_sweep_gather_chunk_mask(st, off, n);
+  return Qnil;
+}
+
 int
 ca_sweep_next_chunk (ca_sweep_state_t *st)
 {
-  int k_op;
   ca_size_t off, n;
+  int tag = 0;
 
   if (st->chunked_state == 0) {
     /* first chunk */
     st->chunk_off     = 0;
     st->chunked_state = 1;
   } else {
-    /* the chunk that just finished is the author's last word on its mask */
     ca_sweep_flush_chunk_mask(st);
     /* advance */
     st->chunk_off += st->chunk_n;
@@ -356,24 +581,16 @@ ca_sweep_next_chunk (ca_sweep_state_t *st)
   if (off + n > st->n_kernel) n = st->n_kernel - off;
   st->chunk_n = n;
 
-  /* set up base[] for the upcoming chunk */
-  for (k_op = 0; k_op < st->n_ops; k_op++) {
-    CArray *ca = st->cx[k_op];
-    if (st->stride[k_op] == 0) {
-      /* scalar: stride 0, base is the single-cell ptr (base_orig) */
-      st->base[k_op] = st->base_orig[k_op];
-    } else if (st->base_orig[k_op]) {
-      /* alias INPUT or OUTPUT: walk through ca->ptr by chunk_off */
-      st->base[k_op] = st->base_orig[k_op] + off * st->stride[k_op];
-    } else {
-      /* non-alias INPUT: per-chunk gather into owned_buf (arena) */
-      ca_chunked_gather(ca, off, n, st->owned_buf[k_op]);
-      st->base[k_op] = st->owned_buf[k_op];
-    }
+  /* The gathers read the operands and so can raise; a raise gives back
+     the walk's resources before it propagates, since the caller's release
+     is never reached. */
+  rb_protect(ca_sweep_next_chunk_body, (VALUE) st, &tag);
+  if (tag) {
+    ca_sweep_abort(st, 1);
+    st->chunked_state = 3;
+    rb_jump_tag(tag);
   }
 
-  if (st->m0) ca_sweep_gather_chunk_mask(st, off, n);
-
   return 1;
 }
 
@@ -381,14 +598,15 @@ void
 ca_sweep_release_chunked (ca_sweep_state_t *st)
 {
   int k_op;
+  volatile VALUE error = Qnil;
+  int tag;
+  /* a walk that raised has already given everything back */
+  if (st->chunked_state == 3) return;
   /* The final chunk has no next_chunk call to flush it, so it is flushed
-   * here.  For INOUT_MASKED forms this is what captures the author's
-   * per-cell m_out writes over that last chunk. */
+   * here. */
   ca_sweep_flush_chunk_mask(st);
-  /* sync OUTPUTs (reverse order, regardless of whether chunk loop ran) */
-  for (k_op = st->n_ops - 1; k_op >= 0; k_op--) {
-    if (st->fsync[k_op] == '1') ca_sync(st->cx[k_op]);
-  }
+  /* sync OUTPUTs (regardless of whether chunk loop ran) */
+  tag = ca_sweep_sync_outputs(st, &error);
   /* detach attached / release arena scratch buffers (reverse order) */
   for (k_op = st->n_ops - 1; k_op >= 0; k_op--) {
     if (st->attached[k_op]) {
@@ -406,13 +624,16 @@ ca_sweep_release_chunked (ca_sweep_state_t *st)
     ca_lazy_arena_release(st->mask_scratch);
     st->mask_scratch = NULL;
   }
+  if (tag) {
+    st->chunked_state = 3;            /* everything has been given back */
+  }
+  ca_sweep_reraise(tag, error);
 }
 
 /* ===== WHOLE_BUFFER function form (rb_ensure-protected) ===== */
 
 typedef struct {
   CArray              *ca;
-  int                  writable;
   ca_with_buffer_body_fn body_fn;
   void                *user_data;
 } ca_with_buffer_ctx_t;
@@ -425,49 +646,15 @@ ca_with_buffer_body_trampoline (VALUE ctx_val)
   return Qnil;
 }
 
-static VALUE
-ca_with_buffer_ensure_trampoline (VALUE ctx_val)
-{
-  ca_with_buffer_ctx_t *ctx = (ca_with_buffer_ctx_t *) ctx_val;
-  if (ctx->writable) ca_sync(ctx->ca);
-  ca_detach(ctx->ca);
-  return Qnil;
-}
-
 void
 rb_ca_call_with_buffer (VALUE r_ca, int writable,
                    ca_with_buffer_body_fn body, void *user_data)
 {
   ca_with_buffer_ctx_t ctx;
   TypedData_Get_Struct(r_ca, CArray, &carray_data_type, ctx.ca);
-  /* CAREFUL: attach outside rb_ensure.  If attach raises, we are not
-   * yet inside the ensure scope, so the ensure trampoline will not
-   * fire on a not-attached ca (which would try to detach and crash). */
-  ca_attach(ctx.ca);
-  ctx.writable  = writable;
   ctx.body_fn   = body;
   ctx.user_data = user_data;
-  rb_ensure(ca_with_buffer_body_trampoline,    (VALUE) &ctx,
-            ca_with_buffer_ensure_trampoline,  (VALUE) &ctx);
+  ca_attach_window(1, &ctx.ca, writable ? "w" : "r",
+                   ca_with_buffer_body_trampoline, (VALUE) &ctx);
 }
 
-void
-ca_sweep_check_same_shape (CArray *ca_in, CArray *ca_out, const char *src_label)
-{
-  int k;
-  const char *lbl = src_label ? src_label : "ca_sweep_check_same_shape";
-  if (ca_in->ndim != ca_out->ndim) {
-    rb_raise(rb_eRuntimeError,
-             "%s: shape mismatch (ndim %d vs %d)",
-             lbl, ca_in->ndim, ca_out->ndim);
-  }
-  for (k = 0; k < ca_in->ndim; k++) {
-    if (ca_in->dim[k] != ca_out->dim[k]) {
-      rb_raise(rb_eRuntimeError,
-               "%s: shape mismatch at dim[%d] (%lld vs %lld)",
-               lbl, k,
-               (long long) ca_in->dim[k],
-               (long long) ca_out->dim[k]);
-    }
-  }
-}