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

Question

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

<div class="onoffswitch">
<TD align="right">Box2</TD>
<input type="checkbox" name="onoffswitch" class="onoffswitch-checkbox"    id="myonoffswitch5" checked>

但是Cython找不到np.float16_t的定义类型，所以我不能只用16替换32。如果我尝试提供另一种占用相同空间的类型，比如np.uint16_t，我得到如下错误：

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

当我谷歌时，我所能找到的只是2011年关于人们试图弄清楚如何支持它的线索......当然必须有一个解决方案吗？

Answer 1

我认为答案是＆＃34;但是，如果你想做任何真正的计算，这是一项合理的工作量。

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

这样做的结果是Cython没有一种简单的方法可以为它需要做的任何计算生成有效的C代码。结果是你可能需要自己写出这个C代码。

有许多级别的复杂性，具体取决于您实际想要做的事情：

1）不要输入任何内容

Cython实际上并不需要你指定任何类型 - 它可以很快地编译Python代码。因此，不要将类型分配给半浮点数组，只需让Python对象使用它们即可。这可能不会非常快，但值得记住它会起作用。

2）要移动数据，您可以view将其作为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的转换功能来改变你的半浮动＆＃34;数据到完全浮动和返回。如果你忘记这样做，你得到的答案将毫无意义。首先将其包含在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)
)