在 Cython 中使用半精度 NumPy 浮点数

Question

我正在尝试从一些 Cython 代码将 float16 数据发送到 Nvidia P100 卡。当我使用 float32 时，我可以在 Cython 中定义我的类型，如下所示：

DTYPE = np.float32
ctypedef np.float32_t DTYPE_t
cdef np.ndarray[DTYPE_t, ndim=2] mat = np.empty((100, 100), dtype=DTYPE)

但 Cython 无法找到 np.float16_t 的定义类型，因此我不能仅将 32 替换为 16。如果我尝试提供另一种占用相同空间量的类型（例如 np.uint16_t），我会收到类似错误:

Does not understand character buffer dtype format string ('e')

当我用谷歌搜索时，我所能找到的只是 2011 年的一条线索，内容是人们试图弄清楚如何支持它……现在肯定有解决方案了吗？

Answer 1

我认为答案是“有点，但如果你想做任何实际的计算，这是一个合理的工作量”。

基本问题是 C 在大多数 PC 上不支持 16 位浮点类型（因为处理器指令不存在）。因此，numpy 所做的就是typedefed 一个 16 位无符号 int 来存储 16 位浮点数，然后编写一组函数将其与支持的浮点数类型相互转换。使用的任何计算np.float16实际上都是在 32 位或 64 位浮点上完成的，但在计算之间数据以 16 位格式存储。

这样做的结果是 Cython 没有一种简单的方法来为它需要执行的任何计算生成有效的 C 代码。结果是您可能需要自己编写此 C 代码。

复杂程度有多个级别，具体取决于您实际想要执行的操作：

1) 不要输入任何内容

Cython 实际上不需要您指定任何类型 - 它可以轻松编译 Python 代码。因此，不要为半浮点数组分配类型，而只需通过 Python 对象来分配类型。这可能不是很快，但值得记住的是它会起作用。

2）要移动数据，您可以将view其作为uint16

如果您只是打乱数据，那么可以定义uint16数组并使用它们将其从一个地方复制到另一个地方。使用 numpyview函数以 Cython 识别的格式获取数据并将其返回。然而，你不能在这种模式下做数学（答案将毫无意义）。

from libc.stdint cimport uint16_t

import numpy as np

def just_move_stuff(x):
    assert x.dtype == np.float16
    # I've used memoryviews by cdef np.ndarray should be fine too
    cdef uint16_t[:] x_as_uint = x.view(np.uint16)

    cdef uint16_t[:] y_as_uint = np.empty(x.shape,dtype=np.uint16)

    for n in range(x_as_uint.shape[0]):
        y_as_uint[n] = x_as_uint[n]

    return np.asarray(y_as_uint).view(dtype=np.float16)

该view函数不进行复制，因此使用起来非常便宜。

3）通过手动转换进行数学计算

如果您想做任何计算，您需要使用 numpy 的转换函数将“半浮点”数据更改为全浮点并返回。如果您忘记这样做，您得到的答案将毫无意义。首先将它们包含在numpy/halffloat.h：

cdef extern from "numpy/halffloat.h":
    ctypedef uint16_t npy_half

    # conversion functions
    float npy_half_to_float(npy_half h);
    npy_half npy_float_to_half(float f);

def do_some_maths(x):
    assert x.dtype == np.float16
    cdef uint16_t[:] x_as_uint = x.view(np.uint16)

    cdef uint16_t[:] y_as_uint = np.empty(x.shape,dtype=np.uint16)

    for n in range(x_as_uint.shape[0]):
        y_as_uint[n] = npy_float_to_half(2*npy_half_to_float(x_as_uint[n]))

    return np.asarray(y_as_uint).view(dtype=np.float16)

此代码要求您链接到 numpy 核心数学库：

from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
from numpy.distutils.misc_util import get_info
info = get_info('npymath')

ext_modules = [Extension("module_name", ["module_name.pyx"],**info)]

setup(
  ext_modules = cythonize(ext_modules)
)