File size: 5,233 Bytes
2166d7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | #include "l1_tensor.h"
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
cs_status tensor_new(tensor *t, size_t rank, const size_t *shape)
{
if (!t || !shape || rank == 0 || rank > CS_TENSOR_MAX_RANK) return CS_ERR_ARG;
size_t size = 1;
for (size_t i = 0; i < rank; i++) {
if (shape[i] == 0) return CS_ERR_ARG;
if (size > SIZE_MAX / shape[i]) return CS_ERR_NOMEM;
size *= shape[i];
}
if (size > SIZE_MAX / sizeof(gl_t)) return CS_ERR_NOMEM;
gl_t *d = calloc(size, sizeof(gl_t));
if (!d) return CS_ERR_NOMEM;
memset(t, 0, sizeof *t);
t->rank = rank;
memcpy(t->shape, shape, rank * sizeof(size_t));
t->d = d;
t->size = size;
return CS_OK;
}
void tensor_free(tensor *t)
{
if (!t) return;
free(t->d);
memset(t, 0, sizeof *t);
}
static int offset_of(const tensor *t, const size_t *idx, size_t *off)
{
size_t o = 0;
for (size_t i = 0; i < t->rank; i++) {
if (idx[i] >= t->shape[i]) return 0;
o = o * t->shape[i] + idx[i];
}
*off = o;
return 1;
}
cs_status tensor_get(const tensor *t, const size_t *idx, gl_t *out)
{
size_t o;
if (!t || !idx || !out || !t->d) return CS_ERR_ARG;
if (!offset_of(t, idx, &o)) return CS_ERR_RANGE;
*out = t->d[o];
return CS_OK;
}
cs_status tensor_set(tensor *t, const size_t *idx, gl_t v)
{
size_t o;
if (!t || !idx || !t->d) return CS_ERR_ARG;
if (!offset_of(t, idx, &o)) return CS_ERR_RANGE;
t->d[o] = gl_reduce(v);
return CS_OK;
}
static int same_shape(const tensor *a, const tensor *b)
{
return a->rank == b->rank && memcmp(a->shape, b->shape, a->rank * sizeof(size_t)) == 0;
}
static cs_status pointwise(const tensor *a, const tensor *b, tensor *out, int mul)
{
if (!a || !b || !out || !a->d || !b->d) return CS_ERR_ARG;
if (!same_shape(a, b)) return CS_ERR_RANGE;
tensor r;
cs_status s = tensor_new(&r, a->rank, a->shape);
if (s != CS_OK) return s;
for (size_t i = 0; i < r.size; i++)
r.d[i] = mul ? gl_mul(a->d[i], b->d[i]) : gl_add(a->d[i], b->d[i]);
*out = r;
return CS_OK;
}
cs_status tensor_add(const tensor *a, const tensor *b, tensor *out) { return pointwise(a, b, out, 0); }
cs_status tensor_hadamard(const tensor *a, const tensor *b, tensor *out) { return pointwise(a, b, out, 1); }
cs_status tensor_outer(const tensor *a, const tensor *b, tensor *out)
{
if (!a || !b || !out || !a->d || !b->d) return CS_ERR_ARG;
if (a->rank + b->rank > CS_TENSOR_MAX_RANK) return CS_ERR_RANGE;
size_t shape[CS_TENSOR_MAX_RANK];
memcpy(shape, a->shape, a->rank * sizeof(size_t));
memcpy(shape + a->rank, b->shape, b->rank * sizeof(size_t));
tensor r;
cs_status s = tensor_new(&r, a->rank + b->rank, shape);
if (s != CS_OK) return s;
for (size_t i = 0; i < a->size; i++)
for (size_t j = 0; j < b->size; j++) r.d[i * b->size + j] = gl_mul(a->d[i], b->d[j]);
*out = r;
return CS_OK;
}
cs_status tensor_contract(const tensor *a, size_t axis_a, const tensor *b, size_t axis_b, tensor *out)
{
if (!a || !b || !out || !a->d || !b->d) return CS_ERR_ARG;
if (axis_a >= a->rank || axis_b >= b->rank) return CS_ERR_RANGE;
if (a->shape[axis_a] != b->shape[axis_b]) return CS_ERR_RANGE;
size_t rr = a->rank + b->rank - 2;
if (rr == 0 || rr > CS_TENSOR_MAX_RANK) return CS_ERR_RANGE;
size_t shape[CS_TENSOR_MAX_RANK], n = 0;
for (size_t i = 0; i < a->rank; i++) if (i != axis_a) shape[n++] = a->shape[i];
for (size_t i = 0; i < b->rank; i++) if (i != axis_b) shape[n++] = b->shape[i];
tensor r;
cs_status s = tensor_new(&r, rr, shape);
if (s != CS_OK) return s;
size_t ext = a->shape[axis_a];
size_t ia[CS_TENSOR_MAX_RANK], ib[CS_TENSOR_MAX_RANK], ir[CS_TENSOR_MAX_RANK] = {0};
for (size_t o = 0; o < r.size; o++) {
/* decode o into the free indices of a then b */
size_t rem = o;
for (size_t i = rr; i-- > 0;) { ir[i] = rem % r.shape[i]; rem /= r.shape[i]; }
size_t p = 0;
for (size_t i = 0; i < a->rank; i++) if (i != axis_a) ia[i] = ir[p++];
for (size_t i = 0; i < b->rank; i++) if (i != axis_b) ib[i] = ir[p++];
gl_t acc = 0;
for (size_t k = 0; k < ext; k++) {
ia[axis_a] = k;
ib[axis_b] = k;
size_t oa = 0, ob = 0;
if (!offset_of(a, ia, &oa) || !offset_of(b, ib, &ob)) {
tensor_free(&r);
return CS_ERR_RANGE;
}
acc = gl_add(acc, gl_mul(a->d[oa], b->d[ob]));
}
r.d[o] = acc;
}
*out = r;
return CS_OK;
}
cs_status tensor_from_pmat(const pmat *m, tensor *t)
{
if (!m || !t || !m->d) return CS_ERR_ARG;
size_t shape[2] = {m->rows, m->cols};
cs_status s = tensor_new(t, 2, shape);
if (s != CS_OK) return s;
memcpy(t->d, m->d, t->size * sizeof(gl_t));
return CS_OK;
}
cs_status tensor_to_pmat(const tensor *t, pmat *m)
{
if (!t || !m || !t->d) return CS_ERR_ARG;
if (t->rank != 2) return CS_ERR_RANGE;
cs_status s = pmat_new(m, t->shape[0], t->shape[1]);
if (s != CS_OK) return s;
memcpy(m->d, t->d, t->size * sizeof(gl_t));
return CS_OK;
}
|