Numba 函数中添加 break 语句导致性能显著下降的原因及解决方案

import numba
import numpy as np
from timeit import timeit

@numba.njit
def count_in_range(arr, min_value, max_value):
    count = 0
    for a in arr:
        if min_value < a < max_value:
            count += 1
    return count

@numba.njit
def count_in_range2(arr, min_value, max_value):
    count = 0
    for a in arr:
        if min_value < a < max_value:
            count += 1
            break  # <---- break here
    return count

rng = np.random.default_rng(0)
arr = rng.random(10 * 1000 * 1000)

# To compare on even conditions, choose the condition that does not terminate early.
min_value = 0.5
max_value = min_value - 1e-10
assert not np.any(np.logical_and(min_value <= arr, arr <= max_value))

n = 100
for f in (count_in_range, count_in_range2):
    f(arr, min_value, max_value)
    elapsed = timeit(lambda: f(arr, min_value, max_value), number=n) / n
    print(f"{f.__name__}: {elapsed * 1000:.3f} ms")

#include <cstdint>
#include <cstdlib>
#include <vector>

int64_t count_in_range(const std::vector<double>& arr, double min_value, double max_value)
{
    int64_t count = 0;

    for(int64_t i=0 ; i<arr.size() ; ++i)
    {
        double a = arr[i];

        if (min_value < a && a < max_value)
        {
            count += 1;
        }
    }

    return count;
}

@numba.njit
def count_in_range_faster(arr, min_value, max_value):
    count = 0
    for i in range(0, arr.size, 16):
        if arr.size - i >= 16:
            # Optimized SIMD-friendly computation of 1 chunk of size 16
            tmp_view = arr[i:i+16]
            for j in range(0, 16):
                if min_value < tmp_view[j] < max_value:
                    count += 1
            if count > 0:
                return 1
        else:
            # Fallback implementation (variable-sized chunk)
            for j in range(i, arr.size):
                if min_value < arr[j] < max_value:
                    count += 1
            if count > 0:
                return 1
    return 0

count_in_range:          7.112 ms
count_in_range2:        35.317 ms
count_in_range_faster:   5.827 ms