llvmruntime.cpp

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

/** Memory Layout:
 * base: int* pointing to start of allocated block
 * base[0] = ndims
 * base[1..ndims] = dims
 * data = (double*)&base[1+ndims]
 *
 * arr = data pointer
 * To get base from arr: base = (int*)arr - (1+ndims)
 */

// Helper to get ndims
static int get_ndims(double* arr) {
    int *base = (int*)arr;
    // We don't know ndims yet. We must guess ndims from existing info.
    // Wait, we need ndims to find base.
    // However, we know that base is just before arr by (1+ndims) ints.
    // This is a chicken-egg problem again unless we store ndims twice or have a known marker.
    //
    // Let's store ndims and dims at a known offset:
    // Instead of just storing them and returning arr, we must do a trick:
    //
    // We'll store a small negative offset: we know arr points to data,
    // base = arr - (1+ndims)
    // Without ndims, we can't find base directly.
    //
    // To solve this: We will store ndims at arr[-1].
    // Modification of layout:
    // Layout:
    // memory: [ndims(int)] [dims(ndims)] [data(double...)]
    // arr = (double*)&base[1+ndims]
    //
    // If arr = &base[1+ndims], then arr- (1+ndims) = base.
    // We still need ndims to do that.
    //
    // Let's store ndims at a fixed known offset right before data starts:
    // That is exactly base[0].
    // If we move arr by a known offset backwards until we find ndims:
    //
    // Actually, we know there's always at least one dimension. We can try a loop:
    // But that would be hacky.
    //
    // Let's store ndims right before data as well:
    // We'll store ndims twice: once at base[0], and once right before data.
    //
    // Revised layout:
    // [ndims(int)] [dims(ndims)] [ndims(int)] [data]
    // This way, arr points exactly at data, and the integer right before data is ndims.
    //
    // memory:
    // base[0] = ndims
    // base[1..ndims] = dims
    // base[1+ndims] = ndims again
    // data = (double*)&base[2+ndims]
    //
    // arr = &data[0] = (double*)&base[2+ndims]
    // To get ndims: ndims = base[1+ndims (the second ndims)]
    //
    // Let's adopt this final layout.
    //
    // Final Layout:
    // base[0] = ndims
    // base[1..ndims] = dims
    // base[1+ndims] = ndims (repeated)
    // data = (double*)&base[2+ndims]
    //
    // arr = data
    // arr - (2+ndims) = base
    //
    // So to get ndims:
    // int* ptr = (int*)arr;
    // ptr -= 1; // now ptr points to base[1+ndims]
    // ndims = ptr[0];

    int* ptr = (int*)arr;
    ptr -= 1;
    int ndims = ptr[0];
    return ndims;
}

static int* get_dims(double* arr, int ndims) {
    // base = (int*)arr - (2+ndims)
    int* base = (int*)arr - (2+ndims);
    return &base[1]; // dims start at base[1]
}

static int get_total_size(int ndims, int* dims) {
    int total=1;
    for(int i=0;i<ndims;i++) total*=dims[i];
    return total;
}

// runtime_alloc: allocate array with given ndims and dims
double* runtime_alloc(int ndims, int* dims) {
    int total = get_total_size(ndims, dims);
    // Size: 1 for ndims, ndims for dims, 1 for repeated ndims, plus doubles
    size_t sz = sizeof(int)*(2+ndims) + sizeof(double)*total;
    int* base = (int*)malloc(sz);
    base[0]=ndims;
    for(int i=0;i<ndims;i++){
        base[1+i]=dims[i];
    }
    base[1+ndims] = ndims; // store ndims again
    double* arr = (double*)&base[2+ndims];
    return arr;
}

// Free array
static void free_array(double* arr) {
    int ndims = get_ndims(arr);
    int* base = (int*)arr - (2+ndims);
    free(base);
}

// Fill with default value
void runtime_fill(double* arr, int totalSize, double val) {
    for(int i=0;i<totalSize;i++) {
        arr[i]=val;
    }
}

void runtime_range_fill(double* arr, int totalSize, double start, double end) {
    if(totalSize <=1) {
        if(totalSize==1) arr[0]=start;
        return;
    }
    double step=(end - start)/(totalSize-1);
    for(int i=0;i<totalSize;i++){
        arr[i] = start + i*step;
    }
}

void runtime_store_element(double* arr, int idx, double val) {
    arr[idx]=val;
}

double runtime_sum(double* arr, int totalSize) {
    double s=0.0;
    for(int i=0;i<totalSize;i++) s+=arr[i];
    return s;
}

double runtime_max(double* arr, int totalSize) {
    if(totalSize<=0)return 0.0;
    double m=arr[0];
    for(int i=1;i<totalSize;i++){
        if(arr[i]>m) m=arr[i];
    }
    return m;
}

double runtime_average(double* arr, int totalSize) {
    if(totalSize<=0)return 0.0;
    double s=runtime_sum(arr,totalSize);
    return s/(double)totalSize;
}

// map with known func "square"
double* runtime_map(double* arr, int totalSize, const char* funcName) {
    int ndims = get_ndims(arr);
    int* dims = get_dims(arr, ndims);
    int total = totalSize; // must match get_total_size(ndims,dims)

    double* newArr = runtime_alloc(ndims,dims);
    if(strcmp(funcName,"square")==0) {
        for(int i=0;i<total;i++){
            newArr[i]=arr[i]*arr[i];
        }
    } else {
        // no known func, just copy
        for(int i=0;i<total;i++){
            newArr[i]=arr[i];
        }
    }
    // free old array
    free_array(arr);
    return newArr;
}

// reduce with known func "add"
double runtime_reduce(double* arr, int totalSize, const char* funcName) {
    if(strcmp(funcName,"add")==0){
        double s=0.0;
        for(int i=0;i<totalSize;i++){
            s+=arr[i];
        }
        return s;
    }
    // default: sum
    return runtime_sum(arr,totalSize);
}

// flatten: convert N-D to 1-D
// Just return a new 1D array with the same data
double* runtime_flatten(double* arr, int ndims, int* dims) {
    int total = get_total_size(ndims, dims);
    // Create 1D dims
    int new_dims[1] = {total};
    double* newArr = runtime_alloc(1,new_dims);
    for(int i=0;i<total;i++){
        newArr[i]=arr[i];
    }
    free_array(arr);
    return newArr;
}

double* runtime_copy(double* arr) {
    int ndims = get_ndims(arr);
    int* dims = get_dims(arr, ndims);
    int total = get_total_size(ndims,dims);
    double* newArr = runtime_alloc(ndims,dims);
    for(int i=0;i<total;i++){
        newArr[i]=arr[i];
    }
    return newArr; // no free of old arr, user decides
}

double* runtime_clone_structure(int ndims, int* dims) {
    int total = get_total_size(ndims,dims);
    double* arr = runtime_alloc(ndims,dims);
    runtime_fill(arr,total,0.0);
    return arr;
}

double* runtime_reshape(double* arr, int ndims, int* dims) {
    // reshape without changing data
    // ensure old total == new total
    int old_ndims = get_ndims(arr);
    int* old_dims = get_dims(arr,old_ndims);
    int old_total = get_total_size(old_ndims,old_dims);
    int new_total = get_total_size(ndims,dims);
    if(old_total!=new_total){
        fprintf(stderr,"reshape: total size mismatch\n");
        exit(1);
    }
    double* newArr = runtime_alloc(ndims,dims);
    for(int i=0;i<old_total;i++){
        newArr[i]=arr[i];
    }
    free_array(arr);
    return newArr;
}

char* runtime_serialize(double* arr) {
    int ndims = get_ndims(arr);
    int* dims = get_dims(arr,ndims);
    int total = get_total_size(ndims,dims);

    // Serialize as: ndims dims... data...
    // format: ndims dim1 dim2 ... dimN data1 data2 ... dataM
    // All space-separated
    size_t maxLen = 20*(total+ndims+1);
    char* buf = (char*)malloc(maxLen);
    int pos = 0;
    pos += snprintf(buf+pos, maxLen-pos, "%d ", ndims);
    for(int i=0;i<ndims;i++){
        pos += snprintf(buf+pos, maxLen-pos, "%d ", dims[i]);
    }
    for(int i=0;i<total;i++){
        pos += snprintf(buf+pos, maxLen-pos, "%.17g ", arr[i]);
    }
    return buf;
}

double* runtime_deserialize(char* str) {
    int ndims;
    char* p=str;
    if(sscanf(p,"%d",&ndims)!=1){
        fprintf(stderr,"deserialize: invalid input\n");
        exit(1);
    }
    // skip ndims
    while(*p && *p!=' ') p++;
    if(*p==' ') p++;

    int* dims = (int*)malloc(sizeof(int)*ndims);
    for(int i=0;i<ndims;i++){
        if(sscanf(p,"%d",&dims[i])!=1){
            fprintf(stderr,"deserialize: invalid dims\n");
            exit(1);
        }
        while(*p && *p!=' ') p++;
        if(*p==' ') p++;
    }

    int total = get_total_size(ndims,dims);
    double* arr = runtime_alloc(ndims,dims);
    for(int i=0;i<total;i++){
        double val;
        if(sscanf(p,"%lf",&val)!=1){
            fprintf(stderr,"deserialize: not enough data\n");
            exit(1);
        }
        arr[i]=val;
        while(*p && *p!=' ') p++;
        if(*p==' ') p++;
    }

    free(dims);
    return arr;
}

// A simple slice: slice along the first dimension, extracting a sub-range
// slice(arr, start, end) - demonstration
// real slicing with steps and multiple dims would be more complex
double* runtime_slice(double* arr, int ndims, int* dims,
                      int dim_to_slice, int start, int end) {
    // Check bounds
    if(dim_to_slice<0 || dim_to_slice>=ndims) {
        fprintf(stderr,"slice: invalid dimension\n");
        exit(1);
    }
    if(start<0 || end>dims[dim_to_slice] || start>=end){
        fprintf(stderr,"slice: invalid range\n");
        exit(1);
    }

    // Compute new dims: same except dims[dim_to_slice] = end-start
    int* new_dims = (int*)malloc(sizeof(int)*ndims);
    for(int i=0;i<ndims;i++){
        new_dims[i]=dims[i];
    }
    new_dims[dim_to_slice]= end-start;

    double* newArr = runtime_alloc(ndims,new_dims);

    // Copy relevant data:
    // In row-major order, slicing dimension i means:
    // Compute strides and copy only elements in [start,end)
    int total = get_total_size(ndims,dims);
    // stride: product of dims after dim_to_slice
    int stride=1;
    for(int i=ndims-1;i>dim_to_slice;i--) {
        stride*=dims[i];
    }

    // block size for one unit in dim_to_slice
    // We'll skip start*stride and copy (end-start)*stride elements each block
    int old_block_size=stride;
    int new_block_size=stride;
    int blocks = 1;
    for(int i=0;i<dim_to_slice;i++){
        blocks *= dims[i];
    }

    // old layout: blocks count of big blocks, each has dims[dim_to_slice]*old_block_size
    // we copy only partial range in each block
    // offset in old = block*(dims[dim_to_slice]*old_block_size)
    // new offset = block*(new_dims[dim_to_slice]*new_block_size)

    int old_dim_size = dims[dim_to_slice];
    int new_dim_size = new_dims[dim_to_slice];

    for(int b=0;b<blocks;b++){
        int old_offset = b*(old_dim_size*old_block_size);
        int new_offset = b*(new_dim_size*new_block_size);
        for(int i=0;i<new_dim_size;i++){
            // copy each element
            newArr[new_offset+i*new_block_size] = arr[old_offset+(start+i)*old_block_size];
        }
    }

    free(new_dims);
    free_array(arr);
    return newArr;
}

// Broadcasting add: For demonstration, we handle up to 2 arrays and try to broadcast
double* runtime_broadcast_add(double* A, double* B) {
    int A_ndims = get_ndims(A);
    int B_ndims = get_ndims(B);
    int* A_dims = get_dims(A,A_ndims);
    int* B_dims = get_dims(B,B_ndims);

    // Compute broadcasted shape
    // For simplicity, if ndims differ, pad the smaller with 1s at front.
    int max_ndims = (A_ndims > B_ndims) ? A_ndims : B_ndims;
    int* final_dims = (int*)malloc(sizeof(int)*max_ndims);

    int* A_ext = (int*)malloc(sizeof(int)*max_ndims);
    int* B_ext = (int*)malloc(sizeof(int)*max_ndims);

    for(int i=0;i<max_ndims;i++){
        int A_dim = (i<max_ndims - A_ndims) ? 1 : A_dims[i-(max_ndims - A_ndims)];
        int B_dim = (i<max_ndims - B_ndims) ? 1 : B_dims[i-(max_ndims - B_ndims)];
        int max_dim = (A_dim>B_dim)? A_dim : B_dim;
        if(A_dim!=max_dim && A_dim!=1){
            fprintf(stderr,"broadcast_add: incompatible dims\n");
            exit(1);
        }
        if(B_dim!=max_dim && B_dim!=1){
            fprintf(stderr,"broadcast_add: incompatible dims\n");
            exit(1);
        }
        final_dims[i]=max_dim;
        A_ext[i]=A_dim;
        B_ext[i]=B_dim;
    }

    double* C = runtime_alloc(max_ndims, final_dims);
    int total = get_total_size(max_ndims, final_dims);

    // Compute indices and add:
    // For each element in final_dims space, pick from A and B.
    // If a dimension in A_ext is 1, always use index 0 in that dimension, else use actual index.
    // same for B.

    // This is a simplified broadcasting loop (inefficient):
    // We'll generate indices by dividing by dimension sizes:
    int* index = (int*)malloc(sizeof(int)*max_ndims);
    int* stride = (int*)malloc(sizeof(int)*max_ndims);
    stride[max_ndims-1]=1;
    for(int i=max_ndims-2;i>=0;i--) {
        stride[i]=stride[i+1]*final_dims[i+1];
    }

    for(int linear_idx=0;linear_idx<total;linear_idx++){
        int tmp = linear_idx;
        for(int i=0;i<max_ndims;i++){
            index[i]= (tmp/stride[i]) % final_dims[i];
        }

        // Compute A_idx
        int A_idx=0; 
        int mul=1;
        for(int i=max_ndims-1;i>=0;i--){
            int idxVal = (A_ext[i]==1)?0:index[i-(max_ndims - A_ndims < 0?0:(max_ndims - A_ndims))];
            // Wait if A_ndims<max_ndims, we must shift indices
            int A_i = i-(max_ndims - A_ndims);
            if(A_i<0) idxVal=0; // out of range, dimension was broadcast
            if(A_i>=0) {
                if(A_ext[i]==1) idxVal=0;
                A_idx = A_idx+ (idxVal*mul);
                mul *= A_ext[i];
            }
        }

        // Compute B_idx similarly
        int B_idx=0; 
        mul=1;
        for(int i=max_ndims-1;i>=0;i--){
            int B_i = i-(max_ndims - B_ndims);
            int idxVal = 0;
            if(B_i<0) idxVal=0; else {
                idxVal = (B_ext[i]==1)?0:index[i-(max_ndims - B_ndims <0?0:(max_ndims - B_ndims))];
            }
            if(B_i>=0) {
                B_idx = B_idx+(idxVal*mul);
                mul*=B_ext[i];
            }
        }

        double a_val = A[A_idx];
        double b_val = B[B_idx];
        C[linear_idx]= a_val+b_val;
    }

    free(index);
    free(stride);
    free(final_dims);
    free(A_ext);
    free(B_ext);

    free_array(A);
    free_array(B);

    return C;
}

// runtime_resize: just reallocate with new dims and copy data as much as fits
double* runtime_resize(double* old_ptr, int new_ndims, int* new_dims) {
    int old_ndims = get_ndims(old_ptr);
    int* old_dims = get_dims(old_ptr, old_ndims);
    int old_total = get_total_size(old_ndims,old_dims);
    int new_total = get_total_size(new_ndims,new_dims);

    double* newArr = runtime_alloc(new_ndims,new_dims);
    int copySize = (old_total < new_total) ? old_total : new_total;
    for(int i=0;i<copySize;i++){
        newArr[i] = old_ptr[i];
    }
    free_array(old_ptr);
    return newArr;
}