#include <assert.h>

#include <gpuarray/elemwise.h>
#include <gpuarray/array.h>
#include <gpuarray/error.h>
#include <gpuarray/kernel.h>
#include <gpuarray/util.h>

#include "private.h"
#include "util/strb.h"

struct _GpuElemwise {
  const char *expr; /* Expression code (to be able to build kernels on-demand) */
  const char *preamble; /* Preamble code */
  gpuelemwise_arg *args; /* Argument descriptors */
  GpuKernel k_contig; /* Contiguous kernel */
  GpuKernel *k_basic; /* Normal basic kernels */
  GpuKernel *k_basic_32; /* 32-bit address basic kernels */
  size_t *dims; /* Preallocated shape buffer for dimension collapsing */
  ssize_t **strides; /* Preallocated strides buffer for dimension collapsing */
  unsigned int nd; /* Current maximum number of dimensions allocated */
  unsigned int n; /* Number of arguments */
  unsigned int narray; /* Number of array arguments */
  int flags; /* Flags for the operation (none at the moment */
};

#define GEN_ADDR32      0x1
#define GEN_CONVERT_F16 0x2

/* This makes sure we have the same value for those flags since we use some shortcuts */
STATIC_ASSERT(GEN_CONVERT_F16 == GE_CONVERT_F16, same_flags_value_elem1);

#define is_array(a) (ISCLR((a).flags, GE_SCALAR))
#define is_output(a) (ISSET((a).flags, GE_WRITE))

static inline int k_initialized(GpuKernel *k) {
  return k->k != NULL;
}

static inline const char *ctype(int typecode) {
  return gpuarray_get_type(typecode)->cluda_name;
}

/* dst has to be zero-initialized on entry */
static int copy_arg(gpuelemwise_arg *dst, gpuelemwise_arg *src) {
  dst->name = strdup(src->name);
  if (dst->name == NULL)
    return -1;

  dst->typecode = src->typecode;
  dst->flags = src->flags;

  return 0;
}

static void clear_arg(gpuelemwise_arg *a) {
  free((void *)a->name);
  a->name = NULL;
}

static gpuelemwise_arg *copy_args(unsigned int n, gpuelemwise_arg *a) {
  gpuelemwise_arg *res = calloc(n, sizeof(gpuelemwise_arg));
  unsigned int i;

  if (res == NULL) return NULL;

  for (i = 0; i < n; i++)
    if (copy_arg(&res[i], &a[i]) != 0)
      goto bail;

  return res;
 bail:
  for (; i > 0; i--) {
    clear_arg(&res[i]);
  }
  return NULL;
}

static void free_args(unsigned int n, gpuelemwise_arg *args) {
  unsigned int i;

  if (args != NULL)
    for (i = 0; i < n; i++)
      clear_arg(&args[i]);
  free(args);
}

#define MUL_NO_OVERFLOW ((size_t)1 << (sizeof(size_t) * 4))

static int reallocaz(void **p, size_t elsz, size_t old, size_t new) {
  char *res;

  assert(old <= new);

  if ((new >= MUL_NO_OVERFLOW || elsz >= MUL_NO_OVERFLOW) &&
      new > 0 && SIZE_MAX / new < elsz) {
    return 1;
  }
  res = realloc(*p, elsz*new);
  if (res == NULL) return 1;
  memset(res + (elsz*old), 0, elsz*(new-old));
  *p = (void *)res;
  return 0;
}

static int ge_grow(GpuElemwise *ge, unsigned int nd) {
  unsigned int i;

  assert(nd > ge->nd);

  if (reallocaz((void **)&ge->k_basic, sizeof(GpuKernel), ge->nd, nd) ||
      reallocaz((void **)&ge->k_basic_32, sizeof(GpuKernel), ge->nd, nd) ||
      reallocaz((void **)&ge->dims, sizeof(size_t), ge->nd, nd))
    return 1;
  for (i = 0; i < ge->narray; i++) {
    if (reallocaz((void **)&ge->strides[i], sizeof(ssize_t), ge->nd, nd))
      return 1;
  }
  ge->nd = nd;
  return 0;
}

static int gen_elemwise_basic_kernel(GpuKernel *k, gpucontext *ctx,
                                     char **err_str,
                                     const char *preamble,
                                     const char *expr,
                                     unsigned int nd, /* Number of dims */
                                     unsigned int n, /* Length of args */
                                     gpuelemwise_arg *args,
                                     int gen_flags) {
  strb sb = STRB_STATIC_INIT;
  unsigned int i, _i, j;
  int *ktypes;
  char *size = "ga_size", *ssize = "ga_ssize";
  unsigned int p;
  int flags = 0;
  int res;

  if (ISSET(gen_flags, GEN_ADDR32)) {
    size = "ga_uint";
    ssize = "ga_int";
  }

  flags |= gpuarray_type_flagsa(n, args);

  p = 1 + nd;
  for (j = 0; j < n; j++) {
    p += ISSET(args[j].flags, GE_SCALAR) ? 1 : (2 + nd);
  }

  ktypes = calloc(p, sizeof(int));
  if (ktypes == NULL)
    return error_sys(ctx->err, "calloc");

  p = 0;

  strb_appends(&sb, "#include \"cluda.h\"\n");
  if (preamble)
    strb_appends(&sb, preamble);
  strb_appends(&sb, "\nKERNEL void elem(const ga_size n, ");
  ktypes[p++] = GA_SIZE;
  for (i = 0; i < nd; i++) {
    strb_appendf(&sb, "const ga_size dim%u, ", i);
    ktypes[p++] = GA_SIZE;
  }
  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      strb_appendf(&sb, "GLOBAL_MEM %s *%s_data, const ga_size %s_offset%s",
                   ctype(args[j].typecode), args[j].name, args[j].name,
                   nd == 0 ? "" : ", ");
      ktypes[p++] = GA_BUFFER;
      ktypes[p++] = GA_SIZE;

      for (i = 0; i < nd; i++) {
        strb_appendf(&sb, "const ga_ssize %s_str_%u%s", args[j].name, i,
                     (i == (nd - 1)) ? "": ", ");
        ktypes[p++] = GA_SSIZE;
      }
    } else {
      strb_appendf(&sb, "%s %s", ctype(args[j].typecode), args[j].name);
      ktypes[p++] = args[j].typecode;
    }
    if (j != (n - 1)) strb_appends(&sb, ", ");
  }
  strb_appendf(&sb, ") {\n"
               "const %s idx = LDIM_0 * GID_0 + LID_0;\n"
               "const %s numThreads = LDIM_0 * GDIM_0;\n"
               "%s i;\n", size, size, size);

  strb_appends(&sb, "for(i = idx; i < n; i += numThreads) {\n");
  if (nd > 0)
    strb_appendf(&sb, "%s ii = i;\n%s pos;\n", size, size);
  for (j = 0; j < n; j++) {
    if (is_array(args[j]))
      strb_appendf(&sb, "%s %s_p = %s_offset;\n",
                   size, args[j].name, args[j].name);
  }
  for (_i = nd; _i > 0; _i--) {
    i = _i - 1;
    if (i > 0)
      strb_appendf(&sb, "pos = ii %% (%s)dim%u;\nii = ii / (%s)dim%u;\n", size, i, size, i);
    else
      strb_appends(&sb, "pos = ii;\n");
    for (j = 0; j < n; j++) {
      if (is_array(args[j]))
        strb_appendf(&sb, "%s_p += pos * (%s)%s_str_%u;\n", args[j].name,
                     ssize, args[j].name, i);
    }
  }
  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      strb_appendf(&sb, "%s %s;", ctype(ISSET(gen_flags, GEN_CONVERT_F16) && args[j].typecode == GA_HALF ?
                                        GA_FLOAT : args[j].typecode), args[j].name);
      if (ISSET(args[j].flags, GE_READ)) {
        if (args[j].typecode == GA_HALF && ISSET(gen_flags, GEN_CONVERT_F16)) {
          strb_appendf(&sb, "%s = ga_half2float(*(GLOBAL_MEM ga_half *)(((GLOBAL_MEM char *)%s_data) + %s_p));\n",
                       args[j].name, args[j].name, args[j].name);
        } else {
          strb_appendf(&sb, "%s = *(GLOBAL_MEM %s *)(((GLOBAL_MEM char *)%s_data) + %s_p);\n",
                       args[j].name, ctype(args[j].typecode), args[j].name, args[j].name);
        }
      }
    }
  }
  strb_appends(&sb, expr);
  strb_appends(&sb, ";\n");
  for (j = 0; j < n; j++) {
    if (is_array(args[j]) && ISSET(args[j].flags, GE_WRITE)) {
      if (args[j].typecode == GA_HALF && ISSET(gen_flags, GEN_CONVERT_F16)) {
        strb_appendf(&sb, "*(GLOBAL_MEM ga_half *)(((GLOBAL_MEM char *)%s_data) + %s_p) = ga_float2half(%s);\n",
                     args[j].name, args[j].name, args[j].name);
      } else {
        strb_appendf(&sb, "*(GLOBAL_MEM %s *)(((GLOBAL_MEM char *)%s_data) + %s_p) = %s;\n",
                     ctype(args[j].typecode), args[j].name, args[j].name, args[j].name);
      }
    }
  }
  strb_appends(&sb, "}\n}\n");
  if (strb_error(&sb)) {
    res = GA_MEMORY_ERROR;
    goto bail;
  }

  res = GpuKernel_init(k, ctx, 1, (const char **)&sb.s, &sb.l, "elem",
                       p, ktypes, flags, err_str);
 bail:
  free(ktypes);
  strb_clear(&sb);
  return res;
}

static ssize_t **strides_array(unsigned int num, unsigned int nd) {
  ssize_t **res = calloc(num, sizeof(ssize_t *));
  unsigned int i;

  if (res == NULL) return NULL;
  for (i = 0; i < num; i++) {
    res[i] = calloc(nd, sizeof(ssize_t));
    if (res[i] == NULL)
      goto bail;
  }

  return res;

 bail:
  for (i = 0; i < num; i++)
    free(res[i]);
  free(res);
  return NULL;
}

static int check_basic(GpuElemwise *ge, void **args, int flags,
                       size_t *_n, unsigned int *_nd, size_t **_dims,
                       ssize_t ***_strides, int *_call32) {
  size_t n;
  gpucontext *ctx = GpuKernel_context(&ge->k_contig);
  GpuArray *a = NULL, *v;
  unsigned int i, j, p, num_arrays = 0, nd = 0, nnd;
  int call32 = 1;
  unsigned int nd_i = 0;
  size_t v_dim_j = 0;

  /* Go through the list and grab some info */
  for (i = 0; i < ge->n; i++) {
    if (is_array(ge->args[i])) {
      nd_i = ((GpuArray *)args[i])->nd;
      if (num_arrays == 0)
        nd = nd_i;
      else if (nd_i != nd) {
        if (flags & GE_PADSHAPE)
          nd = nd_i > nd ? nd_i : nd;
        else
          return error_fmt(ctx->err, GA_VALUE_ERROR, "Arg %u has differing nd = %u", i, nd_i);
      }
      ++num_arrays;
      if (a == NULL && is_output(ge->args[i]))
        a = (GpuArray *)args[i];
    }
  }

  if (a == NULL)
    return error_set(ctx->err, GA_VALUE_ERROR, "No output arrays");

  /* Check if we need to grow the internal buffers */
  if (nd > ge->nd) {
    nnd = ge->nd * 2;
    while (nd > nnd) nnd *= 2;
    if (ge_grow(ge, nnd))
      return error_sys(ctx->err, "ge_grow");
  }

  /* Now we know that all array arguments have at most nd
     dimensions and that the expected output size is the size of a */

  /* And copy their initial values in */
  memcpy(ge->dims, a->dimensions, nd*sizeof(size_t));
  p = 0;
  for (i = 0; i < ge->n; i++) {
    if (is_array(ge->args[i])) {
      /* Left-pad strides with zero on implicitly broadcasted dimensions */
      memset(ge->strides[p], 0, nd*sizeof(ssize_t));
      nd_i = ((GpuArray *)args[i])->nd;
      memcpy((char *)(ge->strides[p]) + (nd - nd_i)*sizeof(ssize_t),
             ((GpuArray *)args[i])->strides, nd_i*sizeof(ssize_t));
      p++;
    }
  }

  /* Check that all arrays are the same size (or broadcast-compatible
     if GE_BROADCAST), adjust strides of broadcastable dimensions and
     check if we can use the 32 bit address version.

     Basically for each dimension go over all the arguments and make
     sure that the dimension size matches. */
  n = 1;
  for (j = 0; j < nd; j++) {
    p = 0;
    for (i = 0; i < ge->n; i++) {
      if (is_array(ge->args[i])) {
        v = (GpuArray *)args[i];
        nd_i = v->nd;
        /* Pad shape with 1 if needed for implicitly broadcasted dimensions
           and shift if needed */
        if (j < nd - nd_i)
          v_dim_j = 1;
        else
          v_dim_j = v->dimensions[j - (nd - nd_i)];
        if (ge->dims[j] != v_dim_j) {
          /* We can't broadcast outputs */
          if (ISCLR(flags, GE_BROADCAST) || is_output(ge->args[i]) ||
              v_dim_j != 1) {
            return error_fmt(ctx->err, GA_VALUE_ERROR, "Mismatched dimension %u for input %u (expected %" SPREFIX "u got %" SPREFIX "u)", j, i, ge->dims[j], v_dim_j);
          }
        }
        /* If the dimension is 1 set the strides to 0 regardless since
           it won't change anything in the non-broadcast case. */
        if (v_dim_j == 1) {
          ge->strides[p][j] = 0;
        }
        call32 &= v->offset < ADDR32_MAX;
        call32 &= (SADDR32_MIN < ge->strides[p][j] &&
                   ge->strides[p][j] < SADDR32_MAX);
        p++;
      } /* is_array() */
    } /* for each arg */
    /* We have the final value in dims[j] */
    n *= ge->dims[j];
  } /* for each dim */

  call32 &= n < ADDR32_MAX;

  if (ISCLR(flags, GE_NOCOLLAPSE) && nd > 1) {
    gpuarray_elemwise_collapse(num_arrays, &nd, ge->dims, ge->strides);
  }

  *_n = n;
  *_nd = nd;
  *_dims = ge->dims;
  *_strides = ge->strides;
  *_call32 = call32;

  return GA_NO_ERROR;
}

static int call_basic(GpuElemwise *ge, void **args, size_t n, unsigned int nd,
                      size_t *dims, ssize_t **strs, int call32) {
  GpuKernel *k;
  size_t ls = 0, gs = 0;
  unsigned int p = 0, i, j, l;
  int err;

  if (nd == 0) return error_set(GpuKernel_context(&ge->k_contig)->err, GA_VALUE_ERROR, "nd == 0");

  if (call32)
    k = &ge->k_basic_32[nd-1];
  else
    k = &ge->k_basic[nd-1];

  if (!k_initialized(k)) {
    err = gen_elemwise_basic_kernel(k, GpuKernel_context(&ge->k_contig), NULL,
                                    ge->preamble, ge->expr, nd, ge->n,
                                    ge->args, ((call32 ? GEN_ADDR32 : 0) |
                                               (ge->flags & GE_CONVERT_F16)));
    if (err != GA_NO_ERROR)
      return err;
  }

  err = GpuKernel_setarg(k, p++, &n);
  if (err != GA_NO_ERROR) goto error_call_basic;

  for (i = 0; i < nd; i++) {
    err = GpuKernel_setarg(k, p++, &dims[i]);
    if (err != GA_NO_ERROR) goto error_call_basic;
  }

  /* l is the number of arrays to date */
  l = 0;
  for (j = 0; j < ge->n; j++) {
    if (is_array(ge->args[j])) {
      GpuArray *v = (GpuArray *)args[j];
      err = GpuKernel_setarg(k, p++, v->data);
      if (err != GA_NO_ERROR) goto error_call_basic;
      err = GpuKernel_setarg(k, p++, &v->offset);
      if (err != GA_NO_ERROR) goto error_call_basic;
      for (i = 0; i < nd; i++) {
        err = GpuKernel_setarg(k, p++, &strs[l][i]);
        if (err != GA_NO_ERROR) goto error_call_basic;
      }
      l++;
    } else {
      err = GpuKernel_setarg(k, p++, args[j]);
      if (err != GA_NO_ERROR) goto error_call_basic;
    }
  }

  err = GpuKernel_sched(k, n, &gs, &ls);
  if (err != GA_NO_ERROR) goto error_call_basic;

  err = GpuKernel_call(k, 1, &gs, &ls, 0, NULL);
 error_call_basic:
  return err;
}

static int gen_elemwise_contig_kernel(GpuKernel *k,
                                      gpucontext *ctx, char **err_str,
                                      const char *preamble,
                                      const char *expr,
                                      unsigned int n,
                                      gpuelemwise_arg *args,
                                      int gen_flags) {
  strb sb = STRB_STATIC_INIT;
  int *ktypes = NULL;
  unsigned int p;
  unsigned int j;
  int flags = 0;
  int res;

  flags |= gpuarray_type_flagsa(n, args);

  p = 1;
  for (j = 0; j < n; j++)
    p += ISSET(args[j].flags, GE_SCALAR) ? 1 : 2;

  ktypes = calloc(p, sizeof(int));
  if (ktypes == NULL) {
    res = error_sys(ctx->err, "calloc");
    goto bail;
  }

  p = 0;

  strb_appends(&sb, "#include \"cluda.h\"\n");
  if (preamble)
    strb_appends(&sb, preamble);
  strb_appends(&sb, "\nKERNEL void elem(const ga_size n, ");
  ktypes[p++] = GA_SIZE;
  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      strb_appendf(&sb, "GLOBAL_MEM %s *%s_p,  const ga_size %s_offset",
                   ctype(args[j].typecode), args[j].name, args[j].name);
      ktypes[p++] = GA_BUFFER;
      ktypes[p++] = GA_SIZE;
    } else {
      strb_appendf(&sb, "%s %s", ctype(args[j].typecode), args[j].name);
      ktypes[p++] = args[j].typecode;
    }
    if (j != (n - 1))
      strb_appends(&sb, ", ");
  }
  strb_appends(&sb, ") {\n"
               "const ga_size idx = LDIM_0 * GID_0 + LID_0;\n"
               "const ga_size numThreads = LDIM_0 * GDIM_0;\n"
               "ga_size i;\n"
               "GLOBAL_MEM char *tmp;\n\n");
  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      strb_appendf(&sb, "tmp = (GLOBAL_MEM char *)%s_p;"
                   "tmp += %s_offset; %s_p = (GLOBAL_MEM %s *)tmp;",
                   args[j].name, args[j].name, args[j].name,
                   ctype(args[j].typecode));
    }
  }

  strb_appends(&sb, "for (i = idx; i < n; i += numThreads) {\n");
  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      strb_appendf(&sb, "%s %s;\n", ctype(ISSET(gen_flags, GEN_CONVERT_F16) && args[j].typecode == GA_HALF ?
                                          GA_FLOAT : args[j].typecode), args[j].name);
      if (ISSET(args[j].flags, GE_READ)) {
        if (args[j].typecode == GA_HALF && ISSET(gen_flags, GEN_CONVERT_F16)) {
          strb_appendf(&sb, "%s = ga_half2float(%s_p[i]);\n", args[j].name, args[j].name);
        } else {
          strb_appendf(&sb, "%s = %s_p[i];\n", args[j].name, args[j].name);
        }
      }
    }
  }
  strb_appends(&sb, expr);
  strb_appends(&sb, ";\n");

  for (j = 0; j < n; j++) {
    if (is_array(args[j])) {
      if (ISSET(args[j].flags, GE_WRITE)) {
        if (args[j].typecode == GA_HALF && ISSET(gen_flags, GEN_CONVERT_F16)) {
          strb_appendf(&sb, "%s_p[i] = ga_float2half(%s);\n", args[j].name, args[j].name);
        } else {
          strb_appendf(&sb, "%s_p[i] = %s;\n", args[j].name, args[j].name);
        }
      }
    }
  }
  strb_appends(&sb, "}\n}\n");

  if (strb_error(&sb)) {
    res = error_set(ctx->err, GA_MISC_ERROR, "Formatting error creating kernel source");
    goto bail;
  }

  res = GpuKernel_init(k, ctx, 1, (const char **)&sb.s, &sb.l, "elem",
                       p, ktypes, flags, err_str);
 bail:
  strb_clear(&sb);
  free(ktypes);
  return res;
}

static int check_contig(GpuElemwise *ge, void **args,
                        size_t *_n, int *contig) {
  GpuArray *a = NULL, *v;
  size_t n = 1;
  unsigned int i, j;
  int c_contig = 1, f_contig = 1;

  for (i = 0; i < ge->n; i++) {
    if (is_array(ge->args[i])) {
      v = (GpuArray *)args[i];
      if (a == NULL) {
        a = v;
        for (j = 0; j < a->nd; j++) n *= a->dimensions[j];
      }
      c_contig &= GpuArray_IS_C_CONTIGUOUS(v);
      f_contig &= GpuArray_IS_F_CONTIGUOUS(v);
      if (a != v) {
        if (a->nd != v->nd)
          return -1; /* We don't check the value of the error code */
        for (j = 0; j < a->nd; j++) {
          if (v->dimensions[j] != a->dimensions[j])
            return -1; /* We don't check the value of the error code */
        }
      }
    }
  }
  *contig = f_contig || c_contig;
  *_n = n;
  return GA_NO_ERROR;
}

static int call_contig(GpuElemwise *ge, void **args, size_t n) {
  GpuArray *a;
  size_t ls = 0, gs = 0;
  unsigned int i, p;
  int err;

  p = 0;
  err = GpuKernel_setarg(&ge->k_contig, p++, &n);
  if (err != GA_NO_ERROR) return err;
  for (i = 0; i < ge->n; i++) {
    if (is_array(ge->args[i])) {
      a = (GpuArray *)args[i];
      err = GpuKernel_setarg(&ge->k_contig, p++, a->data);
      if (err != GA_NO_ERROR) return err;
      err = GpuKernel_setarg(&ge->k_contig, p++, &a->offset);
      if (err != GA_NO_ERROR) return err;
    } else {
      err = GpuKernel_setarg(&ge->k_contig, p++, args[i]);
      if (err != GA_NO_ERROR) return err;
    }
  }
  err = GpuKernel_sched(&ge->k_contig, n, &gs, &ls);
  if (err != GA_NO_ERROR) return err;
  return GpuKernel_call(&ge->k_contig, 1, &gs, &ls, 0, NULL);
}

GpuElemwise *GpuElemwise_new(gpucontext *ctx,
                             const char *preamble, const char *expr,
                             unsigned int n, gpuelemwise_arg *args,
                             unsigned int nd, int flags) {
  GpuElemwise *res;
#ifdef DEBUG
  char *errstr = NULL;
#endif
  unsigned int i;
  int ret;

  res = calloc(1, sizeof(*res));
  if (res == NULL) {
    error_sys(ctx->err, "calloc");
    return NULL;
  }

  res->flags = flags;
  res->nd = 8;
  res->n = n;

  res->expr = strdup(expr);
  if (res->expr == NULL) {
    error_sys(ctx->err, "strdup");
    goto fail;
  }
  if (preamble != NULL) {
    res->preamble = strdup(preamble);
    if (res->preamble == NULL) {
      error_sys(ctx->err, "strdup");
      goto fail;
    }
  }

  res->args = copy_args(n, args);
  if (res->args == NULL) {
    error_sys(ctx->err, "copy_args");
    goto fail;
  }

  /* Count the arrays in the arguements */
  res->narray = 0;
  for (i = 0; i < res->n; i++)
    if (is_array(res->args[i])) res->narray++;

  while (res->nd < nd) res->nd *= 2;
  res->dims = calloc(res->nd, sizeof(size_t));
  if (res->dims == NULL) {
    error_sys(ctx->err, "calloc");
    goto fail;
  }
  res->strides = strides_array(res->narray, res->nd);
  if (res->strides == NULL) {
    error_sys(ctx->err, "strides_array");
    goto fail;
  }
  res->k_basic = calloc(res->nd, sizeof(GpuKernel));
  if (res->k_basic == NULL) {
    error_sys(ctx->err, "calloc");
    goto fail;
  }

  res->k_basic_32 = calloc(res->nd, sizeof(GpuKernel));
  if (res->k_basic_32 == NULL) {
    error_sys(ctx->err, "calloc");
    goto fail;
  }

  ret = gen_elemwise_contig_kernel(&res->k_contig, ctx,
#ifdef DEBUG
                                   &errstr,
#else
                                   NULL,
#endif
                                   res->preamble, res->expr,
                                   res->n, res->args,
                                   (res->flags & GE_CONVERT_F16));
  if (ret != GA_NO_ERROR) {
#ifdef DEBUG
    if (errstr != NULL)
      fprintf(stderr, "%s\n", errstr);
    free(errstr);
#endif
    goto fail;
  }

  if (ISCLR(flags, GE_NOADDR64)) {
    for (i = 0; i < nd; i++) {
      ret = gen_elemwise_basic_kernel(&res->k_basic[i], ctx,
#ifdef DEBUG
                                      &errstr,
#else
                                      NULL,
#endif
                                      res->preamble, res->expr,
                                      i+1, res->n, res->args,
                                      (res->flags & GE_CONVERT_F16));
      if (ret != GA_NO_ERROR) {
#ifdef DEBUG
        if (errstr != NULL)
          fprintf(stderr, "%s\n", errstr);
        free(errstr);
#endif
        goto fail;
      }
    }
  }

  for (i = 0; i < nd; i++) {
    ret = gen_elemwise_basic_kernel(&res->k_basic_32[i], ctx,
#ifdef DEBUG
                                    &errstr,
#else
                                    NULL,
#endif
                                    res->preamble, res->expr,
                                    i+1, res->n, res->args,
                                    GEN_ADDR32 | (res->flags & GE_CONVERT_F16));
    if (ret != GA_NO_ERROR) {
#ifdef DEBUG
      if (errstr != NULL)
        fprintf(stderr, "%s\n", errstr);
      free(errstr);
#endif
      goto fail;
    }
  }

  return res;

 fail:
  GpuElemwise_free(res);
  return NULL;
}

void GpuElemwise_free(GpuElemwise *ge) {
  unsigned int i;
  if (ge->k_basic_32 != NULL)
    for (i = 0; i < ge->nd; i++) {
      if (k_initialized(&ge->k_basic_32[i]))
        GpuKernel_clear(&ge->k_basic_32[i]);
    }
  if (ge->k_basic != NULL)
    for (i = 0; i < ge->nd; i++) {
      if (k_initialized(&ge->k_basic[i]))
        GpuKernel_clear(&ge->k_basic[i]);
    }
  if (ge->strides != NULL)
    for (i = 0; i < ge->narray; i++) {
      free(ge->strides[i]);
    }
  if (k_initialized(&ge->k_contig))
    GpuKernel_clear(&ge->k_contig);
  free(ge->k_basic_32);
  free(ge->k_basic);
  free_args(ge->n, ge->args);
  free((void *)ge->preamble);
  free((void *)ge->expr);
  free(ge->dims);
  free(ge->strides);
  free(ge);
}

int GpuElemwise_call(GpuElemwise *ge, void **args, int flags) {
  size_t n = 0;
  size_t *dims = NULL;
  ssize_t **strides = NULL;
  unsigned int nd = 0;
  int contig = 0;
  int call32 = 0;
  int err;

  err = check_contig(ge, args, &n, &contig);
  if (err == GA_NO_ERROR && contig) {
    if (n == 0) return GA_NO_ERROR;
    return call_contig(ge, args, n);
  }
  err = check_basic(ge, args, flags, &n, &nd, &dims, &strides, &call32);
  if (err == GA_NO_ERROR) {
    if (n == 0) return GA_NO_ERROR;
    return call_basic(ge, args, n, nd, dims, strides, call32);
  }
  return err;
}
