diff --git a/README.md b/README.md index f2a524d..2d42b95 100644 --- a/README.md +++ b/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 diff --git a/check_targets.jai b/check_targets.jai index 8da168c..c5769cb 100644 --- a/check_targets.jai +++ b/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; diff --git a/first.jai b/first.jai index b1b0600..02a7ca5 100644 --- a/first.jai +++ b/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"; diff --git a/modules/OCR/budget.jai b/modules/OCR/budget.jai new file mode 100644 index 0000000..d715187 --- /dev/null +++ b/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 = .{}; +} diff --git a/modules/OCR/detect.jai b/modules/OCR/detect.jai index b5fb184..9a9a3cd 100644 --- a/modules/OCR/detect.jai +++ b/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 { diff --git a/modules/OCR/module.jai b/modules/OCR/module.jai index 203b5cc..b46a5c7 100644 --- a/modules/OCR/module.jai +++ b/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"; diff --git a/modules/OCR/recognize.jai b/modules/OCR/recognize.jai index 90fd218..ae3e4a3 100644 --- a/modules/OCR/recognize.jai +++ b/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)); diff --git a/modules/onnxruntime/session.jai b/modules/onnxruntime/session.jai index 7baf56a..8e1da3c 100644 --- a/modules/onnxruntime/session.jai +++ b/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;