Work with tensors

/* Inspect a model, check a backend, then run it on a raw float tensor.
 *
 * usage: tensors MODEL.slab LICENSE_FILE
 */

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

#include "slab.h"

#define S(str) (uint8_t *)(str), strlen(str)

static int read_file(const char *path, char *buf, size_t cap, size_t *len) {
    FILE *f = fopen(path, "rb");
    if (!f) return 0;
    *len = fread(buf, 1, cap, f);
    fclose(f);
    while (*len && (buf[*len - 1] == '\n' || buf[*len - 1] == '\r')) (*len)--;
    return 1;
}

static void print_tensor(uint8_t *s, int input, int32_t i) {
    uint8_t *name = input ? slab_session_input_name_ptr(s, i) : slab_session_output_name_ptr(s, i);
    size_t name_len = input ? slab_session_input_name_len(s, i) : slab_session_output_name_len(s, i);
    int32_t rank = input ? slab_session_input_rank(s, i) : slab_session_output_rank(s, i);
    printf("%s %.*s:", input ? "input " : "output", (int)name_len, (char *)name);
    for (int32_t a = 0; a < rank; a++)
        printf("%s%lld", a ? "x" : " ", (long long)(input ? slab_session_input_dim(s, i, a) : slab_session_output_dim(s, i, a)));
    printf("\n");
}

int main(int argc, char **argv) {
    if (argc < 3) { fprintf(stderr, "usage: %s MODEL.slab LICENSE_FILE\n", argv[0]); return 2; }

    uint8_t *cfg = slab_config_new();
    char token[4096]; size_t token_len;
    if (!read_file(argv[2], token, sizeof token, &token_len) ||
        !slab_config_set_runtime_license(cfg, (uint8_t *)token, token_len)) return 1;

    uint8_t *s = slab_session_new(cfg);

    /* Inspect reads the model's description without loading its weights. */
    if (!slab_session_inspect(s, S(argv[1]))) {
        fprintf(stderr, "%.*s\n", (int)slab_session_error_message_len(s), (char *)slab_session_error_message_ptr(s));
        return 1;
    }
    for (int32_t i = 0; i < slab_session_input_count(s); i++) print_tensor(s, 1, i);
    for (int32_t i = 0; i < slab_session_output_count(s); i++) print_tensor(s, 0, i);
    printf("estimated memory: %zu bytes\n", slab_session_estimated_memory_bytes(s));
    if (!slab_session_check_backend(s, SLAB_BACKEND_GPU))
        printf("GPU cannot run it: %.*s\n", (int)slab_session_error_message_len(s), (char *)slab_session_error_message_ptr(s));

    if (!slab_session_open(s, S(argv[1]))) return 1;

    /* An all-zero input of the declared size. */
    size_t count = 1;
    for (int32_t a = 0; a < slab_session_input_rank(s, 0); a++) count *= (size_t)slab_session_input_dim(s, 0, a);
    float *x = calloc(count, sizeof(float));

    uint8_t *in = slab_input_new();
    uint8_t *out = slab_output_new();
    /* copy = 0: Slab reads `x` only during run(), so it may be freed after. */
    slab_input_set_f32(in, 0, x, count, 0);
    if (!slab_session_run(s, in, out)) {
        fprintf(stderr, "%.*s\n", (int)slab_session_error_message_len(s), (char *)slab_session_error_message_ptr(s));
        return 1;
    }
    free(x);

    float *y = slab_output_f32_ptr(out, 0);
    size_t n = slab_output_f32_len(out, 0);
    size_t best = 0;
    for (size_t i = 1; i < n; i++) if (y[i] > y[best]) best = i;
    printf("output 0: %zu values, max %.4f at %zu, backend %d\n", n, y[best], best, slab_output_backend_used(out));

    slab_output_free(out);
    slab_input_free(in);
    slab_session_free(s);
    slab_config_free(cfg);
    return 0;
}