使用numba.cuda在GPU上运行Sympy lambdify函数

问题描述

我正在尝试使用numba.cuda在GPU上运行任意sympy lambdify函数。到目前为止，由于numba.jit允许函数返回值，但是numba.cuda.jit不允许这样做（numba.cuda.jit内核无法返回值），因此我在每个步骤中都遇到了错误。这可能是由于我对numba的工作方式有一个基本的误解，但是文档中的示例有些稀疏，因此我尝试对每个给定的示例进行变异，以尝试做我不希望的事情。

我尝试过的例子：

非CUDA jit功能（有效）

import sympy
from sympy.abc import y
import numba

f = sympy.lambdify(y,sympy.sin(y),'math')
g = numba.jit(f)
g(1) #returns 0.8414709848078965

相同代码的CUDA jit示例

import sympy
from sympy.abc import y
from numba import cuda

f = sympy.lambdify(y,'math')
g = cuda.jit(f)
g(1) #error

返回以下内容：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No conversion from float64 to none for '$8return_value.3',defined at None
File "<lambdifygenerated-5>",line 2:
def _lambdifygenerated(y):
    return (sin(y))
    ^

During: typing of assignment at <lambdifygenerated-5> (2)

File "<lambdifygenerated-5>",line 2:
def _lambdifygenerated(y):
    return (sin(y))

进行检查，因为以这种方式创建的cuda函数无法返回值，但似乎暗示numba可以为cuda编译此函数没有问题，只是它无法弄清楚如何处理最后的值。

（天真）明显的答案似乎是创建一个函数来获取该值并将其分配给变量，类似于Numba文档（https://numba.pydata.org/numba-doc/dev/cuda/kernels.html#thread-positioning）中的示例中的操作：

原始代码（有效）：

import sympy
import numpy
from sympy.abc import y
from numba import cuda

f = sympy.lambdify(y,'math')

@cuda.jit
def increment_by_one(an_array):
    # Thread id in a 1D block
    tx = cuda.threadIdx.x
    # Block id in a 1D grid
    ty = cuda.blockIdx.x
    # Block width,i.e. number of threads per block
    bw = cuda.blockDim.x
    # Compute flattened index inside the array
    pos = tx + ty * bw
    if pos < an_array.size:  # Check array boundaries
        an_array[pos] += 1
        
array = numpy.arange(3.)
print(array) #returns [0. 1. 2.]
blockspergrid = 2
threadsperblock = 32
increment_by_one[blockspergrid,threadsperblock](array)
print(array) #returns [1. 2. 3.]

更改了代码（无效）

import sympy
import numpy
from sympy.abc import y
from numba import cuda

f = sympy.lambdify(y,'math')

@cuda.jit
def cuda_f(an_array):
    # Thread id in a 1D block
    tx = cuda.threadIdx.x
    # Block id in a 1D grid
    ty = cuda.blockIdx.x
    # Block width,i.e. number of threads per block
    bw = cuda.blockDim.x
    # Compute flattened index inside the array
    pos = tx + ty * bw
    if pos < an_array.size:  # Check array boundaries
        an_array[pos] = f(an_array[pos])
        
array = numpy.arange(3.)
print(array) #returns [0. 1. 2.]
cuda_f(array)
blockspergrid = 2
threadsperblock = 32
cuda_f[blockspergrid,threadsperblock](array) #error
print(array)

此操作失败，并显示以下代码：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'f': cannot determine Numba type of <class 'function'>

File "<ipython-input-7-9ef7fd8543d7>",line 19:
def cuda_f(an_array):
    <source elided>
    if pos < an_array.size:  # Check array boundaries
        an_array[pos] += f(an_array[pos])
        ^

此错误（“无法确定类的Numba类型”）似乎在我进行GPU验证该功能的所有其他尝试中都普遍存在，包括基于此处所有代码段的尝试：{{ 3}}。

很明显，这意味着我在执行此操作方面做错了。可以做到这一点（使用numba.cuda在GPU上并行运行sympy lambdify函数）吗？

编辑：通过将sympy lambdify函数转换为numba.jit函数，然后在cuda内核中运行它，我取得了一些成功。我使用了以下代码：

import sympy
import numpy
from sympy.abc import y
from numba import cuda
import numba

f = sympy.lambdify(y,'math')
g = numba.jit(f)

@cuda.jit
def sympy_kernel(x,out):
    startx = cuda.grid(1)    
    stridex = cuda.gridsize(1) 

    for i in range(startx,x.shape[0],stridex):
        out[i] = g(x[i])

@numba.jit
def sympy_cpu(x,out):
    for i in range(len(out)):
        out[i] = g(x[i])
        
        
array = numpy.arange(100000000.)
array_device = cuda.to_device(array)
out = numpy.arange(100000000.)
out_device = cuda.to_device(out)
blockspergrid = 64
threadsperblock = 64
%timeit -n5 sympy_kernel[blockspergrid,threadsperblock](array_device,out_device); cuda.synchronize()
%timeit -n5 sympy_cpu(array,out)
out_host = out_device.copy_to_host()
print(out_host)

返回值是：

26.9 ms ± 15.7 ms per loop (mean ± std. dev. of 7 runs,5 loops each)
958 ms ± 16.3 ms per loop (mean ± std. dev. of 7 runs,5 loops each)
[ 0.          0.84147098  0.90929743 ... -0.87103474 -0.05727351
  0.80914472]

解决方法

我正在尝试使用numba.cuda在GPU上运行任意sympy lambdify函数

在当前的开发状态下，Numba尚不支持并且无法做到这一点。 Numba在GPU上仅支持Python语言功能的骨架，如果您不能直接将功能降级为受支持的math函数，那么GPU上将不支持任何外部功能。