// Sessions: one model file each, run on an NVIDIA GPU through CUDA or on the CPU, and running them on
// a tensor. Shared by the modules that run models (OCR, Labels, Segment).

// CUDA sessions are made, first run and released one at a time, process-wide: hold ort_gpu_lock
// around it. Done on two threads at once (onnxruntime 1.30.0, CUDA 13, cuBLASLt 13.8, RTX 5080
// laptop), it crashed with a DEP access violation in cuBLASLt or near it, every time in a debug build;
// one at a time, never. Running sessions on many threads at once is fine. A spin lock, needing no
// setup, as engines can be made before anything else runs; it is held for a second or so at a time.
ort_gpu_lock :: () {
  while atomic_swap(*ort_gpu_held, 1) != 0  sleep_milliseconds(2);
}

ort_gpu_unlock :: () {
  atomic_swap(*ort_gpu_held, 0);
}

#scope_file
ort_gpu_held: s32;
#scope_export

Ort_Session :: struct {
  session: *OrtSession;
  shrink_options: *OrtRunOptions; // for a run that shrinks the arena afterwards (ort_run's shrink); null when the settings did not ask for it
  input_name, output_name: *u8;   // heap C strings
  input_shape: [4] s64;           // as the model declares it: -1 for a size it leaves open
  input_rank: s64;
}

// How a CUDA session uses GPU memory. Its arena takes blocks from CUDA as runs need them, and by
// default never joins or frees them: a model whose input size changes (OCR's) leaves a block of
// every larger size behind, gigabytes of them, unless the arena is shrunk now and then. Shrinking
// after every run costs a third to a half of OCR's speed (CUDA frees and allocates again, waiting for
// the GPU), so OCR shrinks once per call (ocr_process).
Ort_Gpu_Settings :: struct {
  arena_extend_strategy := "kNextPowerOfTwo";   // how much the arena takes when it runs short: twice as much each time; "kSameAsRequested" takes what is asked
  cudnn_max_workspace := false;                 // let cuDNN take as much working memory as it finds best (onnxruntime's own default)
  shrink_arena := false;                        // a run with ort_run's shrink set then gives back the arena's blocks no run is using
}

// A session for the model file at path; with CUDA when gpu is set. output names the output that
// ort_run returns; "" for the first. missing_hint is added when the file is not there, to say how
// to get it. On failure, the reason.
ort_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, settings := Ort_Gpu_Settings.{}, output := "", missing_hint := "") -> Ort_Session, ok: bool, error: string {
  result: Ort_Session;
  api := ort.api;
  options: *OrtSessionOptions;
  ok, error := ort_check(api.CreateSessionOptions(*options));
  if !ok  return result, false, error;
  defer api.ReleaseSessionOptions(options);

  ort_check(api.SetSessionGraphOptimizationLevel(options, .ORT_ENABLE_ALL));
  ort_check(api.SetIntraOpNumThreads(options, intra_op_threads));
  ort_check(api.SetInterOpNumThreads(options, 1));

  if gpu {
    cuda: *OrtCUDAProviderOptionsV2;
    ok, error = ort_check(api.CreateCUDAProviderOptions(*cuda));
    if !ok  return result, false, error;
    defer api.ReleaseCUDAProviderOptions(cuda);
    // Inputs come in many sizes: an exhaustive cuDNN search would be redone for each new shape.
    keys := (*u8).["device_id", "cudnn_conv_algo_search", "arena_extend_strategy", "cudnn_conv_use_max_workspace"];
    values := (*u8).["0", "HEURISTIC", temp_c_string(settings.arena_extend_strategy), temp_c_string(ifx settings.cudnn_max_workspace then "1" else "0")];
    ok, error = ort_check(api.UpdateCUDAProviderOptions(cuda, keys.data, values.data, xx keys.count));
    if !ok  return result, false, error;
    ok, error = ort_check(api.SessionOptionsAppendExecutionProvider_CUDA_V2(options, cuda));
    if !ok  return result, false, error;
  }

  if !file_exists(path)  return result, false, tprint("% is not there%.", path, ifx missing_hint then tprint("; %", missing_hint) else "");
  ok, error = ort_check(api.CreateSession(env, ort_path(path), options, *result.session));
  if !ok  return result, false, error;
  if gpu && settings.shrink_arena {
    // Device 0's arena, as device_id above; the CPU's is left alone.
    ok, error = ort_check(api.CreateRunOptions(*result.shrink_options));
    if ok  ok, error = ort_check(api.AddRunConfigEntry(result.shrink_options, "memory.enable_memory_arena_shrinkage", "gpu:0"));
    if !ok { ort_release_session(*result); return result, false, error; }
  }

  allocator: *OrtAllocator;
  ort_check(api.GetAllocatorWithDefaultOptions(*allocator));
  name: *u8;
  ort_check(api.SessionGetInputName(result.session, 0, allocator, *name));
  result.input_name = to_c_string(to_string(name),, context.default_allocator);
  ort_check(api.AllocatorFree(allocator, name));

  // The output: the one named, else the first.
  output_count: u64;
  ort_check(api.SessionGetOutputCount(result.session, *output_count));
  names: [..] string;
  names.allocator = temp;
  for i: 0..cast(s64) output_count - 1 {
    ort_check(api.SessionGetOutputName(result.session, xx i, allocator, *name));
    array_add(*names, copy_string(to_string(name),, temp));
    ort_check(api.AllocatorFree(allocator, name));
  }
  chosen := ifx names.count then names[0] else "";
  if output {
    found := array_find(names, output);
    if !found { ort_release_session(*result); return result, false, tprint("% has no output %; it has %.", path, output, names); }
    chosen = output;
  }
  result.output_name = to_c_string(chosen,, context.default_allocator);

  input_type: *OrtTypeInfo;
  if ort_check(api.SessionGetInputTypeInfo(result.session, 0, *input_type)) {
    tensor_info: *OrtTensorTypeAndShapeInfo;   // belongs to input_type
    if ort_check(api.CastTypeInfoToTensorInfo(input_type, *tensor_info)) && tensor_info {
      rank: u64;
      ort_check(api.GetDimensionsCount(tensor_info, *rank));
      result.input_rank = min(cast(s64) rank, result.input_shape.count);
      ort_check(api.GetDimensions(tensor_info, result.input_shape.data, xx result.input_rank));
    }
    api.ReleaseTypeInfo(input_type);
  }
  return result, true, "";
}

ort_release_session :: (s: *Ort_Session) {
  if s.session  ort.api.ReleaseSession(s.session);
  if s.shrink_options  ort.api.ReleaseRunOptions(s.shrink_options);
  free(s.input_name,, context.default_allocator);
  free(s.output_name,, context.default_allocator);
  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;
  data: *float;      // the values, for a float output (every model here has one)
  dims: [4] s64;
  dim_count: s64;
}

// Runs the session on one input tensor of shape, bytes long at data, with element type element.
// A session may be run from several threads at once. With shrink, and a session made with
// shrink_arena, the arena gives back the blocks no run is using once this run is done.
ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s64, element: ONNXTensorElementDataType, shape: [] s64, shrink := false) -> Ort_Output, ok: bool, error: string {
  out: Ort_Output;
  api := ort.api;
  input: *OrtValue;
  ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data, xx bytes, shape.data, xx shape.count, element, *input));
  if !ok  return out, false, error;
  defer api.ReleaseValue(input);

  ok, error = ort_check(api.Run(s.session, ifx shrink then s.shrink_options else null, *s.input_name, *input, 1, *s.output_name, 1, *out.value));
  if !ok  return out, false, error;

  raw: *void;
  ort_check(api.GetTensorMutableData(out.value, *raw));
  out.data = cast(*float) raw;
  info: *OrtTensorTypeAndShapeInfo;
  ort_check(api.GetTensorTypeAndShape(out.value, *info));
  count: u64;
  ort_check(api.GetDimensionsCount(info, *count));
  out.dim_count = min(cast(s64) count, out.dims.count);
  ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count));
  api.ReleaseTensorTypeAndShapeInfo(info);
  return out, true, "";
}

ort_run_float :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64, shrink := false) -> Ort_Output, ok: bool, error: string {
  out, ok, error := ort_run(s, memory_info, data.data, data.count * size_of(float), .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, shape, shrink);
  return out, ok, error;
}

// As ort_run_float, for a model with several outputs: those named, into outs (one each, in order),
// each released with ort_release_output. On failure none is set.
ort_run_float_outputs :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64, names: [] string, outs: [] Ort_Output) -> ok: bool, error: string {
  assert(names.count == outs.count);
  api := ort.api;
  input: *OrtValue;
  ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data.data, xx (data.count * size_of(float)), shape.data, xx shape.count, .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, *input));
  if !ok  return false, error;
  defer api.ReleaseValue(input);

  output_names := NewArray(names.count, *u8,, temp);
  for names  output_names[it_index] = temp_c_string(it);
  values := NewArray(names.count, *OrtValue, initialized = true,, temp);
  ok, error = ort_check(api.Run(s.session, null, *s.input_name, *input, 1, output_names.data, xx names.count, values.data));
  if !ok {
    for values  if it  api.ReleaseValue(it);
    return false, error;
  }
  for * out: outs {
    out.* = .{value = values[it_index]};
    raw: *void;
    ort_check(api.GetTensorMutableData(out.value, *raw));
    out.data = cast(*float) raw;
    info: *OrtTensorTypeAndShapeInfo;
    ort_check(api.GetTensorTypeAndShape(out.value, *info));
    count: u64;
    ort_check(api.GetDimensionsCount(info, *count));
    out.dim_count = min(cast(s64) count, out.dims.count);
    ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count));
    api.ReleaseTensorTypeAndShapeInfo(info);
  }
  return true, "";
}

ort_release_output :: (out: *Ort_Output) {
  if out.value  ort.api.ReleaseValue(out.value);
  out.* = .{};
}