c_autograd_02

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

#define MAX_VARS 1000
#define MAX_OPS 1000
#define MAX_INPUTS 2
#define MAX_OUTPUTS 1

typedef struct {
    double val;
    double grad;
    int id;
    int is_active;
} NaiveVar;

typedef enum {
    OP_SUM,
    OP_PROD,
    OP_SOFTPLUS
} OpType;

typedef struct {
    OpType type;
    int n_inputs;
    int input_ids[MAX_INPUTS];
    int output_ids[MAX_OUTPUTS];
} MyOperation;

typedef struct {
    NaiveVar vars[MAX_VARS];
    int var_count;
    MyOperation ops[MAX_OPS];
    int op_count;
} NaiveTape;

NaiveTape create_tape() {
    NaiveTape tape = {
        .var_count = 0,
        .op_count = 0
    };
    return tape;
}

int create_var(NaiveTape* tape, double val)
{
    if (tape->var_count >= MAX_VARS) {
        fprintf(stderr, "Maximum variable limit reached\n");
        exit(1);
    }

    int id = tape->var_count++;
    tape->vars[id] = (NaiveVar){
        .val = val,
        .grad = 0.0,
        .id = id,
        .is_active = 1
    };
    return id;
}

void add_operation(NaiveTape* tape, MyOperation op)
{
    if (tape->op_count >= MAX_OPS) {
        fprintf(stderr, "Maximum operation limit reached\n");
        exit(1);
    }
    tape->ops[tape->op_count++] = op;
}

int tape_sum(NaiveTape* tape, int* var_ids, int n_vars)
{
    double sum_val = 0.0;
    for (int i = 0; i < n_vars; i++) {
        sum_val += tape->vars[var_ids[i]].val;
    }

    int result_id = create_var(tape, sum_val);

    MyOperation op = {
        .type = OP_SUM,
        .n_inputs = n_vars
    };

    memcpy(op.input_ids, var_ids, n_vars * sizeof(int));
    op.output_ids[0] = result_id;

    add_operation(tape, op);
    return result_id;
}

int tape_prod(NaiveTape* tape, int var1_id, int var2_id)
{
    double prod_val = tape->vars[var1_id].val * tape->vars[var2_id].val;
    int result_id = create_var(tape, prod_val);

    MyOperation op = {
        .type = OP_PROD,
        .n_inputs = 2,
        .input_ids = {var1_id, var2_id},
        .output_ids = {result_id}
    };

    add_operation(tape, op);
    return result_id;
}

int tape_softplus(NaiveTape* tape, int var_id)
{
    double softplus_val = log1p(exp(tape->vars[var_id].val));
    int result_id = create_var(tape, softplus_val);

    MyOperation op = {
        .type = OP_SOFTPLUS,
        .n_inputs = 1,
        .input_ids = {var_id},
        .output_ids = {result_id}
    };

    add_operation(tape, op);
    return result_id;
}

void tape_backward(NaiveTape* tape, int var_id)
{
    tape->vars[var_id].grad = 1.0;

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

        switch (op->type) {
            case OP_SUM: {
                double grad = tape->vars[op->output_ids[0]].grad;
                for (int j = 0; j < op->n_inputs; j++) {
                    tape->vars[op->input_ids[j]].grad += grad;
                }
                break;
            }
            case OP_PROD: {
                double grad = tape->vars[op->output_ids[0]].grad;
                NaiveVar* in1 = &tape->vars[op->input_ids[0]];
                NaiveVar* in2 = &tape->vars[op->input_ids[1]];
                in1->grad += in2->val * grad;
                in2->grad += in1->val * grad;
                break;
            }
            case OP_SOFTPLUS: {
                int input_id = op->input_ids[0];
                double exp_val = exp(-tape->vars[input_id].val);
                tape->vars[input_id].grad +=
                    tape->vars[op->output_ids[0]].grad /
                    (1.0 + exp_val);
                break;
            }
        }
    }
}

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

    start = clock();
    for (size_t iter = 0; iter < iterations; iter++) {
        NaiveTape tape = create_tape();

        int x_id = create_var(&tape, (float)rand() / RAND_MAX);
        int y_id = create_var(&tape, (float)rand() / RAND_MAX);

        int var_ids[] = {x_id, y_id};
        int sum_xy_id = tape_sum(&tape, var_ids, 2);
        int prod_sum_xy_id = tape_prod(&tape, sum_xy_id, sum_xy_id);
        int softplus_prod_id = tape_softplus(&tape, prod_sum_xy_id);

        tape_backward(&tape, softplus_prod_id);

        if (iter == iterations - 1) {
            printf("sum_xy->val: %f\n", tape.vars[sum_xy_id].val);
            printf("prod_sum_xy->val: %f\n", tape.vars[prod_sum_xy_id].val);
            printf("softplus_prod->val: %f\n", tape.vars[softplus_prod_id].val);
            printf("sum_xy->grad: %f\n", tape.vars[sum_xy_id].grad);
            printf("x->grad: %f\n", tape.vars[x_id].grad);
            printf("y->grad: %f\n", tape.vars[y_id].grad);
        }
    }

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

    return 0;
}