Pickling a specialized CUDA kernel

example_pickle_cuda.py

from numba import cuda, typeof
import numpy as np
import pickle


@cuda.jit
def foo(x, v):
    x[0] = v


# Specializes the cuda kernel
example_int32_input = np.zeros(shape=1, dtype=np.int32)
int32_input_type = typeof(example_int32_input)
print('typeof(example_int32_input):', int32_input_type)

# Specialize GPU kernel
signature = (int32_input_type, int32_input_type.dtype)
print('signature:', signature)
foo_int32 = foo.compile(signature)


# Specialized GPU kernel is pickle-able.  The PTX code is stored.
dumped = pickle.dumps(foo_int32)
print('Pickled:', dumped)
rebuilt = pickle.loads(dumped)

print('precall', example_int32_input)
rebuilt[1, 1](example_int32_input, 23)
print('postcall', example_int32_input)