main.c

#include <assert.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdbool.h>

#define safetensors_file_size 548105171
#define parameters_raw_size 548090880
#define json_raw_size 14283
#define data_token_count 338024
#define enc_file_size 722883
#define d_vocab 50257
#define d_model 768
#define d_seq 1024
#define d_k 64
#define rsqrt_k 0.125f


#define MATRIX_INDEX(base, i,  d_model) ((base) + ((i) * (d_model) ))


struct token {
    uint32_t offset;
    uint32_t size;
}__attribute__((packed));

//__attribute__((packed))

struct decoder {
    struct token tokens[d_vocab];
    char raw[320827];
}__attribute__((packed));


struct data{
    uint16_t tokens[data_token_count];
};

struct parameters{
    struct {
        float* weight;
    } wte;

    struct {
        float* weight;
    } wpe;
    struct {
        struct {
            float* weight;
            float* bias;
        } ln_1;
        struct {
            struct {
                float* weight;
                float* bias;
            } c_attn;
            struct {
                float* weight;
                float* bias;
            } c_proj
        } attn;
        struct {
            float* weight;
            float* bias;
        } ln_2;
        struct {
            struct {
                float* weight;
                float* bias;
            } c_fc;
            struct {
                float* weight;
                float* bias;
            } c_proj;
        } mlp;
    } h[12];
    struct {
        float* weight;
        float* bias;
    } ln_f;

};

struct activations {

    struct {
        float out[d_seq][d_model];
    } embedding;

    struct { 
        struct {
            float out[d_seq][d_model]
        } ln_1;
        struct {
            struct {
                float out[d_seq][3 * d_model];
            } c_attn;
            struct {
                float out[12][d_seq][d_seq];
            } attn;
            struct {
                float out[12][d_seq][d_seq];
            } s;
            struct {
                float out[d_seq][d_seq];
            } z;
        }attn;
    }h[12];

    
};


static size_t tensor_offset(char* json_raw, char* tensor_name, size_t expected_size){
    char* name = strstr(json_raw, tensor_name);
    assert(name);
    //printf("%.*s\n",20, name );
    char* data_offsets = strstr(name, "data_offsets");
    //printf("%.*s\n",20, data_offsets );
    char *start = data_offsets + 15;
    char *end = strstr(start, ",");
    char temp[32];

    size_t count = end - start;
    memcpy(temp, start, count);
    temp[count] = 0;
    size_t start_offset = (size_t)atoi(temp);

    start = end + 1;
    end = strstr(start, "]");
    count = end - start;
    memcpy(temp, start, count);
    temp[count] = 0;
    size_t end_offset = (size_t)atoi(temp);
    assert(end_offset - start_offset == expected_size * 4);

    // divide by 4 because what we are seeking are floats 
    return start_offset/4;

}

struct fc { 
    const float* in;
    const float* weight;
    const float* bias;
    float* out;
    size_t in_count;
    size_t out_count;
    size_t sample_count;
};

static void fc(const struct fc* fc){

    for (size_t i = 0; i < fc->sample_count; i++){
        const float* in = fc->in + i * fc->in_count;

        float* weight = fc->weight;
        float* weight_end = weight + fc->in_count * fc->out_count;
        float* bias = fc->bias;

        float* out = fc->out + i * fc->out_count;
        float* out_end = out + fc->out_count;
        float* out_reset = out;

        memcpy(out, bias, fc->out_count*sizeof(float));

        while (true) {
            *out += *weight * *in;
            weight++;
            out++;
            if (out == out_end){
                out = out_reset;
                in++;
                if (weight == weight_end){
                    break;
                }
            }
        }

    }

}

int main(void){

    #define align(x) (((size_t)x + 255) & ~(size_t)0xff) 

    size_t offset = 0;
    size_t decoder_offset = offset;
    offset += align(sizeof(struct decoder));
    size_t data_offset = offset;
    offset += align(sizeof(struct data));
    size_t json_raw_offset = offset;
    offset += align(json_raw_size);
    size_t parameters_raw_offset = offset;
    offset += align(parameters_raw_size);
    size_t parameters_offset = offset;
    offset += align(sizeof(struct parameters));
    size_t activations_offset = offset;
    offset += align(sizeof(struct activations));
    #undef align


    char* raw_mem = malloc(offset);
    struct decoder* decoder = (struct decoder*)(raw_mem + decoder_offset);
    struct data* data = (struct data*)(raw_mem + data_offset);
    char* json_raw = (char*)(raw_mem + json_raw_offset);
    float* parameters_raw = (float*)(raw_mem + parameters_raw_offset);
    struct parameters* parameters = (struct parameters*)(raw_mem + parameters_offset);
    struct activations* activations = (struct activations*)(raw_mem + activations_offset);
{
    FILE* f = fopen("ref/enc", "r");
    assert(f);

    unsigned long read = fread(decoder, 1, sizeof(struct decoder), f);
    printf("read is :%lu\n", read);
    printf("struct decoder is :%lu\n", sizeof(struct decoder));
    assert(read == sizeof(struct decoder));

    //for (int i=0; i<d_vocab; i++){
    //    uint32_t o = decoder->tokens[i].offset;
    //    uint32_t s = decoder->tokens[i].size;
    //    printf("offset : %u, size : %u, %.*s\n", o,s,s, decoder->raw + o);
    //}
    fclose(f);
}

{
    FILE* f = fopen("ref/data", "r");
    assert(f);

    unsigned long read = fread(data, 1, data_token_count * sizeof(uint16_t), f);
    assert(read == data_token_count* sizeof(uint16_t));

    //for (int i=0; i<data_token_count; i++){
    //    uint16_t token = data->tokens[i];
    //    uint32_t offset = decoder->tokens[token].offset;
    //    uint32_t size = decoder->tokens[token].size;
    //    printf("%.*s",size, decoder->raw + offset);
    //}

    fclose(f);
}

{
    FILE *f = fopen("model.safetensors", "r");
    assert(f);

    uint64_t json_size;
    unsigned long read = fread(&json_size, 1, 8, f);
    assert(read == 8);

    read = fread(json_raw, 1, json_raw_size, f);
    assert(read == json_raw_size);

    read = fread(parameters_raw, 1, parameters_raw_size, f);
    assert(read == parameters_raw_size);

    parameters->wpe.weight = parameters_raw + tensor_offset(json_raw, "wpe.weight", d_model * d_seq );
    printf("wpe weight is : %d\n", parameters->wpe.weight);
    parameters->wte.weight = parameters_raw + tensor_offset(json_raw, "wte.weight", d_model * d_vocab);
    printf("wte weight is : %d\n", parameters->wte.weight);

    for (int i=0; i<12; i++){
        char temp[64];

        snprintf(temp, sizeof(temp), "h.%u.ln_1.weight", i);
        parameters->h[i].ln_1.weight = parameters_raw + tensor_offset(json_raw, temp, d_model);

        snprintf(temp, sizeof(temp), "h.%u.ln_1.bias", i);
        parameters->h[i].ln_1.bias = parameters_raw + tensor_offset(json_raw, temp, d_model);

        snprintf(temp, sizeof(temp), "h.%u.attn.c_attn.weight", i);
        parameters->h[i].attn.c_attn.weight = parameters_raw + tensor_offset(json_raw, temp, d_model * d_model*3);

        snprintf(temp, sizeof(temp), "h.%u.attn.c_attn.bias", i);
        parameters->h[i].attn.c_attn.bias = parameters_raw + tensor_offset(json_raw, temp, d_model*3);

        snprintf(temp, sizeof(temp), "h.%u.attn.c_proj.weight", i);
        parameters->h[i].attn.c_proj.weight = parameters_raw + tensor_offset(json_raw, temp, d_model * d_model);

        snprintf(temp, sizeof(temp), "h.%u.attn.c_proj.bias", i);
        parameters->h[i].attn.c_proj.bias = parameters_raw + tensor_offset(json_raw, temp, d_model);

        snprintf(temp, sizeof(temp), "h.%u.ln_2.weight", i);
        parameters->h[i].ln_2.weight = parameters_raw + tensor_offset(json_raw, temp, d_model);

        snprintf(temp, sizeof(temp), "h.%u.ln_2.bias", i);
        parameters->h[i].ln_2.bias = parameters_raw + tensor_offset(json_raw, temp, d_model);

        snprintf(temp, sizeof(temp), "h.%u.mlp.c_fc.weight", i);
        parameters->h[i].mlp.c_fc.weight = parameters_raw + tensor_offset(json_raw, temp, d_model * 4 * d_model);

        snprintf(temp, sizeof(temp), "h.%u.mlp.c_fc.bias", i);
        parameters->h[i].mlp.c_fc.bias = parameters_raw + tensor_offset(json_raw, temp,  4 * d_model);

        snprintf(temp, sizeof(temp), "h.%u.mlp.c_proj.weight", i);
        parameters->h[i].mlp.c_proj.weight = parameters_raw + tensor_offset(json_raw, temp, d_model * 4 * d_model);

        snprintf(temp, sizeof(temp), "h.%u.mlp.c_proj.bias", i);
        parameters->h[i].mlp.c_proj.bias = parameters_raw + tensor_offset(json_raw, temp,   d_model);
    }

    parameters->ln_f.weight = parameters_raw + tensor_offset(json_raw, "ln_f.weight",  d_model );
    parameters->ln_f.bias = parameters_raw + tensor_offset(json_raw, "ln_f.bias",  d_model);

    fclose(f);
}

size_t input_size = 64;
uint16_t* input = data->tokens;

for (size_t i = 0; i<input_size; i++){
    uint16_t token = input[i];


    float* wte = parameters->wte.weight + token * d_model;
    //float* wte = MATRIX_INDEX(parameters->wte.weight, token, d_model);
    float* wpe = parameters->wpe.weight + i * d_model;

    float* out = (float*)activations->embedding.out + i * d_model;
    float* out_end = out + d_model;

    // For each input we have a vector embedding
    // For each vector embedding we need to iterate over its values to make the computation ( ie wte + wpe)
    for (; out != out_end; wte++, wpe++, out++) {
        *out = *wte + *wpe;

    }
}
//double total = 0.0;
//for (size_t i = 0; i<input_size * d_model; i++){
//    total += (double)((float*)activations->embedding.out)[i];
//
//}
int layer_i = 0;

for(size_t i =0; i<input_size; i++){
    float* in = (float*)activations->embedding.out + i * d_model;
    float* in_end = in + d_model; 
    float* in_reset = in;

    float mean = 0.0f;
    for (; in != in_end; in++){
        mean += *in;
    }

    mean /= d_model;

    float total_diff_sq = 0.0f;
    for (in = in_reset; in != in_end; in++){
        float diff = *in - mean;
        total_diff_sq += diff * diff;
    }
    float r_stddev = 1.0f / sqrtf(total_diff_sq/d_model);

    float *out = (float*)activations->h[layer_i].ln_1.out + i*d_model;
    float* weight = parameters->h[layer_i].ln_1.weight;
    float* bias = parameters->h[layer_i].ln_1.bias;

    for (in = in_reset; in != in_end; in++, weight++, bias++, out++){
        float in_norm = (*in - mean) * r_stddev;
        *out = in_norm * (*weight) + (*bias);
    }

}



// Filling out c_att (with Q,K and V)
for (size_t i = 0; i < input_size; i++){
    float* in = (float*)activations->h[layer_i].ln_1.out + i * d_model;

    float* weight = parameters->h[layer_i].attn.c_attn.weight;
    float* weight_end = weight + d_model * 3 * d_model;
    float* bias = parameters->h[layer_i].attn.c_attn.bias;

    float* out = (float*)activations->h[layer_i].attn.c_attn.out + i * 3 * d_model;
    float* out_end = out + 3*d_model;
    float* out_reset = out;

    memcpy(out, bias, 3*d_model*sizeof(float));

    while (true) {
        *out += *weight * *in;
        weight++;
        out++;
        if (out == out_end){
            out = out_reset;
            in++;
            if (weight == weight_end){
                break;
            }
        }
    }

}

// the crux of it all, computing self attention 

memset(activations->h[layer_i].attn.z.out,0, sizeof(activations->h[layer_i].attn.z.out));

for (size_t head_i = 0; head_i < 12; head_i++){
    for (size_t q_i = 0; q_i < input_size; q_i++){
        float softmax_max = -INFINITY;

        for (size_t k_i=0; k_i<=q_i; k_i++){
            float* k = (float*)activations->h[layer_i].attn.c_attn.out + k_i * 3 * d_model + d_model + head_i * d_k;

            float* q = (float*)activations->h[layer_i].attn.c_attn.out + q_i * 3 * d_model + head_i * d_k;
            float* q_end = q + d_k;

            size_t k_stride = 3*d_model - d_k;

            float dot = 0.0f;
            for (; q != q_end; q++, k++){
                dot += *q * *k;
            }

            dot *= rsqrt_k;

            activations->h[layer_i].attn.attn.out[head_i][q_i][k_i] = dot;

            if (dot > softmax_max){
                softmax_max = dot;
            }

        }

        float softmax_sum = 0.0f;

        for (size_t k_i=0; k_i <= q_i; k_i++){
            float e = activations->h[layer_i].attn.attn.out[head_i][q_i][k_i];  
            float softmax_exp_i = expf(e - softmax_max);
            activations->h[layer_i].attn.s.out[head_i][q_i][k_i] = softmax_exp_i;

            softmax_sum += softmax_exp_i;
        }

        float r_softmax_sum = 1.0f / softmax_sum;

        for (size_t k_i=0; k_i <= q_i; k_i++){
            activations->h[layer_i].attn.s.out[head_i][q_i][k_i] *= r_softmax_sum;

        }

        for (size_t v_i=0; v_i <= q_i; v_i++){
            float* v = (float*)activations->h[layer_i].attn.c_attn.out + v_i * 3 * d_model + 2 * d_model + head_i * d_k;
            float *v_end = v + d_k;

            float* z = (float*)activations->h[layer_i].attn.z.out + q_i *  d_model + head_i * d_k;

            float factor = activations->h[layer_i].attn.s.out[head_i][q_i][v_i];

            for (; v != v_end; v++, z++){
                *z += *v * factor;
            }
        }
    }
}
double total = 0.0;
for (size_t i = 0; i<input_size  * d_model; i++){
    printf("%d value is : %f\n", i, (double)((float*)activations->h[layer_i].attn.z.out)[i]);
    total += (double)((float*)activations->h[layer_i].attn.z.out)[i];
}
printf("%f\n", total);
}