tools/npu_workspace_probe.cpp

// Standalone feasibility probe; does not change engine allocation or scheduling.
#include <rknn_api.h>
#include <rknn_matmul_api.h>
#include <array>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <stdexcept>
#include <utility>
#include <vector>

void Check(int rc, const char * op) {
    if (rc != 0) throw std::runtime_error(std::string(op) + ": " + std::to_string(rc));
}
struct Experiment {
    std::vector<rknn_matmul_ctx> contexts;
    std::vector<std::pair<rknn_context, rknn_tensor_mem *>> memory;
    ~Experiment() {
        for (auto i = memory.rbegin(); i != memory.rend(); ++i) rknn_destroy_mem(i->first, i->second);
        for (auto i = contexts.rbegin(); i != contexts.rend(); ++i) rknn_matmul_destroy(*i);
    }
    rknn_tensor_mem * Own(rknn_context ctx, rknn_tensor_mem * mem) {
        if (!mem) throw std::runtime_error("memory creation failed");
        memory.emplace_back(ctx, mem); return mem;
    }
};
void Probe(int k, int n, int core) {
        constexpr std::array<int, 8> rows {256, 2, 4, 8, 16, 32, 64, 128};
        for (bool views : {false, true}) {
            Experiment e;
            std::array<rknn_matmul_io_attr, rows.size()> attrs {};
            for (std::size_t i = 0; i < rows.size(); ++i) {
                rknn_matmul_info info {};
                info.M = rows[i]; info.K = k; info.N = n;
                info.type = RKNN_INT4_MM_INT4_TO_INT16;
                info.B_layout = RKNN_MM_LAYOUT_NATIVE;
                info.AC_layout = RKNN_MM_LAYOUT_NATIVE;
                info.B_quant_type = RKNN_QUANT_TYPE_PER_LAYER_SYM;
                info.AC_quant_type = RKNN_QUANT_TYPE_PER_LAYER_SYM;
                info.iommu_domain_id = 2;
                rknn_matmul_ctx ctx = 0;
                Check(rknn_matmul_create(&ctx, &info, &attrs[i]), "create context");
                e.contexts.push_back(ctx);
                Check(rknn_matmul_set_core_mask(ctx, static_cast<rknn_core_mask>(1U << core)), "set core");
            }
            const auto owner = e.contexts[0];
            auto alloc = [&](unsigned size) {
                return e.Own(owner, rknn_create_mem2(owner, size, RKNN_FLAG_MEMORY_CACHEABLE));
            };
            auto * a = alloc(attrs[0].A.size), * b = alloc(attrs[0].B.size), * c = alloc(attrs[0].C.size);
            std::memset(b->virt_addr, 0x11, b->size);
            Check(rknn_mem_sync(owner, b, RKNN_MEMORY_SYNC_TO_DEVICE), "sync B");
            std::array<rknn_tensor_mem *, rows.size()> av {}, cv {};
            for (std::size_t i = 0; i < rows.size(); ++i) {
                auto ctx = e.contexts[i];
                av[i] = views ? e.Own(ctx, rknn_create_mem_from_fd(ctx, a->fd, a->virt_addr, attrs[i].A.size, 0)) : a;
                cv[i] = views ? e.Own(ctx, rknn_create_mem_from_fd(ctx, c->fd, c->virt_addr, attrs[i].C.size, 0)) : c;
                Check(rknn_matmul_set_io_mem(ctx, av[i], &attrs[i].A), "bind shared A");
                Check(rknn_matmul_set_io_mem(ctx, b, &attrs[i].B), "bind B");
                Check(rknn_matmul_set_io_mem(ctx, cv[i], &attrs[i].C), "bind shared C");
                if (views && (av[i]->size != attrs[i].A.size || cv[i]->size != attrs[i].C.size))
                    throw std::runtime_error("FD view did not retain requested synchronization size");
            }
            for (int repeat = 0; repeat < 8; ++repeat) {
                for (std::size_t i : {0, 1, 7, 2, 6, 3, 5, 4, 1}) {
                    auto ctx = e.contexts[i];
                    const int value = 1 + repeat % 2;
                    std::memset(a->virt_addr, value * 0x11, a->size);
                    std::memset(c->virt_addr, 0x55, c->size);
                    Check(rknn_mem_sync(owner, c, RKNN_MEMORY_SYNC_TO_DEVICE), "guard C");
                    Check(rknn_mem_sync(ctx, av[i], RKNN_MEMORY_SYNC_TO_DEVICE), "sync shared A");
                    Check(rknn_matmul_run(ctx), "run");
                    Check(rknn_mem_sync(ctx, cv[i], RKNN_MEMORY_SYNC_FROM_DEVICE), "sync shared C");
                    // Read guards after synchronizing the full owner mapping.
                    Check(rknn_mem_sync(owner, c, RKNN_MEMORY_SYNC_FROM_DEVICE), "sync guard");
                    auto * result = static_cast<std::int16_t *>(c->virt_addr);
                    for (unsigned j = 0; j < c->size / 2; ++j) {
                        const int expected = j < attrs[i].C.size / 2 ? value * k : 0x5555;
                        if (result[j] != expected) throw std::runtime_error("output/guard mismatch");
                    }
                }
            }
            std::cout << "PASS mode=" << (views ? "per-shape-fd-view" : "oversized-owner")
                      << " K=" << k << " N=" << n << " core=" << core
                      << " shapes=8 transitions=72 domain=2 output_and_guards=exact\n";
        }
}
int main() {
    try {
        Probe(64, 128, 0);
        // Real qkvfgb per-core shards and expert down projection dimensions.
        Probe(1536, 3456, 0);
        Probe(1536, 3456, 1);
        Probe(1536, 3392, 2);
        Probe(512, 1536, 1);
    } catch (const std::exception & error) {
        std::cerr << "probe_failed=" << error.what() << '\n'; return 1;
    }
}