Partial update to use the new libgpuarray interface

theano/gpuarray/c_code/topk_dense.cu

@@ -9,14 +9,19 @@ extern "C" __global__ void k_topk_dense(
         // size_t dims_1, ssize_t dims_2, ... , dims_$${NDIM}
         $dstv
         // INPUT_TYPE *dstv
+        $dstv_offset
+        // size_t offset
         $dstv_strides
         // ssize_t dstv_strides_0, ssize_t dstv_strides_1, ... , dstv_strides_$${NDIM}
         $dsti
         // INDEX_TYPE *dsti
+        $dsti_offset
+        // size_t offset
         $dsti_strides
         // ssize_t dsti_strides_0, ssize_t dsti_strides_1, ... , dsti_strides_$${NDIM}
         ssize_t k,
         INPUT_TYPE* src,
+	size_t src_offset
         $src_strides
         // ssize_t src_strides_0, ssize_t src_strides_1, ... , src_strides_$${NDIM}
         size_t size) {
@@ -28,6 +33,9 @@ extern "C" __global__ void k_topk_dense(
     size_t out_idx;
 
     const unsigned char warp_id = idx / GA_WARP_SIZE;
+    dstv = ptr_add(dstv, dstv_offset);
+    dsti = ptr_add(dsti, dsti_offset);
+    src = ptr_add(src, src_offset);
 
     // 0. get the slice for thread block to work on
 

theano/gpuarray/c_code/topk_dense_large.cu

@@ -199,14 +199,19 @@ extern "C" __global__ void KERNEL_NAME(
         // size_t dims_1, ssize_t dims_2, ... , dims_$${NDIM}
         $dstv
         // INPUT_TYPE *dstv
+        $dstv_offset
+        // size_t offset
         $dstv_strides
         // ssize_t dstv_strides_0, ssize_t dstv_strides_1, ... , dstv_strides_$${NDIM}
         $dsti
         // INDEX_TYPE *dsti
+        $dsti_offset
+        // size_t offset
         $dsti_strides
         // ssize_t dsti_strides_0, ssize_t dsti_strides_1, ... , dsti_strides_$${NDIM}
         ssize_t k,
         INPUT_TYPE* src,
+	size_t src_offset
         $src_strides
         // ssize_t src_strides_0, ssize_t src_strides_1, ... , src_strides_$${NDIM}
         size_t size) {
@@ -217,6 +222,9 @@ extern "C" __global__ void KERNEL_NAME(
     k = (order ? k : -k);
     const int idx = threadIdx.x;
     const int warp_id = idx / GA_WARP_SIZE;
+    dstv = ptr_add(dstv, dstv_offset);
+    dsti = ptr_add(dsti, dsti_offset);
+    src = ptr_add(src, src_offset);
 
     // get the slice for thread block to work on
     // size <- the axis to work on

theano/gpuarray/sort.py

@@ -53,7 +53,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
             pygpu.get_include()]
 
     def c_code_cache_version(self):
-        return (2,)
+        return (3,)
 
     def gpu_kernels(self, node, nodename):
         # load kernel source
@@ -83,7 +83,9 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
                 out_t=ga.dtype_to_ctype(self.idx_dtype),
                 dims=''.join('size_t dims_%d, ' % i for i in range(1, ndim)),
                 dstv='INPUT_TYPE *dstv,' if self.return_values else '',
+                dstv_offset='size_t dstv_offset,' if self.return_values else '',
                 dsti='INDEX_TYPE *dsti,' if self.return_indices else '',
+                dsti_offset='size_t dsti_offset,' if self.return_indices else '',
                 dstv_strides=dstv_strides_code if self.return_values else '',
                 dsti_strides=dsti_strides_code if self.return_indices else '',
                 src_strides=src_strides_code,
@@ -100,9 +102,11 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
         param_types = [ga.SIZE] * (ndim - 1)  # dims
         for _ in range(self.return_values + self.return_indices):
             param_types.append(ga.GpuArray)  # dst*
+            param_types.extend([ga.SIZE] * ndim)  # offset
             param_types.extend([ga.SSIZE] * ndim)  # dst*_strides
         param_types.append(ga.SIZE)  # k
         param_types.append(ga.GpuArray)  # src
+        param_types.append(ga.SIZE)  # offset
         param_types.extend([ga.SSIZE] * ndim)  # src_strides
         param_types.append(ga.SIZE)  # size
 
@@ -174,7 +178,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
         prep_output = ''
         if self.return_values:
             def_dvstrides = 'const ssize_t *dvstrides = PyGpuArray_STRIDES(%s)' % yv
-            params_dv = '(void*)((char*)(%s->ga.data) + (%s->ga.offset)),\n' % (yv, yv)
+            params_dv = '(void*)((char*)(%s->ga.data), (%s->ga.offset)),\n' % (yv, yv)
             params_dv += ''.join('(void*)(dvstrides+%d), ' % i for i in reordered_axes)
             prep_output += '''
     if (0 != theano_prep_output(
@@ -187,7 +191,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
 
         if self.return_indices:
             def_distrides = 'const ssize_t *distrides = PyGpuArray_STRIDES(%s)' % yi
-            params_di = '(void*)((char*)(%s->ga.data) + (%s->ga.offset)),\n' % yi
+            params_di = '(void*)((char*)(%s->ga.data), (%s->ga.offset)),\n' % (yi, yi)
             params_di += ''.join('(void*)(distrides+%d), ' % i for i in reordered_axes)
             prep_output += '''
     if (0 != theano_prep_output(
@@ -246,7 +250,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
         %(params_dv)s
         %(params_di)s
         (void*)(&k_),
-        (void*)((char*)(%(x)s->ga.data) + (%(x)s->ga.offset)),
+        (void*)((char*)(%(x)s->ga.data), (%(x)s->ga.offset)),
         %(sstrides)s,
         (void*)(dims+%(axis)d),
     };