From 282b5ad5e8fcb623340f64ad6f3bb107a214b283 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Sun, 23 Aug 2026 16:58:38 -0400 Subject: [PATCH] Optimize MLX threshold and rate kernels --- src/ezmsg/event/kernel_activation.py | 79 ++++++++++++++++++-------- src/ezmsg/event/util/peak_mlx_metal.py | 68 +++++++++++----------- 2 files changed, 88 insertions(+), 59 deletions(-) diff --git a/src/ezmsg/event/kernel_activation.py b/src/ezmsg/event/kernel_activation.py index 84b18fe..1946344 100644 --- a/src/ezmsg/event/kernel_activation.py +++ b/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: diff --git a/src/ezmsg/event/util/peak_mlx_metal.py b/src/ezmsg/event/util/peak_mlx_metal.py index e1b8123..0fe3963 100644 --- a/src/ezmsg/event/util/peak_mlx_metal.py +++ b/src/ezmsg/event/util/peak_mlx_metal.py @@ -73,15 +73,14 @@ 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,)], @@ -89,13 +88,7 @@ def threshold_crossings_mlx_metal( ) 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, )