Numpy中A加v后列级Argmax的高效计算优化问询

我有一个形状为(m, 1)的Numpy数组v，以及一个形状为(m, n)的二维Numpy数组A。需要将v与A的每一列相加，然后计算每一列的argmax值。目前通过广播实现加法后调用np.argmax完成计算，但该操作需重复执行多次，且m的取值在万级范围，希望能加速计算。

已尝试使用Numba并行方案，获得了30%的提速，想进一步优化。另外补充提问：若每次调用时仅v发生变化、A保持固定，是否存在其他优化手段？

当前实现代码 #

v = np.array([[1], [2], [3]])
A = np.array([[2, 3, 4, 5], [6, 7, 8, 9], [10, 11, 12, 13]])

# 将v与A的每一列相加
B = A + v

# 计算每一列的argmax
argmax_cols = np.argmax(B, axis=0)

尝试的Numba并行实现 #

from numba import njit, prange
import numpy as np
import numba as nb

@njit((nb.int64[:])(nb.float64[:], nb.float64[:,:]), parallel=True)
def find_index_numba(v, A):
    
    _, n = A.shape

    indices = np.empty(n, dtype=np.int64)

    for j in prange(n):
        summation = np.add(v, A[:,j])
        indices[j] = np.argmax(summation)

    return indices

一、通用优化（不区分A是否固定） #

1. 消除Numba循环内的临时数组开销 #

当前Numba实现中每次循环都会创建summation临时数组，带来额外内存分配和拷贝开销。可以直接遍历元素追踪最大值及索引，避免临时数组：

@njit((nb.int64[:])(nb.float64[:], nb.float64[:,:]), parallel=True)
def find_index_numba_optimized(v, A):
    m, n = A.shape
    indices = np.empty(n, dtype=np.int64)
    
    for j in prange(n):
        max_val = -np.inf
        max_idx = 0
        for i in range(m):
            current = A[i, j] + v[i]
            if current > max_val:
                max_val = current
                max_idx = i
        indices[j] = max_idx
    return indices

该实现省去了临时数组创建和np.argmax调用的额外开销，能进一步提升并行效率。

2. 降低数据精度减少内存带宽占用 #

若业务场景允许，将v和A的数据类型从float64改为float32，可减少一半内存占用，提升内存访问速度。同时修改Numba的类型签名适配新类型。

3. 简化Numpy向量化操作 #

原Numpy实现可直接合并为一行，部分场景下Numpy会优化内存使用，避免完整存储中间数组B：

argmax_cols = np.argmax(A + v, axis=0)

二、A固定时的特殊优化 #

1. 将A设为Numba全局变量预编译 #

把固定的A设为全局变量，让Numba在编译时就能获取A的形状和内存布局，生成更高效的机器码：

from numba import njit, prange
import numpy as np
import numba as nb

# 预定义固定的A为全局变量
A_global = np.array([[2, 3, 4, 5], [6, 7, 8, 9], [10, 11, 12, 13]])

@njit(nb.int64[:](nb.float64[:]), parallel=True)
def find_index_numba_fixed_A(v):
    m, n = A_global.shape
    indices = np.empty(n, dtype=np.int64)
    
    for j in prange(n):
        max_val = -np.inf
        max_idx = 0
        for i in range(m):
            current = A_global[i, j] + v[i]
            if current > max_val:
                max_val = current
                max_idx = i
        indices[j] = max_idx
    return indices

2. GPU加速（CuPy） #

若有GPU资源，用CuPy替代Numpy可获得数量级提速。将固定的A提前加载到GPU显存，每次仅传递v到GPU计算：

import cupy as cp

# 预加载固定的A到GPU显存
A_gpu = cp.array(A)

def compute_argmax_gpu(v):
    v_gpu = cp.array(v)
    argmax_cols_gpu = cp.argmax(A_gpu + v_gpu, axis=0)
    return cp.asnumpy(argmax_cols_gpu)

3. 使用MKL优化版Numpy #

若使用Intel CPU，切换到MKL优化的Numpy版本（如Anaconda默认Numpy），MKL对广播、argmax等操作有更高效的并行实现，无需修改代码即可获得显著提速。

内容的提问来源于stack exchange，提问作者optimal-br