[0.1.28] gpu-bench: OCR's GPU memory budget in the results

What the engine measured its runs to need (detector MB fixed + per megapixel,
recognizer per 1000 line-pixels), the limit at the end, and budget_waits, the
runs that waited for GPU memory, in the stages.

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
+  threads for the engine's life. Since 0.1.25 the runs on the GPU at once are as many as fit in its
-  calls, at 142 ms a capture against 126 before (gpu-bench, 2026-10-17). The server model has not been
+  memory: each engine measures what its runs need when it starts, and each run waits while the runs in
-  measured again yet.
+  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.26";
+VERSION       :: "0.1.28";
 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.26";
+VERSION :: "0.1.28";
 // 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;

tools/gpu_bench/gpu_bench.jai

@@ -151,6 +151,7 @@ Run :: struct {
   stages: string;          // where the time went, per item
   batch: s64;
   gpu: Gpu_Stats;
+  budget: string;          // OCR on the GPU: what its memory budget measured (modules/OCR/budget.jai), or ""
   trace: [..] string;      // the GPU's memory above the baseline, MB, after each step: load:812 first:1020 warm1:... pass1:.../peak...
   explanation: string;     // why a GPU run is on the CPU
 }
@@ -322,6 +323,10 @@ run_task :: (task: *Task, prefer_gpu: bool, corpus: *Corpus) -> Run {
         if prefer_gpu && !ocr_gpu  run.explanation = copy_string(ocr_why);
         run.provider = sprint("%", ocr_provider(engine.ocr));
         run.batch = engine.ocr.rec_batch;
+        b := engine.ocr.budget;
+        if b.enabled  run.budget = sprint("detector %+% MB/MP, recognizer %+% MB per 1000 line-px", FF(b.det_fixed / MEGABYTE, 0), FF(b.det_per_pixel * 1_000_000 / MEGABYTE, 0),
+                                         FF(b.rec_fixed / MEGABYTE, 0), FF(b.rec_per_column * 1000 / MEGABYTE, 1));
+        else if engine.ocr.provider == .CUDA  run.budget = "off (the arena could not be measured)";
       }
     case .LABELS; #through;
     case .LABELS_TEXT;
@@ -457,7 +462,8 @@ run_task :: (task: *Task, prefer_gpu: bool, corpus: *Corpus) -> Run {
   per := cast(float64) max(run.items * options.repeats, 1) * 1000;
   if task.kind == .OCR {
     s := ocr_stats(engine.ocr);
-    run.stages = sprint("detect:%,orient:%,recognize:%", FF(s.detect_us / per, 2), FF(s.orientation_us / per, 2), FF(s.recognize_us / per, 2));
+    run.stages = sprint("detect:%,orient:%,recognize:%,budget_waits:%", FF(s.detect_us / per, 2), FF(s.orientation_us / per, 2), FF(s.recognize_us / per, 2), s.budget_waits);
+    if engine.ocr.budget.enabled  run.budget = sprint("%; limit % MB at the end", run.budget, FF(engine.ocr.budget.limit / MEGABYTE, 0));
   } else if task.kind == .SEGMENT {
     s := engine.segment.stats;
     run.stages = sprint("prepare:%,model:%,post:%", FF(s.prepare_us / per, 2), FF(s.model_us / per, 2), FF(s.post_us / per, 2));
@@ -501,6 +507,7 @@ print_run :: (task: *Task, run: Run) {
   if run.megapixels  print(", % megapixels/s", FF(run.megapixels / m.median, 1));
   print("\n");
   if run.stages  print("   per %, summed over threads (ms): %\n", run.unit, run.stages);
+  if run.budget  print("   GPU memory budget: %\n", run.budget);
   if run.gpu.samples  print("   GPU: % MB more memory at the peak, % percent busy, % W, % MHz, up to % C\n",
                             FF(run.gpu.peak_mb - run.gpu.baseline_mb, 0), FF(run.gpu.util_sum / run.gpu.samples, 0),
                             FF(run.gpu.power_sum / run.gpu.samples, 0), FF(run.gpu.clock_sum / run.gpu.samples, 0), FF(run.gpu.temperature_max, 0));
@@ -671,6 +678,8 @@ hash_file_cached :: (path: string) -> string {
   return hash;
 }
 
+MEGABYTE :: 1024.0 * 1024;
+
 seconds_since :: (start: Apollo_Time) -> float64 {
   return to_float64_seconds(current_time_monotonic() - start);
 }