Small lazy expressions and reductions cost ~65-145 µs more per call since #741 (Python-side dispatch); float32 arange also slower #791

Since #741 was merged (5ff9823), evaluating a small lazy expression or a full reduction costs 65-145 µs more per call than on main right before the merge (126fa8f), i.e. 1.45-1.75x for 10k-item arrays. The time goes into Python-side dispatch; the native work on the chunk costs the same as before. One other slowdown showed up in the same measurements: blosc2.arange(..., dtype=np.float32) for large arrays. Large full reductions with many threads got slower too, for a different reason; that is #792.

Setup: Ryzen 9 7950X (16 cores / 32 threads, blosc2.nthreads = 28), Linux, Python 3.12, NumPy 2.5.4. Both commits built from source with the same toolchain (conda-forge GCC 15.3, Release) and the same c-blosc2 pin (fa29efc0, the pin of both). A second, independent build of 5ff9823 agrees within ±2%.

import time
import numpy as np
import blosc2

a = blosc2.asarray(np.arange(10_000, dtype=np.float32).reshape(100, 100))
b = blosc2.asarray(np.ones((100, 100), dtype=np.float32))
e = blosc2.lazyexpr("a + b * 2", operands={"a": a, "b": b})

def best_us(f, n=2000, reps=7):
    f()
    best = float("inf")
    for _ in range(reps):
        t0 = time.perf_counter()
        for _ in range(n):
            f()
        best = min(best, (time.perf_counter() - t0) / n)
    return best * 1e6

for name, f in [("(a + b * 2)[:]", lambda: (a + b * 2)[:]), ("e[:] (prebuilt)", lambda: e[:]),
                ("a + b * 2 (build)", lambda: a + b * 2), ("a.sum()", lambda: a.sum()),
                ("(a * b).sum()", lambda: (a * b).sum())]:
    print(f"{name:20s} {best_us(f):6.0f} us")
print(f"{'arange f32 100M':20s} {best_us(lambda: blosc2.arange(0, 100_000_000, dtype=np.float32), n=3) / 1e3:6.1f} ms")

Other small cases (100x100 float32, median of 5 interleaved runs, ±1-3%): (a + b)[:] 137 -> 231 µs, (blosc2.sin(a) + b)[:] 200 -> 343 µs, blosc2.where(a > 5000, a, b)[:] 184 -> 276 µs, (a + b * 2).compute() 223 -> 332 µs, a.mean() 147 -> 218 µs, a.max() 142 -> 211 µs. With blosc2.set_nthreads(1) the numbers are the same. The extra cost per call is constant up to ~100k items; from 1M items on, [:] is much faster than before (1M: 1.48 -> 0.56 ms), while compute() and full reductions stay slower at every size. arange below 1M items, int64/float64 arange at 100M, linspace and a.sum(axis=0) are faster than before (0.5-0.9x).

Where the time goes (prebuilt e[:], 116 -> 201 µs; time.perf_counter_ns wrappers around the functions involved, each step also timed on its own, and switching parts off in one process):

• The native work is unchanged: NDArray._execute_menudet_prefilter() takes ~14.5 µs for the 10k-item chunk, the old _set_pref_expr() + update_data() + remove_prefilter() route took the same.
• LazyExpr.compute() now runs the native-first probes before the chunked path; for small operands they all decline, after ~20-30 µs. Most of it is prepare_native_first() building _native_plan_cache_key() (~11 µs, incl. _graph_configuration()) before its size gate (compressed operands below 2M items are not admitted).
• The typed route re-lowers the expression on every call, also for a prebuilt LazyExpr: compute_typed_prefilter() -> _prepare_typed_map_kernel() -> lower_native_graph() (parse, validate_operands(), _graph_configuration()), ~20-30 µs; the compiled plan itself comes from the cache.
• menudet_prefilter_inputs_eligible() runs twice per call (in compute_typed_prefilter() / compute_typed_reduction(), then again in evaluate_prefilter()), ~12 µs each.
• evaluate_prefilter() adds a Python preamble (descriptor_info(), signature checks) around the scheduling blosc2.uninit() (the old route created a result array of the same cost, so that part is a wash).
• Building a + b * 2: check_dtype() now calls binary_result_dtype() (a NumPy probe with empty arrays) for each binary operation, plus the expression_evaluation() context switches of the expression wrappers.
• Switched off in-process: the three native-first probes (compute_resident_if_eligible, compute_native_prefilter_if_eligible, prepare_native_first) save 9-25 µs; BLOSC2_EXPRESSION_EVALUATOR=legacy saves 8-26 µs on elementwise expressions and 53-67 µs on reductions. With both, an elementwise evaluation is still 55 µs slower than on 126fa8f: many small per-call additions in compute(), fast_eval() and the decorators.

Bisect along #741's first-parent history (each commit built with its own miniexpr and c-blosc2 pins, median of 5 interleaved runs, ±1-2%):

For elementwise expressions there is no single commit: the cost grows in steps over the whole branch. For reductions, 71ce150 is almost the whole step. (The branch base was slower than main at the time of the merge, since main improved in the meantime.)

arange(float32): since 608f9b6 / 2b1fc99 a float ramp is the Menudet kernel float(start) + _flat_idx * float(step), cast to the output dtype. miniexpr's AOT ramp bridge (dsl_menudet_jit.c) only takes a float32 output when the arithmetic is float32 too, so float32 ramps run the generic path: 3.35 ns/item vs 0.65 ns/item for the same kernel with a float64 output (MenudetKernel.evaluate_block, 10M items, one thread). In a separate run at 100M items, float32 arange took 30 -> 44 ms and float64 40 -> 33 ms (28 threads). The values are the same as before (bit-identical for arange(0, 1e8)). A bridge variant for float32 output that computes in double and rounds once would keep these values.

Possible directions (not tried, your call):

• Run the cheap gates of the native-first probes (operand sizes vs MINIMUM_COMPRESSED_*_ITEMS) before building the plan cache key, or skip the probes when all operands are below those sizes.
• Keep the prepared typed map kernel on the LazyExpr (keyed like the native plan cache), so a repeated evaluation does not lower the expression again.
• Check menudet_prefilter_inputs_eligible() once per evaluation.

I have the benchmark and profiling scripts and the raw numbers and can share them. I used Claude Code (Anthropic) to build the commits, run the benchmarks and profiles, and draft this issue; I checked the numbers.

🤖 Generated with Claude Code