Skip to content

Instantly share code, notes, and snippets.

@pentschev

pentschev/cupy_add_broadcast_kernel.py

Created August 9, 2019 20:21
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save pentschev/5946bc72e19b6e64a8b4358f7c73fa9e to your computer and use it in GitHub Desktop.
Save pentschev/5946bc72e19b6e64a8b4358f7c73fa9e to your computer and use it in GitHub Desktop.
Download ZIP
Blog Post - Parallelizing Custom CuPy Kernels with Dask - CuPy Kernel
Raw
cupy_add_broadcast_kernel.py
add_broadcast_kernel = cupy.RawKernel(
r'''
extern "C" __global__
void add_broadcast_kernel(
const float* x, const float* y, float* z,
const int xdim0, const int zdim0)
{
int idx0 = blockIdx.x * blockDim.x + threadIdx.x;
int idx1 = blockIdx.y * blockDim.y + threadIdx.y;
z[idx1 * zdim0 + idx0] = x[idx1 * xdim0 + idx0] + y[idx0];
}
''',
'add_broadcast_kernel'
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment