优化Numba-JIT量子比特移除函数性能及分析工具推荐

问题背景与当前实现 #

你正在优化一个量子态矢量处理的核心函数，该函数负责移除可分量子比特并重组剩余态矢量，是流程中的性能瓶颈（比次慢函数慢2倍），目标是实现约50%的性能提升。函数参数与逻辑如下：

函数参数 #

• psi : npt.NDArray[np.complex128]：长度为$2^{N$的一维复数组，仅前$2}{n_{ ext{qubit}}}$个元素有效，会被原地修改
• nqubit : int：有效量子比特数，取值范围[20,25]
• q : int：待移除的量子比特索引，满足$0 \leq q < n_{ ext{qubit}}$
• atol : float：零值判断容差，默认1e-10

核心逻辑 #

将psi有效部分按q划分为两个不相交子矢量：

• 若两者均非零，选取其一归一化后写入psi前$2^{n_{ ext{qubit}}-1}$位置
• 若其一为零，选取非零子矢量
• 若两者均为零，抛出RuntimeError

当前实现代码 #

import numba as nb
import numpy as np
import numpy.typing as npt
import math

@nb.njit("int32(c16[::1], int32, int32, f8)", parallel=False)
def _remove_qubit_jit(
    psi: npt.NDArray[np.complex128],
    nqubit: int,
    q: int,
    atol: float,
) -> int:
    new_nqubit = nqubit - 1

    n_blocks = 1 << q
    size_block = 1 << nqubit - q  # 2**(nqubits - q)
    size_half_block = size_block >> 1

    # Compute norm of branch 0
    norm2 = 0.0
    shift = 0
    b0 = shift
    for _ in range(n_blocks):
        # If parallelization, set `b0 = b * size_block + shift` with `b` the loop variable to avoid race condition.
        # Parallelization for norm computation is not worth, execution-time controlled by the update loop which can't be parallelized without cache.
        for j in range(size_half_block):
            a = psi[b0 | j]
            a_re = a.real
            a_im = a.imag
            norm2 += a_re * a_re + a_im * a_im
        b0 += size_block

    # If norm of branch 0 is 0, compute norm of branch 1 and set shift to branch 1
    if norm2 <= atol:
        norm2 = 0.0
        shift = size_half_block
        b0 = shift
        for _ in range(n_blocks):
            for j in range(size_half_block):
                a = psi[b0 | j]
                a_re = a.real
                a_im = a.imag
                norm2 += a_re * a_re + a_im * a_im
            b0 += size_block

    if norm2 <= atol:
        raise RuntimeError(f"Attempted to remove qubit {q} from 0-norm statevector.")

    b0 = shift
    k = 0
    inv_norm = 1.0 / math.sqrt(norm2)

    # Update `psi` with selected and normalized elements.
    for _ in range(n_blocks):
        for j in range(size_half_block):
            psi[k] = (
                psi[b0 | j] * inv_norm
            )  # b0 | j equivalent to b0 + j because the active bits of b0 and j don't overlap.
            k += 1
        b0 += size_block

    return new_nqubit

测试情况 #

在Apple M4 24GB设备上，当nq=30、q=nq-1时，函数执行时间为6.657s，且q值越大性能越差。

优化思路 #

1. 优化内存访问模式，缓解内存绑定 #

当q较大时，n_blocks=2^q、size_block=2^(nqubit-q)，内存访问会出现大量跳跃（每次跳转size_block字节），导致缓存命中率极低，成为性能瓶颈。可通过以下方式优化：

• 交换循环层级：将外层块循环与内层元素循环交换，改为先遍历所有块的同一位置元素，再移动到下一个位置，提升缓存局部性。
• 批量内存复制：在最终的psi更新阶段，利用Numba的内存拷贝API（如nb.memcpy）替代逐个元素赋值，减少循环开销。
• 针对大q的特殊处理：当size_half_block=1（即q=nqubit-1）时，直接遍历所有块并批量复制元素，省去内层循环的开销。

2. 提升SIMD计算单元利用率 #

当前代码的范数计算为逐元素操作，未充分利用SIMD向量化能力，可通过以下方式优化：

• 开启FastMath优化：在@nb.njit中添加fastmath=True参数，允许编译器进行更多浮点向量化优化。
• 手动向量化循环：将内层循环改为按固定步长（如4或8，对应SIMD寄存器宽度）批量处理元素，显式利用SIMD指令。例如：

  # 按4个元素为一组处理
  for j in range(0, size_half_block, 4):
      val0 = psi[b0 + j]
      val1 = psi[b0 + j + 1]
      val2 = psi[b0 + j + 2]
      val3 = psi[b0 + j + 3]
      norm2 += (val0.real**2 + val0.imag**2) + (val1.real**2 + val1.imag**2) + \
               (val2.real**2 + val2.imag**2) + (val3.real**2 + val3.imag**2)
  

• 合并范数计算：一次遍历同时计算两个分支的范数，避免两次完整的内存读取，减少内存带宽占用。示例：

  branch0_norm2 = 0.0
  branch1_norm2 = 0.0
  for b in range(n_blocks):
      base = b * size_block
      for j in range(size_half_block):
          val0 = psi[base + j]
          val1 = psi[base + j + size_half_block]
          branch0_norm2 += val0.real**2 + val0.imag**2
          branch1_norm2 += val1.real**2 + val1.imag**2
  # 后续判断分支有效性
  

3. 减少分支与重复计算 #

• 避免双重遍历：如上述合并范数计算的方式，只需一次内存遍历即可获取两个分支的范数，省去分支0无效时的二次遍历开销。
• 简化位运算：将b0 | j替换为b0 + j（因两者位域不重叠，结果等价），减少位运算的微小开销。

4. Numba编译优化 #

• 精确类型注解：使用更精确的Numba类型（如nb.types.Complex128[:]）替代泛型注解，帮助编译器生成更高效的机器码。
• 启用循环展开：通过@nb.njit(loop_unroll=True)开启循环展开优化，减少循环控制的开销。

性能分析工具推荐 #

1. macOS平台：Instruments #

• 使用Time Profiler模板：可查看函数的CPU耗时分布、循环执行时间，定位瓶颈代码。
• 使用Memory Cache模板：分析缓存命中率、内存访问延迟，确认是否为内存绑定问题。

2. Linux平台：perf #

• 记录性能数据：perf record -g python your_test_script.py
• 查看报告：perf report，可查看函数调用栈、CPU周期占比、缓存未命中次数。
• 快速统计缓存情况：perf stat -e cache-misses python your_test_script.py

3. Numba自带工具 #

• 查看编译日志：设置环境变量NUMBA_DEBUG=1运行程序，检查是否开启向量化优化。
• 查看编译信息：使用numba -s your_script.py查看JIT编译后的代码细节。

4. 跨平台工具：Intel VTune #

• 支持ARM架构（Apple M系列），可分析SIMD指令利用率、内存带宽、循环向量化程度，精准定位性能瓶颈。

内容的提问来源于stack exchange，提问作者matulni