File size: 5,218 Bytes
3fd1a35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | // 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;
}
}
|