mcap/modules/OCR/budget.jai

// 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 = .{};
}