[Enhance]: parallelization for large batch sizes

[mmbejani]

Affected Component

ailego

Current Behavior

The current compute_one_to_many_* implementations in zvec::ailego::DistanceBatch process batches sequentially.

• Inner-product computation is vectorized (AVX2) but single-threaded.
• Throughput is limited by a single core when BatchSize is large.

Desired Improvement

For large batch sizes (e.g. hundreds or thousands of vectors), the outer loop over BatchSize / dp_batch becomes embarrassingly parallel and can benefit significantly from multi-core CPUs.

Introduce optional OpenMP parallelization when the batch size exceeds a configurable threshold.

Example strategy:

• Keep current behavior for small batches (to avoid OpenMP overhead).
• Use #pragma omp parallel for over the batch dimension for large batches.
• Guard with #ifdef _OPENMP to preserve portability.

Impact

• Improved throughput for large-scale distance computations.
• No behavior change for small batch sizes.
• Backward compatible when OpenMP is not enabled.

[JalinWang]

Thank you for your interest in Zvec and for such a detailed proposal!

I’d like to provide some context on the current design of ailego:

1. Concurrency Model: Currently, ailego/core is single-threaded by design. We generally handle parallelism at a higher level (e.g., the zvec level). Executing a single query on a single thread allows for better CPU cache utilization. We use manual prefetch instructions to minimize cache misses, which becomes less effective if a single query is distributed across multiple cores due to the unshared L1/L2 cache.
2. Batch Size Constraints: The batch size should ideally be chosen to match the prefetch steps and the CPU cache size. If the batch size is too large, you may see performance degradation due to a rising cache miss rate. We typically don't recommend very large batch sizes; in fact, on CPUs with smaller caches than our baseline (Aliyun g9i), it is often better to decrease it.

Could you share the specific scenario or use case where you found a larger batch size to be necessary?

That said, your point about ailego-level parallelism is valid. For instance, when the number of concurrent queries is lower than the available CPU cores, using OpenMP could help reduce the latency of a single query. However, the trade-off between OpenMP overhead and the degree of parallelism needs to be carefully tested. Also, we are also exploring ways to initialize the batch size dynamically based on the host CPU's specifications.

If you have specific ideas on the implementation of dynamic BS or would like to contribute a PR, we would welcome it!

P.S. We are currently testing AVX-512 one2many support for the inner-product metric, which will be released soon and should provide a noticeable performance boost 😄