perf(rknn): price a ladder engine by what the driver reports

docs/build/rockchip-npu.md

@@ -670,12 +670,18 @@ Three consequences for anyone extending this path:
   before the weight set does.
 - A ladder of engines has to be lazy. Four engines that each load on their own fail at `rknn_init`
   when all four are resident at once, so the ladder loads on demand and evicts. How many rungs share
-  the aperture is a measurement, not a rule: the cost of a context is its compiled length, not its
-  image size, and the rungs co-resident on this board are S=200 x3, 250+250+110 and 110 x4 (all at
-  <= 1.81 GiB) while S=400 coexists only with 110/200 and S=500 with nothing. The accounting for that
-  is shared (`rknn::ApertureBudget`, `include/engine/framework/rknn/engine.h`) and the numbers are
-  not: the per-unit cost and the budget are measurements of one engine on one board, and they live
-  with the sidecar that measured them.
+  the aperture is a measurement, not a rule: the driver reports what a context will occupy (`rknn::
+  engine_memory_mib`: a `RKNN_FLAG_COLLECT_MODEL_INFO_ONLY` init plus `RKNN_QUERY_MEM_SIZE`, ~0.5 s
+  per engine, no network built), and for this estimator that is 786 MiB at S=110, 1129 at 200, 1426
+  at 250, 2615 at 400, 3752 at 500 and 1039 at 900 -- all of them 661-714 MB images, so a context
+  costs its graph and not its image, and the cost is not a function of the compiled length (S=900
+  costs less than S=500). Against a 3900 MiB per-process ceiling, measured by initializing the sets
+  in one process, the sets that admit are 110+200+250 (3341), 110 x4 (3144), 200 x3 (3387),
+  250+250+110 (3638), 110+400 (3401) and 200+400 (3744), while 250+400 (4041), 110+500 (4538) and
+  every set holding an S>=900 engine return -4. The accounting for that is shared
+  (`rknn::ApertureBudget`, `include/engine/framework/rknn/engine.h`) and the numbers are not: the
+  per-engine cost and the budget are measurements of one engine on one board, and they live with the
+  sidecar that measured them.
 - The failure is at init, so it is cheap to detect and cannot corrupt output: `rknn_init` returns -4
   and dmesg names the byte count.
 

include/engine/framework/rknn/engine.h

@@ -27,25 +27,25 @@ std::vector<char> read_engine(const std::string & path, const std::string & labe
 std::string status(const std::string & label, const char * what, int code);
 
 // Aperture accounting for a sidecar that keeps more than one compiled engine resident: the resident
-// set is the sum of the units each loaded engine reserves, and a new engine is pinned only while the
-// whole set stays inside the budget; a unit that cannot share the aperture with the set evicts it and
-// is itself evicted by the next change of shape. What a unit costs and what the board fits are both
-// measurements of one engine on one board, so both are passed in by the model: this type is the
-// accounting, not the policy, and it holds no ceiling of its own.
+// set is the sum of the residency each loaded engine reports, and a new engine is pinned only while
+// the whole set stays inside the budget; a unit that cannot share the aperture with the set evicts it
+// and is itself evicted by the next change of shape. What an engine costs and what the board fits are
+// both measurements of one engine on one board, so both are passed in by the model: this type is the
+// accounting, not the policy, and it holds no ceiling of its own. It counts MiB and nothing else --
+// the model derives that number from the driver's own report (`engine_memory_mib`) or from a proxy it
+// owns, and the accounting stays free of any assumption about how a cost is estimated.
 class ApertureBudget {
 public:
-    ApertureBudget(int64_t mib_per_unit, int64_t budget_mib)
-        : mib_per_unit_(mib_per_unit), budget_mib_(budget_mib) {}
+    explicit ApertureBudget(int64_t budget_mib) : budget_mib_(budget_mib) {}
 
-    // Would adding an engine of `units` compiled length keep the resident set inside the aperture?
-    bool fits(int64_t units) const { return (resident_ + units) * mib_per_unit_ <= budget_mib_; }
-    void reserve(int64_t units) { resident_ += units; }
-    void clear() { resident_ = 0; }
-    int64_t resident() const { return resident_; }
+    // Would adding an engine of `mib` resident megabytes keep the set inside the aperture?
+    bool fits(int64_t mib) const { return resident_mib_ + mib <= budget_mib_; }
+    void reserve(int64_t mib) { resident_mib_ += mib; }
+    void clear() { resident_mib_ = 0; }
+    int64_t resident() const { return resident_mib_; }
 
 private:
-    int64_t resident_ = 0;
-    int64_t mib_per_unit_ = 0;
+    int64_t resident_mib_ = 0;
     int64_t budget_mib_ = 0;
 };
 
@@ -62,6 +62,14 @@ struct Input {
     uint32_t pass_through = 0;
 };
 
+// What a compiled engine will occupy once it is resident, in MiB, read from the driver itself:
+// `rknn_init` with `RKNN_FLAG_COLLECT_MODEL_INFO_ONLY` parses the image without building the network
+// (~0.5 s per engine on RK3588, no NPU memory allocated) and `rknn_query(RKNN_QUERY_MEM_SIZE)` then
+// reports `total_weight_size + total_internal_size` -- the bytes the real init goes on to allocate,
+// less a per-context slack of 30-50 MiB. Returns 0 when the driver cannot say, so a sidecar can keep
+// a proxy for that case.
+int64_t engine_memory_mib(const std::string & path, const std::string & label);
+
 // One compiled engine: a `rknn_context` plus its input/output counts.
 class Engine {
 public:

src/framework/rknn/engine.cpp

@@ -34,6 +34,34 @@ std::string status(const std::string & label, const char * what, int code) {
 
 #if defined(ENGINE_HAVE_RKNN)
 
+int64_t engine_memory_mib(const std::string & path, const std::string & label) {
+    (void)label;
+    std::vector<char> model;
+    try {
+        model = read_engine(path, label);
+    } catch (const std::exception &) {
+        return 0;
+    }
+    rknn_context ctx = 0;
+    if (rknn_init(
+            &ctx, model.data(), static_cast<uint32_t>(model.size()),
+            RKNN_FLAG_COLLECT_MODEL_INFO_ONLY, nullptr) != 0) {
+        return 0;
+    }
+    rknn_mem_size mem{};
+    const int ret = rknn_query(ctx, RKNN_QUERY_MEM_SIZE, &mem, sizeof(mem));
+    rknn_destroy(ctx);
+    if (ret != 0) {
+        return 0;
+    }
+    // `rknn_mem_size` carries the two byte counts as 32-bit, so widen before adding: a 1 GiB weight
+    // set beside a 3.5 GiB workspace is 4.5 GiB, and a 32-bit sum reports 512 MiB.
+    const uint64_t bytes = static_cast<uint64_t>(mem.total_weight_size) + mem.total_internal_size;
+    // Round up: the budget is a ceiling, and a resident set measured at its sum should not be
+    // admitted by fractions of a megabyte.
+    return static_cast<int64_t>((bytes + (1ull << 20) - 1) >> 20);
+}
+
 Engine::~Engine() {
     release();
 }

src/models/cosyvoice3/rknn_flow.cpp

@@ -30,22 +30,23 @@ const char * const kLabel = "CosyVoice3 flow estimator RKNN";
 // unconditioned one, and the ODE loop mixes them outside the engine.
 constexpr int64_t kCfgBatch = 2;
 
-// The aperture accounting of this sidecar, in a linear proxy measured on RK3588. With every context
-// of the ladder resident in one process, the estimator engines occupy 0.27 GiB at S=110, 0.48 at
-// 200, 0.77 at 250, 1.93 at 400, 3.04 at 500 and 4.39 at 900, while all of their files are 661-714
-// MB -- a context costs its compiled length, not its image size. 5 MiB per frame over-estimates the
-// short rungs by ~2x and is within 5% of the long ones, so the budget is a statement about the long
-// rungs with a safety margin on the short ones.
-constexpr int64_t kApertureMiBPerFrame = 5;
+// The aperture accounting of this sidecar, in the driver's own units: each engine's residency is
+// read from the runtime (`rknn::engine_memory_mib`, a model-info-only init plus
+// `RKNN_QUERY_MEM_SIZE`), which on this board reports 786 MiB at S=110, 1129 at 200, 1426 at 250,
+// 2615 at 400, 3752 at 500 and 1039 at 900 -- all of them 661-714 MB files, so a context costs its
+// graph and not its image, and the cost is not a function of the compiled length (S=900 costs less
+// than S=500). The linear proxy this replaces charged 5 MiB per frame, which over-pins the short
+// rungs by up to 11x and under-pins S=500.
+constexpr int64_t kApertureMiBPerFrame = 5;  // only if the driver cannot say
 
-// 2800 MiB of that proxy is the largest resident set this board's I/O aperture fits: the short rungs
-// are genuinely co-resident (S=200 x3, 250+250+110 and 110 x4 all initialize together, at
-// <= 1.81 GiB) and the long ones are not (400 coexists only with 110/200, 500 with nothing);
-// 110+200+250 = 1.52 GiB real (three engines fit 1.81 GiB) and 110+400 = 2.20 GiB (two fit 2.42),
-// while 200+400 = 2.41 GiB is refused by the proxy although it measures just inside the two-engine
-// ceiling -- the proxy under-pins, never over-pins. Measured on this board, not a board property:
-// the accounting is shared (rknn::ApertureBudget), the numbers are this engine's.
-constexpr int64_t kAperturePinBudgetMiB = 2800;
+// 3900 MiB is where this board's driver puts the per-process ceiling, measured by initializing the
+// sets in one process: 110+200+250 (3341 MiB), 110x4 (3144), 200x3 (3387), 250+250+110 (3638),
+// 110+400 (3401) and 200+400 (3744) all initialize together, while 250+400 (4041), 110+500 (4538)
+// and every set holding an S>=900 engine fail at `rknn_init` with -4. The proxy refused 200+400,
+// which this board admits; the measured cost reproduces all twelve outcomes. Measured on this board,
+// not a board property: the accounting is shared (rknn::ApertureBudget), the numbers are this
+// engine's.
+constexpr int64_t kAperturePinBudgetMiB = 3900;
 
 // Graph order, as tools/rknn/prepare_cosyvoice3_rknn_flow.py compiles it.
 const char * const kInputNames[] = {"x", "mask", "mu", "t", "spks", "cond"};
@@ -87,6 +88,9 @@ struct CosyVoice3RknnFlow::Impl {
         int64_t frames = 0;    // the sequence length this engine covers
         int64_t channels = 0;
 #if defined(ENGINE_HAVE_RKNN)
+        // What this engine will occupy once resident, read from the driver at registration; 0 means
+        // the driver could not say, and the linear proxy is used for it.
+        int64_t aperture_mib = 0;
         engine::rknn::Engine engine;
         bool loaded = false;
         std::vector<float> x;
@@ -225,8 +229,8 @@ struct CosyVoice3RknnFlow::Impl {
     std::vector<std::unique_ptr<Slot>> slots;
     int64_t channels = 0;
 #if defined(ENGINE_HAVE_RKNN)
-    // Resident-set accounting for the ladder, in the shared proxy arithmetic.
-    engine::rknn::ApertureBudget budget_{kApertureMiBPerFrame, kAperturePinBudgetMiB};
+    // Resident-set accounting for the ladder, in the driver's own units (see kAperturePinBudgetMiB).
+    engine::rknn::ApertureBudget budget_{kAperturePinBudgetMiB};
 #endif
 
 #if defined(ENGINE_HAVE_RKNN)
@@ -249,6 +253,12 @@ struct CosyVoice3RknnFlow::Impl {
             slot->path = path;
             slot->frames = frames;
             slot->channels = expected_channels;
+            // Ask the driver what this engine will cost before deciding to keep it: a model-info-only
+            // init parses the image without allocating the network (~0.5 s per engine, once at
+            // startup), which is what makes the per-engine table a measurement instead of a guess.
+            slot->aperture_mib = engine::rknn::engine_memory_mib(path, kLabel);
+            debug::trace_log_scalar(
+                "cosyvoice3.flow.dit.rknn.aperture_mib", slot->aperture_mib);
             slots.push_back(std::move(slot));
         }
         std::sort(slots.begin(), slots.end(), [](const auto & a, const auto & b) {
@@ -275,8 +285,10 @@ struct CosyVoice3RknnFlow::Impl {
                 // The rung is kept (and the new one pinned) only while the whole resident set stays
                 // inside the aperture; a rung above the budget can never share it with anything, so
                 // it evicts everything and is itself evicted by the next length change.
-                // ponytail: linear aperture proxy; a per-engine measured table would pin ~2x more.
-                const bool fits = budget_.fits(slot.frames);
+                const int64_t cost_mib = slot.aperture_mib > 0
+                    ? slot.aperture_mib
+                    : slot.frames * kApertureMiBPerFrame;
+                const bool fits = budget_.fits(cost_mib);
                 if (!fits) {
                     for (auto & other : slots) {
                         if (other->loaded) {
@@ -287,9 +299,9 @@ struct CosyVoice3RknnFlow::Impl {
                     budget_.clear();
                 }
                 slot.load(channels);
-                budget_.reserve(slot.frames);
+                budget_.reserve(cost_mib);
                 debug::trace_log_scalar(
-                    "cosyvoice3.flow.dit.rknn.resident_frames", budget_.resident());
+                    "cosyvoice3.flow.dit.rknn.resident_mib", budget_.resident());
             }
             return slot;
         }

tests/unittests/test_rknn_engine.cpp

@@ -67,11 +67,22 @@ void test_status() {
 // The aperture accounting a ladder of engines spends, on the numbers CosyVoice3's flow estimator
 // measured on RK3588 (5 MiB per compiled frame, a 2800 MiB resident budget) and the trace that
 // board stand served: rungs 110/200/250/400/500/900, requests at 118/158/202/192/220/228 frames.
-// The board measured two `rknn_init` calls and a 450-frame resident set for that sequence, so a
-// budget that evicts a short rung it should keep shows up here.
+// The board measured the per-engine residency the driver reports (786 MiB at S=110, 1129 at 200,
+// 1426 at 250, 2615 at 400, 3752 at 500, 1039 at 900) against a 3900 MiB aperture, and the sets that
+// initialize together in one process are the ones the accounting admits -- so a budget that evicts a
+// short rung it should keep, or admits a set the driver refuses, shows up here.
 void test_aperture_budget_ladder() {
     const std::vector<int64_t> rungs = {110, 200, 250, 400, 500, 900};
-    engine::rknn::ApertureBudget budget(5, 2800);
+    const std::vector<int64_t> cost = {786, 1129, 1426, 2615, 3752, 1039};
+    engine::rknn::ApertureBudget budget(3900);
+    const auto cost_of = [&rungs, &cost](int64_t frames) {
+        for (size_t i = 0; i < rungs.size(); ++i) {
+            if (rungs[i] == frames) {
+                return cost[i];
+            }
+        }
+        return int64_t{0};
+    };
     std::vector<bool> loaded(rungs.size(), false);
     int loads = 0;
     for (const int64_t frames : {118, 158, 202, 192, 220, 228}) {
@@ -80,23 +91,24 @@ void test_aperture_budget_ladder() {
                 continue;
             }
             if (!loaded[i]) {
-                if (!budget.fits(rungs[i])) {
+                if (!budget.fits(cost_of(rungs[i]))) {
                     std::fill(loaded.begin(), loaded.end(), false);
                     budget.clear();
                 }
                 loaded[i] = true;
-                budget.reserve(rungs[i]);
+                budget.reserve(cost_of(rungs[i]));
                 ++loads;
             }
             break;
         }
     }
     require_eq(loads, 2, "the served ladder trace loads two engines");
-    require_eq(budget.resident(), int64_t{450}, "the resident set is the two short rungs");
-    require(!budget.fits(500), "a 500-frame rung does not fit beside 200+250");
+    require_eq(budget.resident(), int64_t{2555}, "the resident set is the two short rungs");
+    require(!budget.fits(3752), "S=500 does not fit beside 200+250");
     budget.clear();
-    require(budget.fits(500), "a 500-frame rung fits the aperture alone");
-    require(!budget.fits(900), "the longest rung exceeds the aperture on its own");
+    require(budget.fits(1129 + 2615), "200+400 fits the aperture, as it does on the board");
+    require(!budget.fits(1426 + 2615), "250+400 does not, as it does not on the board");
+    require(budget.fits(1039), "S=900's own cost fits the aperture; it fails at init on its own ceiling");
 }
 
 }  // namespace