test_flexmlrt.cpp

#ifdef _WIN32
    #define NOMINMAX
    #include <windows.h>
#endif

#define FLEXML_LOADER_DEBUG 1
#include "FlexMLClient.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <stdexcept>

#ifdef _WIN32
static bool map_file(const char * path, uint8_t ** buf, size_t * sz) {
    HANDLE hFile = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == INVALID_HANDLE_VALUE) { fprintf(stderr, "open failed: %s\n", path); return false; }
    LARGE_INTEGER fs;
    if (!GetFileSizeEx(hFile, &fs)) { CloseHandle(hFile); return false; }
    HANDLE hMap = CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, fs.QuadPart, NULL);
    if (!hMap) { CloseHandle(hFile); return false; }
    *buf = (uint8_t *)MapViewOfFile(hMap, FILE_MAP_READ, 0, 0, (SIZE_T)fs.QuadPart);
    if (!*buf) { CloseHandle(hMap); CloseHandle(hFile); return false; }
    *sz = (size_t)fs.QuadPart;
    return true;
}
static void unmap_file(uint8_t * buf, size_t) { UnmapViewOfFile(buf); }
#endif

int main(int argc, char ** argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: test_flexmlrt <model_dir> [subgraph_name]\n");
        fprintf(stderr, "  model_dir: directory containing the compiled .rai model\n");
        fprintf(stderr, "  e.g.: test_flexmlrt models vaiml_par_0\n");
        return 1;
    }

    const char * rai_path = argv[1];
    const char * subgraph = argc > 2 ? argv[2] : "vaiml_par_0";

    fprintf(stderr, "=== FlexmlRT standalone test (fbs_buffer mode) ===\n");
    fprintf(stderr, "  rai      : %s\n", rai_path);
    fprintf(stderr, "  subgraph : %s\n", subgraph);

    uint8_t * buf = nullptr;
    size_t    sz  = 0;
    if (!map_file(rai_path, &buf, &sz)) {
        fprintf(stderr, "FAIL: could not mmap %s\n", rai_path);
        return 1;
    }
    fprintf(stderr, "  mmap ok  : %zu bytes\n", sz);

    flexmlrt::client::Options opts;
    opts.modelPath              = rai_path;
    opts.deviceName             = "stx";
    opts.debug                  = true;
    opts.executeMode            = 2;
    opts.extOptions["enable_preemption"] = true;
    opts.extOptions["fbs_buffer"]        = buf;
    opts.extOptions["fbs_buffer_size"]   = sz;
    opts.extOptions["cache_dir"]         = std::string(".");
    opts.subgraphName                    = subgraph;

    fprintf(stderr, "  calling Model() constructor ...\n");
    try {
        auto model = std::make_shared<flexmlrt::client::Model>(opts);
        fprintf(stderr, "  constructor returned, good()=%s\n", model->good() ? "true" : "false");

        if (!model->good()) {
            fprintf(stderr, "FAIL: good() returned false\n");
            return 1;
        }

        auto inputs  = model->getIOTensors("input",  false);
        auto outputs = model->getIOTensors("output", false);
        fprintf(stderr, "  inputs : %zu\n", inputs.size());
        fprintf(stderr, "  outputs: %zu\n", outputs.size());
        auto dump_meta = [](const char * tag, size_t i, flexmlrt::client::ErtTensorType & t) {
            try {
                const auto & m = t.getMetadata();
                std::string shp;
                for (auto d : m.shape) shp += std::to_string(d) + ",";
                fprintf(stderr, "    %s[%zu] idx=%d size=%zu type=%s fmt=%s shape=[%s]\n",
                        tag, i, t.idx, m.size,
                        flexmlrt::client::toString(m.type).c_str(),
                        flexmlrt::client::toString(m.format).c_str(),
                        shp.c_str());
            } catch (const std::exception & e) {
                fprintf(stderr, "    %s[%zu] idx=%d (getMetadata failed: %s)\n", tag, i, t.idx, e.what());
            }
        };
        for (size_t i = 0; i < inputs.size();  ++i) dump_meta("in",  i, inputs[i]);
        for (size_t i = 0; i < outputs.size(); ++i) dump_meta("out", i, outputs[i]);

        // --- forward with dummy buffers matching the COMPILED model metadata ---
        // in[0] = bool mask  [1,188]      (188 bytes) -- length encoded as a frame mask
        // in[1] = audio fp32 [1,128,1498] (766976 bytes)
        // out[0]= fp32       [1,1024,188] (770048 bytes)
        const int M = 128, T = 1498, D = 1024, Tp = 188;
        const int n_mel_real = 1101;
        const int subsampl   = 8;
        const int n_valid    = (n_mel_real + subsampl - 1) / subsampl; // ceil -> 138

        std::vector<float>   audio((size_t)M * T, 0.1f);
        std::vector<uint8_t> mask((size_t)Tp, 0);   // 1 byte per bool
        for (int i = 0; i < Tp; ++i) mask[i] = (i < n_valid) ? 1 : 0; // true = valid frame
        std::vector<float>   out_buf((size_t)D * Tp, 0.0f);

        if (inputs.size() < 2 || outputs.empty()) {
            fprintf(stderr, "FAIL: unexpected IO counts\n");
            return 1;
        }

        inputs[0].data  = mask.data();   // bool mask
        inputs[1].data  = audio.data();  // audio_signal
        outputs[0].data = out_buf.data();

        fprintf(stderr, "  calling forward (n_valid=%d/%d) ...\n", n_valid, Tp);
        model->forward(inputs, outputs);
        fprintf(stderr, "  forward returned\n");

        // print a few output values to confirm it actually wrote something
        fprintf(stderr, "  out[0..4] = %f %f %f %f %f\n",
                out_buf[0], out_buf[1], out_buf[2], out_buf[3], out_buf[4]);

        fprintf(stderr, "SUCCESS\n");
    } catch (const std::exception & e) {
        fprintf(stderr, "FAIL exception: %s\n", e.what());
        return 1;
    }

    return 0;
}