Can't get same values as numpy elementwise matrix multiplication using numba

Here is how you could debug this.

Consider a smaller and simplified example with:

• reduced array sizes, e.g. (2, 3, 1) (so you could actually print the values and be able to read them)
• simple and deterministic contents, e.g. "all ones" (to compare across runs)
• additional kernel arguments for debugging

from __future__ import (division, print_function)

import numpy as np
from numba import cuda

M = 2
N = 3
P = 1

threadsperblock = 1
blockspergrid = (M * N * P + (threadsperblock - 1)) // threadsperblock


@cuda.jit
def mult_gpu_3d(a, b, c, grid_ran, grid_multed):
    grid = cuda.grid(3)
    x, y, z = grid

    grid_ran[x] = 1

    if (x < c.shape[0]) and (y < c.shape[1]) and (z < c.shape[2]):
        grid_multed[x] = 1
        c[grid] = a[grid] * b[grid]


if __name__ == '__main__':
    A = np.ones((M, N, P), np.int32)
    B = np.ones((M, N, P), np.int32)

    A_gpu = cuda.to_device(A)
    B_gpu = cuda.to_device(B)
    C_gpu = cuda.to_device(np.zeros_like(A))

    # Tells whether thread at index i have ran
    grid_ran = cuda.to_device(np.zeros([blockspergrid], np.int32))

    # Tells whether thread at index i have performed multiplication
    grid_multed = cuda.to_device(np.zeros(blockspergrid, np.int32))

    mult_gpu_3d[blockspergrid, threadsperblock](
        A_gpu, B_gpu, C_gpu, grid_ran, grid_multed)

    print("grid_ran.shape    : ", grid_ran.shape)
    print("grid_multed.shape : ", grid_multed.shape)
    print("C_gpu.shape       : ", C_gpu.shape)

    print("grid_ran          : ", grid_ran.copy_to_host())
    print("grid_multed       : ", grid_multed.copy_to_host())

    C = C_gpu.copy_to_host()
    print("C transpose flat  : ", C.T.flatten())
    print("C                 : \n", C)

Output:

grid_ran.shape    :  (6,)
grid_multed.shape :  (6,)
C_gpu.shape       :  (2, 3, 1)
grid_ran          :  [1 1 1 1 1 1]
grid_multed       :  [1 1 0 0 0 0]
C transpose flat  :  [1 1 0 0 0 0]
C                 : 
 [[[1]
  [0]
  [0]]

 [[1]
  [0]
  [0]]]

You can see that the device grid shape does not correspond to the shape of the arrays: the grid is flat (M*N*P), while arrays are all 3-dimensional (M, N, P). That is, first dimension of the grid has indices in range 0..M*N*P-1 (0..5, totaling 6 values in my example), while first dimension of the array is only in 0..M-1 (0..1, totaling 2 values in my example). This mistake typically leads do out-of-bounds access, but you have protected your kernel with a conditional which cuts down the offending threads:

if (x <= c.shape[0])

This line does not allow threads with indices above M-1 (1 in my example) to run (well, sort of [1]), that is why no values are written and you get many zeros in the resulting array.

Possible solutions:

• In general, you could use multidimensional kernel grid configuration, i.e. a 3D vector for blockspergrid instead of a scalar [2].
• In particular, as elementwise multiplication is a map operation and does not depend on array shapes, you could flatten all 3 arrays to 1D arrays, run your kernel as is on 1D grid, then reshape the result back [3], [4].

References:

• [1] How to understand “All threads in a warp execute the same instruction at the same time.” in GPU?
• [2] Understanding CUDA grid dimensions, block dimensions and threads organization (simple explanation)
• [3] numpy.ndarray.flatten
• [4] numpy.ravel