PyCUDA: syntax for function that calls a function

Question

When using a function from a SourceModule that depends on another function in the SourceModule, how do I pass it in the function call, i.e. what is "???" in the following code:

import numpy
import pycuda.autoinit
import pycuda.driver as drv
from pycuda.compiler import SourceModule

mod = SourceModule("""
__global__ void make_square(float *in_array, float *out_array)
{
  int i;
  int N = 5;
  for (i=0; i<N; i++)
  {
    out_array[i] = pow(in_array[i],2);
  }
}
__global__ void make_square_add_one(float *in_array, float *out_array, void make_square(float *, float *))
{
  int N = 5;
  make_square(in_array,out_array);
  for (int i=0; i<N; i++)
    out_array[i] = out_array[i] + 1;
}
""")

make_square = mod.get_function("make_square")
make_square_add_one = mod.get_function("make_square_add_one")
in_array = numpy.array([1.,2.,3.,4.,5.]).astype(numpy.float32)
out_array = numpy.zeros_like(in_array).astype(numpy.float32)
make_square_add_one(drv.In(in_array), drv.Out(out_array), ??? , block = (1,1,1), grid = (1,1))

Thanks for any information.

Answer 1

In the traditional CUDA execution model, __global__ functions are kernels, and they can't be passed as arguments to other kernels and they can't be called by other kernels. It looks like make_square should really be a device function, something like:

__device__ void make_square(float *in_array, float *out_array)
{
  int i;
  for (i=0; i<5; i++)
  {
    out_array[i] = pow(in_array[i],2);
  }
}

which is then called from the running kernel as:

__global__ void make_square_add_one(float *in_array, float *out_array)
{
  int N = 5;
  make_square(in_array,out_array);
  for (int i=0; i<N; i++)
    out_array[i] = out_array[i] + 1;
}

It is worth noting that this kernel is entirely serial and pretty much orthogonal to how CUDA kernels would be expected to be written.