Extending Theano with a GPU Op

Note

This covers the gpuarray back-end for the GPU.

This tutorial covers how to extend Theano with an op that offers a GPUimplementation. It assumes you are familiar with how to write newTheano ops. If that is not the case you should probably follow theCreating a new Op: Python implementation and Extending Theano with a C Op sections beforecontinuing on.

Writing a new GPU op can be done in Python for some simple tasks, butwill usually done in C to access the complete API and avoid paying theoverhead of a Python function call.

Dealing With the Context

One of the major differences with GPU ops is that they require acontext (a.k.a. device) to execute. Most of the time you can inferthe context to run on from your inputs. There is a way for the userto transfer things between contexts and to tag certain variables fortransfer. It might also be the case that your inputs are not all fromthe same context and you would have to choose which one to run on.

In order to support all of those options and have a consistentinterface, theano.gpuarray.basic_ops.infer_context_name() waswritten. An example usage is below:

  1. def make_node(self, a, b, c):
  2.     ctx = infer_context_name(a, b, c)
  3.     a = as_gpuarray_variable(a, ctx)
  4.     b = as_gpuarray_variable(b, ctx)
  5.     c = as_gpuarray_variable(c, ctx)
  6.     return Apply(self, [a, b, c], [a.type()])

In this example the Op takes three inputs, all on the GPU. In caseone or more of your inputs is not supposed to be on the GPU, youshould not pass it to infer_context_name() or callas_gpuarray_variable() on it.

Also note that theano.gpuarray.basic_ops.as_gpuarray_variable()takes context_name as a mandatory parameter. This is because it’snot enough to know you want the value to be on the GPU, you also wantto know which GPU to put it on. In almost all cases, you can pass inthe return value of infer_context_name() there.

If you also need the context during runtime (for example to allocatethe output), you can use the context of one of your inputs to knowwhich one to use. Here is another example:

  1. def perform(self, node, inputs, output_storage):
  2.     A, B = inputs
  3.     C, = output_storage
  4.     C[0] = pygpu.empty([A.shape[0], B.shape[1]], dtype=A.dtype, A.context)
  5.     pygpu.blas.gemm(1, A, B, 0, C, overwrite_c=True)

Finally if you require the context before perform, such as duringmake_thunk() to initialize kernels and such, you can access thecontext of your inputs through the type of the variables:

  1. def make_thunk(self, node, storage_map, compute_map, no_recycling):
  2.     ctx = node.inputs[0].type.context

Note that GpuArrayType objects also have a context_nameattribute which is the symbolic equivalent of context. It can’tbe used for calls to pygpu or libgpuarray, but it should be used fortheano operations and variables.

The last place where you might need the context is in the Cinitialization code. For that you will have to use the params. The params type should betheano.gpuarray.type.gpu_context_type and the params objectshould be a context object from one of your input variables:

  1. def get_params(self, node):
  2.     return node.inputs[0].type.context

If you don’t have any input variables on the GPU you can follow thethe example of GpuFromHost or GpuEye. This is not a case that youshould encounter often, so it will not be covered further.

Defining New Kernels

If your op needs to do some transformation on the data, chances arethat you will need to write a new kernel. The best way to do this isto leverage GpuKernelBase (or CGpuKernelBase if you want to use theCOp functionality).

For plain GpuKernelBase, you have to define amethod called gpu_kernels which returns a list of Kernel objects. You can define as manykernels as you want for a single op. An example would look likethis:

  1. def gpu_kernels(self, node, name):
  2.     code = """
  3. KERNEL void k(GLOBAL_MEM ga_double *a, ga_size n, ga_size m) {
  4.     ga_size nb = n < m ? n : m;
  5.     for (ga_size i = LID_0; i < nb; i += LDIM_0) {
  6.         a[i*m + i] = 1;
  7.     }
  8. }"""
  9.     return [Kernel(
  10.             code=code, name="k",
  11.             params=[gpuarray.GpuArray, gpuarray.SIZE, gpuarray.SIZE],
  12.             flags=Kernel.get_flags('float64'))]

If you want to use COp, then you should use CGpuKernelBaseinstead. It adds a new section to the parsed files whose tag iskernels. Inside that section you can define some kernels with#kernel name:params:flags.

Here name is the name of the kernel function in the followingcode, params is a comma-separeted list of numpy typecode names.There are three exceptions for size_t which should be noted assize, ssize_t which should be noted as ssize and a pointerwhich should be noted as *.

flags is a |-separated list of C kernel flag values (can beempty). The same kernel definition as above would look like this withCGpuKernelBase:

  1. #section kernels
  2.  
  3. #kernel k : *, size, size : GA_USE_DOUBLE
  4.  
  5. KERNEL void k(GLOBAL_MEM ga_double *a, ga_size n, ga_size m) {
  6.     ga_size nb = n < m ? n : m;
  7.     for (ga_size i = LID_0; i < nb; i += LDIM_0) {
  8.     a[i*m + i] = 1;
  9.     }
  10. }

The second method is to handle the kernel compilation and cache onyour own. This is not recommended because there are lots of detailsto pay attention to that can cripple your performance if not doneright, which GpuKernelBase handles for you. But if you really want togo this way, then you can look up the C API for kernels inlibgpuarray.

In any case you will need to call your compiled kernel with some data,in most cases in your c_code() method. This is done by usingthe provided wrapper function. An example calling the above kernelwould be:

  1. size_t ls, gs;
  2. size_t dims[2];
  3.  
  4. // ...
  5.  
  6. ls = 1;
  7. gs = 256;
  8. err = k_call(1, &gs, &ls, 0, input->ga.data, dims[0], dims[1]);
  9.  
  10. // ...

The name of the wrapper function depends on the name you passed toKernel() when you declared it (or the name in your #kernel_statement). It defaults to _name + ‘_call’.

For other operations in the C code you should refer to thelibgpuarray documentation.

A Complete Example

This is a complete example using both approches for a implementationof the Eye operation.

GpuKernelBase

Python File

  1. class GpuEye(GpuKernelBase, Op):
  2.     """
  3.     Eye for GPU.
  4.  
  5.     """
  6.     __props__ = ('dtype', 'context_name')
  7.     _f16_ok = True
  8.  
  9.     def __init__(self, dtype=None, context_name=None):
  10.         if dtype is None:
  11.             dtype = config.floatX
  12.         self.dtype = dtype
  13.         self.context_name = context_name
  14.  
  15.     def get_params(self, node):
  16.         return get_context(self.context_name)
  17.  
  18.     def make_node(self, n, m, k):
  19.         n = tensor.as_tensor_variable(n)
  20.         m = tensor.as_tensor_variable(m)
  21.         k = tensor.as_tensor_variable(k)
  22.         assert n.ndim == 0
  23.         assert m.ndim == 0
  24.         assert k.ndim == 0
  25.         otype = GpuArrayType(dtype=self.dtype,
  26.                              broadcastable=(False, False),
  27.                              context_name=self.context_name)
  28.  
  29.         # k != 0 isn't implemented on the GPU yet.
  30.         assert tensor.get_scalar_constant_value(k) == 0
  31.         return Apply(self, [n, m], [otype()])
  32.  
  33.     def infer_shape(self, node, in_shapes):
  34.         out_shape = [node.inputs[0], node.inputs[1]]
  35.         return [out_shape]
  36.  
  37.     def grad(self, inp, grads):
  38.         return [grad_undefined(self, i, inp[i])
  39.                 for i in xrange(3)]
  40.  
  41.     def gpu_kernels(self, node, name):
  42.         code = """
  43. KERNEL void eye(GLOBAL_MEM %(ctype)s *a, ga_size n, ga_size m) {
  44.     ga_size nb = n < m ? n : m;
  45.     for (ga_size i = LID_0; i < nb; i += LDIM_0) {
  46.         a[i*m + i] = %(write_a)s(1);
  47.     }
  48. }""" % dict(ctype=pygpu.gpuarray.dtype_to_ctype(self.dtype),
  49.             name=name, write_a=write_w(self.dtype))
  50.         return [Kernel(
  51.                 code=code, name="eye",
  52.                 params=[gpuarray.GpuArray, gpuarray.SIZE, gpuarray.SIZE],
  53.                 flags=Kernel.get_flags(self.dtype),
  54.                 objvar='k_eye_' + name)]
  55.  
  56.     def c_code(self, node, name, inp, out, sub):
  57.         n, m = inp
  58.         z, = out
  59.         fail = sub['fail']
  60.         ctx = sub['params']
  61.         typecode = pygpu.gpuarray.dtype_to_typecode(self.dtype)
  62.         sync = bool(config.gpuarray.sync)
  63.         kname = self.gpu_kernels(node, name)[0].objvar
  64.         s = """
  65.         size_t dims[2] = {0, 0};
  66.         size_t ls, gs;
  67.         int err;
  68.  
  69.         dims[0] = ((dtype_%(n)s*)PyArray_DATA(%(n)s))[0];
  70.         dims[1] = ((dtype_%(m)s*)PyArray_DATA(%(m)s))[0];
  71.         Py_CLEAR(%(z)s);
  72.  
  73.         %(z)s = pygpu_zeros(2, dims,
  74.                             %(typecode)s,
  75.                             GA_C_ORDER,
  76.                             %(ctx)s, Py_None);
  77.         if (%(z)s == NULL) {
  78.             %(fail)s
  79.         }
  80.  
  81.         ls = 1;
  82.         gs = 256;
  83.         err = eye_call(1, &gs, &ls, 0, %(z)s->ga.data, dims[0], dims[1]);
  84.         if (err != GA_NO_ERROR) {
  85.             PyErr_Format(PyExc_RuntimeError,
  86.                          "gpuarray error: kEye: %%s. n%%lu, m=%%lu.",
  87.                          GpuKernel_error(&%(kname)s, err),
  88.                          (unsigned long)dims[0], (unsigned long)dims[1]);
  89.             %(fail)s;
  90.         }
  91.  
  92.         if(%(sync)d)
  93.             GpuArray_sync(&%(z)s->ga);
  94.         """ % locals()
  95.  
  96.         return s
  97.  
  98.     def c_code_cache_version(self):
  99.         return (6,)

CGpuKernelBase

Python File

  1. class GpuEye(CGpuKernelBase, Op):
  2.     """
  3.     Eye for GPU.
  4.  
  5.     """
  6.     __props__ = ('dtype', 'context_name')
  7.     _f16_ok = True
  8.  
  9.     def __init__(self, dtype=None, context_name=None):
  10.         if dtype is None:
  11.             dtype = config.floatX
  12.         self.dtype = dtype
  13.         self.context_name = context_name
  14.         CGpuKernelBase.__init__(self, ['tstgpueye.c'],
  15.                                 'APPLY_SPECIFIC(tstgpueye)')
  16.  
  17.     def get_params(self, node):
  18.         return get_context(self.context_name)
  19.  
  20.     def c_headers(self):
  21.         return ['<gpuarray/types.h>', '<gpuarray/kernel.h>']
  22.  
  23.     def make_node(self, n, m):
  24.         n = tensor.as_tensor_variable(n)
  25.         m = tensor.as_tensor_variable(m)
  26.         assert n.ndim == 0
  27.         assert m.ndim == 0
  28.         otype = GpuArrayType(dtype=self.dtype,
  29.                              broadcastable=(False, False),
  30.                              context_name=self.context_name)
  31.  
  32.         return Apply(self, [n, m], [otype()])
  33.  
  34.     def infer_shape(self, node, in_shapes):
  35.         out_shape = [node.inputs[0], node.inputs[1]]
  36.         return [out_shape]
  37.  
  38.     def grad(self, inp, grads):
  39.         return [grad_undefined(self, i, inp[i])
  40.                 for i in xrange(2)]
  41.  
  42.     def get_op_params(self):
  43.         return [('TYPECODE', str(dtype_to_typecode(self.dtype)))]

tstgpueye.c

  1. #section kernels
  2.  
  3. #kernel eye : *, size, size :
  4. /* The eye name will be used to generate supporting objects.  The only
  5.    you probably need to care about is the kernel object which will be
  6.    named 'k_' + <the name above> (k_eye in this case).  This name also
  7.    has to match the kernel function name below.
  8.  */
  9.  
  10. KERNEL void eye(GLOBAL_MEM DTYPE_o0 *a, ga_size n, ga_size m) {
  11.   ga_size nb = n < m ? n : m;
  12.   for (ga_size i = LID_0; i < nb; i += LDIM_0) {
  13.     a[i*m + i] = 1;
  14.   }
  15. }
  16.  
  17. #section support_code_struct
  18.  
  19. int APPLY_SPECIFIC(tstgpueye)(PyArrayObject *n, PyArrayObject *m,
  20.                               PyGpuArrayObject **z, PyGpuContextObject *ctx) {
  21.   size_t dims[2] = {0, 0};
  22.   size_t ls, gs;
  23.   void *args[3];
  24.   int err;
  25.  
  26.   dims[0] = ((DTYPE_INPUT_0 *)PyArray_DATA(n))[0];
  27.   dims[1] = ((DTYPE_INPUT_1 *)PyArray_DATA(m))[0];
  28.  
  29.   Py_XDECREF(*z);
  30.   *z = pygpu_zeros(2, dims,
  31.                    TYPECODE,
  32.                    GA_C_ORDER,
  33.                    ctx, Py_None);
  34.   if (*z == NULL)
  35.     return -1;
  36.  
  37.   ls = 1;
  38.   gs = 256;
  39.   /* The eye_call name comes from the kernel declaration above. */
  40.   err = eye_call(1, &gs, &ls, 0, (*z)->ga.data, dims[0], dims[1]);
  41.   if (err != GA_NO_ERROR) {
  42.     PyErr_Format(PyExc_RuntimeError,
  43.                  "gpuarray error: kEye: %s. n%lu, m=%lu.",
  44.                  GpuKernel_error(&k_eye, err),
  45.                  (unsigned long)dims[0], (unsigned long)dims[1]);
  46.     return -1;
  47.   }
  48.   return 0;
  49. }

Wrapping Exisiting Libraries

PyCUDA

For things in PyCUDA (or things wrapped with PyCUDA), we usually needto create a PyCUDA context. This can be done with the followingcode:

  1. with gpuarray_cuda_context:
  2.     pycuda_context = pycuda.driver.Context.attach()

If you don’t need to create a context, because the library doesn’trequire it, you can also just use the pygpu context and a _with_statement like above for all your code which will make the context thecurrent context on the cuda stack.

GpuArray objects are compatible with PyCUDA and will expose thenecessary interface so that they can be used in most things. Onenotable exception is PyCUDA kernels which require native objects. Ifyou need to convert a pygpu GpuArray to a PyCUDA GPUArray, this codeshould do the trick:

  1. assert pygpu_array.flags['IS_C_CONTIGUOUS']
  2. pycuda_array = pycuda.gpuarray.GPUArray(pygpu_array.shape,
  3.                                         pygpu_array.dtype,
  4.                                         base=pygpu_array,
  5.                                         gpudata=(pygpu_array.gpudata +
  6.                                                  pygpu_array.offset))

As long as the computations happen on the NULL stream there are nospecial considerations to watch for with regards to synchronization.Otherwise, you will have to make sure that you synchronize the pygpuobjects by calling the .sync() method before scheduling any work andsynchronize with the work that happends in the library after all thework is scheduled.