Merge pull request #298 from abergeron/offset_take1

Add offset handling for indices and result in take1()

Author: lamblin

src/gpuarray_array.c

@@ -336,7 +336,7 @@ static int gen_take1_kernel(GpuKernel *k, gpucontext *ctx, char **err_str,
   int flags = GA_USE_CLUDA;
   int res;
 
-  nargs = 7 + 2 * v->nd;
+  nargs = 9 + 2 * v->nd;
 
   atypes = calloc(nargs, sizeof(int));
   if (atypes == NULL)
@@ -351,34 +351,41 @@ static int gen_take1_kernel(GpuKernel *k, gpucontext *ctx, char **err_str,
   }
 
   apos = 0;
-  strb_appendf(&sb, "KERNEL void take1(GLOBAL_MEM %s *r, "
-               "GLOBAL_MEM const %s *v, ga_size off,",
+  strb_appendf(&sb, "KERNEL void take1(GLOBAL_MEM %s *r, ga_size r_off, "
+               "GLOBAL_MEM const %s *v, ga_size v_off,",
                gpuarray_get_type(a->typecode)->cluda_name,
                gpuarray_get_type(v->typecode)->cluda_name);
   atypes[apos++] = GA_BUFFER;
+  atypes[apos++] = GA_SIZE;
   atypes[apos++] = GA_BUFFER;
   atypes[apos++] = GA_SIZE;
   for (i = 0; i < v->nd; i++) {
     strb_appendf(&sb, " ga_ssize s%u, ga_size d%u,", i, i);
     atypes[apos++] = GA_SSIZE;
     atypes[apos++] = GA_SIZE;
   }
-  strb_appendf(&sb, " GLOBAL_MEM const %s *ind, ga_size n0, ga_size n1,"
-               " GLOBAL_MEM int* err) {\n",
+  strb_appendf(&sb, " GLOBAL_MEM const %s *ind, ga_size i_off, "
+               "ga_size n0, ga_size n1, GLOBAL_MEM int* err) {\n",
                gpuarray_get_type(ind->typecode)->cluda_name);
   atypes[apos++] = GA_BUFFER;
   atypes[apos++] = GA_SIZE;
   atypes[apos++] = GA_SIZE;
+  atypes[apos++] = GA_SIZE;
   atypes[apos++] = GA_BUFFER;
   assert(apos == nargs);
   strb_appendf(&sb, "  const %s idx0 = LDIM_0 * GID_0 + LID_0;\n"
                "  const %s numThreads0 = LDIM_0 * GDIM_0;\n"
                "  const %s idx1 = LDIM_1 * GID_1 + LID_1;\n"
                "  const %s numThreads1 = LDIM_1 * GDIM_1;\n"
                "  %s i0, i1;\n", sz, sz, sz, sz, sz);
+  strb_appends(&sb, "  if (idx0 >= n0 || idx1 >= n1) return;\n");
+  strb_appendf(&sb, "  r = (GLOBAL_MEM %s *)(((char *)r) + r_off);\n"
+               "  ind = (GLOBAL_MEM %s *)(((char *)ind) + i_off);\n",
+               gpuarray_get_type(a->typecode)->cluda_name,
+               gpuarray_get_type(ind->typecode)->cluda_name);
   strb_appendf(&sb, "  for (i0 = idx0; i0 < n0; i0 += numThreads0) {\n"
                "    %s ii0 = ind[i0];\n"
-               "    %s pos0 = off;\n"
+               "    %s pos0 = v_off;\n"
                "    if (ii0 < 0) ii0 += d0;\n"
                "    if ((ii0 < 0) || (ii0 >= d0)) {\n"
                "      *err = -1;\n"
@@ -500,13 +507,16 @@ int GpuArray_take1(GpuArray *a, const GpuArray *v, const GpuArray *i,
 
   argp = 0;
   GpuKernel_setarg(&k, argp++, a->data);
+  GpuKernel_setarg(&k, argp++, (void *)&a->offset);
   GpuKernel_setarg(&k, argp++, v->data);
+  /* The cast is to avoid a warning about const */
   GpuKernel_setarg(&k, argp++, (void *)&v->offset);
   for (j = 0; j < v->nd; j++) {
     GpuKernel_setarg(&k, argp++, &v->strides[j]);
     GpuKernel_setarg(&k, argp++, &v->dimensions[j]);
   }
   GpuKernel_setarg(&k, argp++, i->data);
+  GpuKernel_setarg(&k, argp++, (void *)&i->offset);
   GpuKernel_setarg(&k, argp++, &n[0]);
   GpuKernel_setarg(&k, argp++, &n[1]);
   GpuKernel_setarg(&k, argp++, errbuf);
@@ -1083,11 +1093,15 @@ int GpuArray_fdump(FILE *fd, const GpuArray *a) {
     case GA_UINT:
       fprintf(fd, "%u", *(unsigned int *)p);
       break;
+    case GA_LONG:
+      fprintf(fd, "%lld", (long long)*(int64_t *)p);
+      break;
     case GA_SSIZE:
       fprintf(fd, "%" SPREFIX "d", *(ssize_t *)p);
       break;
     default:
       free(buf);
+      fprintf(fd, "<unsupported data type %d>\n", a->typecode);
       return GA_UNSUPPORTED_ERROR;
     }
     s -= gpuarray_get_elsize(a->typecode);

tests/check_array.c

@@ -21,12 +21,12 @@ START_TEST(test_take1_ok) {
   GpuArray v;
   GpuArray vidx;
   GpuArray vres;
-  static const uint32_t data[24] = { 0,  1,  2,  3,  4,  5,
-                                     6,  7,  8,  9, 10, 11,
-                                    12, 13, 14, 15, 16, 17,
-                                    18, 19, 20, 21, 22, 23};
+  const uint32_t data[24] = { 0,  1,  2,  3,  4,  5,
+                              6,  7,  8,  9, 10, 11,
+                              12, 13, 14, 15, 16, 17,
+                              18, 19, 20, 21, 22, 23};
   uint32_t buf[12 * 24];
-  static const size_t data_dims[1] = {24};
+  const size_t data_dims[1] = {24};
   ssize_t indexes[12];
   size_t dims[3];
 
@@ -243,12 +243,40 @@ START_TEST(test_take1_ok) {
 }
 END_TEST
 
+START_TEST(test_take1_offset) {
+  const uint32_t data[24] = {0, 1, 2, 3};
+  const size_t data_dims[1] = {4};
+  const size_t out_dims[1] = {2};
+  const uint32_t idx[4] = {20, 3, 3, 2};
+  GpuArray v;
+  GpuArray i;
+  GpuArray r;
+
+  ga_assert_ok(GpuArray_empty(&v, ctx, GA_UINT, 1, data_dims, GA_C_ORDER));
+  ga_assert_ok(GpuArray_write(&v, data, sizeof(data)));
+
+  ga_assert_ok(GpuArray_empty(&i, ctx, GA_UINT, 1, data_dims, GA_C_ORDER));
+  ga_assert_ok(GpuArray_write(&i, idx, sizeof(idx)));
+
+  ga_assert_ok(GpuArray_empty(&r, ctx, GA_UINT, 1, out_dims, GA_C_ORDER));
+
+  /* Fake subtensor for offset */
+  i.offset = 8;
+  i.dimensions[0] = 2;
+
+  ga_assert_ok(GpuArray_take1(&r, &v, &i, 1));
+  /* The actual results are not important, this is just to check that
+     we don't trigger the out of bounds check */
+}
+END_TEST
+
 Suite *get_suite(void) {
   Suite *s = suite_create("array");
   TCase *tc = tcase_create("take1");
   tcase_add_checked_fixture(tc, setup, teardown);
   tcase_set_timeout(tc, 8.0);
   tcase_add_test(tc, test_take1_ok);
+  tcase_add_test(tc, test_take1_offset);
   suite_add_tcase(s, tc);
   return s;
 }