Blog Post - Parallelizing Custom CuPy Kernels with Dask - Kernel Dispatcher

dispatch_add_broadcast.py

def dispatch_add_broadcast(x, y):
    block_size = (32, 32)
    grid_size = (x.shape[1] // block_size[1], x.shape[0] // block_size[0])

    z = cupy.empty(x.shape, x.dtype)

    xdim0 = x.strides[0] // x.strides[1]
    zdim0 = z.strides[0] // z.strides[1]

    add_broadcast_kernel(grid_size, block_size, (x, y, z, xdim0, zdim0))
    return z