c_autograd_01

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

typedef struct {
    double val;
    double grad;
} NaiveVar;

typedef enum {
    OP_SUM,
    OP_PROD,
    OP_SOFTMAX,
    OP_SOFTPLUS
} OpType;

typedef struct {
    OpType type;
    int n_inputs;
    NaiveVar *inputs[8];
    NaiveVar *outputs[8];
} MyOperation;

typedef struct {
    int n_ops;
    int max_ops;
    MyOperation *ops;  // Changed from MyOperation** to MyOperation*
    int ops_capacity;
} NaiveTape;

// Use a memory pool for NaiveVar allocations
#define POOL_SIZE 1000
static NaiveVar var_pool[POOL_SIZE];
static int pool_index = 0;

inline NaiveVar* naive_var(double val) {
    NaiveVar* var = &var_pool[pool_index++];
    var->val = val;
    var->grad = 0.0;
    return var;
}

NaiveTape* naive_tape_create(int max_ops) {
    NaiveTape *tape = (NaiveTape*) malloc(sizeof(NaiveTape));
    tape->n_ops = 0;
    tape->max_ops = max_ops;
    tape->ops = (MyOperation*) calloc(max_ops, sizeof(MyOperation));
    tape->ops_capacity = max_ops;
    return tape;
}

inline void naive_tape_add_op(NaiveTape *tape, MyOperation *op) {
    if (tape->n_ops >= tape->max_ops) {
        tape->max_ops *= 2;
        tape->ops = (MyOperation*) realloc(tape->ops, tape->max_ops * sizeof(MyOperation));
    }
    memcpy(&tape->ops[tape->n_ops++], op, sizeof(MyOperation));
}

inline NaiveVar* naive_tape_sum(NaiveTape *tape, NaiveVar **vars, int n_vars) {
    double sum_val = 0.0;
    #pragma unroll
    for (int i = 0; i < n_vars; i++) {
        sum_val += vars[i]->val;
    }

    NaiveVar *res = naive_var(sum_val);
    MyOperation op;
    op.type = OP_SUM;
    op.n_inputs = n_vars;

    memcpy(op.inputs, vars, n_vars * sizeof(NaiveVar*));
    op.outputs[0] = res;
    naive_tape_add_op(tape, &op);
    return res;
}

inline NaiveVar* naive_tape_prod(NaiveTape *tape, NaiveVar *var1, NaiveVar *var2) {
    NaiveVar *res = naive_var(var1->val * var2->val);
    MyOperation op;
    op.type = OP_PROD;
    op.n_inputs = 2;
    op.inputs[0] = var1;
    op.inputs[1] = var2;
    op.outputs[0] = res;
    naive_tape_add_op(tape, &op);
    return res;
}

inline NaiveVar* naive_tape_softplus(NaiveTape *tape, NaiveVar *var) {
    NaiveVar *res = naive_var(log1p(exp(var->val)));
    MyOperation op;
    op.type = OP_SOFTPLUS;
    op.n_inputs = 1;
    op.inputs[0] = var;
    op.outputs[0] = res;
    naive_tape_add_op(tape, &op);
    return res;
}

void naive_tape_backward(NaiveTape *tape, NaiveVar *var) {
    var->grad = 1.0;

    MyOperation *op;
    for (int i = tape->n_ops - 1; i >= 0; i--) {
        op = &tape->ops[i];

        switch (op->type) {
            case OP_SUM: {
                double grad = op->outputs[0]->grad;
                #pragma unroll
                for (int j = 0; j < op->n_inputs; j++) {
                    op->inputs[j]->grad += grad;
                }
                break;
            }
            case OP_PROD: {
                double grad = op->outputs[0]->grad;
                NaiveVar *in1 = op->inputs[0];
                NaiveVar *in2 = op->inputs[1];
                in1->grad += in2->val * grad;
                in2->grad += in1->val * grad;
                break;
            }
            case OP_SOFTPLUS:
                op->inputs[0]->grad += op->outputs[0]->grad / (1.0 + exp(-op->inputs[0]->val));
                break;
        }
    }
}

void naive_tape_free(NaiveTape *tape) {
    free(tape->ops);
    free(tape);
}

int main() {
    clock_t start, end;
    double elapsed_time;
    size_t iterations = 1000000;

    start = clock();
    for (size_t i = 0; i < iterations; i++) {
        pool_index = 0;  // Reset pool index for each iteration

        NaiveVar *x = naive_var(1.0);
        NaiveVar *y = naive_var(2.0);
        NaiveTape *tape = naive_tape_create(10);

        NaiveVar *vars[] = {x, y};
        NaiveVar *sum_xy = naive_tape_sum(tape, vars, 2);
        NaiveVar *prod_sum_xy = naive_tape_prod(tape, sum_xy, sum_xy);
        NaiveVar *softplus_prod = naive_tape_softplus(tape, prod_sum_xy);

        naive_tape_backward(tape, softplus_prod);

        if (i == iterations - 1) {
            printf("sum_xy->val: %f\n", sum_xy->val);
            printf("prod_sum_xy->val: %f\n", prod_sum_xy->val);
            printf("softplus_prod->val: %f\n", softplus_prod->val);
            printf("sum_xy->grad: %f\n", sum_xy->grad);
            printf("x->grad: %f\n", x->grad);
            printf("y->grad: %f\n", y->grad);
        }

        naive_tape_free(tape);
    }
    end = clock();
    elapsed_time = ((double) (end - start)) * 1000 / CLOCKS_PER_SEC;
    printf("\nElapsed time: %f ms\n", elapsed_time);

    return 0;
}