utensil

#include "utensil.h"

#include "cifar10_loader.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// MobileNetV1 (Howard et al.), scaled down for CIFAR-10 (32x32, 10 classes)
// instead of the ImageNet table (224x224, 1000 classes, channels up to 1024) --
// that table doesn't fit a 32x32 input and would be wildly over-parameterized
// for a 10-class problem. Keeps the architecture's defining idea (depthwise
// separable conv: a per-channel spatial conv + a 1x1 pointwise conv, each
// followed by BN+ReLU6, replacing one dense conv) and its downsampling
// pattern (stride-2 in a depthwise conv, never in the pointwise conv), just
// right-sized: stem -> 6 depthwise-separable blocks (channels 32-256, three
// stride-2 steps: 32x32 -> 16x16 -> 8x8 -> 4x4) -> GAP -> linear head. No
// data augmentation -- a comparison scaffold, not a from-scratch SOTA run.

// gradient tensors are borrowed from the optimizer's own grads[]; wired once
// in main() after ut_adam_alloc, then read directly during backward.
typedef struct {
  ut_tensor *dw_w, *dw_b, *dw_bn_w, *dw_bn_b, *pw_w, *pw_b, *pw_bn_w, *pw_bn_b;
} block_grads_t;

typedef struct {
  ut_dwconv2d dw;      // depthwise: one 3x3 filter per input channel
  ut_batchnorm2d dw_bn;
  ut_conv2d pw;        // pointwise: 1x1 conv, mixes channels
  ut_batchnorm2d pw_bn;
  block_grads_t g;
} block_t;

typedef struct {
  ut_dwconv2d_cache dw_c;
  ut_batchnorm2d_cache dw_bn_c;
  ut_conv2d_cache pw_c;
  ut_batchnorm2d_cache pw_bn_c;
  ut_tensor *dw_relu_in, *pw_relu_in;  // kept for backward's two ReLU6es
} block_cache_t;

typedef struct {
  ut_conv2d stem_conv;
  ut_batchnorm2d stem_bn;
  block_t blocks[6];
  ut_linear fc;
  ut_tensor *stem_cw, *stem_cb, *stem_bw, *stem_bb, *fc_w, *fc_b;  // borrowed, like block_grads_t
} mobilenet_t;

typedef struct {
  ut_conv2d_cache stem_c;
  ut_batchnorm2d_cache stem_bn_c;
  ut_tensor* stem_relu_in;
  block_cache_t block_c[6];
  ut_tensor* pooled;  // kept for the fc layer's backward
  int pool_h, pool_w;
} mobilenet_cache_t;

static block_t block_alloc(int in_c, int out_c, int stride, ut_dev dev) {
  block_t b = {0};
  b.dw = ut_dwconv2d_alloc(in_c, 3, 3, stride, 1, true, dev);
  b.dw_bn = ut_batchnorm2d_alloc(in_c, dev);
  b.pw = ut_conv2d_alloc(in_c, out_c, 1, 1, 1, 0, true, dev);
  b.pw_bn = ut_batchnorm2d_alloc(out_c, dev);
  return b;
}

static mobilenet_t mobilenet_alloc(ut_dev dev) {
  mobilenet_t m = {0};
  m.stem_conv = ut_conv2d_alloc(3, 32, 3, 3, 1, 1, true, dev);
  m.stem_bn = ut_batchnorm2d_alloc(32, dev);
  m.blocks[0] = block_alloc(32, 64, 1, dev);   // 32x32
  m.blocks[1] = block_alloc(64, 128, 2, dev);  // 32x32 -> 16x16
  m.blocks[2] = block_alloc(128, 128, 1, dev); // 16x16
  m.blocks[3] = block_alloc(128, 256, 2, dev); // 16x16 -> 8x8
  m.blocks[4] = block_alloc(256, 256, 1, dev); // 8x8
  m.blocks[5] = block_alloc(256, 256, 2, dev); // 8x8 -> 4x4
  m.fc = ut_linear_alloc(256, 10, true, dev);
  return m;
}

// appends this block's param tensors to out[] (for ut_adam_alloc) and returns
// the count -- always 8, no conditional (unlike a residual block's shortcut)
static int block_params(block_t* b, ut_tensor** out) {
  int n = 0;
  out[n++] = b->dw.weight, out[n++] = b->dw.bias;
  out[n++] = b->dw_bn.weight, out[n++] = b->dw_bn.bias;
  out[n++] = b->pw.weight, out[n++] = b->pw.bias;
  out[n++] = b->pw_bn.weight, out[n++] = b->pw_bn.bias;
  return n;
}

// mirrors block_params' exact order to wire each gradient pointer to its param
static void block_wire_grads(block_t* b, ut_tensor** g) {
  int n = 0;
  b->g.dw_w = g[n++], b->g.dw_b = g[n++];
  b->g.dw_bn_w = g[n++], b->g.dw_bn_b = g[n++];
  b->g.pw_w = g[n++], b->g.pw_b = g[n++];
  b->g.pw_bn_w = g[n++], b->g.pw_bn_b = g[n++];
}

static ut_tensor* block_forward(block_t* b, ut_tensor* x, bool training, block_cache_t* c) {
  ut_tensor* d1 = ut_dwconv2d_forward(&b->dw, x, c ? &c->dw_c : NULL);
  ut_tensor* d2 = ut_batchnorm2d_forward(&b->dw_bn, d1, training, c ? &c->dw_bn_c : NULL);
  ut_free(d1);
  ut_tensor* d3 = ut_relu6(d2);
  if (c) c->dw_relu_in = d2; else ut_free(d2);

  ut_tensor* p1 = ut_conv2d_forward(&b->pw, d3, c ? &c->pw_c : NULL);
  ut_free(d3);
  ut_tensor* p2 = ut_batchnorm2d_forward(&b->pw_bn, p1, training, c ? &c->pw_bn_c : NULL);
  ut_free(p1);
  ut_tensor* out = ut_relu6(p2);
  if (c) c->pw_relu_in = p2; else ut_free(p2);
  return out;
}

static ut_tensor* block_backward(block_t* b, block_cache_t* c, ut_tensor* dout) {
  ut_tensor* dp2 = ut_relu6_backward(dout, c->pw_relu_in);
  ut_tensor* dp1 = ut_batchnorm2d_backward(&b->pw_bn, &c->pw_bn_c, dp2, b->g.pw_bn_w, b->g.pw_bn_b);
  ut_free(dp2);
  ut_tensor* dd3 = ut_conv2d_backward(&b->pw, &c->pw_c, dp1, b->g.pw_w, b->g.pw_b);
  ut_free(dp1);
  ut_tensor* dd2 = ut_relu6_backward(dd3, c->dw_relu_in);
  ut_free(dd3);
  ut_tensor* dd1 = ut_batchnorm2d_backward(&b->dw_bn, &c->dw_bn_c, dd2, b->g.dw_bn_w, b->g.dw_bn_b);
  ut_free(dd2);
  ut_tensor* dx = ut_dwconv2d_backward(&b->dw, &c->dw_c, dd1, b->g.dw_w, b->g.dw_b);
  ut_free(dd1);
  return dx;
}

static void block_cache_free(block_cache_t* c) {
  ut_dwconv2d_cache_free(&c->dw_c);
  ut_batchnorm2d_cache_free(&c->dw_bn_c);
  ut_conv2d_cache_free(&c->pw_c);
  ut_batchnorm2d_cache_free(&c->pw_bn_c);
  ut_free_all(c->dw_relu_in, c->pw_relu_in);
}

static void block_free(block_t* b) {
  ut_dwconv2d_free(&b->dw);
  ut_batchnorm2d_free(&b->dw_bn);
  ut_conv2d_free(&b->pw);
  ut_batchnorm2d_free(&b->pw_bn);
}

static ut_tensor* mobilenet_forward(mobilenet_t* m, ut_tensor* x, bool training,
                                    mobilenet_cache_t* c) {
  ut_tensor* s = ut_conv2d_forward(&m->stem_conv, x, c ? &c->stem_c : NULL);
  ut_tensor* sbn = ut_batchnorm2d_forward(&m->stem_bn, s, training, c ? &c->stem_bn_c : NULL);
  ut_free(s);
  ut_tensor* h = ut_relu6(sbn);
  if (c) c->stem_relu_in = sbn; else ut_free(sbn);

  for (int i = 0; i < 6; i++) {
    ut_tensor* next = block_forward(&m->blocks[i], h, training, c ? &c->block_c[i] : NULL);
    ut_free(h);
    h = next;
  }

  int ph = h->shape.shape[2], pw = h->shape.shape[3];
  ut_tensor* pooled = ut_global_avgpool2d(h);
  ut_free(h);
  if (c) c->pooled = pooled, c->pool_h = ph, c->pool_w = pw;
  ut_tensor* logits = ut_linear_forward(&m->fc, pooled);
  if (!c) ut_free(pooled);
  return logits;
}

static ut_tensor* mobilenet_backward(mobilenet_t* m, mobilenet_cache_t* c, ut_tensor* dlogits) {
  ut_tensor* dpooled = ut_linear_backward(&m->fc, c->pooled, dlogits, m->fc_w, m->fc_b);
  ut_tensor* dh = ut_global_avgpool2d_backward(dpooled, c->pool_h, c->pool_w);
  ut_free(dpooled);

  for (int i = 5; i >= 0; i--) {
    ut_tensor* dprev = block_backward(&m->blocks[i], &c->block_c[i], dh);
    ut_free(dh);
    dh = dprev;
  }

  ut_tensor* dsbn = ut_relu6_backward(dh, c->stem_relu_in);
  ut_free(dh);
  ut_tensor* ds = ut_batchnorm2d_backward(&m->stem_bn, &c->stem_bn_c, dsbn, m->stem_bw, m->stem_bb);
  ut_free(dsbn);
  ut_tensor* dx = ut_conv2d_backward(&m->stem_conv, &c->stem_c, ds, m->stem_cw, m->stem_cb);
  ut_free(ds);
  return dx;
}

static void mobilenet_cache_free(mobilenet_cache_t* c) {
  ut_conv2d_cache_free(&c->stem_c);
  ut_batchnorm2d_cache_free(&c->stem_bn_c);
  ut_free(c->stem_relu_in);
  for (int i = 0; i < 6; i++) block_cache_free(&c->block_c[i]);
  ut_free(c->pooled);
}

static void mobilenet_free(mobilenet_t* m) {
  ut_conv2d_free(&m->stem_conv);
  ut_batchnorm2d_free(&m->stem_bn);
  for (int i = 0; i < 6; i++) block_free(&m->blocks[i]);
  ut_linear_free(&m->fc);
}

int main(void) {
  srand(42);

  ut_dev dev = UT_METAL;  // flip to UT_CPU to compare against a (CPU-only) PyTorch script

  printf("Loading CIFAR-10…\n");
  const char* train_files[5] = {"cifar/data_batch_1.bin", "cifar/data_batch_2.bin",
                                "cifar/data_batch_3.bin", "cifar/data_batch_4.bin",
                                "cifar/data_batch_5.bin"};
  const char* test_files[1] = {"cifar/test_batch.bin"};
  cifar10_t train = cifar10_load(train_files, 5);
  cifar10_t test = cifar10_load(test_files, 1);
  printf("train: %d  test: %d\n\n", train.n, test.n);

  mobilenet_t m = mobilenet_alloc(dev);

  ut_tensor* params[64];
  int np = 0;
  params[np++] = m.stem_conv.weight, params[np++] = m.stem_conv.bias;
  params[np++] = m.stem_bn.weight, params[np++] = m.stem_bn.bias;
  int block_off[6];
  for (int i = 0; i < 6; i++) {
    block_off[i] = np;
    np += block_params(&m.blocks[i], params + np);
  }
  int fc_off = np;
  params[np++] = m.fc.weight, params[np++] = m.fc.bias;

  ut_adam opt = ut_adam_alloc(params, np, 1e-3f, 0.9f, 0.999f, 1e-8f, 1e-4f);

  // wire each layer's gradient pointers to the optimizer's own grads[], in the
  // same order the params[] above was built
  m.stem_cw = opt.grads[0], m.stem_cb = opt.grads[1];
  m.stem_bw = opt.grads[2], m.stem_bb = opt.grads[3];
  for (int i = 0; i < 6; i++) block_wire_grads(&m.blocks[i], opt.grads + block_off[i]);
  m.fc_w = opt.grads[fc_off], m.fc_b = opt.grads[fc_off + 1];

  int B = 64, epochs = 20, batches = train.n / B;
  ut_tensor* grad_logits = ut_alloc(2, (int[]){B, 10}, dev);
  int* idx = malloc((size_t)train.n * sizeof(int));
  for (int i = 0; i < train.n; i++) idx[i] = i;

  for (int ep = 0; ep < epochs; ep++) {
    cifar10_shuffle(idx, train.n);
    clock_t t0 = clock();
    float loss_sum = 0;
    int correct = 0;

    for (int bi = 0; bi < batches; bi++) {
      float bx[B * CIFAR10_PIX];
      int bl[B];
      for (int i = 0; i < B; i++) {
        int ii = idx[bi * B + i];
        memcpy(bx + i * CIFAR10_PIX, train.imgs + ii * CIFAR10_PIX, CIFAR10_PIX * sizeof(float));
        bl[i] = train.labels[ii];
      }
      ut_tensor* x = ut_from_data(4, (int[]){B, 3, 32, 32}, bx, dev);

      mobilenet_cache_t c;
      ut_tensor* logits = mobilenet_forward(&m, x, true, &c);
      float loss = ut_cross_entropy(logits, bl, grad_logits);
      ut_sync_cpu(logits);
      for (int i = 0; i < B; i++)
        if (cifar10_argmax(logits, i) == bl[i]) correct++;

      ut_tensor* dx = mobilenet_backward(&m, &c, grad_logits);
      ut_adam_step(&opt, 5.0f);
      loss_sum += loss;

      mobilenet_cache_free(&c);
      ut_free_all(x, logits, dx);
    }

    float secs = (float)(clock() - t0) / (float)CLOCKS_PER_SEC;

    int test_ok = 0;
    for (int i = 0; i < test.n; i += B) {
      int nb = (i + B <= test.n) ? B : (test.n - i);
      ut_tensor* tx = ut_from_data(4, (int[]){nb, 3, 32, 32}, test.imgs + i * CIFAR10_PIX, dev);
      ut_tensor* tl = mobilenet_forward(&m, tx, false, NULL);
      ut_sync_cpu(tl);
      for (int j = 0; j < nb; j++)
        if (cifar10_argmax(tl, j) == test.labels[i + j]) test_ok++;
      ut_free_all(tx, tl);
    }

    printf("epoch %2d  loss %7.4f  train %5.1f%%  test %5.1f%%  %6.1fs\n", ep + 1,
           loss_sum / (float)batches, 100.f * (float)correct / (float)(batches * B),
           100.f * (float)test_ok / (float)test.n, secs);
  }

  free(idx);
  ut_free(grad_logits);
  ut_adam_free(&opt);
  mobilenet_free(&m);
  cifar10_free(&train);
  cifar10_free(&test);
  return 0;
}