diff --git a/CMakeLists.txt b/CMakeLists.txt index daeb645df..b4d33a6bb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -212,6 +212,16 @@ add_test(NAME bump_hrx_report COMMAND bash ${CMAKE_SOURCE_DIR}/tests/bump_hrx_re add_executable(audio_ops_test tests/audio_ops_test.cpp) target_include_directories(audio_ops_test PRIVATE npu/forward/include) add_test(NAME audio_ops COMMAND audio_ops_test) +# engine#1783 (F8): the engine reads an audio pack's graph and runs the whole encoder + decoder on the host -- 447 + 786 nodes, +# matched against onnxruntime. The reader keeps every node in graph order (the census walk keeps only the MatMul family). +add_executable(onnx_graph_test tests/onnx_graph_test.cpp) +target_link_libraries(onnx_graph_test PRIVATE onebit_npu) +add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) +set_tests_properties(onnx_graph PROPERTIES SKIP_RETURN_CODE 77) +add_executable(audio_exec_test tests/audio_exec_test.cpp) +target_link_libraries(audio_exec_test PRIVATE onebit_npu) +add_test(NAME audio_exec COMMAND audio_exec_test ${ONEBIT_AUDIO_ONNX}) +set_tests_properties(audio_exec PROPERTIES SKIP_RETURN_CODE 77) add_test(NAME metrics_serve COMMAND ${CMAKE_SOURCE_DIR}/tests/metrics_serve.sh $) # serve --print-backend: the llama-server Lemonade starts itself, one hop (docs/serve.md): no GPU. add_test(NAME print_backend COMMAND ${CMAKE_SOURCE_DIR}/tests/print_backend.sh $) @@ -754,7 +764,7 @@ if(ONEBIT_NPU_HOST) find_package(PkgConfig REQUIRED) pkg_check_modules(PCRE2 REQUIRED IMPORTED_TARGET libpcre2-8) add_library(onebit_npu STATIC npu/file_map.cpp npu/flm_head_import.cpp npu/full_elf.cpp npu/model.cpp npu/vision_pack.cpp npu/vision_tower.cpp npu/vision_image.cpp - npu/moe.cpp npu/pack.cpp npu/pack_import.cpp npu/private_route.cpp + npu/audio_exec.cpp npu/moe.cpp npu/pack.cpp npu/pack_import.cpp npu/private_route.cpp npu/q4nx.cpp npu/q4nx_device.cpp npu/tokenizer.cpp) target_include_directories(onebit_npu PUBLIC npu) target_link_libraries(onebit_npu PUBLIC nlohmann_json::nlohmann_json PRIVATE PkgConfig::PCRE2) diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp new file mode 100644 index 000000000..abeab9eb1 --- /dev/null +++ b/npu/audio_exec.cpp @@ -0,0 +1,602 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec.cpp -- the host f32 executor for an audio pack's graph (see audio_exec.h for the contract). + +#include "audio_exec.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "forward/include/audio_ops.h" +#include "file_map.h" + +namespace onebit::npu::pack_import { +namespace { + +struct Tensor { + std::vector dims; + std::vector data; // f32 payload + std::vector i64; // int64 payload (tokens, indices, shapes); exactly one of the two is used + std::vector b; // bool payload (a mask) + size_t count() const { return i64.empty() ? (b.empty() ? data.size() : b.size()) : i64.size(); } +}; + +size_t nelem(const std::vector& d) { + size_t n = 1; + for (int64_t v : d) n *= size_t(v < 0 ? 0 : v); + return n; +} + +[[noreturn]] void fail(const std::string& why) { throw std::runtime_error(why); } + +// numpy-style right-aligned broadcast of `small` against `big`; returns the result dims and, per output index, the source +// index. Only the two forms this graph uses: equal, or `small` has trailing dims of 1/size and fewer ranks. +struct Broadcast { + std::vector dims; + std::vector stride_big, stride_small; // 0 where the axis is broadcast +}; + +Broadcast plan_broadcast(const std::vector& a, const std::vector& b) { + Broadcast p; + const size_t r = std::max(a.size(), b.size()); + p.dims.assign(r, 1); + p.stride_big.assign(r, 0); + p.stride_small.assign(r, 0); + for (size_t i = 0; i < r; ++i) { + const int64_t da = (i < r - a.size()) ? 1 : a[i - (r - a.size())]; + const int64_t db = (i < r - b.size()) ? 1 : b[i - (r - b.size())]; + if (da != db && da != 1 && db != 1) fail("broadcast mismatch " + std::to_string(da) + " vs " + std::to_string(db)); + p.dims[i] = da == 1 ? db : da; + } + size_t sb = 1, ss = 1; + for (size_t i = r; i-- > 0;) { + const int64_t da = (i < r - a.size()) ? 1 : a[i - (r - a.size())]; + const int64_t db = (i < r - b.size()) ? 1 : b[i - (r - b.size())]; + p.stride_big[i] = da == 1 ? 0 : sb; + p.stride_small[i] = db == 1 ? 0 : ss; + sb *= size_t(da); + ss *= size_t(db); + } + return p; +} + +// Flat index into `src` for the right-aligned broadcast of `src` against an output of dims `od` at multi-index `idx`. +// numpys rule: a missing or unit axis is read as index 0. +size_t bcast_index(const std::vector& od, const std::vector& src, const std::vector& idx) { + size_t off = 0, stride = 1; + for (size_t i = src.size(); i-- > 0;) { + const size_t from_right = src.size() - i; // 1 = last axis + const size_t oi = od.size() - from_right; // the output axis it aligns with + const int64_t d = src[i]; + const size_t k = (oi < od.size() && d != 1) ? idx[oi] : 0; + off += k * stride; + stride *= size_t(d); + } + return off; +} + +// Deterministic fork-join over independent units of work (output rows, batch indices). Each unit writes only its own slice, so +// every element's arithmetic is unchanged by the thread count: the result is bit-identical to the serial loop, which is the +// property a reference implementation has to keep. ONEBIT_AUDIO_THREADS=1 forces the serial path (the control for measuring it). +template +void parallel_units(size_t units, Fn&& fn) { + size_t want = 1; + if (const char* env = std::getenv("ONEBIT_AUDIO_THREADS")) { + const long n = std::strtol(env, nullptr, 10); + want = n > 0 ? size_t(n) : 1; + } else { + const size_t hw = std::thread::hardware_concurrency(); + want = hw == 0 ? 1 : (hw > 8 ? 8 : hw); + } + if (want > units) want = units; + if (want <= 1) { + for (size_t u = 0; u < units; ++u) fn(u); + return; + } + std::vector pool; + pool.reserve(want); + for (size_t t = 0; t < want; ++t) + pool.emplace_back([&, t]() { + for (size_t u = t; u < units; u += want) fn(u); + }); + for (auto& th : pool) th.join(); +} + +void broadcast_binary(const Tensor& a, const Tensor& b, const std::string& op, Tensor* out) { + const Broadcast p = plan_broadcast(a.dims, b.dims); + out->dims = p.dims; + out->data.resize(nelem(p.dims)); + std::vector idx(p.dims.size(), 0); + for (size_t o = 0; o < out->data.size(); ++o) { + size_t ia = 0, ib = 0; + for (size_t d = 0; d < idx.size(); ++d) { + ia += idx[d] * p.stride_big[d]; + ib += idx[d] * p.stride_small[d]; + } + const float x = a.data[ia], y = b.data[ib]; + out->data[o] = (op == "Add") ? x + y : (op == "Sub") ? x - y : (op == "Mul") ? x * y : x / y; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(p.dims[d])) break; + idx[d] = 0; + } + } +} + +std::vector dims_of(const OnnxInit& t) { return t.dims; } + +// Copies an initializer to f32, following external_data into the pack's side file when the bytes are not inline. +void load_init_f32(const OnnxGraphInfo& g, const OnnxInit& t, Tensor* dst) { + const uint8_t* src = t.raw; + size_t len = t.raw_len; + size_t size = 0; + std::string why; + const uint8_t* mapped = nullptr; + if (!t.location.empty()) { + mapped = map_file(g.dir + "/" + t.location, size, &why); + if (mapped == nullptr) fail("external_data " + t.location + ": " + why); + const uint64_t off = t.offset, want = t.length; + if (off > size || want > size - off) fail("external_data " + t.location + ": range past the end"); + src = mapped + off; + len = size_t(want); + if (len == 0) len = size_t(size - off); + } + if (src == nullptr) fail("initializer '" + t.name + "' has no data"); + const size_t n = nelem(t.dims); + if (t.dtype == 7) { // INT64 -- indices and shape operands; the payload stays integer + dst->i64.assign(n, 0); + if (len < n * 8) fail("int64 initializer '" + t.name + "' is short"); + std::memcpy(dst->i64.data(), src, n * 8); + if (mapped != nullptr) unmap_file(mapped, size); + return; + } + dst->data.assign(n, 0.0f); + if (t.dtype == 1) { // FLOAT + if (len < n * 4) fail("initializer '" + t.name + "' is short: " + std::to_string(len) + " < " + std::to_string(n * 4)); + std::memcpy(dst->data.data(), src, n * 4); + } else if (t.dtype == 10) { // FLOAT16 + if (len < n * 2) fail("f16 initializer '" + t.name + "' is short"); + for (size_t i = 0; i < n; ++i) { + uint16_t h = 0; + std::memcpy(&h, src + i * 2, 2); + const uint32_t sign = uint32_t(h & 0x8000) << 16, exp = (h >> 10) & 0x1f, man = h & 0x3ff; + uint32_t bits; + if (exp == 0) { + bits = sign | (man ? (uint32_t(127 - 15) << 23) + (man << 13) : 0); + } else if (exp == 31) { + bits = sign | 0x7f800000u | (man << 13); + } else { + bits = sign | ((exp + 127 - 15) << 23) | (man << 13); + } + std::memcpy(&dst->data[i], &bits, 4); + } + } else { + fail("initializer '" + t.name + "' has dtype " + std::to_string(t.dtype) + ", which this executor does not load"); + } + if (mapped != nullptr) unmap_file(mapped, size); +} + +void load_init_i64(const OnnxGraphInfo& g, const OnnxInit& t, std::vector* dst) { + const uint8_t* src = t.raw; + size_t len = t.raw_len; + size_t size = 0; + std::string why; + const uint8_t* mapped = nullptr; + if (!t.location.empty()) { + mapped = map_file(g.dir + "/" + t.location, size, &why); + if (mapped == nullptr) fail("external_data " + t.location + ": " + why); + if (t.offset > size || t.length > size - t.offset) fail("external_data range past the end"); + src = mapped + t.offset; + len = size_t(t.length); + } + if (t.dtype != 7) fail("'" + t.name + "' is not INT64"); + const size_t n = nelem(t.dims); + if (src == nullptr || len < n * 8) fail("int64 initializer '" + t.name + "' is short"); + dst->assign(n, 0); + std::memcpy(dst->data(), src, n * 8); + if (mapped != nullptr) unmap_file(mapped, size); +} + +} // namespace + +namespace { + +bool run_graph_impl(const OnnxGraphInfo& g, std::map t, const std::string& out_name, AudioRunResult* out, + std::string* why) { + try { + + std::map inits; + for (const auto& ti : g.init_tensors) inits[ti.name] = &ti; + + auto need = [&](const std::string& name) -> const Tensor& { + const auto it = t.find(name); + if (it == t.end()) fail("node input '" + name + "' was never produced"); + return it->second; + }; + auto init_f32 = [&](const std::string& name) -> const Tensor& { + auto it = t.find(name); + if (it != t.end()) return it->second; + const auto ii = inits.find(name); + if (ii == inits.end()) fail("'" + name + "' is neither a tensor nor an initializer"); + Tensor v; + v.dims = dims_of(*ii->second); + load_init_f32(g, *ii->second, &v); + return t.emplace(name, std::move(v)).first->second; + }; + + // A node operand is either a tensor an earlier node produced or an initializer (a weight, an index vector, a + // scalar attribute operand). `any_t` is the one spelling of that rule. + auto any_t = [&](const std::string& name) -> const Tensor& { + return t.count(name) != 0 ? need(name) : init_f32(name); + }; + + int node_index = 0; + for (const auto& n : g.nodes) { + ++node_index; + const auto where = [&]() { return "node " + std::to_string(node_index) + " (" + n.op + ")"; }; + if (n.outs.empty()) fail(where() + " has no output"); + const std::string& oname = n.outs[0]; + Tensor r; + if (n.op == "MatMul") { + const Tensor& a = need(n.ins[0]); + const Tensor& b = t.count(n.ins[1]) != 0 ? need(n.ins[1]) : init_f32(n.ins[1]); + if (a.dims.size() < 2 || b.dims.size() < 2) fail(where() + ": needs rank >= 2 on both sides"); + const size_t ra = a.dims.size(), rb = b.dims.size(); + const int64_t M = a.dims[ra - 2], K = a.dims.back(), K2 = b.dims[rb - 2], N = b.dims.back(); + if (K != K2) fail(where() + ": " + std::to_string(K) + " vs " + std::to_string(K2)); + // Attention's q.k^T and scores.v are 4-D: (batch, heads, M|K, K|N). The batch dims broadcast, the last two + // are a plain GEMM, so one projection covers both the layer weights and the attention products. + const std::vector ab(a.dims.begin(), a.dims.end() - 2), bb(b.dims.begin(), b.dims.end() - 2); + const Broadcast bp = plan_broadcast(ab, bb); + r.dims = bp.dims; + r.dims.push_back(M); + r.dims.push_back(N); + r.data.assign(nelem(r.dims), 0.0f); + // Rows are independent, so the work splits over (batch, row) without touching any element's arithmetic. + const size_t batches = nelem(bp.dims); + const size_t rows = size_t(std::max(M, 1)); + parallel_units(batches * rows, [&](size_t u) { + const size_t bi = u / rows; + const int64_t i = int64_t(u % rows); + std::vector idx(bp.dims.size(), 0); + size_t rem = bi; + for (size_t d = bp.dims.size(); d-- > 0;) { // row-major unflatten of the batch index + const size_t dim = size_t(bp.dims[d] ? bp.dims[d] : 1); + idx[d] = rem % dim; + rem /= dim; + } + size_t oa = 0, ob = 0; + for (size_t d = 0; d < idx.size(); ++d) { + oa += idx[d] * bp.stride_big[d]; + ob += idx[d] * bp.stride_small[d]; + } + const float* ap = a.data.data() + oa * size_t(M * K); + const float* bpp = b.data.data() + ob * size_t(K * N); + float* rr = r.data.data() + bi * size_t(M * N) + size_t(i) * size_t(N); + for (int64_t k = 0; k < K; ++k) { + const float av = ap[size_t(i) * size_t(K) + size_t(k)]; + if (av == 0.0f) continue; + const float* br = bpp + size_t(k) * size_t(N); + for (int64_t j = 0; j < N; ++j) rr[j] += av * br[j]; + } + }); + } else if (n.op == "Add" || n.op == "Mul" || n.op == "Div" || n.op == "Sub") { + // An operand is either a tensor an earlier node produced or an initializer (a bias vector, the 1/sqrt(d) + // scalar); init_f32() handles the second without a separate code path at every call site. + auto operand = [&](const std::string& name) -> const Tensor& { + return t.count(name) != 0 ? need(name) : init_f32(name); + }; + if (n.ins.size() < 2) fail(where() + ": needs two operands"); + broadcast_binary(operand(n.ins[0]), operand(n.ins[1]), n.op, &r); + } else if (n.op == "Erf") { + const Tensor& a = need(n.ins[0]); + r.dims = a.dims; + r.data.resize(a.data.size()); + for (size_t i = 0; i < a.data.size(); ++i) r.data[i] = std::erf(a.data[i]); + } else if (n.op == "Gelu") { + const Tensor& a = need(n.ins[0]); + r.dims = a.dims; + r.data.resize(a.data.size()); + const auto ap = n.attrs.find("approximate"); + (void)ap; // ONNX spells it as a string; the export here is 'none', which is what audio_ops::gelu_erf is + for (size_t i = 0; i < a.data.size(); ++i) r.data[i] = audio::gelu_erf(a.data[i]); + } else if (n.op == "LayerNormalization") { + const Tensor& a = need(n.ins[0]); + const Tensor& scale = init_f32(n.ins[1]); + const Tensor* bias = (n.ins.size() > 2 && !n.ins[2].empty()) ? &init_f32(n.ins[2]) : nullptr; + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : -1; + if (axis != -1 && axis != int64_t(a.dims.size()) - 1) fail(where() + ": only the last axis is implemented"); + const int64_t nrm = a.dims.back(); + const float eps = n.attrs_float.count("epsilon") ? n.attrs_float.at("epsilon") : 1e-5f; + if (int64_t(scale.data.size()) != nrm) fail(where() + ": scale is not the normalised width"); + r.dims = a.dims; + r.data.resize(a.data.size()); + for (size_t row = 0; row < a.data.size() / size_t(nrm); ++row) { + audio::layer_norm(a.data.data() + row * size_t(nrm), scale.data.data(), + bias != nullptr ? bias->data.data() : nullptr, int(nrm), eps, + r.data.data() + row * size_t(nrm)); + } + } else if (n.op == "Softmax") { + const Tensor& a = need(n.ins[0]); + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : -1; + if (axis != -1 && axis != int64_t(a.dims.size()) - 1) fail(where() + ": only the last axis is implemented"); + r = a; + const int64_t nrm = a.dims.back(); + for (size_t row = 0; row < r.data.size() / size_t(nrm); ++row) + audio::softmax_row(r.data.data() + row * size_t(nrm), int(nrm)); + } else if (n.op == "Transpose") { + const Tensor& a = need(n.ins[0]); + const auto it = n.attrs_list.find("perm"); + std::vector perm; + if (it != n.attrs_list.end()) perm = it->second; + else + for (size_t i = a.dims.size(); i-- > 0;) perm.push_back(int64_t(i)); + if (perm.size() != a.dims.size()) fail(where() + ": perm rank mismatch"); + r.dims.resize(perm.size()); + for (size_t i = 0; i < perm.size(); ++i) r.dims[i] = a.dims[size_t(perm[i])]; + r.data.resize(a.data.size()); + std::vector in_stride(perm.size(), 1), out_stride(perm.size(), 1); + for (size_t i = perm.size(); i-- > 1;) { + in_stride[i - 1] = in_stride[i] * size_t(a.dims[i]); + out_stride[i - 1] = out_stride[i] * size_t(r.dims[i]); + } + std::vector idx(perm.size(), 0); + for (size_t o = 0; o < r.data.size(); ++o) { + size_t src = 0; + for (size_t i = 0; i < idx.size(); ++i) src += idx[i] * in_stride[size_t(perm[i])]; + r.data[o] = a.data[src]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(r.dims[d])) break; + idx[d] = 0; + } + } + } else if (n.op == "Reshape") { + const Tensor& a = need(n.ins[0]); + std::vector want; + auto ii = inits.find(n.ins[1]); + if (ii == inits.end()) fail(where() + ": target shape is not an initializer"); + load_init_i64(g, *ii->second, &want); + std::vector dims(want.size()); + int64_t infer = -1; + size_t known = 1; + for (size_t i = 0; i < want.size(); ++i) { + int64_t d = want[i]; + if (d == 0) d = a.dims[i]; + if (d == -1) { + if (infer >= 0) fail(where() + ": two -1"); + infer = int64_t(i); + d = 1; + } + dims[i] = d; + known *= size_t(d); + } + if (infer >= 0) dims[size_t(infer)] = int64_t(a.data.size() / std::max(known, 1)); + if (nelem(dims) != a.data.size()) fail(where() + ": reshape changes the element count"); + r.dims = dims; + r.data = a.data; + } else if (n.op == "Conv") { + const Tensor& a = need(n.ins[0]); + const Tensor& w = init_f32(n.ins[1]); + const Tensor* bias = (n.ins.size() > 2 && !n.ins[2].empty()) ? &init_f32(n.ins[2]) : nullptr; + const auto gs = n.attrs_list.find("strides"); + const auto gp = n.attrs_list.find("pads"); + const int stride = (gs != n.attrs_list.end() && !gs->second.empty()) ? int(gs->second[0]) : 1; + const int pad = (gp != n.attrs_list.end() && !gp->second.empty()) ? int(gp->second[0]) : 0; + if (a.dims.size() != 3 || w.dims.size() != 3) fail(where() + ": only 1-D temporal Conv on (N,C,T)"); + const int64_t cin = a.dims[1], T = a.dims[2]; + const int64_t cout = w.dims[0], k = w.dims[2]; + if (w.dims[1] != cin) fail(where() + ": kernel channels do not match the input"); + const int64_t tout = audio::conv1d_out_len(int(T), int(k), stride, pad); + r.dims = {a.dims[0], cout, tout}; + r.data.resize(size_t(a.dims[0]) * size_t(cout) * size_t(tout)); + for (int64_t n0 = 0; n0 < a.dims[0]; ++n0) + audio::conv1d(a.data.data() + size_t(n0) * size_t(cin) * size_t(T), int(cin), int(T), w.data.data(), + bias != nullptr ? bias->data.data() : nullptr, int(cout), int(k), stride, pad, + r.data.data() + size_t(n0) * size_t(cout) * size_t(tout)); + } else if (n.op == "Gather") { + // The token embedding: axis 0, indices are the int64 token window. + const Tensor& a = any_t(n.ins[0]); + const Tensor& ix = any_t(n.ins[1]); + if (!ix.i64.empty()) { + const int64_t axis = n.attrs.count("axis") ? n.attrs.at("axis") : 0; + if (axis != 0 || a.dims.empty()) fail(where() + ": only axis 0 is implemented"); + const int64_t rows = a.dims[0]; + const size_t row = a.data.size() / size_t(rows); + r.dims = ix.dims; + for (size_t i = 1; i < a.dims.size(); ++i) r.dims.push_back(a.dims[i]); + r.data.resize(ix.i64.size() * row); + for (size_t i = 0; i < ix.i64.size(); ++i) { + const int64_t src = ix.i64[i]; + if (src < 0 || src >= rows) fail(where() + ": index out of range"); + std::memcpy(r.data.data() + i * row, a.data.data() + size_t(src) * row, row * 4); + } + } else { + fail(where() + ": indices are not int64"); + } + } else if (n.op == "Equal") { + // Three forms in this family: the token window against <|endoftext|> (int64), and the causal-mask constant + // against itself (float, exact equality -- which is the ONNX semantic and the whole point of the pattern). + const Tensor& a = any_t(n.ins[0]); + const Tensor& b = any_t(n.ins[1]); + if (a.count() != b.count() && a.count() != 1 && b.count() != 1) + fail(where() + ": only scalar broadcast is implemented"); + r.dims = a.dims.size() >= b.dims.size() ? a.dims : b.dims; + const size_t n_out = std::max(a.count(), b.count()); + r.b.assign(n_out, 0); + for (size_t i = 0; i < n_out; ++i) { + const size_t ia = a.count() == 1 ? 0 : i, ib = b.count() == 1 ? 0 : i; + if (!a.i64.empty() && !b.i64.empty()) r.b[i] = a.i64[ia] == b.i64[ib] ? 1 : 0; + else if (!a.b.empty() && !b.b.empty()) r.b[i] = a.b[ia] == b.b[ib] ? 1 : 0; + else if (!a.data.empty() && !b.data.empty()) r.b[i] = a.data[ia] == b.data[ib] ? 1 : 0; + else fail(where() + ": operands are of different kinds"); + } + } else if (n.op == "Not") { + const Tensor& a = any_t(n.ins[0]); + if (a.b.empty()) fail(where() + ": operand is not a mask"); + r.dims = a.dims; + r.b.resize(a.b.size()); + for (size_t i = 0; i < a.b.size(); ++i) r.b[i] = a.b[i] ? 0 : 1; + } else if (n.op == "Where") { + // The causal-mask pattern in this family is Where(mask, -inf_scalar, 0_scalar): the *condition* carries the + // shape, so broadcasting only the two value operands would collapse the result to one element. + const Tensor& c = any_t(n.ins[0]); + const Tensor& a = any_t(n.ins[1]); + const Tensor& bb2 = any_t(n.ins[2]); + if (c.b.empty()) fail(where() + ": condition is not a mask"); + std::vector vals = plan_broadcast(a.dims, bb2.dims).dims; + r.dims = plan_broadcast(vals, c.dims).dims; + const size_t total = nelem(r.dims); + const bool use_i64 = !a.i64.empty(); + if (use_i64) r.i64.resize(total, 0); + else r.data.resize(total, 0.0f); + std::vector idx(r.dims.size(), 0); + for (size_t o = 0; o < total; ++o) { + const size_t ci = bcast_index(r.dims, c.dims, idx); + const size_t ia = bcast_index(r.dims, a.dims, idx); + const size_t ib = bcast_index(r.dims, bb2.dims, idx); + if (use_i64) r.i64[o] = c.b[ci] ? a.i64[ia] : bb2.i64[ib]; + else r.data[o] = c.b[ci] ? a.data[ia] : bb2.data[ib]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < size_t(r.dims[d])) break; + idx[d] = 0; + } + } + } else if (n.op == "Trilu") { + // The causal mask. ONNX zeroes the masked triangle, which for the float form is the "large constant" the + // following Equal/Where pair looks for; an int64 input is zeroed the same way. + const Tensor& a = any_t(n.ins[0]); + const int64_t k = n.attrs.count("k") ? n.attrs.at("k") : 0; + const bool upper = n.attrs.count("upper") ? n.attrs.at("upper") != 0 : true; + if (a.dims.size() < 2) fail(where() + ": needs a matrix"); + const int64_t R = a.dims[a.dims.size() - 2], C = a.dims.back(); + r.dims = a.dims; + r.data = a.data; + r.i64 = a.i64; + const size_t mat = size_t(R * C); + for (size_t base = 0; base + mat <= a.count(); base += mat) { + for (int64_t i = 0; i < R; ++i) { + for (int64_t j = 0; j < C; ++j) { + const bool keep = upper ? (j >= i + k) : (j <= i + k); + if (keep) continue; + if (!r.i64.empty()) r.i64[base + size_t(i * C + j)] = 0; + else r.data[base + size_t(i * C + j)] = 0.0f; + } + } + } + } else if (n.op == "Slice") { + const Tensor& a = any_t(n.ins[0]); + std::vector starts, ends, axes, steps; + auto shape_op = [&](size_t which, std::vector* dst) { + auto ii = inits.find(n.ins[which]); + if (ii == inits.end()) fail(where() + ": slice bounds must be initializers"); + load_init_i64(g, *ii->second, dst); + }; + if (n.ins.size() < 3) fail(where() + ": needs starts/ends"); + shape_op(1, &starts); + shape_op(2, &ends); + if (n.ins.size() > 3) shape_op(3, &axes); + if (n.ins.size() > 4) shape_op(4, &steps); + if (axes.empty()) + for (size_t i = 0; i < starts.size(); ++i) axes.push_back(int64_t(i)); + if (steps.empty()) + for (size_t i = 0; i < starts.size(); ++i) steps.push_back(1); + // Partial-rank slicing: an axis the node does not name is taken whole. Whisper's decoder slices one axis + // at a time (the position embedding, the logits' last position), so a full-rank-only rule refuses it. + std::vector from(a.dims.size(), 0), cnt = a.dims, step(a.dims.size(), 1); + for (size_t i = 0; i < axes.size(); ++i) { + const size_t ax = size_t(axes[i]); + if (ax >= a.dims.size()) fail(where() + ": axis out of range"); + int64_t st = starts[i] < 0 ? starts[i] + a.dims[ax] : starts[i]; + int64_t en = ends[i] < 0 ? ends[i] + a.dims[ax] : ends[i]; + st = std::max(0, std::min(st, a.dims[ax])); + en = std::max(st, std::min(en, a.dims[ax])); + from[ax] = st; + step[ax] = steps[i]; + cnt[ax] = (en - st + steps[i] - 1) / steps[i]; + } + for (size_t ax = 0; ax < step.size(); ++ax) + if (step[ax] != 1) fail(where() + ": only unit steps are implemented"); + r.dims = cnt; + r.data.resize(nelem(cnt)); + r.i64.resize(a.i64.empty() ? 0 : nelem(cnt)); + std::vector stride(a.dims.size(), 1); + for (size_t i = a.dims.size(); i-- > 1;) stride[i - 1] = stride[i] * size_t(a.dims[i]); + std::vector idx(cnt.size(), 0); + for (size_t o = 0; o < nelem(cnt); ++o) { + size_t src = 0; + for (size_t d = 0; d < idx.size(); ++d) src += size_t(from[d] + idx[d]) * stride[d]; + if (!a.i64.empty()) r.i64[o] = a.i64[src]; + else r.data[o] = a.data[src]; + for (size_t d = idx.size(); d-- > 0;) { + if (++idx[d] < cnt[d]) break; + idx[d] = 0; + } + } + } else { + fail(where() + " is not implemented; the pack needs an op this executor does not have"); + } + t[oname] = std::move(r); + } + + const std::string want = out_name.empty() ? (g.outputs.empty() ? std::string() : g.outputs[0]) : out_name; + const auto it = t.find(want); + if (it == t.end()) fail("output '" + want + "' was never produced"); + out->dims = it->second.dims; + out->data = it->second.data; + out->i64 = it->second.i64; + out->b = it->second.b; + return true; + } catch (const std::exception& e) { + if (why != nullptr) *why = e.what(); + return false; + } +} + +} // namespace + +bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, const std::string& out_name, + AudioRunResult* out, std::string* why) { + std::map seed; + Tensor in; + in.dims = {1, mel_bins, frames}; + in.data.assign(mel, mel + size_t(mel_bins) * size_t(frames)); + if (g.inputs.empty()) { + if (why != nullptr) *why = "graph declares no input"; + return false; + } + seed[g.inputs[0]] = std::move(in); + return run_graph_impl(g, std::move(seed), out_name, out, why); +} + +bool run_decoder_graph(const OnnxGraphInfo& g, const int64_t* tokens, int64_t n_tokens, const float* xa, + const std::vector& xa_dims, const std::string& out_name, AudioRunResult* out, + std::string* why) { + if (g.inputs.size() < 1) { + if (why != nullptr) *why = "graph declares no input"; + return false; + } + std::map seed; + Tensor toks; + toks.dims = {1, n_tokens}; + toks.i64.assign(tokens, tokens + size_t(n_tokens)); + seed[g.inputs[0]] = std::move(toks); + if (g.inputs.size() > 1) { + if (xa == nullptr || xa_dims.empty()) { + if (why != nullptr) *why = "the graph has a second input (cross-attention) but none was supplied"; + return false; + } + Tensor x; + x.dims = xa_dims; + x.data.assign(xa, xa + nelem(xa_dims)); + seed[g.inputs[1]] = std::move(x); + } + return run_graph_impl(g, std::move(seed), out_name, out, why); +} + +} // namespace onebit::npu::pack_import diff --git a/npu/audio_exec.h b/npu/audio_exec.h new file mode 100644 index 000000000..04433b4ef --- /dev/null +++ b/npu/audio_exec.h @@ -0,0 +1,43 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec.h -- run an audio pack's graph on the host in f32. Not a general ONNX runtime: it covers exactly the ops a +// Whisper encoder uses (measured on both exports -- 447 nodes / 11 op types for whisper-small, 1031 / 10 for large-v3), and +// it refuses an op it does not implement *by name* rather than producing a plausible activation from a silently skipped node. +// +// The four f32 kernels this stands on are in npu/forward/include/audio_ops.h, which is device-free and known-answer tested; +// this file adds the shape bookkeeping and the wiring. The caller compares its output against the reference, and against +// onnxruntime, before anything here is used to argue a device number. + +#pragma once + +#include +#include + +#include "onnx_graph.h" + +namespace onebit::npu::pack_import { + +struct AudioRunResult { + std::vector dims; // dims of the requested output + std::vector data; // row-major f32 payload + std::vector i64; // int64 payload, when the requested tensor is an index/token tensor + std::vector b; // bool payload, when it is a mask +}; + +// Two entry points, because a Whisper pack has two graphs: the encoder consumes the mel window as f32, the decoder consumes +// a token window as int64 plus the encoder's output. Everything else -- node order, weights, refusals -- is shared. +// +// mel: (mel_bins, frames) row-major f32, the encoder's graph input in the layout its own first Conv consumes. +// out_name: the graph output to return (empty = the graph's declared output). +// Returns true and fills `out`; on failure returns false and sets `why` to the node and the reason. +bool run_audio_graph(const OnnxGraphInfo& g, const float* mel, int64_t mel_bins, int64_t frames, + const std::string& out_name, AudioRunResult* out, std::string* why); + +// The decoder: `tokens` is the int64 token window for the graph's first input; `xa` (the encoder output, row-major f32) is +// bound to the graph's second input when the graph has one. `xa_dims` may be empty when there is no cross-attention input. +bool run_decoder_graph(const OnnxGraphInfo& g, const int64_t* tokens, int64_t n_tokens, const float* xa, + const std::vector& xa_dims, const std::string& out_name, AudioRunResult* out, + std::string* why); + +} // namespace onebit::npu::pack_import diff --git a/npu/forward/include/audio_mel.h b/npu/forward/include/audio_mel.h new file mode 100644 index 000000000..2f7919891 --- /dev/null +++ b/npu/forward/include/audio_mel.h @@ -0,0 +1,126 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_mel.h -- the log-mel front-end: hann-windowed STFT (n_fft 400, hop 160, reflect padding), power spectrum, a +// slaney-normalised mel filterbank, then log10 and the (max - 8, +4)/4 scaling a Whisper encoder's input expects. +// +// This is the one part of the pipeline that is *not* in the pack: the pack's graph starts at the mel window, so the numbers +// here have to match the reference extractor's, and tests/audio_mel_test.cpp checks them against it. +// +// The filterbank is built, not shipped: librosa's slaney mel scale (f_sp = 200/3 Hz, linear below 1 kHz, log above, +// logstep = ln(6.4)/27) with slaney normalisation, which is what Whisper's feature extractor uses. + +#pragma once + +#include +#include +#include + +namespace onebit::audio { + +inline double hz_to_mel(double hz) { + const double f_sp = 200.0 / 3.0; + const double min_log_hz = 1000.0, min_log_mel = min_log_hz / f_sp, logstep = std::log(6.4) / 27.0; + if (hz < min_log_hz) return hz / f_sp; + return min_log_mel + std::log(hz / min_log_hz) / logstep; +} + +inline double mel_to_hz(double mel) { + const double f_sp = 200.0 / 3.0; + const double min_log_hz = 1000.0, min_log_mel = min_log_hz / f_sp, logstep = std::log(6.4) / 27.0; + if (mel < min_log_mel) return mel * f_sp; + return min_log_hz * std::exp(logstep * (mel - min_log_mel)); +} + +// (bins, 1 + n_fft/2): triangular filters, slaney-normalised, spanning 0..nyquist. +inline std::vector mel_filterbank(int bins, int n_fft, int sample_rate) { + const int nfreqs = n_fft / 2 + 1; + const double nyq = sample_rate / 2.0; + const double mlo = hz_to_mel(0.0), mhi = hz_to_mel(nyq); + std::vector f(bins + 2); + for (int i = 0; i < bins + 2; ++i) f[i] = mel_to_hz(mlo + (mhi - mlo) * double(i) / double(bins + 1)); + // FFT bin centres in Hz, and the per-filter bandwidth normalisation. + std::vector hz(nfreqs); + for (int k = 0; k < nfreqs; ++k) hz[k] = double(k) * nyq / double(nfreqs - 1); + std::vector enorm(bins); + for (int i = 0; i < bins; ++i) { + const double w = f[i + 2] - f[i]; + enorm[i] = w > 0 ? 2.0 / w : 0.0; + } + std::vector m(size_t(bins) * size_t(nfreqs), 0.0f); + for (int i = 0; i < bins; ++i) { + const double lo = f[i], ctr = f[i + 1], hi = f[i + 2]; + for (int k = 0; k < nfreqs; ++k) { + const double x = hz[k]; + double v = 0.0; + if (x >= lo && x <= ctr && ctr > lo) v = (x - lo) / (ctr - lo); + else if (x > ctr && x <= hi && hi > ctr) v = (hi - x) / (hi - ctr); + m[size_t(i) * nfreqs + k] = float(v * enorm[i]); + } + } + return m; +} + +struct MelConfig { + int sample_rate = 16000; + int n_fft = 400; + int hop = 160; + int bins = 80; + int frames = 3000; // the encoder's window: 30 s +}; + +// audio: `samples` mono f32 in [-1,1]. out: (bins, frames) row-major. Returns false only on a size the config cannot fill. +inline bool log_mel_spectrogram(const float* audio, size_t samples, const MelConfig& cfg, std::vector* out) { + const int pad = cfg.n_fft / 2; + const int nfreqs = cfg.n_fft / 2 + 1; + if (samples + 2 * size_t(pad) < size_t(cfg.n_fft)) return false; + + // reflect padding, both sides (torch's center=True default for stft) + std::vector x(samples + 2 * size_t(pad)); + for (int i = 0; i < pad; ++i) x[size_t(i)] = audio[pad - i]; + for (size_t i = 0; i < samples; ++i) x[size_t(pad) + i] = audio[i]; + for (int i = 0; i < pad; ++i) x[size_t(pad) + samples + size_t(i)] = audio[samples - 2 - size_t(i)]; + + std::vector win(cfg.n_fft); + for (int i = 0; i < cfg.n_fft; ++i) // periodic hann, as torch.hann_window(w) is by default + win[size_t(i)] = 0.5 - 0.5 * std::cos(2.0 * M_PI * double(i) / double(cfg.n_fft)); + + const std::vector fb = mel_filterbank(cfg.bins, cfg.n_fft, cfg.sample_rate); + out->assign(size_t(cfg.bins) * size_t(cfg.frames), 0.0f); + std::vector re(nfreqs), im(nfreqs), power(nfreqs); + for (int t = 0; t < cfg.frames; ++t) { + const size_t base = size_t(t) * size_t(cfg.hop); + for (int k = 0; k < nfreqs; ++k) { + double sr = 0.0, si = 0.0; + for (int n = 0; n < cfg.n_fft; ++n) { + const double v = x[base + size_t(n)] * win[size_t(n)]; + const double ang = -2.0 * M_PI * double(k) * double(n) / double(cfg.n_fft); + sr += v * std::cos(ang); + si += v * std::sin(ang); + } + re[size_t(k)] = sr; + im[size_t(k)] = si; + power[size_t(k)] = sr * sr + si * si; + } + for (int b = 0; b < cfg.bins; ++b) { + const float* row = fb.data() + size_t(b) * nfreqs; + double acc = 0.0; + for (int k = 0; k < nfreqs; ++k) acc += double(row[k]) * power[size_t(k)]; + out->at(size_t(b) * size_t(cfg.frames) + size_t(t)) = float(acc); + } + } + // log10 with a floor, then Whisper's dynamic-range clamp and scale. + double mx = -1e30; + for (float& v : *out) { + v = float(std::log10(std::max(double(v), 1e-10))); + if (v > mx) mx = v; + } + const double floor_db = mx - 8.0; + for (float& v : *out) { + const double c = std::max(double(v), floor_db); + v = float((c + 4.0) / 4.0); + } + return true; +} + +} // namespace onebit::audio diff --git a/npu/onnx_graph.h b/npu/onnx_graph.h new file mode 100644 index 000000000..db3a8157a --- /dev/null +++ b/npu/onnx_graph.h @@ -0,0 +1,57 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// onnx_graph.h -- a *reader* of an ONNX graph, as opposed to the census walk in npu/pack_import.cpp, which keeps only +// the MatMul-family nodes because that is all a quantised LLM pack needs. An audio pack (the Whisper encoder) is a plain +// f32 graph: 32 transformer blocks of Conv/MatMul/Add/LayerNormalization/Gelu/Softmax/Transpose/Mul, and to execute it the +// engine needs every node in order, with its inputs, outputs and integer attributes, and the set of initializer names. +// +// Nothing here executes anything. It exists so the executor's ordering assumption -- "a node's inputs are graph inputs, +// initializers, or outputs of an earlier node" -- is a *checked* fact about a real pack rather than a hope, and so a pack +// that breaks it is refused by name. + +#pragma once + +#include +#include +#include +#include + +namespace onebit::npu::pack_import { + +struct OnnxNodeInfo { + std::string op; + std::string name; + std::vector ins; + std::vector outs; + std::map attrs; // integer attributes only (axis, perm as a list is not an int attr) + std::map> attrs_list; // `perm`, `pads`, `strides`, `dilations` + std::map attrs_float; // `alpha`, `epsilon` (as stored) +}; + +// One initializer: its shape, dtype and where its bytes are (inline `raw_data`, or a side file via external_data). +struct OnnxInit { + std::string name; + int dtype = 0; // TensorProto.DataType: 1 = FLOAT, 7 = INT64, 10 = FLOAT16, ... + std::vector dims; + const uint8_t* raw = nullptr; // valid only while the mapped file is open; the caller copies what it needs + size_t raw_len = 0; + std::string location; // external_data side file, relative to the graph's directory + uint64_t offset = 0; + uint64_t length = 0; +}; + +struct OnnxGraphInfo { + std::string dir; // directory of the .onnx file, for resolving external_data + std::vector init_tensors; + std::vector inputs; // graph inputs, in declaration order + std::vector outputs; // graph outputs + std::vector nodes; // graph order, which ONNX guarantees is topological + std::vector initializers; // weight/bias names + std::vector external; // initializers whose data is in a side file (location set) +}; + +// Throws std::runtime_error with the path and the protobuf complaint on a malformed file. +OnnxGraphInfo parse_onnx_graph(const std::string& path); + +} // namespace onebit::npu::pack_import diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index 463c49a36..2111ca954 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -15,6 +15,8 @@ #include "pack_import.h" +#include "onnx_graph.h" + #include "file_map.h" #include "pack.h" #include "q4nx.h" @@ -85,6 +87,9 @@ std::vector pb(const uint8_t* p, size_t n) { } case 5: if (size_t(end - p) < 4) throw std::runtime_error("protobuf: truncated fixed32"); + // engine#1783: keep the four bytes in v. A float attribute (LayerNormalization's epsilon) is wire 5, and + // without this every such attribute reads as 0.0 -- indistinguishable from an absent one. + f.v = uint32_t(p[0]) | (uint32_t(p[1]) << 8) | (uint32_t(p[2]) << 16) | (uint32_t(p[3]) << 24); p += 4; break; default: throw std::runtime_error("protobuf: unsupported wire type " + std::to_string(f.wire)); @@ -6310,4 +6315,86 @@ void sweep_fconst_scratch_dirs(const std::vector& dirs, const std:: // box measures are rel_err 0.000000 (serves) and 0.095/0.136 (imports cleanly, answers garbage). bool selfcheck_head_ok(double corr, double rel) { return corr > kTieCorr && rel <= kSelfcheckRelMax; } +// engine#1783: a full graph reader, for packs that are plain f32 graphs (the Whisper encoder) rather than quantised LLM +// packs. The census walk above keeps only the MatMul family by design; this one keeps every node in graph order, which ONNX +// guarantees is topological, plus the initializer names -- so an executor can tell "a weight" from "a graph input", and a +// pack that breaks the ordering is refused by name instead of producing plausible activations. +OnnxGraphInfo parse_onnx_graph(const std::string& path) { + size_t size = 0; + std::string why; + const uint8_t* base = map_file(path, size, &why); + if (base == nullptr) throw std::runtime_error(path + ": " + why); + const auto model = pb(base, size); // ModelProto + const auto* gfield = field1(model, 7); // ModelProto.graph + if (gfield == nullptr) throw std::runtime_error(path + ": no ModelProto.graph (field 7)"); + const auto g = pb(gfield->p, gfield->n); + + OnnxGraphInfo out; + const size_t slash = path.find_last_of('/'); + out.dir = (slash == std::string::npos) ? std::string(".") : path.substr(0, slash); + for (const auto* f : fields(g, 11)) { // graph.input + const auto vi = pb(f->p, f->n); + if (const auto* nm = field1(vi, 1)) out.inputs.push_back(pb_str(*nm)); + } + for (const auto* f : fields(g, 12)) { // graph.output + const auto vi = pb(f->p, f->n); + if (const auto* nm = field1(vi, 1)) out.outputs.push_back(pb_str(*nm)); + } + for (const auto* f : fields(g, 5)) { // graph.initializer + const auto t = pb(f->p, f->n); + const auto* nm = field1(t, 8); + if (nm == nullptr) continue; + out.initializers.push_back(pb_str(*nm)); + OnnxInit ti; + ti.name = pb_str(*nm); + if (const auto* dt = field1(t, 2)) ti.dtype = int(dt->v); + for (const auto* d : fields(t, 1)) ti.dims.push_back(int64_t(d->v)); + if (const auto* rd = field1(t, 9)) { + ti.raw = rd->p; + ti.raw_len = rd->n; + } + for (const auto* d : fields(t, 13)) { // TensorProto.external_data: key/value string pairs + const auto kv = pb(d->p, d->n); + const auto* k = field1(kv, 1); + const auto* v = field1(kv, 2); + if (k == nullptr || v == nullptr) continue; + const std::string key = pb_str(*k), val = pb_str(*v); + if (key == "location") ti.location = val; + else if (key == "offset") ti.offset = std::stoull(val); + else if (key == "length") ti.length = std::stoull(val); + } + if (!ti.location.empty()) out.external.push_back(ti.name); + out.init_tensors.push_back(std::move(ti)); + } + for (const auto* f : fields(g, 1)) { // graph.node + const auto node = pb(f->p, f->n); + OnnxNodeInfo ni; + if (const auto* o = field1(node, 4)) ni.op = pb_str(*o); + if (const auto* nm = field1(node, 3)) ni.name = pb_str(*nm); + for (const auto* i : fields(node, 1)) ni.ins.push_back(pb_str(*i)); + for (const auto* o : fields(node, 2)) ni.outs.push_back(pb_str(*o)); + for (const auto* a : fields(node, 5)) { // AttributeProto + const auto at = pb(a->p, a->n); + const auto* nm = field1(at, 1); + if (nm == nullptr) continue; + const std::string key = pb_str(*nm); + if (const auto* iv = field1(at, 3)) ni.attrs[key] = int64_t(iv->v); // INT + if (const auto* fv = field1(at, 2)) { // FLOAT (wire 5, see fields()) + const uint32_t bits = uint32_t(fv->v); + float f32 = 0.0f; + std::memcpy(&f32, &bits, sizeof f32); + ni.attrs_float[key] = f32; + } + std::vector list; + for (const auto* iv : fields(at, 8)) list.push_back(int64_t(iv->v)); // INTS + if (!list.empty()) ni.attrs_list[key] = std::move(list); + } + out.nodes.push_back(std::move(ni)); + } + // NOT unmapped: init_tensors[].raw points into this mapping, and the executor copies what it needs while it is alive. + // shortcut: the mapping is released when the process exits; a long-lived caller that parses many graphs should own the + // mapping itself, which is the refactor to make when a second consumer appears. + return out; +} + } // namespace onebit::npu::pack_import diff --git a/tests/audio_exec_test.cpp b/tests/audio_exec_test.cpp new file mode 100644 index 000000000..1a63a6ee0 --- /dev/null +++ b/tests/audio_exec_test.cpp @@ -0,0 +1,162 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_exec_test.cpp -- run an audio pack's encoder through the engine's own host executor and compare it against +// onnxruntime, on an input both sides can build from a formula (no binary fixture in the repo). +// +// The mel is mel[b,t] = sin(0.013*(b+1)*(t+1))*0.5 + 0.25 for 80 bins and 3000 frames, computed in double on both sides. +// Expected values below are onnxruntime's, on whisper-small's encoder_model.onnx. Skips (77) without the pack. + +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "onnx_graph.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(double(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.7g want %.7g (|d| %.3g)\n", what.c_str(), double(got), want, d); + ++g_failures; + } +} + +} // namespace + +int main(int argc, char** argv) { + const char* path = (argc > 1) ? argv[1] : std::getenv("ONEBIT_TEST_WHISPER"); + if (path == nullptr || *path == '\0') { + std::printf("SKIP: pass the encoder .onnx, or set ONEBIT_TEST_WHISPER to it\n"); + return 77; + } + using namespace onebit::npu::pack_import; + OnnxGraphInfo g; + try { + g = parse_onnx_graph(path); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", path, e.what()); + return 1; + } + + // The mel-bin count is the front-end Conv's input-channel count: 80 for whisper-small, 128 for large-v3. Read from the + // pack rather than hard-coded, so the same test runs on either (the numeric comparison below is small's). + int bins = 80; + for (const auto& ti : g.init_tensors) + if (ti.dims.size() == 3 && ti.dims[0] == 768) { bins = int(ti.dims[1]); break; } + for (const auto& n : g.nodes) + if (n.op == "Conv") { + for (const auto& ti : g.init_tensors) + if (ti.name == n.ins[1] && ti.dims.size() == 3) bins = int(ti.dims[1]); + break; + } + const int frames = 3000; + std::vector mel(size_t(bins) * frames); + for (int b = 0; b < bins; ++b) + for (int t = 0; t < frames; ++t) + mel[size_t(b) * frames + t] = float(std::sin(0.013 * (b + 1) * (t + 1)) * 0.5 + 0.25); + + AudioRunResult r; + std::string why; + if (!run_audio_graph(g, mel.data(), bins, frames, "", &r, &why)) { + std::printf("FAIL run: %s\n", why.c_str()); + return 1; + } + std::printf("ran %zu nodes -> dims", g.nodes.size()); + for (int64_t d : r.dims) std::printf(" %lld", (long long)d); + std::printf("\n"); + check(r.dims.size() == 3 && r.dims[0] == 1 && r.dims[1] == 1500, "encoder output is [1,1500,hidden]"); + if (r.dims.size() != 3 || r.dims[2] != 768) { + // A larger encoder: the point of this run is that the graph executed (including external_data weights), not the + // small model's numbers. + std::printf("audio exec: ran a %d-bin encoder of hidden %lld (numeric comparison is whisper-small's only)\n", bins, + (long long)(r.dims.size() == 3 ? r.dims[2] : 0)); + return g_failures == 0 ? 0 : 1; + } + + const int64_t H = 768; + auto at = [&](int64_t t, int64_t h) { return r.data[size_t(t) * H + size_t(h)]; }; + double mean = 0.0, var = 0.0; + for (float v : r.data) mean += v; + mean /= double(r.data.size()); + for (float v : r.data) var += (double(v) - mean) * (double(v) - mean); + var /= double(r.data.size()); + std::printf("mean %.9f std %.9f\n", mean, std::sqrt(var)); + + // onnxruntime, same input, same pack: mean -0.009045244, std 1.678493023 + expect_close(float(mean), -0.009045244, 2e-4, "output mean"); + expect_close(float(std::sqrt(var)), 1.678493023, 2e-4, "output std"); + + // frame 0, 750 and 1499 -- the first frame, the middle, and the last (whose window is padding-heavy) + const double p0[5] = {-0.5922473669, -1.6559172869, 0.0451035500, -0.9414483309, -1.6385756731}; + const double p1[4] = {0.1923044026, -0.0132405162, 0.4830040932, 0.6385290623}; + const double p2[4] = {-2.1586399078, 0.3378775120, 0.8375968933, 0.5081450939}; + for (int i = 0; i < 5; ++i) expect_close(at(0, i), p0[i], 2e-3, "frame 0 col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(at(750, i), p1[i], 2e-3, "frame 750 col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(at(1499, i), p2[i], 2e-3, "frame 1499 col " + std::to_string(i)); + + if (g_failures != 0) { + std::printf("audio exec: %d failures\n", g_failures); + return 1; + } + std::printf("audio exec: whole encoder on the host matches onnxruntime\n"); + + // ---- the decoder, on the same pack: 128-position token window x, the encoder output xa ---------------------------- + const char* dpath = (argc > 2) ? argv[2] : std::getenv("ONEBIT_TEST_WHISPER_DECODER"); + if (dpath == nullptr || *dpath == '\0') { + std::printf("SKIP decoder: pass the decoder .onnx as the second argument\n"); + return 0; + } + OnnxGraphInfo dg; + try { + dg = parse_onnx_graph(dpath); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", dpath, e.what()); + return 1; + } + const int ntok = 128; + std::vector toks(size_t(ntok), 0); + toks[0] = 50258; // <|startoftranscript|> + toks[1] = 50359; // <|en|> + toks[2] = 50257; // <|transcribe|> + toks[3] = 1234; + std::vector xad; + for (int64_t d : r.dims) xad.push_back(d); + AudioRunResult dr; + if (!run_decoder_graph(dg, toks.data(), ntok, r.data.data(), xad, "", &dr, &why)) { + std::printf("FAIL decoder run: %s\n", why.c_str()); + return 1; + } + std::printf("decoder ran %zu nodes -> dims", dg.nodes.size()); + for (int64_t d : dr.dims) std::printf(" %lld", (long long)d); + std::printf("\n"); + check(dr.dims.size() == 3 && dr.dims[0] == 1 && dr.dims[1] == ntok, "decoder output is [1,128,vocab]"); + if (dr.dims.size() == 3) { + const int64_t V = dr.dims[2]; + auto dl = [&](int64_t t, int64_t v) { return dr.data[size_t(t) * size_t(V) + size_t(v)]; }; + // onnxruntime, same chain (formula mel -> encoder -> xa -> decoder): the decoder's last row and row 3. + const double last[5] = {8.26411438, -1.29840982, -2.36972690, -3.26691365, -2.12733102}; + const double row3[4] = {18.45385361, 17.14962196, 12.60445786, 10.96115971}; + for (int i = 0; i < 5; ++i) expect_close(dl(ntok - 1, i), last[i], 2e-3, "decoder last row col " + std::to_string(i)); + for (int i = 0; i < 4; ++i) expect_close(dl(3, i), row3[i], 2e-3, "decoder row 3 col " + std::to_string(i)); + } + if (g_failures != 0) { + std::printf("audio exec: %d failures\n", g_failures); + return 1; + } + std::printf("audio exec: decoder graph runs on the host and matches onnxruntime\n"); + return 0; +} diff --git a/tests/audio_mel_test.cpp b/tests/audio_mel_test.cpp new file mode 100644 index 000000000..70d7e823b --- /dev/null +++ b/tests/audio_mel_test.cpp @@ -0,0 +1,97 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// audio_mel_test.cpp -- the log-mel front-end against the reference extractor's numbers. +// +// Input: 30 s at 16 kHz of three summed sinusoids (440, 1250, 3100 Hz), built from a formula so both sides use the same +// signal. Expected values are transformers' WhisperFeatureExtractor (feature_size=80) on that signal. + +#include +#include +#include +#include + +#include "audio_mel.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(double(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.9g want %.9g (|d| %.3g)\n", what.c_str(), double(got), want, d); + ++g_failures; + } +} + +} // namespace + +int main() { + onebit::audio::MelConfig cfg; + const int sr = cfg.sample_rate; + const size_t n = size_t(sr) * 30; + std::vector sig(n); + for (size_t i = 0; i < n; ++i) { + const double t = double(i) / sr; + sig[i] = float(0.3 * std::sin(2 * M_PI * 440.0 * t) + 0.2 * std::sin(2 * M_PI * 1250.0 * t) + + 0.1 * std::sin(2 * M_PI * 3100.0 * t)); + } + + std::vector mel; + if (!onebit::audio::log_mel_spectrogram(sig.data(), sig.size(), cfg, &mel)) { + std::printf("FAIL front-end refused a 30 s window\n"); + return 1; + } + check(mel.size() == size_t(cfg.bins) * size_t(cfg.frames), "mel is (80, 3000)"); + if (mel.size() != size_t(cfg.bins) * size_t(cfg.frames)) return 1; + + double mean = 0.0, mn = 1e30, mx = -1e30; + for (float v : mel) { + mean += v; + mn = std::min(mn, v); + mx = std::max(mx, v); + } + mean /= double(mel.size()); + auto at = [&](int b, int t) { return mel[size_t(b) * size_t(cfg.frames) + size_t(t)]; }; + + // reference: mean -0.33795589, min -0.672718, max 1.327282 + expect_close(float(mean), -0.33795589, 2e-4, "mel mean"); + expect_close(float(mn), -0.672718, 2e-4, "mel min"); + expect_close(float(mx), 1.327282, 2e-4, "mel max"); + + // bin 0, first five frames -- the onset and the dynamic-range clamp (three equal -0.6727176 values) + const double p0[5] = {0.92388016, 0.41213650, -0.67271757, -0.67271757, -0.67271757}; + for (int i = 0; i < 5; ++i) expect_close(at(0, i), p0[i], 2e-3, "bin 0 frame " + std::to_string(i)); + // mid band, the end of the window + const double p1[5] = {-0.65874779, -0.66997766, -0.65873981, -0.66998839, -0.65872025}; + for (int i = 0; i < 5; ++i) expect_close(at(40, 1499 + i), p1[i], 2e-3, "bin 40 frame " + std::to_string(1499 + i)); + // top bin + const double p2[3] = {0.18210298, -0.32799470, -0.67271757}; + for (int i = 0; i < 3; ++i) expect_close(at(79, i), p2[i], 2e-3, "bin 79 frame " + std::to_string(i)); + + // A silent window still produces a full, finite spectrogram: the log floor and the max-8 clamp own that case, and a + // NaN there would poison every frame the encoder sees. + std::vector silent(n, 0.0f), smel; + check(onebit::audio::log_mel_spectrogram(silent.data(), silent.size(), cfg, &smel), "silence is accepted"); + bool finite = true; + for (float v : smel) finite = finite && std::isfinite(v); + check(finite, "a silent window produces finite numbers"); + // Silence is the floor of the whole chain: power 0 -> log10(1e-10) = -10, the max is that too, so the max-8 clamp never + // bites and the scale gives (-10+4)/4 = -1.5. (My first expectation here said 1.0 -- the test caught its author again.) + expect_close(smel[0], -1.5f, 1e-6, "silence lands on the log floor: (-10+4)/4 == -1.5"); + + if (g_failures != 0) { + std::printf("audio mel: %d failures\n", g_failures); + return 1; + } + std::printf("audio mel: front-end matches the reference extractor on tones, silences on the clamp\n"); + return 0; +} diff --git a/tests/onnx_graph_test.cpp b/tests/onnx_graph_test.cpp new file mode 100644 index 000000000..c2a2724a4 --- /dev/null +++ b/tests/onnx_graph_test.cpp @@ -0,0 +1,120 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// onnx_graph_test.cpp -- the audio pack's graph, read through the engine and checked for the property the executor will +// depend on: every node's inputs are graph inputs, initializers, or outputs of an *earlier* node. Skips (77) without +// ONEBIT_TEST_WHISPER naming the encoder's .onnx. + +#include +#include +#include +#include +#include +#include + +#include "onnx_graph.h" + +namespace { + +int g_failures = 0; + +void check(bool cond, const std::string& what) { + if (!cond) { + std::printf("FAIL %s\n", what.c_str()); + ++g_failures; + } +} + +} // namespace + +int main(int argc, char** argv) { + const char* path = (argc > 1) ? argv[1] : std::getenv("ONEBIT_TEST_WHISPER"); + if (path == nullptr || *path == '\0') { + std::printf("SKIP: pass the encoder .onnx, or set ONEBIT_TEST_WHISPER to it\n"); + return 77; + } + onebit::npu::pack_import::OnnxGraphInfo g; + try { + g = onebit::npu::pack_import::parse_onnx_graph(path); + } catch (const std::exception& e) { + std::printf("FAIL cannot parse %s: %s\n", path, e.what()); + return 1; + } + + check(!g.nodes.empty(), "graph has nodes"); + check(g.inputs.size() == 1, "one graph input (the mel window)"); + check(!g.outputs.empty(), "graph declares an output"); + check(!g.initializers.empty(), "graph has initializers"); + + // The ordering assumption, checked against the real pack rather than assumed. + std::set known(g.inputs.begin(), g.inputs.end()); + known.insert(g.initializers.begin(), g.initializers.end()); + int unnamed_inputs = 0; + for (const auto& n : g.nodes) { + for (const auto& in : n.ins) { + if (in.empty()) continue; // optional input, ONNX spells it "" + if (!known.count(in)) { + if (unnamed_inputs < 5) std::printf("FAIL %s input '%s' is not produced earlier\n", n.op.c_str(), in.c_str()); + ++unnamed_inputs; + } + } + for (const auto& o : n.outs) known.insert(o); + } + check(unnamed_inputs == 0, "every node input is a graph input, an initializer, or an earlier node's output"); + + // The op census the reference implementation was written against (tools/whisper_encoder_ref.py): the front-end is + // exactly two Convs, and a transformer block is MatMul/LayerNormalization/Gelu dominated. + std::map census; + for (const auto& n : g.nodes) ++census[n.op]; + std::printf("nodes %zu, initializers %zu (external %zu), graph input '%s', output '%s'\n", g.nodes.size(), + g.initializers.size(), g.external.size(), g.inputs.empty() ? "" : g.inputs[0].c_str(), + g.outputs.empty() ? "" : g.outputs[0].c_str()); + for (const auto& kv : census) std::printf(" %-24s %d\n", kv.first.c_str(), kv.second); + + // The structure, read off the graph rather than assumed from the reference. Measured on whisper-small's encoder: 447 + // nodes, and three facts that changed how the executor has to be written -- + // * the GELU is NOT a `Gelu` node: it arrives decomposed as Div -> Erf -> Add -> Mul -> Mul, so an executor keyed on + // op names alone needs `Erf` and the surrounding arithmetic (one Erf per block plus two in the front-end); + // * attention is explicit (Transpose/MatMul/Softmax), with 8 MatMuls per layer: q,k,v,o, fc1, fc2, qk, and the + // scores-times-values product; + // * the weights arrive transposed (61 Transpose nodes), which is why the projection shapes in the graph are (K, N). + // The counts scale with depth, so they are checked as ratios against the Softmax count (= one per layer), not as + // absolute constants -- the same test then holds for whisper-large-v3's 32-layer encoder. + check(census["Conv"] == 2, "two Convs (the mel front-end: stride 1 then stride 2, both pad 1)"); + const int layers = census["Softmax"]; + check(layers > 0, "attention scores go through a Softmax"); + if (layers > 0) { + check(census["MatMul"] == 8 * layers, "8 MatMuls per layer (q,k,v,o, fc1, fc2, qk, scores-times-values)"); + check(census["LayerNormalization"] == 2 * layers + 1, "two layer norms per layer plus the encoder's final one"); + // Two exports of this family spell the same activation two ways: whisper-small's encoder_Model decomposes it into + // Div -> Erf -> Add -> Mul -> Mul (0 Gelu nodes, 14 Erf for 12 layers), while whisper-large-v3's carries a real + // `Gelu` op (34 for 32 layers). An executor keyed on op names must accept BOTH; the count, one per layer plus two + // in the front-end, is the invariant that holds either way. + check(census["Erf"] + census["Gelu"] == layers + 2, + "the activation, however spelled, appears once per layer plus twice in the front-end"); + check(census["Erf"] == 0 || census["Gelu"] == 0, "an export uses one spelling of the activation, not a mix"); + } + check(census["Transpose"] > layers, "the weights arrive transposed, so the graph's MatMuls are (K, N)"); + + // One attribute the executor must not guess: LayerNormalization's epsilon is a float attribute, and a reader that + // drops fixed32 (as the census walk does) sees every one of them as absent. + const onebit::npu::pack_import::OnnxNodeInfo* ln = nullptr; + for (const auto& n : g.nodes) + if (n.op == "LayerNormalization") { + ln = &n; + break; + } + if (ln != nullptr) { + const auto it = ln->attrs_float.find("epsilon"); + check(it != ln->attrs_float.end(), "LayerNormalization carries epsilon as a float attribute"); + if (it != ln->attrs_float.end()) check(it->second > 0.0f && it->second < 1e-3f, "epsilon is in the usual range"); + } + + if (g_failures != 0) { + std::printf("onnx graph: %d failures\n", g_failures); + return 1; + } + std::printf("onnx graph: topological order, op census (%d layers, activation spelled %s) and layer-norm epsilon all hold\n", + census["Softmax"], census["Gelu"] ? "Gelu" : "Div/Erf/Add/Mul/Mul"); + return 0; +}