结构化数组的Numba签名

Question

Numba的文档没有给出接受结构化数组的函数的签名的任何示例。我尝试了几种方法，但都被Numba(和Pylance)拒绝了。

import numba as nb
import numpy as np
PairSpec = [("x", np.float32), ("y", np.float32)]
Pair = np.dtype(PairSpec)
NumbaPair = nb.from_dtype(Pair)
# BUG None of this works
# @nb.jit(np.float32(Pair[:]))
# @nb.jit(np.float32(NumbaPair[:]))
@nb.jit
def sum(pairs):
pair = pairs[0]
return pair.x + pair.y

pairs = np.array([(2, 3)], dtype=PairSpec)
print(sum(pairs))

如何给一个函数签名，采取结构化数组?

Answer 1

正确的签名是nb.float32(NumbaPair[:])。注意这里使用的是nb.float32而不是np.float32。另外请注意，结构的数组(AoS)通常比数组的结构(SoA)效率更低。对于坐标来说尤其如此，因为大多数字段通常是被读取的，而AoS阻止了任何有效的向量化(而现代x86-64处理器通常每个周期和每个内核可以计算~16个float32值，而标量值只能计算2个)。