ndtypes 0.2.0dev5 → 0.2.0dev6

data/ext/ruby_ndtypes/ndtypes/libndtypes/util.c

@@ -95,21 +95,46 @@ ndt_asprintf(ndt_context_t *ctx, const char *fmt, ...)
 /*                       Type functions (unstable API)                       */
 /*****************************************************************************/
 
+/* Number of elements, currently only for ndarrays. Return -1 on error. */
+int64_t
+ndt_nelem(const ndt_t *t)
+{
+    NDT_STATIC_CONTEXT(ctx);
+    ndt_ndarray_t x;
+    int64_t n;
+
+    if (ndt_as_ndarray(&x, t, &ctx) < 0) {
+        ndt_err_clear(&ctx);
+        return -1;
+    }
+
+    n = 1;
+    for (int i = 0; i < x.ndim; i++) {
+        n *= x.shape[i];
+    }
+
+    return n;
+}
+
 /* Return the next type in a dimension chain.  Undefined for non-dimensions. */
 static const ndt_t *
-next_dim(const ndt_t *a)
+next_dim(const ndt_t *t)
 {
-    assert(a->ndim > 0);
+    assert(t->ndim > 0);
 
-    switch (a->tag) {
+    switch (t->tag) {
     case FixedDim:
-        return a->FixedDim.type;
+        return t->FixedDim.type;
     case VarDim:
-        return a->VarDim.type;
+        return t->VarDim.type;
+    case VarDimElem:
+        return t->VarDimElem.type;
     case SymbolicDim:
-        return a->SymbolicDim.type;
+        return t->SymbolicDim.type;
     case EllipsisDim:
-        return a->EllipsisDim.type;
+        return t->EllipsisDim.type;
+    case Array:
+        return t->Array.type;
     default:
         /* NOT REACHED: tags should be exhaustive. */
         ndt_internal_error("invalid value");
@@ -126,16 +151,19 @@ ndt_dtype(const ndt_t *t)
     return t;
 }
 
-const ndt_t *
-ndt_dim_at(const ndt_t *t, int n)
+int
+ndt_logical_ndim(const ndt_t *t)
 {
-    assert(n <= t->ndim);
+    int ndim = t->ndim;
 
-    for (int i = 0; i < n; i++) {
-        t = next_dim(t);
+    while (t->tag == VarDim || t->tag == VarDimElem) {
+        if (t->tag == VarDimElem) {
+            --ndim;
+        }
+        t = t->VarDim.type;
     }
 
-    return t;
+    return ndim;
 }
 
 const ndt_t *
@@ -170,6 +198,75 @@ ndt_dims_dtype(const ndt_t *dims[NDT_MAX_DIM], const ndt_t **dtype, const ndt_t
     return n;
 }
 
+static const ndt_t *
+compress(const ndt_t *t)
+{
+    while (t->tag == VarDimElem) {
+        t = t->VarDimElem.type;
+    }
+
+    return t;
+}
+
+/*
+ * Return the next logical type in a dimension chain.  Undefined for
+ * non-dimensions.
+ */
+static const ndt_t *
+next_logical_dim(const ndt_t *t)
+{
+    switch (t->tag) {
+    case FixedDim:
+        return t->FixedDim.type;
+    case VarDim:
+        return compress(t->VarDim.type);
+    case VarDimElem:
+        return next_logical_dim(compress(t));
+    case SymbolicDim:
+        return t->SymbolicDim.type;
+    case EllipsisDim:
+        return t->EllipsisDim.type;
+    default:
+        return t;
+    }
+}
+
+const ndt_t *
+ndt_logical_dim_at(const ndt_t *t, int n)
+{
+    for (int i = 0; i < n; i++) {
+        t = next_logical_dim(t);
+    }
+
+    return t;
+}
+
+const ndt_t *
+ndt_copy_contiguous_at(const ndt_t *t, int n, const ndt_t *dtype, ndt_context_t *ctx)
+{
+    const ndt_t *u;
+
+    if (!ndt_is_ndarray(t)) {
+        ndt_err_format(ctx, NDT_NotImplementedError,
+            "partial copies are currently restricted to fixed dimensions");
+        return NULL;
+    }
+
+    if (n < 0 || n > t->ndim) {
+        ndt_err_format(ctx, NDT_ValueError, "n out of bounds");
+        return NULL;
+    }
+
+    u = ndt_logical_dim_at(t, n);
+
+    if (dtype != NULL) {
+        return ndt_copy_contiguous_dtype(u, dtype, 0, ctx);
+    }
+    else {
+        return ndt_copy_contiguous(u, 0, ctx);
+    }
+}
+
 int
 ndt_as_ndarray(ndt_ndarray_t *a, const ndt_t *t, ndt_context_t *ctx)
 {
@@ -205,9 +302,89 @@ ndt_as_ndarray(ndt_ndarray_t *a, const ndt_t *t, ndt_context_t *ctx)
     return n;
 }
 
+static const ndt_t *
+_ndt_transpose(const ndt_ndarray_t *a, const int p[], const ndt_t *dtype,
+               ndt_context_t *ctx)
+{
+    const ndt_t *t;
+
+    t = dtype;
+    ndt_incref(t);
+
+    for (int i = a->ndim-1; i >= 0; i--) {
+        const ndt_t *u = ndt_fixed_dim(t, a->shape[p[i]], a->steps[p[i]], ctx);
+        ndt_decref(t);
+        if (u == NULL) {
+            return NULL;
+        }
+        t = u;
+    }
+
+    return t;
+}
+
+const ndt_t *
+ndt_transpose(const ndt_t *t, const int *p, int ndim, ndt_context_t *ctx)
+{
+    ndt_ndarray_t a;
+    int permute[NDT_MAX_DIM];
+
+    if (ndt_as_ndarray(&a, t, ctx) < 0) {
+        return NULL;
+    }
+
+    if (p == NULL) {
+        if (ndim != 0) {
+            ndt_err_format(ctx, NDT_RuntimeError,
+                "internal error: ndim != 0 but no permutations given");
+            return NULL;
+        }
+
+        ndim = a.ndim;
+        if (ndim <= 1) {
+            ndt_incref(t);
+            return t;
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            permute[i] = ndim-i-1;
+        }
+        p = permute;
+    }
+    else {
+        if (ndim != a.ndim) {
+            ndt_err_format(ctx, NDT_ValueError,
+                "number of permutations != ndim");
+            return NULL;
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            permute[i] = 0;
+        }
+
+        for (int i = 0; i < ndim; i++) {
+            if (p[i] < 0 || p[i] >= ndim) {
+                ndt_err_format(ctx, NDT_ValueError,
+                    "permutation with value %d out of bounds", p[i]);
+                return NULL;
+            }
+
+            if (permute[p[i]] != 0) {
+                ndt_err_format(ctx, NDT_ValueError,
+                    "duplicate permutation index=%d", p[i]);
+                return NULL;
+            }
+
+            permute[p[i]] = 1;
+        }
+    }
+
+    return _ndt_transpose(&a, p, ndt_dtype(t), ctx);
+}
+
 /* Unoptimized hash function for experimenting. */
 ndt_ssize_t
-ndt_hash(ndt_t *t, ndt_context_t *ctx)
+ndt_hash(const ndt_t *t, ndt_context_t *ctx)
 {
     unsigned char *s, *cp;
     size_t len;
@@ -242,11 +419,11 @@ ndt_hash(ndt_t *t, ndt_context_t *ctx)
 
 const ndt_apply_spec_t ndt_apply_spec_empty = {
   .flags = 0U,
-  .nout = 0,
-  .nbroadcast = 0,
   .outer_dims = 0,
-  .out = {NULL},
-  .broadcast = {NULL}
+  .nin = 0,
+  .nout = 0,
+  .nargs = 0,
+  .types = {NULL}
 };
 
 ndt_apply_spec_t *
@@ -259,9 +436,10 @@ ndt_apply_spec_new(ndt_context_t *ctx)
         return ndt_memory_error(ctx);
     }
     spec->flags = 0U;
-    spec->nout = 0;
-    spec->nbroadcast = 0;
     spec->outer_dims = 0;
+    spec->nin = 0;
+    spec->nout = 0;
+    spec->nargs = 0;
 
     return spec;
 }
@@ -269,24 +447,20 @@ ndt_apply_spec_new(ndt_context_t *ctx)
 void
 ndt_apply_spec_clear(ndt_apply_spec_t *spec)
 {
-    int i;
-
     if (spec == NULL) {
         return;
     }
 
-    for (i = 0; i < spec->nbroadcast; i++) {
-        ndt_del(spec->broadcast[i]);
-    }
-
-    for (i = 0; i < spec->nout; i++) {
-        ndt_del(spec->out[i]);
+    for (int i = 0; i < spec->nin+spec->nout; i++) {
+        ndt_decref(spec->types[i]);
+        spec->types[i] = NULL;
     }
 
     spec->flags = 0U;
-    spec->nout = 0;
-    spec->nbroadcast = 0;
     spec->outer_dims = 0;
+    spec->nin = 0;
+    spec->nout = 0;
+    spec->nargs = 0;
 }
 
 void
@@ -300,6 +474,15 @@ ndt_apply_spec_del(ndt_apply_spec_t *spec)
     ndt_free(spec);
 }
 
+#define X (NDT_INNER_XND)
+#define S (NDT_INNER_STRIDED)
+#define C (NDT_INNER_C)
+#define F (NDT_INNER_F)
+
+#define ES (NDT_EXT_STRIDED)
+#define EC (NDT_EXT_C)
+#define EZ (NDT_EXT_ZERO)
+
 /* This function is used in places where it is _really_ convenient not to
    have to deal with deallocating an error message. */
 const char *
@@ -307,38 +490,35 @@ ndt_apply_flags_as_string(const ndt_apply_spec_t *spec)
 {
     switch (spec->flags) {
     case 0: return "None";
-    case NDT_C: return "C";
-    case NDT_FORTRAN: return "Fortran";
-    case NDT_FORTRAN|NDT_C: return "C|Fortran";
-    case NDT_STRIDED: return "Strided";
-    case NDT_STRIDED|NDT_C: return "C|Strided";
-    case NDT_STRIDED|NDT_FORTRAN: return "Fortran|Strided";
-    case NDT_STRIDED|NDT_FORTRAN|NDT_C: return "C|Fortran|Strided";
-    case NDT_XND: return "Xnd";
-    case NDT_XND|NDT_C: return "C|Xnd";
-    case NDT_XND|NDT_FORTRAN: return "Fortran|Xnd";
-    case NDT_XND|NDT_FORTRAN|NDT_C: return "C|Fortran|Xnd";
-    case NDT_XND|NDT_STRIDED: return "Strided|Xnd";
-    case NDT_XND|NDT_STRIDED|NDT_C: return "C|Strided|Xnd";
-    case NDT_XND|NDT_STRIDED|NDT_FORTRAN: return "Fortran|Strided|Xnd";
-    case NDT_XND|NDT_STRIDED|NDT_FORTRAN|NDT_C: return "C|Fortran|Strided|Xnd";
-    case NDT_ELEMWISE_1D: return "Elemwise1D";
-    case NDT_ELEMWISE_1D|NDT_C: return "Elemwise1D|C";
-    case NDT_ELEMWISE_1D|NDT_FORTRAN: return "Elemwise1D|Fortran";
-    case NDT_ELEMWISE_1D|NDT_FORTRAN|NDT_C: return "Elemwise1D|C|Fortran";
-    case NDT_ELEMWISE_1D|NDT_STRIDED: return "Elemwise1D|Strided";
-    case NDT_ELEMWISE_1D|NDT_STRIDED|NDT_C: return "Elemwise1D|C|Strided";
-    case NDT_ELEMWISE_1D|NDT_STRIDED|NDT_FORTRAN: return "Elemwise1D|Fortran|Strided";
-    case NDT_ELEMWISE_1D|NDT_STRIDED|NDT_FORTRAN|NDT_C: return "Elemwise1D|C|Fortran|Strided";
-    case NDT_ELEMWISE_1D|NDT_XND: return "Elemwise1D|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_C: return "Elemwise1D|C|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_FORTRAN: return "Elemwise1D|Fortran|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_FORTRAN|NDT_C: return "Elemwise1D|C|Fortran|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_STRIDED: return "Elemwise1D|Strided|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_STRIDED|NDT_C: return "Elemwise1D|C|Strided|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_STRIDED|NDT_FORTRAN: return "Elemwise1D|Fortran|Strided|Xnd";
-    case NDT_ELEMWISE_1D|NDT_XND|NDT_STRIDED|NDT_FORTRAN|NDT_C: return "Elemwise1D|C|Fortran|Strided|Xnd";
-    default: return "error: invalid combination of spec->flags";
+    case X: return "Xnd";
+
+    case S|X: return "Strided|Xnd";
+    case C|S|X: return "C|Strided|Xnd";
+    case F|S|X: return "Fortran|Strided|Xnd";
+    case C|F|S|X: return "C|Fortran|Strided|Xnd";
+
+    case ES|S|X: return "OptS|Strided|Xnd";
+    case ES|C|S|X: return "OptS|C|Strided|Xnd";
+    case ES|F|S|X: return "OptS|Fortran|Strided|Xnd";
+    case ES|C|F|S|X: return "OptS|C|Fortran|Strided|Xnd";
+    case EC|ES|C|S|X: return "OptC|OptS|C|Strided|Xnd";
+    case EC|ES|C|F|S|X: return "OptC|OptS|C|Fortran|Strided|Xnd";
+    case EZ|EC|ES|C|F|S|X: return "OptZ|OptC|OptS|C|Fortran|Strided|Xnd";
+    case EZ|EC|ES|C|S|X: return "OptZ|OptC|OptS|C|Strided|Xnd";
+    case EZ|ES|C|S|X: return "OptZ|OptS|C|Strided|Xnd";
+    case EZ|ES|C|F|S|X: return "OptZ|OptS|C|Fortran|Strided|Xnd";
+
+    default:
+        if (spec->flags & NDT_EXT_ZERO) fprintf(stderr, "EZ ");
+        if (spec->flags & NDT_EXT_C) fprintf(stderr, "EC ");
+        if (spec->flags & NDT_EXT_STRIDED) fprintf(stderr, "ES ");
+        if (spec->flags & NDT_INNER_C) fprintf(stderr, "C ");
+        if (spec->flags & NDT_INNER_F) fprintf(stderr, "F ");
+        if (spec->flags & NDT_INNER_STRIDED) fprintf(stderr, "S ");
+        if (spec->flags & NDT_INNER_XND) fprintf(stderr, "X ");
+        fprintf(stderr, "\n");
+
+        return "unknown flags";
     }
 }
 
@@ -362,16 +542,21 @@ ndt_meta_new(ndt_context_t *ctx)
 }
 
 void
-ndt_meta_del(ndt_meta_t *m)
+ndt_meta_clear(ndt_meta_t *m)
 {
     if (m == NULL) {
         return;
     }
 
     for (int i = 0; i < m->ndims; i++) {
-        ndt_free(m->offsets[i]);
+        ndt_decref_offsets(m->offsets[i]);
     }
+}
 
+void
+ndt_meta_del(ndt_meta_t *m)
+{
+    ndt_meta_clear(m);
     ndt_free(m);
 }
 
@@ -380,95 +565,118 @@ ndt_meta_del(ndt_meta_t *m)
 /*                     Optimized kernel strategy (unstable API)              */
 /*****************************************************************************/
 
-static bool
-all_inner_1D_contiguous(const ndt_t *types[], int n)
+static uint32_t
+check_c(uint32_t flags, ndt_ndarray_t *x, int outer)
 {
-    for (int i = 0; i < n; i++) {
-        if (types[i]->ndim == 0) {
-            return false;
+    int inner = x->ndim-outer;
+    int64_t *shape = x->shape+outer;
+    int64_t *steps = x->steps+outer;
+    int64_t step = 1;
+
+    for (int i = inner-1; i >= 0; i--) {
+        if (shape[i] > 1 && steps[i] != step) {
+            return flags & ~(NDT_INNER_C|NDT_EXT_C|NDT_EXT_ZERO);
         }
-        if (!ndt_is_c_contiguous(ndt_dim_at(types[i], types[i]->ndim-1))) {
-            return false;
+        step *= shape[i];
+    }
+
+    if (x->ndim == inner) {
+        return flags & ~(NDT_EXT_C|NDT_EXT_ZERO);
+    }
+    else if (x->ndim > inner) {
+        if (shape[-1] > 1 && steps[-1] != step) {
+            flags &= ~NDT_EXT_C;
+            if (steps[-1] != 0) {
+                flags &= ~NDT_EXT_ZERO;
+            }
         }
     }
 
-    return true;
+    return flags;
 }
 
-static bool
-all_inner_c_contiguous(const ndt_t *types[], int n, int outer)
+static uint32_t
+check_f(uint32_t flags, ndt_ndarray_t *x, int outer)
 {
-    for (int i = 0; i < n; i++) {
-        if (!ndt_is_c_contiguous(ndt_dim_at(types[i], outer))) {
-            return false;
+    int inner = x->ndim-outer;
+    int64_t *shape = x->shape+outer;
+    int64_t *steps = x->steps+outer;
+    int64_t step = 1;
+
+    for (int i = 0; i < inner; i++) {
+        if (shape[i] > 1 && steps[i] != step) {
+            return flags & ~NDT_INNER_F;
         }
+        step *= shape[i];
     }
 
-    return true;
+    return flags;
 }
 
-static bool
-all_inner_f_contiguous(const ndt_t *types[], int n, int outer)
+static uint32_t
+check_strided(uint32_t flags, int outer)
 {
-    for (int i = 0; i < n; i++) {
-        if (!ndt_is_f_contiguous(ndt_dim_at(types[i], outer))) {
-            return false;
-        }
+    if (outer == 0) {
+        return flags & ~NDT_EXT_STRIDED;
     }
 
-    return true;
+    return flags;
 }
 
-static bool
-all_ndarray(const ndt_t *types[], int n)
+static uint32_t
+check_var(uint32_t flags, const ndt_t *t, int outer)
 {
-    for (int i = 0; i < n; i++) {
-        if (!ndt_is_ndarray(types[i])) {
-            return false;
-        }
-        if (ndt_subtree_is_optional(types[i])) {
-            return false;
-        }
-    }
-
-    return true;
+#if 0
+    /*
+     * This currently does not handle a corner case with zeros in a shape.
+     * The loop in gumath is safe, however, since it stops at shape==0.
+     */
+    if (ndt_logical_ndim(t) == outer) {
+        return flags & NDT_INNER_XND;
+    }
+
+    return 0U;
+#endif
+    (void)t;
+    (void)outer;
+    return flags & NDT_INNER_XND;
 }
 
-void
-ndt_select_kernel_strategy(ndt_apply_spec_t *spec, const ndt_t *sig,
-                           const ndt_t *in[], int nin)
+static uint32_t
+select_flags(const ndt_t *types[], int n, int outer, ndt_context_t *ctx)
 {
-    const ndt_t **out = (const ndt_t **)spec->out;
-    bool in_inner_c, in_inner_f, out_inner_c;
+    uint32_t flags = NDT_SPEC_FLAGS_ALL;
+    ndt_ndarray_t x;
 
-    assert(sig->tag == Function);
-    assert(spec->flags == 0);
+    for (int i = 0; i < n; i++) {
+        const ndt_t *t = types[i];
 
-    if (spec->nbroadcast > 0) {
-        in = (const ndt_t **)spec->broadcast;
-        nin = spec->nbroadcast;
+        if (ndt_as_ndarray(&x, t, ctx) < 0) { /* var dimension */
+            ndt_err_clear(ctx);
+            if (t->tag == VarDim || t->tag == VarDimElem) {
+                flags = check_var(flags, t, outer);
+            }
+        }
+        else {
+            if (outer > t->ndim) {
+                ndt_err_format(ctx, NDT_RuntimeError,
+                               "number of outer dimensions greater than ndim");
+                return UINT32_MAX;
+            }
+
+            flags = check_strided(flags, outer);
+            flags = check_c(flags, &x, outer);
+            flags = check_f(flags, &x, outer);
+        }
     }
 
-    in_inner_c = all_inner_c_contiguous(in, nin, spec->outer_dims);
-    in_inner_f = all_inner_f_contiguous(in, nin, spec->outer_dims);
-    out_inner_c = all_inner_c_contiguous(out, spec->nout, spec->outer_dims);
-
-    spec->flags = NDT_XND;
+    return flags;
+}
 
-    if (sig->Function.elemwise) {
-        if (all_inner_1D_contiguous(in, nin) &&
-            all_inner_1D_contiguous(out, spec->nout)) {
-            spec->flags |= NDT_ELEMWISE_1D;
-        }
-    }
-    if (in_inner_c && out_inner_c) {
-        spec->flags |= NDT_C;
-    }
-    if (in_inner_f && out_inner_c) {
-        spec->flags |= NDT_FORTRAN;
-    }
+int
+ndt_select_kernel_strategy(ndt_apply_spec_t *spec, ndt_context_t *ctx)
+{
+    spec->flags = select_flags(spec->types, spec->nargs, spec->outer_dims, ctx);
 
-    if (all_ndarray(in, nin) && all_ndarray(out, spec->nout)) {
-        spec->flags |= NDT_STRIDED;
-    }
+    return spec->flags == UINT32_MAX ? -1 : 0;
 }