Расчет Numbapro Jit дает неверный результат

Question

У меня есть кусок кода, который использует Numbapro для написания простого ядра, чтобы выстроить в квадрат содержимое двух массивов размером 41724, сложить их вместе и сохранить в другом массиве. Все массивы имеют одинаковый размер и являются float32. Код ниже:

import numpy as np
from numba import *
from numbapro import cuda

@cuda.jit('void(float32[:],float32[:],float32[:])')
def square_add(a,b,c):
    tx = cuda.threadIdx.x
    bx = cuda.blockIdx.x
    bw = cuda.blockDim.x

    i = tx + bx * bw

    #Since the length of a is 41724 and the total
    #threads is 41*1024 = 41984, this check is necessary
    if (i>len(a)):
            return
    else:
            c[i] = a[i]*a[i] + b[i]*b[i]


a = np.array(range(0,41724),dtype = np.float32)
b = np.array(range(41724,83448),dtype=np.float32)
c = np.zeros(shape=(1,41724),dtype=np.float32)

d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_c = cuda.to_device(c,copy=False)

#Launch the kernel; Gridsize = (1,41),Blocksize=(1,1024)
square_add[(1,41),(1,1024)](d_a,d_b,d_c)

c = d_c.copy_to_host()
print c
print len(c[0])

Значения, которые я получаю, когда я печатаю результат операции (массив c), полностью отличаются от тех, которые я делаю точно так же в терминале Python. Я не знаю, что я делаю не так здесь.

Answer 1

Здесь есть две проблемы.

Во-первых, вы указываете размер блока и сетки для запуска ядра CUDA, что несовместимо со схемой индексации, которую вы выбрали для использования в ядре.

Это:

square_add[(1,41),(1,1024)](d_a,d_b,d_c)

запускает двумерную сетку, где все потоки имеют одинаковые размеры блоков и нитей в x и меняются только по y. Это подразумевает, что

tx = cuda.threadIdx.x
bx = cuda.blockIdx.x
bw = cuda.blockDim.x

i = tx + bx * bw

даст i=0 для каждой темы. Если вы измените запуск ядра на это:

square_add[(41,1),(1024,1)](d_a,d_b,d_c)

вы обнаружите, что в индексации будет работать правильно.

Второе, что c был объявлен как двумерный массив, но подпись функции ядра была объявлена ​​как одномерный массив. При некоторых обстоятельствах среда выполнения numbapro должна обнаружить это и вызвать ошибку.

Я смог заставить ваш пример работать правильно, вот так:

import numpy as np
from numba import *
from numbapro import cuda

@cuda.jit('void(float32[:],float32[:],float32[:,:])')
def square_add(a,b,c):
    tx = cuda.threadIdx.x
    bx = cuda.blockIdx.x
    bw = cuda.blockDim.x

    i = tx + bx * bw

    if (i<len(a)):
        c[0,i] = a[i]*a[i] + b[i]*b[i]

a = np.array(range(0,41724),dtype=np.float32)
b = np.array(range(41724,83448),dtype=np.float32)
c = np.zeros(shape=(1,41724),dtype=np.float32)

d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_c = cuda.to_device(c, copy=False)

square_add[(41,1),(1024,1)](d_a,d_b,d_c)

c = d_c.copy_to_host()
print(c)
print(c.shape)

[Обратите внимание, что я использую Python 3, поэтому здесь используются операторы печати нового стиля]

$ ipython numbatest.py 
numbapro:1: ImportWarning: The numbapro package is deprecated in favour of the accelerate package. Please update your code to use equivalent functions from accelerate.
[[  1.74089216e+09   1.74097562e+09   1.74105907e+09 ...,   8.70371021e+09
    8.70396006e+09   8.70421094e+09]]
(1, 41724)