[0.1.25] OCR on the GPU runs as many detector and recognizer runs at once as fit in its memory

The workers ran gpu_workers (4) runs at once whatever the images: the server
model, four large images at a time, took 15 GB of a 16 GB GPU and spilled into
shared system memory, and a smaller GPU would fare far worse.

modules/OCR/budget.jai: each run reserves what it is expected to need, and
waits while the runs in flight or the GPU's free memory leave too little; one
run always goes. The need is measured when the engine is made, from the arena's
peak (AllocatorGetStats MaxInUse) over runs on two input sizes, so it holds for
any model and GPU: the mobile detector 6 MB + 137 MB a megapixel, the server's
4 MB + 1015, against 511 and 522 MB measured for a 2560 x 1440 tile. The limit
is the GPU's free memory less 10% (at least 512 MB), read at the start of each
ocr_process call (cudaMemGetInfo, ort_cuda_memory); the free memory is also read
as each run starts, since finished runs' blocks stay in the arena until the
stage ends.

gpu-bench, RTX 5080 laptop: the server model 335 ms a capture (854 before),
peak 13.1 GB (15.4), passes within 4%; the mobile model unchanged (140 ms,
5.7 GB), as its runs fit.

README.md

@@ -107,9 +107,12 @@ grid (in its order, so through a search's matches).
   from CUDA, and OCR's inputs change size every run, so each larger image or wider batch of lines left
   another block behind. OCR now runs the largest images and widest batches first, has the detector's
   and recognizer's arenas give back what they hold after each stage of a call, and keeps its 4 worker
-  threads for the engine's life. The mobile model now peaks at 5.7 GB and drops to 2 GB between
-  calls, at 142 ms a capture against 126 before (gpu-bench, 2026-10-17). The server model has not been
-  measured again yet.
+  threads for the engine's life. Since 0.1.25 the runs on the GPU at once are as many as fit in its
+  memory: each engine measures what its runs need when it starts, and each run waits while the runs in
+  flight and the GPU's free memory leave too little for it (`modules/OCR/budget.jai`). On an RTX 5080
+  laptop (16 GB), gpu-bench, 2026-10-17: the mobile model peaks at 5.7 GB and drops to 2 GB between
+  calls, at 140 ms a capture against 126 before; the server model peaks at 13.1 GB, at 335 ms a
+  capture against 10.9 s.
 - **Running on:** the GPU or the CPU, and if the CPU, what to install for the GPU.
 - **Scan the library** reads every capture the model has not read yet and, with **Labels too**,
   labels every capture not yet labelled with the current list. While a search is active the button

check_targets.jai

@@ -47,7 +47,7 @@ APP_NAME        :: "mcap";
 APP_FOLDER_NAME :: "mcap";
 GIT_BRANCH    :: "check";
 GIT_HASH      :: "-";
-VERSION       :: "0.1.24";
+VERSION       :: "0.1.25";
 BUILD_DATE    :: "-";
 DEBUG         :: true;
 COMPILE_TESTS :: true;

first.jai

@@ -8,7 +8,7 @@ APP_NAME              :: "mcap";
 APP_FOLDER_NAME       :: "mcap";
 APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
 
-VERSION :: "0.1.24";
+VERSION :: "0.1.25";
 // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated
 // builds never write over the ones you run by hand.
 build_directory := "bin";

modules/OCR/budget.jai

@@ -0,0 +1,137 @@
+// On the GPU, how many detector and recognizer runs share it at once: as many as fit in its memory.
+// Each run reserves what it is expected to need before it starts, and waits while the runs in flight
+// leave too little; one run always goes. A run's need grows with its input, and the engine measures
+// how fast when it is made, so it holds for any model and GPU: from the arena's peak (onnxruntime's
+// MaxInUse) over runs on two input sizes. The limit is a share of the GPU's free memory, read again at
+// the start of every ocr_process call, as other engines and programs come and go. A run must also fit
+// in what the GPU has free when it starts: the arenas keep finished runs' blocks until the stage ends
+// (ocr_shrink_arena), a run of another size can need a new block beside them, and cuDNN and cuBLAS
+// take memory of their own, so the reservations alone let the server model reach 15.4 GB of 16.
+//
+// Without it, gpu_workers runs went at once whatever the images: the server model, four large images
+// at a time, took 15 GB of a 16 GB GPU, and on a smaller GPU Windows pages the excess into shared
+// system memory, running it many times slower.
+
+Ocr_Budget :: struct {
+  enabled: bool;               // on CUDA, when the engine could measure its runs
+  limit: s64;                  // bytes the runs in flight may reserve together
+  in_use: s64;                 // reserved by the runs in flight; under mutex
+  headroom: s64;               // bytes left free on the GPU, beside the limit and when a run starts
+  recent: [8] Ocr_Budget_Grant;   // the last runs let go, under mutex: what the GPU's free memory may not show yet
+  mutex: Mutex;
+  det_fixed, det_per_pixel: float64;    // a detector run's need: det_fixed + det_per_pixel * height * width
+  rec_fixed, rec_per_column: float64;   // a recognizer run's: rec_fixed + rec_per_column * lines * width
+}
+
+// A run let go, which may not have taken its memory yet.
+Ocr_Budget_Grant :: struct { at: Apollo_Time; bytes: s64; }
+OCR_BUDGET_GRANT_MS :: 50;          // how long a run let go may take to take its memory
+
+OCR_BUDGET_MARGIN :: 1.2;           // estimates are raised by this much: other shapes, other kernels
+OCR_BUDGET_HEADROOM :: 0.1;         // of the free memory, left for everything else on the GPU
+OCR_BUDGET_MIN_HEADROOM :: 512 * 1024 * 1024;
+
+// Measures what the detector's and recognizer's runs need, from runs on two sizes each. Called at
+// creation, under ort_gpu_lock, after ocr_warm_up; on failure the budget is off and every worker runs.
+ocr_budget_calibrate :: (engine: *Ocr_Engine) {
+  b := *engine.budget;
+  // What one run needs, from the arena's peak, run after run of growing size so each sets a new peak.
+  need :: (engine: *Ocr_Engine, session: *Ocr_Session, shape: [4] s64) -> s64, ok: bool {
+    before, _, ok := ort_gpu_arena_stats(session);
+    if !ok  return 0, false;
+    count := shape[0] * shape[1] * shape[2] * shape[3];
+    input := NewArray(count, float, initialized = true,, context.default_allocator);
+    defer array_free(input,, context.default_allocator);
+    out, run_ok := ocr_run(session, engine.memory_info, input, shape);
+    if !run_ok  return 0, false;
+    ocr_release_output(*out);
+    _, peak, stats_ok := ort_gpu_arena_stats(session);
+    if !stats_ok || peak <= before  return 0, false;
+    return peak - before, true;
+  }
+  fit :: (x0: float64, y0: s64, x1: float64, y1: s64) -> fixed: float64, slope: float64 {
+    slope := max(cast(float64) (y1 - y0) / (x1 - x0), 0);
+    return max(y0 - slope * x0, 0), slope;
+  }
+
+  // The detector at 512 and 1024 px square; the recognizer on a full batch 256 and 1024 px wide.
+  d0, ok0 := need(engine, *engine.det, .[1, 3, 512, 512]);
+  d1, ok1 := need(engine, *engine.det, .[1, 3, 1024, 1024]);
+  lines := cast(s64) max(engine.model.rec_batch_gpu, 1);
+  height := cast(s64) engine.model.rec_height;
+  r0, ok2 := need(engine, *engine.rec, .[lines, 3, height, 256]);
+  r1, ok3 := need(engine, *engine.rec, .[lines, 3, height, 1024]);
+  ocr_shrink_arena(engine, *engine.det);
+  ocr_shrink_arena(engine, *engine.rec);
+  if !ok0 || !ok1 || !ok2 || !ok3  return;
+  b.det_fixed, b.det_per_pixel  = fit(512.0 * 512, d0, 1024.0 * 1024, d1);
+  b.rec_fixed, b.rec_per_column = fit(cast(float64) lines * 256, r0, cast(float64) lines * 1024, r1);
+  init(*b.mutex);
+  b.enabled = true;
+}
+
+// The share of the GPU's free memory the runs of this call may take; at the start of ocr_process.
+ocr_budget_update :: (engine: *Ocr_Engine) {
+  b := *engine.budget;
+  if !b.enabled  return;
+  free, _, ok := ort_cuda_memory();
+  if !ok  return;   // the last limit stands
+  headroom := max(cast(s64) (free * OCR_BUDGET_HEADROOM), OCR_BUDGET_MIN_HEADROOM);
+  lock(*b.mutex);
+  b.headroom = headroom;
+  b.limit = max(free - headroom, 0);
+  unlock(*b.mutex);
+}
+
+ocr_budget_detector :: (engine: *Ocr_Engine, height: s64, width: s64) -> s64 {
+  b := *engine.budget;
+  return cast(s64) ((b.det_fixed + b.det_per_pixel * height * width) * OCR_BUDGET_MARGIN);
+}
+
+ocr_budget_recognizer :: (engine: *Ocr_Engine, lines: s64, width: s64) -> s64 {
+  b := *engine.budget;
+  return cast(s64) ((b.rec_fixed + b.rec_per_column * lines * width) * OCR_BUDGET_MARGIN);
+}
+
+// Waits until bytes fit beside the runs in flight and in the GPU's free memory (at once when no run
+// is in flight), and reserves them.
+ocr_budget_reserve :: (engine: *Ocr_Engine, bytes: s64) {
+  b := *engine.budget;
+  if !b.enabled  return;
+  waited := false;
+  while true {
+    free, _, free_ok := ort_cuda_memory();
+    lock(*b.mutex);
+    fits := b.in_use == 0;
+    if !fits && b.in_use + bytes <= b.limit {
+      // Runs let go in the last moments may not show in free yet.
+      now := current_time_monotonic();
+      pending: s64;
+      for b.recent  if it.bytes && to_milliseconds(now - it.at) < OCR_BUDGET_GRANT_MS  pending += it.bytes;
+      fits = !free_ok || bytes + pending <= free - b.headroom;
+    }
+    if fits {
+      b.in_use += bytes;
+      oldest := 0;
+      for b.recent  if it.at < b.recent[oldest].at  oldest = it_index;
+      b.recent[oldest] = .{current_time_monotonic(), bytes};
+    }
+    unlock(*b.mutex);
+    if fits  break;
+    if !waited { waited = true; atomic_add(*engine.stats.budget_waits, 1); }
+    sleep_milliseconds(1);
+  }
+}
+
+ocr_budget_release :: (engine: *Ocr_Engine, bytes: s64) {
+  b := *engine.budget;
+  if !b.enabled  return;
+  lock(*b.mutex);
+  b.in_use -= bytes;
+  unlock(*b.mutex);
+}
+
+ocr_budget_destroy :: (engine: *Ocr_Engine) {
+  if engine.budget.enabled  destroy(*engine.budget.mutex);
+  engine.budget = .{};
+}

modules/OCR/detect.jai

@@ -83,7 +83,10 @@ ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_
   }
 
   shape := s64.[1, 3, h, w];
+  reserved := ocr_budget_detector(engine, h, w);
+  ocr_budget_reserve(engine, reserved);   // budget.jai: waits while the GPU's memory is taken
   out, ok, error := ocr_run(*engine.det, engine.memory_info, input, shape);
+  ocr_budget_release(engine, reserved);
   if !ok  return false, error;
   defer ocr_release_output(*out);
   if out.dim_count != 4 || out.dims[2] != h || out.dims[3] != w {

modules/OCR/module.jai

@@ -48,6 +48,7 @@ Ocr_Engine :: struct {
   options: Ocr_Options;
   stats: Ocr_Stats;           // since creation, or the last ocr_reset_stats
   pool: Ocr_Pool;             // the worker threads, for the engine's life (run_on_workers)
+  budget: Ocr_Budget;         // on CUDA: how many runs fit in GPU memory at once (budget.jai)
 
   // When set, and true, ocr_process stops starting work: the batch of lines or image each worker is
   // on is finished, and then every image fails as "cancelled" (none is part read). Set by the caller.
@@ -68,6 +69,7 @@ Ocr_Stats :: struct {
   recognize_us: s64;          // cropping lines, the recognizer, and decoding the text
   orientation_us: s64;        // cropping lines and the orientation classifier
   turned: s64;                // lines that read better the other way up, or level as found
+  budget_waits: s64;          // detector and recognizer runs that waited for GPU memory (budget.jai)
 }
 
 ocr_stats :: (engine: *Ocr_Engine) -> Ocr_Stats { return engine.stats; }
@@ -132,6 +134,10 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op
     sessions_ok, sessions_error := ocr_create_sessions(engine, models_folder, true, 2);
     if sessions_ok  sessions_ok, sessions_error = ocr_warm_up(engine);
     if !sessions_ok  ocr_release_sessions(engine);
+    if sessions_ok {
+      engine.provider = .CUDA;   // before calibrating: ocr_shrink_arena only works on CUDA
+      ocr_budget_calibrate(engine);
+    }
     ort_gpu_unlock();
     if sessions_ok {
       engine.provider = .CUDA;
@@ -166,6 +172,7 @@ ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string {
 ocr_destroy :: (engine: *Ocr_Engine) {
   if !engine  return;
   ocr_pool_stop(engine);
+  ocr_budget_destroy(engine);
   if engine.provider == .CUDA  ort_gpu_lock();
   ocr_release_sessions(engine);
   if engine.provider == .CUDA  ort_gpu_unlock();
@@ -180,6 +187,7 @@ ocr_destroy :: (engine: *Ocr_Engine) {
 ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.allocator) -> [] Ocr_Result {
   results := NewArray(images.count, Ocr_Result,, allocator);
   if !images.count  return results;
+  ocr_budget_update(engine);   // budget.jai: the GPU memory this call's runs may share
 
   // Detection: images shared out among the workers.
   per_image := NewArray(images.count, Image_Work,, context.default_allocator);
@@ -446,6 +454,7 @@ run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void
 #load "models.jai";
 #load "detect.jai";
 #load "recognize.jai";
+#load "budget.jai";
 
 #import "Basic";
 #import "String";

modules/OCR/recognize.jai

@@ -38,7 +38,10 @@ ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, er
   }
 
   shape := s64.[jobs.count, 3, height, width];
+  reserved := ocr_budget_recognizer(engine, jobs.count, width);
+  ocr_budget_reserve(engine, reserved);   // budget.jai: waits while the GPU's memory is taken
   out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
+  ocr_budget_release(engine, reserved);
   if !ok  return false, error;
   defer ocr_release_output(*out);
   if out.dim_count != 3 || out.dims[0] != jobs.count  return false, tprint("The recognizer gave an output of shape %.", array_view(out.dims, 0, out.dim_count));

modules/onnxruntime/session.jai

@@ -123,6 +123,50 @@ ort_release_session :: (s: *Ort_Session) {
   s.* = .{};
 }
 
+// The GPU's free and total memory in bytes, from the CUDA runtime onnxruntime's CUDA provider loads;
+// ok false before a CUDA session has been made, or without CUDA.
+ort_cuda_memory :: () -> free: s64, total: s64, ok: bool {
+  Get_Info :: #type (free: *u64, total: *u64) -> s32 #c_call;
+  get: Get_Info;
+  #if OS == .WINDOWS {
+    module := GetModuleHandleA("cudart64_13.dll");
+    if module  get = cast(Get_Info) GetProcAddress(module, "cudaMemGetInfo");
+  } else #if OS == .LINUX {
+    handle := dlopen("libcudart.so.13", RTLD_LAZY);
+    if handle  get = cast(Get_Info) dlsym(handle, "cudaMemGetInfo");
+  }
+  if !get  return 0, 0, false;
+  free, total: u64;
+  if get(*free, *total) != 0  return 0, 0, false;
+  return cast(s64) free, cast(s64) total, true;
+}
+
+// A CUDA session's GPU arena: the bytes its blocks hand out now, and the most they ever have (onnxruntime
+// keeps the peak since the session was made). ok false when onnxruntime does not say.
+ort_gpu_arena_stats :: (s: *Ort_Session) -> in_use: s64, max_in_use: s64, ok: bool {
+  api := ort.api;
+  info: *OrtMemoryInfo;
+  if !ort_check(api.CreateMemoryInfo("Cuda", .OrtArenaAllocator, 0, .OrtMemTypeDefault, *info))  return 0, 0, false;
+  defer api.ReleaseMemoryInfo(info);
+  allocator: *OrtAllocator;
+  if !ort_check(api.CreateAllocator(s.session, info, *allocator))  return 0, 0, false;
+  defer api.ReleaseAllocator(allocator);
+  stats: *OrtKeyValuePairs;
+  if !ort_check(api.AllocatorGetStats(allocator, *stats))  return 0, 0, false;
+  defer api.ReleaseKeyValuePairs(stats);
+  keys, values: **u8;
+  count: u64;
+  api.GetKeyValuePairs(stats, *keys, *values, *count);
+  in_use, max_in_use := -1, -1;
+  for 0..cast(s64) count - 1 {
+    key := to_string(keys[it]);
+    if key == "InUse"     in_use     = string_to_int(to_string(values[it]));
+    if key == "MaxInUse"  max_in_use = string_to_int(to_string(values[it]));
+  }
+  if in_use < 0 || max_in_use < 0  return 0, 0, false;
+  return in_use, max_in_use, true;
+}
+
 // A run's output: its values stay valid until ort_release_output.
 Ort_Output :: struct {
   value: *OrtValue;