DEV Community

Cover image for The Model Is a Variable: Compiling a Neural Network Into Your C++ Binary
Eugene Ostroukhov
Eugene Ostroukhov

Posted on • Originally published at uchenml.tech

The Model Is a Variable: Compiling a Neural Network Into Your C++ Binary

Overview

I strongly believe software frameworks cannot be designed in the abstract.
UchenML came out of my experiments integrating machine learning into edge
applications, and out of what I want to build next — distributed, heterogeneous
intelligence.

The first demo, dots, puts a trained neural network in your
browser tab as your opponent. No inference server, no API call — a 418 KB
WebAssembly binary and a weights file, served as static assets. (It still
carries introspection hooks for a future demo; stripped, it gets smaller.)

Behind it: 2,760,322 parameters packed to fp16, and a 128-channel residual trunk
running eleven 3×3 convolutions across the 24×24 board for every position it
evaluates — up to 128 positions per move, all in the browser tab. That model was
trained with UchenML and deployed with UchenML. It is almost certainly bigger
than it needs to be — I spent the training budget on making it stronger, not on
making it smaller.

I am not currently able to publish the source code. There is a
publicly available snapshot, about a year
old. I plan to publish the current codebase, but I cannot commit to a timeline
yet.

C++ ML Framework

I chose C++ for ease of integration and portability. With Python, the model
lives in a separate world, and reaching it from the application that needs it
means writing and maintaining glue. I want the model inside one software
development workflow: same repository, same build, same review, same tests.

The longer-term bet is heterogeneous, distributed intelligence: not one large
model behind an API, but many small specialized ones running where the data
already is — in the browser tab, inside the app, or in a cloud service. That
only works if a model is cheap to embed anywhere, which makes the deployment
target the design constraint rather than an afterthought.

UchenML is C++20, built with Bazel, and tested on Visual C++, GCC, and Clang. CI
runs the suite on Linux, macOS, Windows, and under Emscripten for WebAssembly.
The compute backend is hand-written portable SIMD on Google's Highway, so a
single kernel source covers AVX-512, NEON, and WASM SIMD with no
architecture-specific branch in the library. Small, fast, and easy to drop into
an existing C++ project is the whole design brief.

The model is a variable — I usually make it constexpr — declared in a header.
Inference, training, and everything else are templated on that variable, so the
same information is never stated twice. You never describe a layer's input type
when the previous layer already declared its output; the compiler works that
out, along with parameter counts and scratch buffer sizes. There are extra
facilities for introspection and for advanced use cases, but the core is a
simple composition of layers.

The framework is split into clearly defined modules, and you only pay for what
you use — convolution, RNN, and attention are separate targets, and depending on
one does not drag in the others. Training mirrors that structure in a parallel
tree, and deliberately stays out of the runtime one.

Defining a Model

Models compose with |:

#include "uchen/layers.h"
#include "uchen/linear.h"

constexpr auto model = uchen::layers::FloatModel<1>
                     | uchen::layers::Linear<2>
                     | uchen::layers::Relu
                     | uchen::layers::Linear<1>;

// (1×2 weights + 2 biases) + (2×1 weights + 1 bias)
static_assert(model.all_parameters_count() == 7);
Enter fullscreen mode Exit fullscreen mode

Linear<2> declares only its output width; the input width is deduced from
whatever it is piped onto. That is why the same descriptor composes anywhere in
a stack without restating shapes — and why getting it wrong is a compile error.

Because layers are values, a block of architecture is a named constant. This is
the residual block the dots network is built from — two padded 3×3 convolutions
on one branch, identity on the other, summed and activated — with C channels
over an S×S grid:

template <size_t C, size_t S>
inline constexpr auto kResBlock =
    uchen::layers::Fork<2>
  | uchen::layers::Parallel(uchen::layers::Conv2d<C, 3, 3, 1, 1>
                              | uchen::layers::Relu
                              | uchen::layers::Conv2d<C, 3, 3, 1, 1>,
                            uchen::Layer<>())
  | uchen::layers::Join
  | uchen::layers::Reduce<uchen::PlusOp, 2>
  | uchen::layers::Relu
  | uchen::layers::Reshape<uchen::convolution::ConvolutionInput<C, S, S>>;
Enter fullscreen mode Exit fullscreen mode

kResBlock<64, 16> and kResBlock<96, 24> are different types with different
parameter counts, both resolved at compile time. A trunk is a constexpr
pipeline of these; multi-head outputs come from Fork and Parallel the same
way, with each head's loss supplied separately at training time.

The next post will dig into real-world model definitions and the more advanced
composition they need.

Inference

The model variable is a functor, so evaluating it is a call, and a
ModelParameters object supplies the weights:

uchen::ModelParameters parameters(&model, weights);  // std::span<const float>
auto output = model(input, parameters);
Enter fullscreen mode Exit fullscreen mode

An optional third argument preallocates the scratch space. Its size is another
compile-time property of the model, and passing it is what makes the forward
pass allocation-free:

// The context is usually too large for the stack, so it goes on the heap. It
// can equally come from a memory-mapped region, an arena, or anywhere else.
auto context = std::make_unique<uchen::ModelContext<Model>>(parameters);
auto output = model(input, parameters, *context);
Enter fullscreen mode Exit fullscreen mode

One context can be reused across calls. Reusing it is not thread-safe, but the
model is a pure function of its input and parameters, so a context per thread
runs in parallel — and the parameters, read-only throughout, are shared between
them.

The parameter count is a property of the type too, so the weight buffer is a
fixed-size array sized by the compiler — no allocation, no length field to get
wrong:

static constexpr size_t kParamCount = uchen::ModelParameters<Model>::P;
std::array<float, kParamCount> weights{};
Enter fullscreen mode Exit fullscreen mode

The on-disk format is that same flat array of floats, so weights can be memory
mapped rather than parsed. In the browser build the array is filled straight
from the ArrayBuffer the page fetched: the demo ships its weights fp16-packed
— 5.5 MB on the wire, widened to floats on load, since WASM offers no fp16 SIMD
— and binds them with one span. There is no model loader, no operator
registry, no interpreter warm-up. Once the bytes have arrived, the network is
ready.

Training a Model

Training uses the same model definition, with uchen/training/ added. Fitting a
piecewise-linear function, end to end:

#include "uchen/training/training.h"

uchen::training::TrainingData<uchen::Vector<float, 1>,
                              uchen::Vector<float, 1>> data(samples.begin(),
                                                            samples.end());

uchen::training::Training training(&model,
                                   {&model, kInitialWeights},
                                   uchen::training::SgdOptimizer<Model>{});

while (training.Loss(data) > 1e-4f) {
  training = training.Generation(data, 0.001f).next;
}

uchen::ModelParameters parameters = training.parameters();
Enter fullscreen mode Exit fullscreen mode

TrainingData is a collection of input-output pairs. It also handles splitting
a dataset into batches and reserving validation records.

Generation does not mutate anything — it returns the next training state,
which makes a step cheap to inspect, discard, or checkpoint, and makes the loop
read as a fold over generations. Optimizers are pluggable: SgdOptimizer
(stochastic gradient descent), momentum SGD, and AdamOptimizer come out of the
box, alongside Kaiming-He initialization, per-layer gradient statistics, and
multi-threaded batch gradients.

The loss function is deduced from the model's output type, and the framework
provides a few common ones. Pass your own to the Training constructor instead
and its gradient is computed automatically through the chain.

What comes out is a flat float array — and that is the entire deployment
format, the same bytes the browser build maps into its std::array. Training
and serving cannot disagree, because there is nothing for them to disagree
about.

Training is multi-threaded, but it does not by itself scale to real-world
runs. The next post will discuss the primitives it is built from, and how I used
them to run a larger, more reliable training pipeline with Bazel as the executor.

Extending the Framework

Adding a layer means writing a type the composition machinery recognizes. The
runtime contract is small:

class MyLayer {
 public:
  using input_t  = uchen::Vector<float, 8>;
  using output_t = uchen::Vector<float, 4>;
  static constexpr size_t parameter_count = 36;
  using scratch_area_t = std::array<float, 4>;

  output_t operator()(const input_t& input,
                      const uchen::Parameters& parameters,
                      memory::LayerContext<scratch_area_t>* context) const;
};
Enter fullscreen mode Exit fullscreen mode

scratch_area_t comes from the model context, and the framework controls its
lifetime. The layer uses it for its output and for any intermediate buffers it
needs.

operator() is reached through a dispatcher that asks the compiler which of
those signatures your layer actually defines, so a layer takes only what it
needs — the input alone, the input and its parameters, or those plus a scratch
context. Composite layers own no storage; they slice the model's single flat
parameter blob at compile-time offsets, which is why a Fork of four heads
still costs exactly one allocation-free array.

A pipeable descriptor turns the type into something that composes:

inline constexpr uchen::Layer<MyLayerDesc> MyLayer;
Enter fullscreen mode Exit fullscreen mode

The descriptor's job is to map "the model so far" onto a concrete layer type,
binding the input shape the caller never had to write down. It hooks in through
an ADL-found StackLayer(model, desc). It is also where optimizations hook in —
kernel fusion, or eliding a layer that is a no-op at the given input shape.

Gradients are a separate, optional surface. A layer becomes trainable when an
overload of ComputeGradients for it is visible:

auto ComputeGradients(const MyLayer& layer, const MyLayer::input_t& input,
                      const Grad& output_gradient,
                      const uchen::Parameters& parameters,
                      std::span<float, MyLayer::parameter_count> gradients,
                      const void* scratch);
Enter fullscreen mode Exit fullscreen mode

Every layer in the framework — convolution, RMSNorm, pooling, RNN, softmax —
arrived through exactly this path. There is no privileged built-in set.

Conclusion

UchenML is still very much a work in progress, but it has reached the point
where I can hand agentic tools a model and let them build and train it — which
frees me to spend more time on the application side. The feature list is long.
Among the things I want to add:

  • Integration with common ML tools such as TensorBoard, for visualizing and monitoring training runs.
  • ONNX support, for importing models to and from other frameworks.
  • More backends. The current one was designed for WebAssembly first, so it is single-threaded and fp32. I have clear ideas about a Metal backend, and that project will be a lot of fun.
  • A generalized training pipeline, lifted out of the dots project. It has proven robust, surviving power loss and a full disk. Adding distributed training should be straightforward, since Bazel already underpins it.

I plan to post every few weeks. Next up is a deep dive into the dots project and
how UchenML held up in the real world, followed by articles on the optimization
techniques inside the framework, kernel fusion among them. More demos are in the
works too.

Top comments (0)