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

# 基准测试代码
def run_benchmark():
    rng = np.random.default_rng(0)
    arr = rng.random(10 * 1000 * 1000)

    # 选择一个不触发早期退出的条件，以确保公平比较循环遍历整个数组的情况
    min_value = 0.5
    max_value = min_value - 1e-10 # 确保范围为空，不会触发if条件
    assert not np.any(np.logical_and(min_value <= arr, arr <= max_value))

    n = 100
    print("--- 初始基准测试 ---")
    for f in (count_in_range, count_in_range2):
        f(arr, min_value, max_value) # 预热JIT
        elapsed = timeit(lambda: f(arr, min_value, max_value), number=n) / n
        print(f"{f.__name__}: {elapsed * 1000:.3f} ms")

# run_benchmark()

count_in_range: 3.351 ms
count_in_range2: 42.312 ms

# ... (Numba函数定义同上) ...

def partition(arr, threshold):
    """将数组元素分为小于阈值和大于等于阈值两部分，并拼接。"""
    less = arr[arr < threshold]
    more = arr[~(arr < threshold)]
    return np.concatenate((less, more))

def partition_with_error(arr, threshold, error_rate):
    """在分区的基础上引入错误率，打乱部分元素以增加分支预测难度。"""
    less = arr[arr < threshold]
    more = arr[~(arr < threshold)]
    # 引入错误，将一部分小于阈值的元素混入大于阈值的部分，反之亦然
    less_error, less_correct = np.split(less, [int(len(less) * error_rate)])
    more_error, more_correct = np.split(more, [int(len(more) * error_rate)])
    mostly_less = np.concatenate((less_correct, more_error))
    mostly_more = np.concatenate((more_correct, less_error))
    rng = np.random.default_rng(0)
    rng.shuffle(mostly_less)
    rng.shuffle(mostly_more)
    out = np.concatenate((mostly_less, mostly_more))
    assert np.array_equal(np.sort(out), np.sort(arr)) # 确保元素不变
    return out

def bench(f, arr, min_value, max_value, n=10, info=""):
    f(arr, min_value, max_value) # 预热JIT
    elapsed = timeit(lambda: f(arr, min_value, max_value), number=n) / n
    print(f"{f.__name__}: {elapsed * 1000:.3f} ms, min_value: {min_value:.1f}, {info}")

def main_benchmark():
    rng = np.random.default_rng(0)
    arr = rng.random(10 * 1000 * 1000)
    thresholds = np.linspace(0, 1, 11)

    print("\n# --- 随机数据 ---")
    for min_value in thresholds:
        bench(
            count_in_range2,
            arr,
            min_value=min_value,
            max_value=min_value - 1e-10, # 确保范围为空
        )

    print("\n# --- 分区数据（仍是随机的）---")
    for min_value in thresholds:
        bench(
            count_in_range2,
            partition(arr, threshold=min_value),
            min_value=min_value,
            max_value=min_value - 1e-10,
        )

    print("\n# --- 带有概率错误的已分区数据 ---")
    for ratio in thresholds:
        bench(
            count_in_range2,
            partition_with_error(arr, threshold=0.5, error_rate=ratio),
            min_value=0.5,
            max_value=0.5 - 1e-10, # 确保范围为空
            info=f"error: {ratio:.0%}",
        )

# main_benchmark()

@numba.njit
def count_in_range_faster(arr, min_value, max_value):
    """
    通过分块处理优化，实现类似提前退出但支持向量化的查找。
    返回1如果找到，0如果未找到。
    """
    count = 0
    # 设定一个块大小，例如16，这是常见的SIMD寄存器宽度（双精度浮点数）
    chunk_size = 16 

    for i in range(0, arr.size, chunk_size):
        # 处理完整的块
        if arr.size - i >= chunk_size:
            # 创建一个视图来处理当前块，LLVM可以对这种固定大小的循环进行向量化
            tmp_view = arr[i : i + chunk_size]
            for j in range(chunk_size): # 循环固定次数
                if min_value < tmp_view[j] < max_value:
                    count += 1
            if count > 0: # 检查当前块是否找到，如果找到则可以提前返回
                return 1
        else:
            # 处理剩余的、不足一个完整块的元素
            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     <----------