Numba并行计算中如何安全修改NumPy数组元素

• 方案1：预分配固定存储区间，按循环下标隔离写入
  你的循环总次数是固定的100次，完全可以避开动态append操作，提前分配好存储所有索引的数组空间，每个并行迭代的线程只写入自己专属的内存区间，从根源上避免跨线程写冲突：
  • 如果get_indice()每次返回的索引长度固定，直接预分配二维数组存储所有轮次的索引，循环结束后拉平数组统一更新原数组即可，参考实现：

    import numpy as np
    from numba import njit, prange
    
    @njit(parallel=True)
    def myfunction(arr):
        # 替换per_idx_len为get_indice()实际返回的固定索引长度
        per_idx_len = 5
        all_indices = np.empty((100, per_idx_len), dtype=np.int64)
        for i in prange(100):
            indices = get_indice()
            all_indices[i] = indices
        # 并行循环结束后单线程批量更新，无竞态
        arr[all_indices.ravel()] += 1
    

  • 如果get_indice()每次返回的索引长度不固定，先串行执行一次100轮循环，统计每轮返回的索引长度，通过前缀和计算每轮i对应的写入起止偏移，再预分配总长度匹配的一维索引数组，并行阶段每个迭代只往自己的偏移区间写入数据，同样是线程安全的。
• 方案2：使用线程本地存储隔离动态列表
  并行场景下list.append()不安全的核心原因是多个线程同时修改同一个列表的内部元数据，你可以让每个执行线程维护自己私有的索引列表，并行阶段线程只操作自己的私有列表，不会产生冲突。等所有并行任务执行完成后，单线程合并所有线程的私有列表得到完整索引数组，再统一更新原数组即可。
• 方案3：直接使用原子操作省略索引收集步骤
  如果你只是做+=1这类简单的累加更新，完全不需要绕路收集索引，直接调用Numba内置的原子操作就能自动处理并行写冲突，代码更简洁，性能通常也比收集索引的方案更好，参考实现：

  import numpy as np
  from numba import njit, prange
  from numba.np import atomic
  
  @njit(parallel=True)
  def myfunction(arr):
      for i in prange(100):
          indices = get_indice()
          for idx in indices:
              atomic.add(arr, idx, 1)
  

  注意：原子操作仅支持加减、赋值、位运算这类基础内存操作，如果你的更新逻辑包含复杂的条件判断、多值联动写入，优先选择前两种方案。

• 方案4：线程私有副本+最终归并
  如果待更新的数组体积不大，可以在并行阶段不对原数组做任何写入，每个线程维护自己的局部计数副本，并行循环结束后把所有线程的局部副本累加合并到原数组上，这种方式逻辑最简单，完全不需要考虑索引存储的问题，适合索引重复率高、更新逻辑复杂的场景。

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