← Commits · de5af5ff
de5af5ff1f6270503285eb364be8a4eed0b832ac
diff --git a/Makefile b/Makefile
index d6beb5d..ec84f5a 100644
--- a/Makefile
+++ b/Makefile
@@ -1,7 +1,13 @@
CFLAGS ?= -std=c99 -Wall -Wextra -g
+LDFLAGS ?= -lm
+
+METAL := $(shell test -d /System/Library/Frameworks/Metal.framework && echo 1 || echo 0)
+ifeq ($(METAL),1)
+ LDFLAGS += -framework Metal -framework MetalPerformanceShaders -framework Foundation -framework Accelerate
+endif
all:
- $(CC) $(CFLAGS) test.c -o test
+ $(CC) $(CFLAGS) $(LDFLAGS) test.c -o test
./test
.PHONY: all
diff --git a/test.c b/test.c
index 487c52e..961d613 100644
--- a/test.c
+++ b/test.c
@@ -2,6 +2,16 @@
#include "utensil.h"
+#define assert_eq(actual, expected, tol) \
+ do { \
+ if (fabs(actual - expected) > tol) { \
+ fprintf(stderr, "FAIL(%d): got %f expected %f\n", __LINE__, (actual), (expected)); \
+ abort(); \
+ } \
+ } while (0)
+#define assert_data(actual, expected, tol) \
+ for (int i = 0; i < (actual)->shape.nelem; i++) assert_eq((actual)->data[i], (expected)[i], tol)
+
static void test_shape(void) {
ut_shape s = ut_shape_new(3, (int[]){2, 3, 4});
assert(s.nelem == 24);
@@ -13,7 +23,7 @@ static void test_shape(void) {
}
static void test_lifetime(void) {
- ut_tensor* a = ut_alloc(2, (int[]){4, 4});
+ ut_tensor* a = ut_alloc(2, (int[]){4, 4}, UT_CPU);
a->data[0] = 42.f;
ut_tensor* v = ut_view(a, 1, (int[]){16});
@@ -30,7 +40,7 @@ static void test_lifetime(void) {
}
void test_reshape(void) {
- ut_tensor* a = ut_randn(2, (int[]){2, 3}, 0.f, 1.f);
+ ut_tensor* a = ut_randn(2, (int[]){2, 3}, 0.f, 1.f, UT_CPU);
ut_tensor* v = ut_view(a, 1, (int[]){6});
for (int i = 0; i < 6; i++) assert(v->data[i] == a->data[i]);
ut_free(v);
@@ -38,29 +48,80 @@ void test_reshape(void) {
}
void test_elementwise(void) {
- ut_tensor* a = ut_from_data(1, (int[]){5}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f});
- ut_tensor* b = ut_from_data(1, (int[]){5}, (float[]){5.f, 4.f, 3.f, 2.f, 1.f});
-
- ut_tensor* c = ut_add(a, b);
- for (int i = 0; i < 5; i++) assert(c->data[i] == a->data[i] + b->data[i]);
-
- ut_tensor* d = ut_mul(a, b);
- for (int i = 0; i < 5; i++) assert(d->data[i] == a->data[i] * b->data[i]);
+ ut_tensor* a = ut_from_data(1, (int[]){5}, (float[]){1.f, 2.f, 3.f, 4.f, 5.f}, UT_METAL);
+ ut_tensor* b = ut_from_data(1, (int[]){5}, (float[]){5.f, 4.f, 3.f, 2.f, 1.f}, UT_METAL);
+ // unary
+ {
+ ut_tensor* c = ut_neg(a);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){-1.f, -2.f, -3.f, -4.f, -5.f}), 1e-6f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_exp(a);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){2.7182817f, 7.389056f, 20.085537f, 54.59815f, 148.41316f}), 1e-4f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_sigmoid(a);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){0.7310586f, 0.880797f, 0.9525741f, 0.9820138f, 0.9933071f}), 1e-4f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_tanh(a);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){0.7615942f, 0.9640276f, 0.9950547f, 0.9993293f, 0.9999092f}), 1e-4f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_relu(a);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){1.f, 2.f, 3.f, 4.f, 5.f}), 1e-6f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_scale(a, 2.f);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){2.f, 4.f, 6.f, 8.f, 10.f}), 1e-6f);
+ ut_free(c);
+ }
- ut_tensor* e = ut_neg(a);
- for (int i = 0; i < 5; i++) assert(e->data[i] == -a->data[i]);
+ // binary
+ {
+ ut_tensor* c = ut_add(a, b);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){6.f, 6.f, 6.f, 6.f, 6.f}), 1e-6f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_sub(a, b);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){-4.f, -2.f, 0.f, 2.f, 4.f}), 1e-6f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_mul(a, b);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){5.f, 8.f, 9.f, 8.f, 5.f}), 1e-6f);
+ ut_free(c);
+ }
+ {
+ ut_tensor* c = ut_div(a, b);
+ ut_sync_cpu(c);
+ assert_data(c, ((float[]){0.2f, 0.5f, 1.f, 2.f, 5.f}), 1e-6f);
+ ut_free(c);
+ }
ut_free(a);
ut_free(b);
- ut_free(c);
- ut_free(d);
- ut_free(e);
}
int main() {
test_shape();
test_lifetime();
test_reshape();
- test_elementwise();
+ test_elementwise();
return 0;
}
diff --git a/utensil.h b/utensil.h
index aa5172f..b500d48 100644
--- a/utensil.h
+++ b/utensil.h
@@ -2,25 +2,243 @@
#define UTENSIL_H
#include <math.h>
+#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#define UT_MAX_DIMS 4
+typedef enum { UT_CPU, UT_METAL } ut_dev;
+
typedef struct ut_shape {
- int ndim;
- int nelem;
- int shape[UT_MAX_DIMS];
- int stride[UT_MAX_DIMS];
+ int ndim; // number of dimensions
+ int nelem; // total number of elements
+ int shape[UT_MAX_DIMS]; // size of each dimension
+ int stride[UT_MAX_DIMS]; // stride for each dimension
} ut_shape;
typedef struct ut_tensor {
- ut_shape shape;
- float* data;
- int rc;
- struct ut_tensor* owner;
+ ut_shape shape; // shape of the tensor
+ float* data; // pointer to the data buffer
+ void* gpu_buf; /* MTLBuffer */
+ int rc; // reference count for memory management
+ struct ut_tensor* owner; // if this is a view, points to the owner tensor
+ ut_dev dev; // device where the tensor is allocated
+ bool dirty_cpu; // GPU is newer than CPU mirror
+ bool dirty_gpu; // CPU is newer than GPU buffer
} ut_tensor;
+// =========================================================
+// Metal context management
+// =========================================================
+#ifdef __APPLE__
+#include <Accelerate/Accelerate.h>
+#include <objc/message.h>
+#include <objc/runtime.h>
+// Type-safe objc_msgSend wrappers
+static inline void* _m0(void* o, const char* s) {
+ return ((id (*)(id, SEL))objc_msgSend)((id)o, sel_getUid(s));
+}
+static inline void* _m1(void* o, const char* s, void* a) {
+ return ((id (*)(id, SEL, id))objc_msgSend)((id)o, sel_getUid(s), (id)a);
+}
+static inline void* _m1s(void* o, const char* s, const char* a) {
+ return ((id (*)(id, SEL, const char*))objc_msgSend)((id)o, sel_getUid(s), a);
+}
+static inline void* _m2ll(void* o, const char* s, long a, long b) {
+ return ((id (*)(id, SEL, long, long))objc_msgSend)((id)o, sel_getUid(s), a, b);
+}
+static inline void* _m4l(void* o, const char* s, long a, long b, long c, long d) {
+ return ((id (*)(id, SEL, long, long, long, long))objc_msgSend)((id)o, sel_getUid(s), a, b, c, d);
+}
+static inline void _v0(void* o, const char* s) {
+ ((void (*)(id, SEL))objc_msgSend)((id)o, sel_getUid(s));
+}
+static inline void _v1(void* o, const char* s, void* a) {
+ ((void (*)(id, SEL, id))objc_msgSend)((id)o, sel_getUid(s), (id)a);
+}
+static inline void* _p0(void* o, const char* s) {
+ return ((void* (*)(id, SEL))objc_msgSend)((id)o, sel_getUid(s));
+}
+static inline unsigned long _l0(void* o, const char* s) {
+ return ((unsigned long (*)(id, SEL))objc_msgSend)((id)o, sel_getUid(s));
+}
+static inline const char* _c0(void* o, const char* s) {
+ return ((const char* (*)(id, SEL))objc_msgSend)((id)o, sel_getUid(s));
+}
+
+static const char* _mtl_src =
+ "#include <metal_stdlib>\nusing namespace metal;\n"
+ /* unary */
+ "kernel void uneg(device const float*i,device float*o,constant int&n,"
+ "uint idx[[thread_position_in_grid]]){o[idx]=-i[idx];}\n"
+ "kernel void urelu(device const float*i,device float*o,constant int&n,"
+ "uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=max(i[idx],0.f);}\n"
+ "kernel void usig(device const float*i,device float*o,constant int&n,"
+ "uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=1.f/(1.f+exp(-i[idx]));}\n"
+ "kernel void utanh(device const float*i,device float*o,constant int&n,"
+ "uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=tanh(i[idx]);}\n"
+ "kernel void uexp(device const float*i,device float*o,constant int&n,"
+ "uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=exp(i[idx]);}\n"
+
+ /* binary */
+ "kernel void badd(device const float*a,device const float*b,device float*o,"
+ "constant int&n,uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=a[idx]+b[idx];}\n"
+ "kernel void bsub(device const float*a,device const float*b,device float*o,"
+ "constant int&n,uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=a[idx]-b[idx];}\n"
+ "kernel void bmul(device const float*a,device const float*b,device float*o,"
+ "constant int&n,uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=a[idx]*b[idx];}\n"
+ "kernel void bdiv(device const float*a,device const float*b,device float*o,"
+ "constant int&n,uint idx[[thread_position_in_grid]]){if((int)idx<n)o[idx]=a[idx]/b[idx];}\n"
+
+ /* scalar multiply: args = {float s, int n} packed as 8 bytes */
+ "struct ScaleArgs{ float s; int n; };\n"
+ "kernel void bscale(device const float*a,device float*o,"
+ "constant ScaleArgs&args,uint idx[[thread_position_in_grid]])"
+ "{if((int)idx<args.n)o[idx]=a[idx]*args.s;}\n"
+ "";
+
+#define _MTL_MAX_PL 32 // maximum number of pipeline states
+typedef struct {
+ void *device, *queue, *library, *pl[_MTL_MAX_PL];
+ const char* pl_name[_MTL_MAX_PL];
+ int pl_count;
+ void* pending_cmd; /* open command buffer for batching; NULL if none */
+} _mtl_ctx_t;
+
+typedef struct {
+ unsigned long w, h, d;
+} _msize_t;
+
+static void* _mtl_init(void) {
+ _mtl_ctx_t* c = calloc(1, sizeof(_mtl_ctx_t));
+ extern id MTLCreateSystemDefaultDevice(void);
+ c->device = MTLCreateSystemDefaultDevice();
+ if (!c->device) {
+ free(c);
+ return NULL;
+ }
+ c->queue = _m0(c->device, "newCommandQueue");
+ void* src = _m1s(objc_getClass("NSString"), "stringWithUTF8String:", _mtl_src);
+ void* opts = _m0(objc_getClass("MTLCompileOptions"), "new");
+ void* lib = ((id (*)(id, SEL, id, id, void*))objc_msgSend)(
+ (id)c->device, sel_getUid("newLibraryWithSource:options:error:"), (id)src, (id)opts, NULL);
+ _v0(opts, "release");
+ if (!lib) {
+ _v0(c->queue, "release");
+ _v0(c->device, "release");
+ free(c);
+ return NULL;
+ }
+ c->library = lib;
+ return c;
+}
+
+static void _mtl_free(void* ctx) {
+ if (!ctx) return;
+ _mtl_ctx_t* c = (_mtl_ctx_t*)ctx;
+ for (int i = 0; i < c->pl_count; i++) _v0(c->pl[i], "release");
+ _v0(c->library, "release");
+ _v0(c->queue, "release");
+ _v0(c->device, "release");
+ free(c);
+}
+
+static void _mtl_flush(_mtl_ctx_t* c) {
+ if (!c->pending_cmd) return;
+ _v0(c->pending_cmd, "commit");
+ _v0(c->pending_cmd, "waitUntilCompleted");
+ _v0(c->pending_cmd, "release");
+ c->pending_cmd = NULL;
+}
+
+static void* _mtl_get_pl(_mtl_ctx_t* c, const char* name) {
+ for (int i = 0; i < c->pl_count; i++)
+ if (c->pl_name[i] && !strcmp(c->pl_name[i], name)) return c->pl[i];
+ void* fn =
+ _m1(c->library,
+ "newFunctionWithName:", _m1s(objc_getClass("NSString"), "stringWithUTF8String:", name));
+ if (!fn) {
+ fprintf(stderr, "utensil: Metal fn '%s' not found\n", name);
+ return NULL;
+ }
+ void* ps = ((id (*)(id, SEL, id, void*))objc_msgSend)(
+ (id)c->device, sel_getUid("newComputePipelineStateWithFunction:error:"), (id)fn, NULL);
+ _v0(fn, "release");
+ if (c->pl_count < _MTL_MAX_PL) {
+ c->pl[c->pl_count] = ps;
+ c->pl_name[c->pl_count] = name;
+ c->pl_count++;
+ }
+ return ps;
+}
+
+/* Get or create the batched command buffer. */
+static void* _mtl_get_cmd(_mtl_ctx_t* c) {
+ if (!c->pending_cmd) {
+ void* cb = _m0(c->queue, "commandBuffer");
+ c->pending_cmd = ((id (*)(id, SEL))objc_msgSend)((id)cb, sel_getUid("retain"));
+ }
+ return c->pending_cmd;
+}
+
+static void _mtl_dispatch(_mtl_ctx_t* c, const char* kern, void** bufs, int nbufs,
+ const void* bytes, int blen, int n) {
+ void* cb = _mtl_get_cmd(c);
+ void* enc = _m0(cb, "computeCommandEncoder");
+ void* ps = _mtl_get_pl(c, kern);
+ if (!ps) {
+ _v0(enc, "endEncoding");
+ return;
+ }
+ _v1(enc, "setComputePipelineState:", ps);
+ for (int i = 0; i < nbufs; i++)
+ ((void (*)(id, SEL, id, long, long))objc_msgSend)(
+ (id)enc, sel_getUid("setBuffer:offset:atIndex:"), (id)bufs[i], 0L, (long)i);
+ if (bytes)
+ ((void (*)(id, SEL, const void*, long, long))objc_msgSend)(
+ (id)enc, sel_getUid("setBytes:length:atIndex:"), bytes, (long)blen, (long)nbufs);
+ unsigned long tg = _l0(ps, "maxTotalThreadsPerThreadgroup");
+ unsigned long tgs = ((unsigned long)n + tg - 1) / tg;
+ _msize_t grp = {tgs, 1, 1}, thr = {tg, 1, 1};
+ ((void (*)(id, SEL, _msize_t, _msize_t))objc_msgSend)(
+ (id)enc, sel_getUid("dispatchThreadgroups:threadsPerThreadgroup:"), grp, thr);
+ _v0(enc, "endEncoding");
+}
+
+static void* _mtl_buf_alloc(_mtl_ctx_t* c, size_t bytes) {
+ return _m2ll(c->device, "newBufferWithLength:options:", (long)bytes, 0L);
+}
+static void _mtl_buf_free(void* b) { _v0(b, "release"); }
+static void _mtl_buf_write(void* b, const float* src, size_t n) {
+ memcpy(_p0(b, "contents"), src, n * sizeof(float));
+}
+static void _mtl_buf_read(void* b, float* dst, size_t n) {
+ memcpy(dst, _p0(b, "contents"), n * sizeof(float));
+}
+
+static void* _g_metal = NULL;
+static int _g_metal_init = 0;
+
+void* ut_metal_ctx(void) {
+ if (!_g_metal_init) _g_metal_init = 1, _g_metal = _mtl_init();
+ return _g_metal;
+}
+
+#if defined(__APPLE__) && defined(__GNUC__)
+__attribute__((destructor)) static void _ut_metal_cleanup(void) {
+ if (_g_metal) {
+ _mtl_free(_g_metal);
+ _g_metal = NULL;
+ }
+}
+
+#else
+void* ut_metal_ctx(void) { return NULL; }
+#endif
+
+#endif
+
// =========================================================
// Allocation and lifetime management
// =========================================================
@@ -37,52 +255,123 @@ int ut_index(ut_shape s, const int* idx) {
return flat;
}
-ut_tensor* ut_alloc(int ndim, int* dim) {
+void ut_sync_cpu(ut_tensor* t) {
+ if (t->owner) { // Delegate to owner
+ ut_sync_cpu(t->owner);
+ t->data = t->owner->data;
+ return;
+ }
+ if (!t->dirty_cpu) return;
+ // flush any pending GPU work before reading back to CPU
+ _mtl_ctx_t* _mc_s = (_mtl_ctx_t*)ut_metal_ctx();
+ if (_mc_s) _mtl_flush(_mc_s);
+ if (!t->data) t->data = malloc((size_t)t->shape.nelem * sizeof(float));
+ _mtl_buf_read(t->gpu_buf, t->data, (size_t)t->shape.nelem);
+ t->dirty_cpu = false;
+}
+
+void ut_sync_gpu(ut_tensor* t) {
+ if (!t->dirty_gpu || !t->gpu_buf) return;
+ _mtl_buf_write(t->gpu_buf, t->data, (size_t)t->shape.nelem);
+ t->dirty_gpu = false;
+}
+
+void ut_to_device(ut_tensor* t, ut_dev dev) {
+ if (t->owner) { // Delegate device move to owner; update our alias pointers
+ ut_to_device(t->owner, dev);
+ t->data = t->owner->data;
+ t->gpu_buf = t->owner->gpu_buf;
+ t->dev = t->owner->dev;
+ return;
+ }
+ if (t->dev == dev) {
+ dev == UT_CPU ? ut_sync_cpu(t) : ut_sync_gpu(t);
+ return;
+ }
+ if (dev == UT_METAL) {
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (!mc) return;
+ if (!t->gpu_buf) t->gpu_buf = _mtl_buf_alloc(mc, (size_t)t->shape.nelem * sizeof(float));
+ _mtl_buf_write(t->gpu_buf, t->data, (size_t)t->shape.nelem);
+ t->dev = UT_METAL;
+ t->dirty_cpu = t->dirty_gpu = false;
+ } else {
+ ut_sync_cpu(t);
+ _mtl_buf_free(t->gpu_buf);
+ t->gpu_buf = NULL;
+ t->dev = UT_CPU;
+ t->dirty_gpu = false;
+ }
+}
+
+ut_tensor* ut_alloc(int ndim, const int* dim, ut_dev dev) {
ut_tensor* t = (ut_tensor*)malloc(sizeof(ut_tensor));
- *t = (struct ut_tensor){.shape = ut_shape_new(ndim, dim), .owner = NULL, .rc = 1};
+ *t = (struct ut_tensor){.shape = ut_shape_new(ndim, dim), .dev = dev, .owner = NULL, .rc = 1};
t->data = (float*)malloc(t->shape.nelem * sizeof(float));
+ if (dev == UT_METAL) {
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ if (mc)
+ t->gpu_buf = _mtl_buf_alloc(mc, (size_t)t->shape.nelem * sizeof(float));
+ else
+ t->dev = UT_CPU;
+ }
return t;
}
-ut_tensor* ut_randn(int ndim, int* dim, float mean, float stddev) {
- ut_tensor* t = ut_alloc(ndim, dim);
+ut_tensor* ut_randn(int ndim, int* dim, float mean, float stddev, ut_dev dev) {
+ ut_tensor* t = ut_alloc(ndim, dim, UT_CPU);
for (int i = 0; i < t->shape.nelem; i++) {
// Box-Muller transform to generate normally distributed random numbers
- float u1 = (float)rand() / RAND_MAX;
- float u2 = (float)rand() / RAND_MAX;
+ float u1 = (float)rand() / (float)RAND_MAX;
+ float u2 = (float)rand() / (float)RAND_MAX;
float z0 = sqrtf(-2.0f * logf(u1)) * cosf(6.28f * u2);
t->data[i] = z0 * stddev + mean;
}
+ if (dev == UT_METAL) ut_to_device(t, UT_METAL);
return t;
}
-ut_tensor* ut_from_data(int ndim, const int* dim, const float* data) {
- ut_tensor* t = ut_alloc(ndim, dim);
+ut_tensor* ut_from_data(int ndim, const int* dim, const float* data, ut_dev dev) {
+ ut_tensor* t = ut_alloc(ndim, dim, UT_CPU);
memcpy(t->data, data, (size_t)t->shape.nelem * sizeof(float));
+ if (dev == UT_METAL) ut_to_device(t, UT_METAL);
return t;
}
ut_tensor* ut_clone(ut_tensor* t) {
- ut_tensor* c = ut_alloc(t->shape.ndim, t->shape.shape);
+ ut_sync_cpu(t);
+ ut_tensor* c = ut_alloc(t->shape.ndim, t->shape.shape, UT_CPU);
memcpy(c->data, t->data, (size_t)t->shape.nelem * sizeof(float));
+ if (t->dev == UT_METAL) ut_to_device(c, UT_METAL);
return c;
}
void ut_free(ut_tensor* t) {
if (!--t->rc) return;
- if (!t->owner) free(t->data);
+ if (t->owner)
+ ut_free(t->owner);
+ else {
+ free(t->data);
+ if (t->gpu_buf) _mtl_buf_free(t->gpu_buf);
+ }
free(t);
}
+ut_tensor* ut_retain(ut_tensor* t) { return t->rc++, t; }
+
ut_tensor* ut_view(ut_tensor* t, int ndim, const int* dim) {
+ ut_sync_cpu(t->owner ? t->owner : t);
ut_shape ns = ut_shape_new(ndim, dim);
ut_tensor* v = (ut_tensor*)malloc(sizeof(ut_tensor));
- *v = (struct ut_tensor){.shape = ns, .data = t->data, .owner = t->owner ? t->owner : t};
+ *v = (struct ut_tensor){.shape = ns,
+ .data = t->data,
+ .gpu_buf = t->gpu_buf,
+ .rc = 1,
+ .dev = t->dev,
+ .owner = ut_retain(t)};
return v;
}
-ut_tensor* ut_retain(ut_tensor* t) { return t->rc++, t; }
-
// =========================================================
// Elementwise operations
// =========================================================
@@ -113,30 +402,67 @@ static void ew_mul(float* out, const float* a, const float* b, int n) {
static void ew_div(float* out, const float* a, const float* b, int n) {
for (int i = 0; i < n; i++) out[i] = a[i] / b[i];
}
-static ut_tensor* ew_unary(ut_tensor* a, void (*fn)(float*, const float*, int)) {
- ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
- fn(out->data, a->data, a->shape.nelem);
+static ut_tensor* ew_unary(ut_tensor* a, const char* kern, void (*fn)(float*, const float*, int)) {
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, a->dev);
+ if (a->dev == UT_METAL && mc) {
+ ut_sync_gpu(a);
+ void* bufs[2] = {a->gpu_buf, out->gpu_buf};
+ int n = a->shape.nelem;
+ _mtl_dispatch(mc, kern, bufs, 2, &n, sizeof(int), n);
+ out->dirty_cpu = true;
+
+ } else {
+ ut_sync_cpu(a);
+ fn(out->data, a->data, a->shape.nelem);
+ }
return out;
}
-static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b,
+static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b, const char* kern,
void (*fn)(float*, const float*, const float*, int)) {
- ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
- fn(out->data, a->data, b->data, a->shape.nelem);
+ ut_dev dev = (a->dev == UT_METAL || b->dev == UT_METAL) ? UT_METAL : UT_CPU;
+ ut_to_device(a, dev);
+ ut_to_device(b, dev);
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, dev);
+ if (dev == UT_METAL && mc) {
+ void* bufs[3] = {a->gpu_buf, b->gpu_buf, out->gpu_buf};
+ int n = a->shape.nelem;
+ _mtl_dispatch(mc, kern, bufs, 3, &n, sizeof(int), n);
+ out->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(a);
+ ut_sync_cpu(b);
+ fn(out->data, a->data, b->data, a->shape.nelem);
+ }
return out;
}
-ut_tensor* ut_neg(ut_tensor* a) { return ew_unary(a, ew_neg); }
-ut_tensor* ut_exp(ut_tensor* a) { return ew_unary(a, ew_exp); }
-ut_tensor* ut_sigmoid(ut_tensor* a) { return ew_unary(a, ew_sigmoid); }
-ut_tensor* ut_tanh(ut_tensor* a) { return ew_unary(a, ew_tanh); }
-ut_tensor* ut_relu(ut_tensor* a) { return ew_unary(a, ew_relu); }
-ut_tensor* ut_add(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_add); }
-ut_tensor* ut_sub(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_sub); }
-ut_tensor* ut_mul(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_mul); }
-ut_tensor* ut_div(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_div); }
+ut_tensor* ut_neg(ut_tensor* a) { return ew_unary(a, "uneg", ew_neg); }
+ut_tensor* ut_exp(ut_tensor* a) { return ew_unary(a, "uexp", ew_exp); }
+ut_tensor* ut_sigmoid(ut_tensor* a) { return ew_unary(a, "usig", ew_sigmoid); }
+ut_tensor* ut_tanh(ut_tensor* a) { return ew_unary(a, "utanh", ew_tanh); }
+ut_tensor* ut_relu(ut_tensor* a) { return ew_unary(a, "urelu", ew_relu); }
+ut_tensor* ut_add(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, "badd", ew_add); }
+ut_tensor* ut_sub(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, "bsub", ew_sub); }
+ut_tensor* ut_mul(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, "bmul", ew_mul); }
+ut_tensor* ut_div(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, "bdiv", ew_div); }
ut_tensor* ut_scale(ut_tensor* a, float s) {
- ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
- for (int i = 0; i < a->shape.nelem; i++) out->data[i] = a->data[i] * s;
+ _mtl_ctx_t* mc = (_mtl_ctx_t*)ut_metal_ctx();
+ ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape, a->dev);
+ if (a->dev == UT_METAL && mc) {
+ ut_sync_gpu(a);
+ void* bufs[2] = {a->gpu_buf, out->gpu_buf};
+ struct {
+ float s;
+ int n;
+ } args = {s, a->shape.nelem};
+ _mtl_dispatch(mc, "bscale", bufs, 2, &args, (int)sizeof(args), a->shape.nelem);
+ out->dirty_cpu = true;
+ } else {
+ ut_sync_cpu(a);
+ for (int i = 0; i < a->shape.nelem; i++) out->data[i] = a->data[i] * s;
+ }
return out;
}