Work with tensors

// Inspect a model, check it fits a backend, then run it on raw tensors.
//
// Use this path for models whose inputs are not images, or when the app
// prepares the input tensor itself.

import "stdlib/io" as io;
import "stdlib/env" as env;
import "stdlib/fs" as fs;
import "stdlib/result" as result;
import "stdlib/option" as option;
import "stdlib/vec" as vec;
import "ffi/api" as slab;

fn print_shape(ref s: slab::Session, idx: i32, input: bool) {
    for a in 0..6 {
        let d = if input { s.input_dim(idx, a) } else { s.output_dim(idx, a) };
        if d <= 0i64 { break; }
        io::print(if a == 0 { " " } else { "x" });
        io::print("${d}");
    }
    io::println("");
}

fn main() -> i32 {
    var cfg = slab::Config::new();
    cfg.set_models_dir("models");
    guard let option::Option::Some(path) = env::var("SLAB_RUNTIME_LICENSE_FILE") else { return 1; };
    guard let result::Result::Ok(token) = fs::read_to_string(path.view()) else { return 1; };
    if !cfg.set_runtime_license(token.view()) { return 1; }

    var session = slab::Session::new(cfg);

    // Inspect reads the model's description without loading its weights.
    if !session.inspect("classifier.slab") {
        io::eprintln("Cannot inspect: ${session.error_message()}");
        return 1;
    }
    for i in 0..session.input_count() {
        io::print("input  ${session.input_name(i)}:");
        print_shape(session, i, true);
    }
    for i in 0..session.output_count() {
        io::print("output ${session.output_name(i)}:");
        print_shape(session, i, false);
    }
    io::println("estimated memory: ${session.estimated_memory_bytes()} bytes");
    if !session.check_backend(slab::BACKEND_GPU) {
        io::println("GPU cannot run it: ${session.error_message()}");
    }

    if !session.open("classifier.slab") {
        io::eprintln("Cannot open: ${session.error_message()}");
        return 1;
    }

    // Build an input tensor of the declared size (here, all zeros).
    var count: usize = 1;
    for a in 0..6 {
        let d = session.input_dim(0, a);
        if d <= 0i64 { break; }
        count = count * (d as usize);
    }
    var data = vec::with_capacity::[f32](count);
    for i in 0..(count as i32) { let _ = data.append(0.0f32); }
    guard let option::Option::Some(ptr) = data.at_ptr(0) else { return 1; };

    var inp = slab::Input::new();
    var out = slab::Output::new();
    // copy = false: Slab reads the buffer during run() only.
    if !inp.set_named_f32(session.input_name(0), ptr, count, false) {
        io::eprintln("Input rejected");
        return 1;
    }
    if !session.run(inp, out) {
        io::eprintln("Run failed: ${session.error_message()}");
        return 1;
    }

    let y = out.f32_ptr(0);
    io::println("output 0: ${out.f32_len(0)} values, first ${y[0]}, on ${slab::backend_name(out.backend_used())}");
    return 0;
}