Optimize MLX kernel reuse and dense rate reductions

src/ezmsg/event/kernel_activation.py

@@ -400,10 +400,11 @@ def _process_events(self, message: AxisArray) -> AxisArray:
     def _process_dense_count_sum(self, message: AxisArray) -> AxisArray:
         """Fast path: dense input + COUNT kernel + SUM aggregation.
 
-        Bins are summed using cumulative-sum arithmetic in the input's array
-        namespace, so accelerator-resident inputs (MLX, CuPy) stay on device.
+        Bins are summed in the input's array namespace, so accelerator-resident
+        inputs (MLX, CuPy) stay on device. Uniform completed bins use direct
+        reshape reductions; irregular fractional bins use cumulative sums.
         Carry-over for the partial bin spanning chunk boundaries is held in
-        ``state.activation`` (numpy) and shuttled across boundaries.
+        ``state.dense_carry`` in the active array namespace.
         """
         xp = get_namespace(message.data)
         data = message.data
@@ -450,32 +451,64 @@ def _process_dense_count_sum(self, message: AxisArray) -> AxisArray:
         bin_end_samples = np.asarray(step.cut_points, dtype=np.int64) - n_carry_before
         bin_start_samples = np.concatenate(([np.int64(0)], bin_end_samples[:-1]))
 
-        # Cumulative sum, prepended with zeros so cumsum_padded[k] = sum(contrib[:k]).
-        # Use cumsum (in both numpy and mlx); numpy via array_api_compat also exposes
-        # the standard `cumulative_sum`, but mlx does not.
-        cumsum = xp.cumsum(contrib, axis=0)
-        zero_row = xp.zeros((1,) + feature_shape, dtype=cumsum.dtype)
-        cumsum_padded = xp.concat((zero_row, cumsum), axis=0)
-
-        end_idx = xp.asarray(bin_end_samples)
-        start_idx = xp.asarray(bin_start_samples)
-        bin_sums = xp.take(cumsum_padded, end_idx, axis=0) - xp.take(cumsum_padded, start_idx, axis=0)
-
-        # Add carry-over from the previous chunk's partial bin into bin 0.
-        overflow_pad_first = overflow_xp[None, ...]
-        if n_bins > 1:
-            overflow_pad_rest = xp.zeros((n_bins - 1,) + feature_shape, dtype=bin_sums.dtype)
-            overflow_pad = xp.concat((overflow_pad_first, overflow_pad_rest), axis=0)
+        last_bin_end = int(bin_end_samples[-1])
+        bin_widths_after_first = np.diff(bin_end_samples)
+        uniform_complete_bins = n_bins == 1 or np.all(bin_widths_after_first == bin_widths_after_first[0])
+        # For a stream of sub-bin chunks, the occasional call that closes only
+        # the carried bin is faster through the existing cumsum path. Direct
+        # reduction pays off for a carry-free bin or when it can amortise setup
+        # across multiple completed bins.
+        use_direct_reduction = uniform_complete_bins and (n_carry_before == 0 or n_bins > 1)
+
+        if use_direct_reduction:
+            # The first bin may complete a partial bin carried from the previous
+            # chunk. Every later bin is wholly inside this chunk; when those bin
+            # widths are uniform, reshape + reduction avoids materialising a
+            # full-size cumulative-sum array and two gather-index arrays.
+            first_bin_end = int(bin_end_samples[0])
+            first_bin_sum = overflow_xp + (
+                xp.sum(contrib[:first_bin_end], axis=0)
+                if first_bin_end > 0
+                else xp.zeros(feature_shape, dtype=xp.float32)
+            )
+            if n_bins == 1:
+                output = first_bin_sum[None, ...]
+            else:
+                bin_width = int(bin_widths_after_first[0])
+                complete = xp.reshape(
+                    contrib[first_bin_end:last_bin_end],
+                    (n_bins - 1, bin_width) + feature_shape,
+                )
+                later_bin_sums = xp.sum(complete, axis=1)
+                output = xp.concat((first_bin_sum[None, ...], later_bin_sums), axis=0)
         else:
-            overflow_pad = overflow_pad_first
-        output = bin_sums + overflow_pad
+            # Fractional schedules can produce nonuniform interior bin widths;
+            # this path also handles the cheaper single carried-bin completion.
+            # Cumulative sums handle arbitrary cut points without a Python loop.
+            # Use cumsum (in both numpy and mlx); numpy via array_api_compat also
+            # exposes `cumulative_sum`, but mlx does not.
+            cumsum = xp.cumsum(contrib, axis=0)
+            zero_row = xp.zeros((1,) + feature_shape, dtype=cumsum.dtype)
+            cumsum_padded = xp.concat((zero_row, cumsum), axis=0)
+
+            end_idx = xp.asarray(bin_end_samples)
+            start_idx = xp.asarray(bin_start_samples)
+            bin_sums = xp.take(cumsum_padded, end_idx, axis=0) - xp.take(cumsum_padded, start_idx, axis=0)
+
+            # Add carry-over from the previous chunk's partial bin into bin 0.
+            overflow_pad_first = overflow_xp[None, ...]
+            if n_bins > 1:
+                overflow_pad_rest = xp.zeros((n_bins - 1,) + feature_shape, dtype=bin_sums.dtype)
+                overflow_pad = xp.concat((overflow_pad_first, overflow_pad_rest), axis=0)
+            else:
+                overflow_pad = overflow_pad_first
+            output = bin_sums + overflow_pad
 
         # New carry-over: events past the last complete bin remain in the partial bin.
-        last_bin_end = int(bin_end_samples[-1])
         if last_bin_end < n_samples:
             new_overflow = xp.sum(contrib[last_bin_end:], axis=0)
         else:
-            new_overflow = xp.zeros(feature_shape, dtype=cumsum.dtype)
+            new_overflow = xp.zeros(feature_shape, dtype=contrib.dtype)
         self._state.dense_carry = new_overflow
 
         if self.settings.rate_normalize:

src/ezmsg/event/util/peak_mlx_metal.py

@@ -73,29 +73,22 @@ def threshold_crossings_mlx_metal(
     elapsed_flat = elapsed.astype(mx.int32).reshape(n_channels)
     n_words = (n_samples + 31) // 32
 
+    # Keep values that vary by stream or chunk out of Metal templates. The
+    # kernels obtain dimensions from generated shape metadata and receive the
+    # remaining values here, so jittered chunk lengths reuse one compilation.
     params = mx.array([float(threshold)], dtype=mx.float32)
+    metadata = mx.array([n_samples, refrac_width], dtype=mx.int32)
 
     crossing_words, last_sample = _crossing_words_kernel(
         inputs=[x_flat, prev_sample_flat, params],
-        template=[
-            ("N_SAMPLES", n_samples),
-            ("N_CHANNELS", n_channels),
-            ("N_WORDS", n_words),
-        ],
         grid=(n_words, n_channels, 1),
         threadgroup=(1, 1, 1),
         output_shapes=[(n_words, n_channels), (n_channels,)],
         output_dtypes=[mx.uint32, mx.float32],
     )
 
     events_flat, elapsed_out = _refractory_dense_kernel(
-        inputs=[crossing_words, elapsed_flat],
-        template=[
-            ("N_SAMPLES", n_samples),
-            ("N_CHANNELS", n_channels),
-            ("N_WORDS", n_words),
-            ("REFRAC_WIDTH", refrac_width),
-        ],
+        inputs=[crossing_words, elapsed_flat, metadata],
         grid=(n_channels, 1, 1),
         threadgroup=(1, 1, 1),
         output_shapes=[(n_samples, n_channels), (n_channels,)],
@@ -113,77 +106,80 @@ def threshold_crossings_mlx_metal(
 _CROSSING_WORDS_KERNEL_SOURCE = r"""
     uint word = thread_position_in_grid.x;
     uint ch = thread_position_in_grid.y;
-    if (word >= N_WORDS || ch >= N_CHANNELS) {
+    uint n_samples = x_in_shape[0];
+    uint n_channels = x_in_shape[1];
+    uint n_words = (n_samples + 31) / 32;
+    if (word >= n_words || ch >= n_channels) {
         return;
     }
 
     float threshold = params[0];
     uint start = word * 32;
-    float ref_sample = start == 0 ? prev_sample_in[ch] : x_in[(start - 1) * N_CHANNELS + ch];
+    float ref_sample = start == 0 ? prev_sample_in[ch] : x_in[(start - 1) * n_channels + ch];
     uint prev = threshold >= 0.0f ? (ref_sample >= threshold) : (ref_sample <= threshold);
 
     uint bits = 0;
     for (uint bit = 0; bit < 32; ++bit) {
         uint t = start + bit;
-        if (t >= N_SAMPLES) {
+        if (t >= n_samples) {
             break;
         }
 
-        float sample = x_in[t * N_CHANNELS + ch];
+        float sample = x_in[t * n_channels + ch];
         uint over = threshold >= 0.0f ? (sample >= threshold) : (sample <= threshold);
         if (over && !prev) {
             bits |= (1u << bit);
         }
         prev = over;
     }
-    crossing_words_out[word * N_CHANNELS + ch] = bits;
+    crossing_words_out[word * n_channels + ch] = bits;
 
-    if (word == N_WORDS - 1) {
-        last_sample_out[ch] = x_in[(N_SAMPLES - 1) * N_CHANNELS + ch];
+    if (word == n_words - 1) {
+        last_sample_out[ch] = x_in[(n_samples - 1) * n_channels + ch];
     }
 """
 
 
 _REFRACTORY_DENSE_KERNEL_SOURCE = r"""
     uint ch = thread_position_in_grid.x;
-    if (ch >= N_CHANNELS) {
+    uint n_words = crossing_words_in_shape[0];
+    uint n_channels = crossing_words_in_shape[1];
+    uint n_samples = metadata[0];
+    int refrac_width = metadata[1];
+    if (ch >= n_channels) {
         return;
     }
 
-    for (uint t = 0; t < N_SAMPLES; ++t) {
-        events_out[t * N_CHANNELS + ch] = 0;
+    for (uint t = 0; t < n_samples; ++t) {
+        events_out[t * n_channels + ch] = 0;
     }
 
     int elapsed = elapsed_in[ch];
     int last_t = -1;
 
-    for (uint word = 0; word < N_WORDS; ++word) {
-        uint bits = crossing_words_in[word * N_CHANNELS + ch];
+    for (uint word = 0; word < n_words; ++word) {
+        uint bits = crossing_words_in[word * n_channels + ch];
         while (bits != 0) {
-            uint bit = 0;
-            uint mask = 1u;
-            while ((bits & mask) == 0u) {
-                bit += 1;
-                mask <<= 1;
-            }
-            bits &= ~mask;
+            // Locate and clear the least-significant crossing in constant time.
+            uint bit = ctz(bits);
+            bits &= bits - 1;
 
             uint t = word * 32 + bit;
-            if (t >= N_SAMPLES) {
+            if (t >= n_samples) {
                 break;
             }
 
             elapsed += int(t) - last_t;
             last_t = int(t);
 
-            if (REFRAC_WIDTH <= 2 || elapsed > REFRAC_WIDTH) {
-                events_out[t * N_CHANNELS + ch] = 1;
+            if (refrac_width <= 2 || elapsed > refrac_width) {
+                events_out[t * n_channels + ch] = 1;
                 elapsed = 0;
             }
         }
     }
 
-    elapsed += int(N_SAMPLES) - 1 - last_t;
+    elapsed += int(n_samples) - 1 - last_t;
     elapsed_out[ch] = elapsed;
 """
 
@@ -198,7 +194,7 @@ def threshold_crossings_mlx_metal(
 
 _refractory_dense_kernel = mx.fast.metal_kernel(
     name="threshold_refractory_dense",
-    input_names=["crossing_words_in", "elapsed_in"],
+    input_names=["crossing_words_in", "elapsed_in", "metadata"],
     output_names=["events_out", "elapsed_out"],
     source=_REFRACTORY_DENSE_KERNEL_SOURCE,
 )