ggml_graph.hpp - mozsearch

Enable keyboard shortcuts

#pragma once

#include <functional>

#include <vector>

struct ggml_context;

struct ggml_tensor;

namespace pk {

// One-shot CPU graph runner.

//

// Allocates a ggml context of `mem_bytes` (no_alloc=false, so tensor data is

// stored inline in the context buffer). Calls `build(ctx)` which must create

// input tensors, write their data, build the computation graph, and return the

// output tensor. Then the forward graph for that output is built, executed on

// CPU with `n_threads` threads, and `ggml_nelements(output)` f32 values are

// copied into `out`. Returns true on success, false on any failure.

bool run_graph(size_t mem_bytes, int n_threads,

               const std::function<ggml_tensor*(ggml_context*)>& build,

               std::vector<float>& out);

// Process-global override for the ggml compute thread count.

//

// `run_graph` reads this override and, when it has been explicitly set to a

// positive value, uses it in place of the per-call `n_threads` argument that

// the components pass. This lets a single `--threads N` switch control EVERY

// graph computation (the encoder is the bulk) without threading a thread-count

// parameter through every component.

//

// The default (0 == "unset") means: honor whatever `n_threads` each caller

// passes. So existing behavior — and the test suite, which never sets this —

// is unchanged. Setting it back to 0 clears the override.

void set_num_threads(int n);

int  num_threads();  // current override (0 == unset)

// Gallocr buffer size (bytes) reserved for the most recent single-backend (CPU)

// run_graph compute. Used by tests to assert that banded attention memory scales

// O(T*window), not O(T^2). Reflects the high-water mark of the persistent

// gallocr, so query it after a fresh run at the size of interest.

size_t last_graph_alloc_bytes();

class Backend;

// The process-global persistent Backend (created lazily on first use). Exposed

// so the weight-realization path can give the loader's tensors a backend buffer

// on the SAME CPU backend that graphs run on.

Backend& global_backend();

// Free the process-global backend. Call once at program exit (after all model

// objects are destroyed) so GPU backends release device memory while the driver

// is still alive — otherwise static destruction frees it after the CUDA atexit

// handler and aborts. A later global_backend() call recreates it.

void shutdown_backend();

} // namespace pk