如何注册np.float128为有效Numba类型以实现数组求和？

需要编写Numba函数对np.float128数组求和（环境中该类型为80位精度，转为标准float128也可接受），要求快速且无精度损失，返回np.float128类型结果。

已完成基础的类型映射定义，但不清楚如何替换bfloat16的加法实现以适配float128，调用现有Numba求和函数时会报错NumbaValueError: Unsupported array dtype: float128。

补充性能背景：

• float64数组场景下，启用fastmath=True的Numba求和比np.sum快数倍；
• 转为np.float128后，np.sum性能下降，但Numba无法直接处理该类型数组。

1. 完善float128的Numba类型适配 #

仅定义类型与dtype的映射不足以让Numba支持该类型，必须为np.float128实现基础运算（如加法）的内联函数，结合x86平台的硬件指令（因为np.float128在x86上对应80位扩展精度的long double）完成适配。

2. 替换bfloat16_add为float128加法实现 #

以下是适配float128的完整代码，包含类型定义、加法内联函数及操作符重载：

import numpy as np
from numba.core.types.scalars import Number
from numba.np import numpy_support
from numba.core import ir, types
from numba.core.extending import intrinsic, overload
import operator

# 定义float128类型
class Float128(Number):
    def __init__(self, *args, **kws):
        super().__init__(name='float128')

float128_type = Float128()
# 双向映射Numba类型与numpy dtype
numpy_support.FROM_DTYPE[np.dtype(np.float128)] = float128_type
numpy_support.TO_DTYPE[float128_type] = np.dtype(np.float128)

# 实现float128加法的内联汇编函数
@intrinsic
def float128_add(typingctx, a, b):
    # 定义函数签名：两个float128输入，返回float128
    sig = float128_type(float128_type, float128_type)
    
    def codegen(context, builder, sig, args):
        # x86平台的80位扩展精度对应ir.FloatType(80)
        f80 = ir.FloatType(80)
        func_type = ir.FunctionType(f80, [f80, f80])
        # 内联汇编使用FPU指令完成加法，保证精度
        asm_code = """
            fldt $1
            fldt $2
            faddp %st(1), %st(0)
            fstpt $0
        """
        # 约束说明：=t表示返回值存在FPU栈顶，t表示输入参数在FPU栈顶
        asm = ir.InlineAsm(func_type, asm_code, "=t,t,t")
        return builder.call(asm, args)
    
    return sig, codegen

# 重载加法操作符，让Numba识别float128的加法
@overload(operator.add)
def overload_float128_add(a, b):
    if isinstance(a, Float128) and isinstance(b, Float128):
        def impl(a, b):
            return float128_add(a, b)
        return impl

3. 编写支持float128的Numba求和函数 #

import numba as nb

@nb.njit(cache=True)
def fast_float128_sum(arr):
    # 初始化求和变量为float128类型的0
    s = np.float128(0.0)
    for v in arr:
        s += v
    return s

关键注意事项 #

• x86平台上np.float128对应硬件的80位扩展精度，内联汇编直接使用FPU指令，确保精度无损失；
• 若需要乘法、减法等其他运算，需参照加法的实现方式编写对应的内联函数并重载操作符；
• 若启用fastmath=True，需注意部分优化可能破坏80位精度的保留，严格要求无精度损失时建议关闭该选项。

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