diff --git a/CMakeLists.txt b/CMakeLists.txt index aec5e5659..ec80f2751 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -804,7 +804,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) @@ -1378,6 +1378,28 @@ endif() # q4nx_pack.c is listed here rather than reached through onebit_npu: the integer host arm's checks # (engine#1808) call q4nx_host_gemm_i8*() directly, and those live in that file. add_executable(npu_pack_test tests/npu_pack_test.cpp npu/forward/src/q4nx_pack.c npu/forward/src/q4nx_dequant.c) +# The Whisper encoder front-end's four f32 host kernels: no container, no device, so it runs wherever this file is built +# (the same reason #1830/#1863 put their arms in a binary CI already compiles). Values come from the encoder reference. +add_executable(audio_ops_test tests/audio_ops_test.cpp) +target_include_directories(audio_ops_test PRIVATE npu/forward/include) +# The audio pack's graph, read through the engine: pins the topological-order assumption the executor needs and the op +# census the host reference was written against. Skips (77) without a path. +add_executable(onnx_graph_test tests/onnx_graph_test.cpp) +target_link_libraries(onnx_graph_test PRIVATE onebit_npu) +add_executable(whisper_transcribe tools/whisper_transcribe.cpp) +target_include_directories(whisper_transcribe PRIVATE npu/forward/include) +target_link_libraries(whisper_transcribe PRIVATE onebit_npu) +add_executable(audio_mel_test tests/audio_mel_test.cpp) +target_include_directories(audio_mel_test PRIVATE npu/forward/include) +add_executable(audio_exec_test tests/audio_exec_test.cpp) +target_link_libraries(audio_exec_test PRIVATE onebit_npu) +# The executor's fork-join (npu/audio_parallel.h): MatMul bit-identical across ONEBIT_AUDIO_THREADS, and a throwing unit +# is a caught exception rather than std::terminate. The graph is built in memory, so: no pack, no device. +add_executable(audio_parallel_test tests/audio_parallel_test.cpp) +target_link_libraries(audio_parallel_test PRIVATE onebit_npu Threads::Threads) +# ci.yml's build list is not ours to edit without `workflow` scope, so these reach CI the way ops_dump_test does: as build +# dependencies of a target it already builds (tools/ci/ci_targets_check.py follows this edge). +add_dependencies(onebit audio_ops_test audio_mel_test onnx_graph_test audio_exec_test audio_parallel_test) # A DFlash draft package as a model source (engine#1780): config + container -> a checked model whose tensors decode # one at a time. The loader is compiled here and not into onebit_npu because it needs the q4_1 tile decoder, which # lives with the forward (npu/forward/src/q4nx_dequant.c); it moves there with the draft's forward. Synthetic package, @@ -1729,6 +1751,14 @@ endif() if(ONEBIT_NPU_MODEL AND ONEBIT_NPU_BO_DUMP) add_test(NAME npu_pack_model COMMAND npu_pack_test ${ONEBIT_NPU_MODEL} ${ONEBIT_NPU_BO_DUMP}) endif() +add_test(NAME audio_ops COMMAND audio_ops_test) +add_test(NAME audio_mel COMMAND audio_mel_test) +set_tests_properties(audio_ops PROPERTIES SKIP_RETURN_CODE 77) +add_test(NAME onnx_graph COMMAND onnx_graph_test ${ONEBIT_AUDIO_ONNX}) +add_test(NAME audio_exec COMMAND audio_exec_test ${ONEBIT_AUDIO_ONNX}) +add_test(NAME audio_parallel COMMAND audio_parallel_test) +set_tests_properties(audio_exec PROPERTIES SKIP_RETURN_CODE 77) +set_tests_properties(onnx_graph PROPERTIES SKIP_RETURN_CODE 77) # The private NPU route registry (npu/private_route.h), without an add-on. add_executable(npu_private_route_test tests/npu_private_route_test.cpp) target_link_libraries(npu_private_route_test PRIVATE onebit_npu) diff --git a/npu/audio_exec.cpp b/npu/audio_exec.cpp new file mode 100644 index 000000000..edb394ad7 --- /dev/null +++ b/npu/audio_exec.cpp @@ -0,0 +1,587 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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 "audio_parallel.h" + +#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; +} + +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..df7e099ce --- /dev/null +++ b/npu/audio_exec.h @@ -0,0 +1,55 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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/audio_parallel.h b/npu/audio_parallel.h new file mode 100644 index 000000000..30ebfc05c --- /dev/null +++ b/npu/audio_parallel.h @@ -0,0 +1,81 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// audio_parallel.h -- the fork-join the host audio executor splits independent work over (engine#1931). +// +// Deterministic by construction: unit u always runs fn(u) exactly once, and a unit writes only its own slice of the +// output, so every element's arithmetic is unchanged by the thread count and the result is bit-identical to the serial +// loop. That is the property a reference implementation has to keep, and it is asserted, not assumed: +// tests/audio_parallel_test.cpp compares the executor's MatMul byte for byte across thread counts. +// +// A unit that throws does not take the process down. A std::thread whose function lets an exception escape calls +// std::terminate, so each worker catches, every worker is joined, and the first exception (lowest worker index, which +// makes the choice repeatable) is rethrown on the calling thread, where the executor turns it into a refusal with a +// reason like every other failure. +// +// Header-only and free of the executor's types, so the test needs no pack and no device. + +#pragma once + +#include +#include +#include +#include +#include + +namespace onebit::npu::pack_import { + +// How many threads the audio executor may use. ONEBIT_AUDIO_THREADS=N sets it; anything that is not a positive number +// means 1 (serial, the control for measuring the split), so a typo cannot turn into "as many as the parse produced". +// Unset: the hardware's count, at most 8. +inline size_t audio_thread_count() { + if (const char* env = std::getenv("ONEBIT_AUDIO_THREADS")) { + const long n = std::strtol(env, nullptr, 10); + return n > 0 ? size_t(n) : 1; + } + const size_t hw = std::thread::hardware_concurrency(); + return hw == 0 ? 1 : (hw > 8 ? 8 : hw); +} + +// Run fn(0) .. fn(units-1), each exactly once, on at most `want` threads. Worker t takes units t, t+want, t+2*want, ... +template +void parallel_units(size_t units, size_t want, Fn&& fn) { + if (want > units) want = units; + if (want <= 1) { + for (size_t u = 0; u < units; ++u) fn(u); + return; + } + std::vector failed(want); + std::vector pool; + pool.reserve(want); + for (size_t t = 0; t < want; ++t) + pool.emplace_back([&, t]() { + try { + for (size_t u = t; u < units; u += want) fn(u); + } catch (...) { + failed[t] = std::current_exception(); + } + }); + for (auto& th : pool) th.join(); + for (const auto& ep : failed) + if (ep) std::rethrow_exception(ep); +} + +template +void parallel_units(size_t units, Fn&& fn) { + parallel_units(units, audio_thread_count(), static_cast(fn)); +} + +} // 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..bb0b2d8f8 --- /dev/null +++ b/npu/forward/include/audio_mel.h @@ -0,0 +1,138 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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/forward/include/audio_ops.h b/npu/forward/include/audio_ops.h new file mode 100644 index 000000000..8fd3dfef6 --- /dev/null +++ b/npu/forward/include/audio_ops.h @@ -0,0 +1,98 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// audio_ops.h -- the four f32 host kernels a Whisper encoder front-end is made of: a temporal conv with stride and zero +// padding, layer norm, erf-GELU and a row softmax. Header-only, no dependencies. +// +// Why these and not a graph executor: the encoder's arithmetic that is easy to get subtly wrong is here -- the mel +// front-end is two strided 3-tap convs ([1280,128,3] stride 1 then [1280,1280,3] stride 2, both pad 1), and the GELU is +// the *exact* erf form, which numpy does not have (the reference builds it with A&S 7.1.26, |err| < 1.5e-7). Every number +// in tests/audio_ops_test.cpp was computed in double from the same formulas the reference uses, so a change here that +// moves a value is visible before it reaches the device. + +#pragma once + +#include +#include + +namespace onebit::audio { + +// (T + 2*pad - k) / stride + 1, floor -- the ONNX Conv default for a valid (non-SAME) convolution. +inline int conv1d_out_len(int T, int k, int stride, int pad) { + const int span = T + 2 * pad - k; + return span < 0 ? 0 : span / stride + 1; +} + +// x: (cin, T) row-major. w: (cout, cin, k) row-major -- the layout the ONNX initializer carries, so a caller reading +// weights through the pack's own loader needs no transpose. b: (cout,) or null. out: (cout, tout) row-major. +inline void conv1d(const float* x, int cin, int T, const float* w, const float* b, int cout, int k, int stride, int pad, + float* out) { + const int tout = conv1d_out_len(T, k, stride, pad); + for (int co = 0; co < cout; ++co) { + for (int t = 0; t < tout; ++t) { + float acc = (b != nullptr) ? b[co] : 0.0f; + for (int ci = 0; ci < cin; ++ci) { + const float* xr = x + static_cast(ci) * T; + const float* wr = w + (static_cast(co) * cin + ci) * k; + for (int j = 0; j < k; ++j) { + const int src = t * stride + j - pad; + if (src >= 0 && src < T) acc += wr[j] * xr[src]; + } + } + out[static_cast(co) * tout + t] = acc; + } + } +} + +// y = (x - mean) / sqrt(var + eps) * gamma + beta, biased variance (divide by n), one row. +inline void layer_norm(const float* x, const float* gamma, const float* beta, int n, float eps, float* out) { + double sum = 0.0; + for (int i = 0; i < n; ++i) sum += x[i]; + const double mean = sum / n; + double var = 0.0; + for (int i = 0; i < n; ++i) { + const double d = x[i] - mean; + var += d * d; + } + var /= n; + const double inv = 1.0 / std::sqrt(var + eps); + for (int i = 0; i < n; ++i) { + const float g = (gamma != nullptr) ? gamma[i] : 1.0f; + const float b = (beta != nullptr) ? beta[i] : 0.0f; + out[i] = static_cast((x[i] - mean) * inv) * g + b; + } +} + +// The exact form: 0.5*x*(1+erf(x/sqrt(2))). `approximate=none` in the export, so tanh is not a drop-in. +inline float gelu_erf(float x) { + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752f)); +} + +// In place, one row: numerically stable (subtract the row max) so a logit of 100 does not overflow. +inline void softmax_row(float* x, int n) { + if (n <= 0) return; + float mx = x[0]; + for (int i = 1; i < n; ++i) + if (x[i] > mx) mx = x[i]; + double sum = 0.0; + for (int i = 0; i < n; ++i) { + x[i] = std::exp(x[i] - mx); + sum += x[i]; + } + const double inv = 1.0 / sum; + for (int i = 0; i < n; ++i) x[i] = static_cast(x[i] * inv); +} + +} // namespace onebit::audio diff --git a/npu/onnx_graph.h b/npu/onnx_graph.h new file mode 100644 index 000000000..87acc99cd --- /dev/null +++ b/npu/onnx_graph.h @@ -0,0 +1,69 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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 165d7ee39..25aeeafb5 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#1931: 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)); @@ -6519,6 +6524,88 @@ 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#1931: 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; +} + const char* known_bad_architecture(const std::string& model_type) { if (model_type == "gemma2") return "gemma-2-2b hybrid packs import cleanly and answer garbage on this route (`myselfGetBytes...`, two packs, " diff --git a/tests/audio_exec_test.cpp b/tests/audio_exec_test.cpp new file mode 100644 index 000000000..0b923244a --- /dev/null +++ b/tests/audio_exec_test.cpp @@ -0,0 +1,174 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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..e7c368443 --- /dev/null +++ b/tests/audio_mel_test.cpp @@ -0,0 +1,109 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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/audio_ops_test.cpp b/tests/audio_ops_test.cpp new file mode 100644 index 000000000..9107a8c53 --- /dev/null +++ b/tests/audio_ops_test.cpp @@ -0,0 +1,149 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// audio_ops_test.cpp -- known-answer checks for npu/forward/include/audio_ops.h. No container, no device, no pack: every +// expected value was computed in double in Python from the same formulas the Whisper encoder reference uses +// (tools/whisper_encoder_ref.py), so this binary runs wherever CI builds, which is the point. + +#include +#include +#include +#include + +#include "audio_ops.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_eq(int got, int want, const std::string& what) { + if (got != want) { + std::printf("FAIL %s: got %d want %d\n", what.c_str(), got, want); + ++g_failures; + } +} + +void expect_close(float got, double want, double tol, const std::string& what) { + const double d = std::fabs(static_cast(got) - want); + if (!(d <= tol)) { + std::printf("FAIL %s: got %.10g want %.10g (|d| %.3g > %.3g)\n", what.c_str(), static_cast(got), want, d, + tol); + ++g_failures; + } +} + +// 1. The front-end conv, at the shapes the encoder actually uses: two strided 3-tap convs, pad 1 on both ends. +void conv1d_arms() { + // cin 2, T 5, cout 3, k 3, stride 2, pad 1 -> tout 3 + const std::vector x = {0.5f, -1.0f, 0.25f, 2.0f, -0.75f, // channel 0 + 1.5f, 0.0f, -0.5f, 1.0f, 0.25f}; // channel 1 + std::vector w(3 * 2 * 3); + for (int i = 0; i < 18; ++i) w[i] = static_cast((i + 1) * 0.0625); + const std::vector b = {0.125f, -0.25f, 0.5f}; + + expect_eq(onebit::audio::conv1d_out_len(5, 3, 2, 1), 3, "conv1d_out_len T5 k3 s2 p1"); + expect_eq(onebit::audio::conv1d_out_len(3000, 3, 2, 1), 1500, "conv1d_out_len 3000 -> 1500"); + + std::vector out(3 * 3); + onebit::audio::conv1d(x.data(), 2, 5, w.data(), b.data(), 3, 3, 2, 1, out.data()); + + // double, same formulas: [0.46875, 0.6875, 0.484375, 0.46875, 0.96875, 1.046875, 1.59375, 2.375, 2.734375] + const double want[9] = {0.46875, 0.6875, 0.484375, 0.46875, 0.96875, 1.046875, 1.59375, 2.375, 2.734375}; + for (int i = 0; i < 9; ++i) expect_close(out[i], want[i], 1e-6, "conv1d value " + std::to_string(i)); + + // edge case: stride 1, pad 0, k 3 -- the second conv's arithmetic without the padding, and a 1-tap full-window + // where tout * stride + k - pad == T exactly (no tail row may be emitted past the input) + std::vector out2(2); + onebit::audio::conv1d(x.data(), 1, 5, w.data(), nullptr, 1, 3, 1, 0, out2.data()); + expect_close(out2[0], 0.03076171875 * 0.0 + 0.0625 * 0.5 + 0.125 * -1.0 + 0.1875 * 0.25, 1e-6, "conv1d no bias"); + expect_eq(onebit::audio::conv1d_out_len(5, 3, 1, 0), 3, "conv1d_out_len T5 k3 s1 p0"); +} + +// 2. Layer norm: biased variance, eps added under the root, gamma/beta applied after the normalise. +void layernorm_arm() { + const std::vector v = {0.25f, -1.5f, 2.0f, 0.75f}; + const std::vector g = {1.0f, 0.5f, -1.0f, 2.0f}; + const std::vector be = {0.0f, 1.0f, 0.25f, -0.5f}; + std::vector out(4); + onebit::audio::layer_norm(v.data(), g.data(), be.data(), 4, 1e-5f, out.data()); + // double: [-0.0995034038, 0.2537244718, -1.0435442489, 0.0970204226] + const double want[4] = {-0.0995034038, 0.2537244718, -1.0435442489, 0.0970204226}; + for (int i = 0; i < 4; ++i) expect_close(out[i], want[i], 1e-6, "layer_norm value " + std::to_string(i)); + + // A constant row normalises to exactly beta whatever eps is: the variance is 0, so the numerator is 0 too. + const std::vector flat = {2.0f, 2.0f, 2.0f, 2.0f}; + onebit::audio::layer_norm(flat.data(), nullptr, nullptr, 4, 1e-5f, out.data()); + for (int i = 0; i < 4; ++i) expect_close(out[i], 0.0, 1e-7, "layer_norm flat " + std::to_string(i)); + // ... and with no gamma/beta the default is identity scale, zero shift. + std::vector out2(4); + const std::vector v2 = {1.0f, 2.0f, 3.0f, 4.0f}; + onebit::audio::layer_norm(v2.data(), nullptr, nullptr, 4, 1e-5f, out2.data()); + expect_close(out2[0], -1.34163542, 1e-6, "layer_norm null gamma/beta"); +} + +// 3. erf-GELU -- the export is approximate=none, so tanh must not pass this. +void gelu_arm() { + // double: -0.0455002639, 0.3457312306, 2.9959503059 + expect_close(onebit::audio::gelu_erf(-2.0f), -0.0455002639, 1e-6, "gelu(-2)"); + expect_close(onebit::audio::gelu_erf(0.5f), 0.3457312306, 1e-6, "gelu(0.5)"); + expect_close(onebit::audio::gelu_erf(3.0f), 2.9959503059, 1e-6, "gelu(3)"); + + // The two properties that separate erf-GELU from tanh-GELU: GELU(x) - GELU(-x) == x exactly, and the tanh form is + // 1.7e-4 away at x = 2 -- far outside the 1e-6 this test runs at, so a tanh substitute cannot pass silently. + expect_close(onebit::audio::gelu_erf(2.0f) - onebit::audio::gelu_erf(-2.0f), 2.0, 1e-6, "gelu odd part"); + const float tanh_form = 0.5f * 2.0f * (1.0f + std::tanh(0.7978845608f * (2.0f + 0.044715f * 8.0f))); + check(std::fabs(tanh_form - onebit::audio::gelu_erf(2.0f)) > 1e-5, "tanh GELU is distinguishable from erf GELU"); +} + +// 4. Softmax: normalised, row-max stable, and order-preserving. +void softmax_arm() { + std::vector s = {1.0f, -2.0f, 0.5f, 0.25f}; + onebit::audio::softmax_row(s.data(), 4); + // double: [0.4697737513, 0.0233886579, 0.2849321833, 0.2219054075] + const double want[4] = {0.4697737513, 0.0233886579, 0.2849321833, 0.2219054075}; + for (int i = 0; i < 4; ++i) expect_close(s[i], want[i], 1e-6, "softmax value " + std::to_string(i)); + double sum = 0.0; + for (float v : s) sum += v; + expect_close(static_cast(sum), 1.0, 1e-6, "softmax sums to 1"); + + // +1000 on every logit must not change the result: the encoder's attention logits reach this range after the + // 1/sqrt(d) scale is dropped on a long window. + std::vector big = {1001.0f, 998.0f, 1000.5f, 1000.25f}; + onebit::audio::softmax_row(big.data(), 4); + for (int i = 0; i < 4; ++i) expect_close(big[i], want[i], 1e-6, "softmax shift-invariance " + std::to_string(i)); +} + +} // namespace + +int main() { + conv1d_arms(); + layernorm_arm(); + gelu_arm(); + softmax_arm(); + if (g_failures != 0) { + std::printf("audio ops: %d failures\n", g_failures); + return 1; + } + std::printf("audio ops: conv1d (strided + padded + edge), layer_norm (biased var, null affine), gelu (erf), softmax " + "(stable, shift-invariant) -- all pass\n"); + return 0; +} diff --git a/tests/audio_parallel_test.cpp b/tests/audio_parallel_test.cpp new file mode 100644 index 000000000..93b71a551 --- /dev/null +++ b/tests/audio_parallel_test.cpp @@ -0,0 +1,198 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// audio_parallel_test.cpp -- the two claims npu/audio_parallel.h makes, as assertions (engine#1931 review): +// +// 1. the executor's MatMul is BIT-identical whatever ONEBIT_AUDIO_THREADS says. A tolerance cannot see this, and the +// pack-gated comparison against onnxruntime runs at one thread count, so here the same MatMul runs through +// run_audio_graph() at 1, 2, 3, 7, 8 and 64 threads and unset, and the outputs are compared with memcmp; +// 2. a unit that throws is an exception on the calling thread after every worker has been joined, not std::terminate. +// +// The graph is built in memory (one MatMul against an inline initializer), so this needs no pack, no file and no device. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "audio_parallel.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; + } +} + +using namespace onebit::npu::pack_import; + +// y = x . W for x (1, M, K) and W (K, N): M rows, which is what the executor splits over. +std::vector matmul_at(const char* threads, int64_t M, int64_t K, int64_t N, const std::vector& x, + const std::vector& w) { + if (threads != nullptr) setenv("ONEBIT_AUDIO_THREADS", threads, 1); + else unsetenv("ONEBIT_AUDIO_THREADS"); + OnnxGraphInfo g; + g.dir = "."; + g.inputs = {"x"}; + g.outputs = {"y"}; + OnnxInit wi; + wi.name = "W"; + wi.dtype = 1; + wi.dims = {K, N}; + wi.raw = reinterpret_cast(w.data()); + wi.raw_len = w.size() * sizeof(float); + g.initializers = {"W"}; + g.init_tensors.push_back(wi); + OnnxNodeInfo n; + n.op = "MatMul"; + n.name = "mm"; + n.ins = {"x", "W"}; + n.outs = {"y"}; + g.nodes.push_back(n); + AudioRunResult out; + std::string why; + if (!run_audio_graph(g, x.data(), M, K, "y", &out, &why)) { + std::printf("FAIL MatMul at ONEBIT_AUDIO_THREADS=%s refused: %s\n", threads ? threads : "(unset)", why.c_str()); + ++g_failures; + return {}; + } + check(out.dims == std::vector({1, M, N}), "MatMul output is (1, M, N)"); + return out.data; +} + +void test_bit_identity() { + // 37 rows: divisible by none of the thread counts below, so the last worker's share is short at every one of them. + const int64_t M = 37, K = 53, N = 29; + std::vector x(size_t(M * K)), w(size_t(K * N)); + for (size_t i = 0; i < x.size(); ++i) x[i] = float(std::sin(0.37 * double(i + 1)) * 3.0); + for (size_t i = 0; i < w.size(); ++i) w[i] = float(std::cos(0.11 * double(i + 3)) * 0.7 + 1e-3 * double(i % 7)); + for (size_t i = 0; i < x.size(); i += 5) x[i] = 0.0f; // the zero shortcut is on the path + for (int64_t k = 0; k < K; ++k) x[size_t(11 * K + k)] = 0.0f; // and one whole row of zeros + + const std::vector serial = matmul_at("1", M, K, N, x, w); + check(serial.size() == size_t(M * N), "the serial MatMul produced M*N values"); + if (serial.size() != size_t(M * N)) return; + + // The serial result is the arithmetic the split must not change: an ordered sum over k, per element, in float. + bool same_as_plain_loop = true; + for (int64_t i = 0; i < M && same_as_plain_loop; ++i) + for (int64_t j = 0; j < N; ++j) { + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) { + const float av = x[size_t(i * K + k)]; + if (av == 0.0f) continue; + acc += av * w[size_t(k * N + j)]; + } + if (std::memcmp(&acc, &serial[size_t(i * N + j)], sizeof acc) != 0) { + same_as_plain_loop = false; + break; + } + } + check(same_as_plain_loop, "the serial MatMul is the ordered per-element sum, bit for bit"); + bool zero_row = true; + for (int64_t j = 0; j < N; ++j) zero_row = zero_row && serial[size_t(11 * N + j)] == 0.0f; + check(zero_row, "a row of zeros gives a row of zeros"); + + for (const char* threads : {"2", "3", "7", "8", "64", static_cast(nullptr)}) { + const std::vector got = matmul_at(threads, M, K, N, x, w); + const std::string name = threads ? threads : "(unset)"; + check(got.size() == serial.size() && std::memcmp(got.data(), serial.data(), serial.size() * sizeof(float)) == 0, + "MatMul at ONEBIT_AUDIO_THREADS=" + name + " is bit-identical to the serial result"); + } + // What is not a thread count means serial, and the answer is still the same one. + for (const char* threads : {"0", "-4", "many", ""}) { + setenv("ONEBIT_AUDIO_THREADS", threads, 1); + check(audio_thread_count() == 1, std::string("ONEBIT_AUDIO_THREADS='") + threads + "' means one thread"); + const std::vector got = matmul_at(threads, M, K, N, x, w); + check(got.size() == serial.size() && std::memcmp(got.data(), serial.data(), serial.size() * sizeof(float)) == 0, + std::string("and MatMul at ONEBIT_AUDIO_THREADS='") + threads + "' is the serial result"); + } + unsetenv("ONEBIT_AUDIO_THREADS"); +} + +void test_every_unit_once() { + for (size_t units : {size_t(0), size_t(1), size_t(2), size_t(7), size_t(37), size_t(1000)}) + for (size_t want : {size_t(0), size_t(1), size_t(2), size_t(3), size_t(8), size_t(64), size_t(5000)}) { + std::vector> seen(units); + for (auto& s : seen) s = 0; + parallel_units(units, want, [&](size_t u) { seen[u].fetch_add(1); }); + bool once = true; + for (auto& s : seen) once = once && s.load() == 1; + check(once, "every one of " + std::to_string(units) + " units runs exactly once on " + std::to_string(want) + + " thread(s)"); + } +} + +void test_a_throwing_unit() { + // If the exception escaped a worker, the process would call std::terminate here and the test would die by signal. + for (size_t want : {size_t(1), size_t(2), size_t(4), size_t(8)}) { + std::atomic ran{0}; + bool caught = false; + std::string what; + try { + parallel_units(40, want, [&](size_t u) { + ran.fetch_add(1); + if (u == 5) throw std::runtime_error("unit 5 refused"); + }); + } catch (const std::runtime_error& e) { + caught = true; + what = e.what(); + } + check(caught && what == "unit 5 refused", + "a unit that throws on " + std::to_string(want) + " thread(s) is an exception the caller catches"); + // the other workers ran to the end and were joined before the rethrow: only the throwing worker stopped early + if (want > 1) { + const int stopped_worker_left = int((40 - 1 - 5) / want); // units after 5 in unit 5's worker + check(ran.load() == 40 - stopped_worker_left, + "the other workers finish and are joined before the caller sees it (" + std::to_string(want) + " threads)"); + } + } + // Two workers throw: the caller gets the lower worker's, every time. + for (int rep = 0; rep < 20; ++rep) { + std::string what; + try { + parallel_units(8, 4, [&](size_t u) { + if (u == 1) throw std::runtime_error("worker 1"); + if (u == 3) throw std::runtime_error("worker 3"); + }); + } catch (const std::runtime_error& e) { + what = e.what(); + } + if (what != "worker 1") { + check(false, "when two workers throw, the lower worker's exception is the one rethrown (got '" + what + "')"); + break; + } + } +} + +} // namespace + +int main() { + test_bit_identity(); + test_every_unit_once(); + test_a_throwing_unit(); + if (g_failures == 0) std::printf("audio_parallel: ok\n"); + return g_failures == 0 ? 0 : 1; +} diff --git a/tests/onnx_graph_test.cpp b/tests/onnx_graph_test.cpp new file mode 100644 index 000000000..2d6234dc5 --- /dev/null +++ b/tests/onnx_graph_test.cpp @@ -0,0 +1,132 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// 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; +} diff --git a/tools/whisper_transcribe.cpp b/tools/whisper_transcribe.cpp new file mode 100644 index 000000000..119a17131 --- /dev/null +++ b/tools/whisper_transcribe.cpp @@ -0,0 +1,203 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// whisper_transcribe.cpp -- the whole host pipeline: wav -> log-mel -> encoder -> greedy decode -> token ids. +// +// ./whisper_transcribe [expected-ids-csv ...] +// +// A tool rather than a test because one forward pass of each graph is ~26 s at these shapes and a sentence is 25 of them +// (see PR #1931); `ctest` would time out for a reason that has nothing to do with correctness. The numbers it prints are +// the ones a test would assert, so `expected-ids-csv` turns it into a check when run by hand. +// +// The pack on disk carries no tokenizer, so this stops at ids: `400 370 452 ...` is the sentence, and comparing ids is a +// stricter check than comparing text (two ids can share a detokenisation, and neither can share an id). + +#include +#include +#include +#include +#include +#include +#include + +#include "audio_exec.h" +#include "audio_mel.h" +#include "onnx_graph.h" + +namespace { + +std::vector read_wav_16k_mono(const char* path, std::string* why) { + FILE* f = std::fopen(path, "rb"); + if (f == nullptr) { + *why = std::string("cannot open ") + path; + return {}; + } + std::vector raw; + uint8_t buf[65536]; + size_t n; + while ((n = std::fread(buf, 1, sizeof buf, f)) > 0) raw.insert(raw.end(), buf, buf + n); + std::fclose(f); + if (raw.size() < 44 || std::memcmp(raw.data(), "RIFF", 4) != 0 || std::memcmp(raw.data() + 8, "WAVE", 4) != 0) { + *why = "not a RIFF/WAVE file"; + return {}; + } + size_t pos = 12; + uint16_t channels = 0, bits = 0; + uint32_t rate = 0; + const uint8_t* pcm = nullptr; + size_t pcm_len = 0; + while (pos + 8 <= raw.size()) { + const uint32_t len = uint32_t(raw[pos + 4]) | (uint32_t(raw[pos + 5]) << 8) | (uint32_t(raw[pos + 6]) << 16) | + (uint32_t(raw[pos + 7]) << 24); + const size_t body = pos + 8; + if (std::memcmp(raw.data() + pos, "fmt ", 4) == 0 && body + 16 <= raw.size()) { + channels = uint16_t(raw[body + 2] | (raw[body + 3] << 8)); + rate = uint32_t(raw[body + 4]) | (uint32_t(raw[body + 5]) << 8) | (uint32_t(raw[body + 6]) << 16) | + (uint32_t(raw[body + 7]) << 24); + bits = uint16_t(raw[body + 14] | (raw[body + 15] << 8)); + } else if (std::memcmp(raw.data() + pos, "data", 4) == 0) { + pcm = raw.data() + body; + pcm_len = std::min(len, raw.size() - body); + } + pos = body + len + (len & 1); + } + if (pcm == nullptr || channels == 0) { + *why = "no data chunk"; + return {}; + } + if (rate != 16000) { + *why = "expected 16 kHz, got " + std::to_string(rate); + return {}; + } + if (bits != 16) { + *why = "expected 16-bit PCM, got " + std::to_string(bits) + "-bit"; + return {}; + } + const size_t frames = pcm_len / (2 * size_t(channels)); + std::vector out(frames); + for (size_t i = 0; i < frames; ++i) { + int32_t acc = 0; + for (uint16_t c = 0; c < channels; ++c) { + const uint8_t* s = pcm + (i * channels + c) * 2; + acc += int16_t(uint16_t(s[0] | (s[1] << 8))); + } + out[i] = float(double(acc) / double(channels) / 32768.0); + } + return out; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc < 4) { + std::fprintf(stderr, "usage: %s [expected-id ...]\n", argv[0]); + return 2; + } + using namespace onebit::npu::pack_import; + + std::string why; + // The spans the last line reports, each named for what it contains and nothing else: wav (reading the file), mel + // (the log-mel front-end), load (parsing both graphs), encoder, decode. `wall` is their sum, from before the file + // is opened. + const auto t_begin = std::chrono::steady_clock::now(); + std::vector audio = read_wav_16k_mono(argv[3], &why); + if (audio.empty()) { + std::fprintf(stderr, "wav: %s\n", why.c_str()); + return 1; + } + // Whisper's window is a fixed 30 s: pad with silence, do not scale. + onebit::audio::MelConfig cfg; + const size_t want = size_t(cfg.sample_rate) * size_t(cfg.frames) * size_t(cfg.hop) / size_t(cfg.frames); + audio.resize(std::max(audio.size(), want), 0.0f); + + const auto t_start = std::chrono::steady_clock::now(); + std::vector mel; + if (!onebit::audio::log_mel_spectrogram(audio.data(), audio.size(), cfg, &mel)) { + std::fprintf(stderr, "mel: refused\n"); + return 1; + } + const auto t_mel = std::chrono::steady_clock::now(); + OnnxGraphInfo enc, dec; + try { + enc = parse_onnx_graph(argv[1]); + dec = parse_onnx_graph(argv[2]); + } catch (const std::exception& e) { + std::fprintf(stderr, "graph: %s\n", e.what()); + return 1; + } + const auto t_load = std::chrono::steady_clock::now(); + AudioRunResult xa; + if (!run_audio_graph(enc, mel.data(), cfg.bins, cfg.frames, "", &xa, &why)) { + std::fprintf(stderr, "encoder: %s\n", why.c_str()); + return 1; + } + const auto t_enc = std::chrono::steady_clock::now(); + std::printf("encoder: ["); + for (int64_t d : xa.dims) std::printf(" %lld", (long long)d); + std::printf(" ]\n"); + + // <|startoftranscript|>, <|en|>, <|transcribe|>, <|notimestamps|> -- then greedy, reading the logits of the last real + // prompt position (HF pads the 448/128-token window with <|endoftext|>, so the prompt is a prefix, not a suffix). + const int64_t kSot = 50258, kEn = 50259, kTranscribe = 50359, kNoTimestamps = 50363, kEot = 50257; + const int window = 128; + std::vector seq = {kSot, kEn, kTranscribe, kNoTimestamps}; + std::vector ids; + for (int step = 0; step < 60; ++step) { + std::vector toks(size_t(window), kEot); + for (size_t i = 0; i < seq.size() && i < size_t(window); ++i) toks[i] = seq[i]; + AudioRunResult y; + if (!run_decoder_graph(dec, toks.data(), window, xa.data.data(), xa.dims, "", &y, &why)) { + std::fprintf(stderr, "decoder at step %d: %s\n", step, why.c_str()); + return 1; + } + const int64_t vocab = y.dims.empty() ? 0 : y.dims.back(); + const size_t row = (seq.size() - 1) * size_t(vocab); + int64_t best = 0; + for (int64_t v = 1; v < vocab; ++v) + if (y.data[row + size_t(v)] > y.data[row + size_t(best)]) best = v; + if (best == kEot || seq.size() >= size_t(window)) break; + ids.push_back(best); + seq.push_back(best); + } + const auto t_dec = std::chrono::steady_clock::now(); + const auto secs = [](std::chrono::steady_clock::time_point a, std::chrono::steady_clock::time_point b) { + return std::chrono::duration(b - a).count(); + }; + std::printf("time: wav %.3f s, mel %.3f s, load %.3f s, encoder %.3f s, decode %.3f s for %zu token(s) -> %.3f s/token, " + "wall %.3f s\n", + secs(t_begin, t_start), secs(t_start, t_mel), secs(t_mel, t_load), secs(t_load, t_enc), secs(t_enc, t_dec), + ids.size(), ids.empty() ? 0.0 : secs(t_enc, t_dec) / double(ids.size()), secs(t_begin, t_dec)); + std::printf("ids:"); + for (int64_t v : ids) std::printf(" %lld", (long long)v); + std::printf("\n"); + + if (argc > 4) { + int bad = 0; + const int n = argc - 4; + if (int(ids.size()) != n) { + std::printf("FAIL length: got %zu want %d\n", ids.size(), n); + ++bad; + } else { + for (int i = 0; i < n; ++i) + if (ids[size_t(i)] != std::atoll(argv[4 + i])) { + std::printf("FAIL id[%d]: got %lld want %s\n", i, (long long)ids[size_t(i)], argv[4 + i]); + ++bad; + } + } + if (bad == 0) std::printf("transcript ids match the reference decode\n"); + return bad == 0 ? 0 : 1; + } + return 0; +}