# ZML Concepts

## Model lifecycle

ZML is an inference stack that helps running Machine Learning (ML) models, and
particulary Neural Networks (NN).

The lifecycle of a model is implemented in the following steps:

1. Open the model file and read the shapes of the weights, but don't load the weights 
   yet.

2. Using the loaded shapes and optional metadata, instantiate a model struct
   with `Tensor`s, representing the shape and layout of each layer of the NN.

3. Compile the model struct and it's `forward` function into an accelerator
   specific executable. The `forward` function describes the mathematical
   operations corresponding to the model inference.

4. Load the model weights from disk, onto the accelerator memory.

5. Load some user inputs, and copy them to the accelerator.

6. Call the executable using the weights and the user inputs.

7. Fetch the returned model output from accelerator into host memory, and
   finally present it to the user.

8. When all user inputs have been processed, free the executable resources, weights and
   inputs.


**Some details:**

Note that the compilation and weight loading steps are both bottlenecks to your
model startup time, but they can be done in parallel thanks to Zig's `std.Io` interface.

The accelerator is typically a GPU, but can be another chip, or even the CPU
itself, churning vector instructions.


## Tensor Bros.

In ZML, we leverage Zig's static type system to differentiate between a few
concepts, hence we not only have a `Tensor` to work with, like other ML
frameworks, but also `Buffer`, `Slice`, and `Shape`.

Let's explain all that.

* `Shape`: _describes_ a multi-dimension array.
    - `Shape.init(.{16}, .f32)` represents a vector of 16 floats of 32 bits
      precision.
    - `Shape.init(.{512, 1024}, .f16)` represents a matrix of `512*1024` floats
      of 16 bits precision, i.e. a `[512][1024]f16` array.

    A `Shape` is only **metadata**, it doesn't point to or own any memory. The
    `Shape` struct can also represent a regular number, aka a scalar:
    `Shape.init(.{}, .i32)` represents a 32-bit signed integer.

* `Slice`: _is_ the combination of a `Shape` and raw bytes (bytes that are **on the CPU**).
  - can own the underlying memory - but can also accomodate non-owned memory.

* `Buffer`: _is_ a multi-dimension array, whose memory is allocated **on an
  accelerator**.
    - contains a handle that the ZML runtime can use to convert it into a
      physical address, but there is no guarantee this address is visible from
      the CPU.
    - can be created from a `Slice` by calling `Buffer.fromSlice(...)`. 

* `Tensor`: is a mathematical object representing an intermediary result of a
  computation or an input to an executable (including weights).
  - is basically a `Shape` with an attached MLIR value when in the context of a compilation.


## The model struct

The model struct is the Zig code that describes your Neural Network (NN).
Let's look a the following model architecture:

![Multilayer perceptrons](https://raw.githubusercontent.com/zml/zml.github.io/refs/heads/main/docs-assets/perceptron.png)

This is how we can describe it in a Zig struct:

```zig
const Model = struct {
    input_layer: zml.Tensor,
    output_layer: zml.Tensor,

    pub fn forward(self: Model, input: zml.Tensor) zml.Tensor {
      const hidden = self.input_layer.matmul(input);
      const output = self.output_layer.matmul(hidden);
      return output;
    }
}
```

NNs are generally seen as a composition of smaller NNs, which are split into
layers. ZML makes it easy to mirror this structure in your code.

```zig
const Model = struct {
    input_layer: MyOtherLayer,
    output_layer: MyLastLayer,

    pub fn forward(self: Model, input: zml.Tensor) zml.Tensor {
      const hidden = self.input_layer.forward(input);
      const output = self.output_layer.forward(hidden);
      return output;
    }
}
```

`zml.nn` module provides a number of well-known layers to more easily bootstrap
models.

Since the `Model` struct contains `Tensor`s, it is only ever useful during the
compilation stage, but not during inference. If we want to represent the model
with actual `Buffer`s, we can use the `zml.Bufferize(Model)`, which is a mirror
struct of `Model` but with a `Buffer` replacing every `Tensor`.

## Strong type checking

Let's look at the model life cycle again, but this time annotated with the
corresponding types.

1. Open the model file and read the shapes of the weights -> `zml.Shape`

2. Instantiate a model struct -> `Model` struct (with `zml.Tensor` inside)

3. Compile the model struct and its `forward` function into an executable.
   `foward` is a `Tensor -> Tensor` function, executable is a
   `zml.Executable`

4. Load the model weights from disk, onto accelerator memory ->
   `zml.Bufferized(Model)` struct (with `zml.Buffer` inside)

5. Load some user inputs (custom struct), encode them into arrays of numbers
   (`zml.Slice`), and copy them to the accelerator (`zml.Buffer`).

6. Call the executable on the user inputs. `module.call` accepts `zml.Buffer`
   arguments and returns `zml.Buffer`

7. Return the model output (`zml.Buffer`) to the host (`zml.Slice`),
   decode it (custom struct) and finally return to the user.