python - 在 Cython 中使用半精度 NumPy float

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 的定义类型，所以我不能只用 16 替换 32。如果我尝试提供另一种占用相同空间量的类型，例如 np.uint16_t，我得到如下错误:

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

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

Answer 1

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

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

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

有多种复杂程度，具体取决于您实际想要做什么:

1) 不要输入任何内容

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

2) 要移动数据，您可以将其查看为uint16

如果您只是打乱数据，则可以定义 uint16 数组并使用它们将其从一个位置复制到另一个位置。使用 numpy view 函数以 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)

此代码要求您link against the numpy core math library :

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)
)