ONNX models on the GPU (experimental)

This chapter is experimental. [OnnxGpu("…")] takes an image in a buffer and returns the result in another buffer.

The example model

Model file	filter.onnx, in the same folder as Sharpen.cs
Input	image — a 64×64 image with 3 components (shape [1,3,64,64])
Output	out — an image of the same size with 3 components (shape [1,3,64,64])
Layers	Four convolutions. The channel count in between grows 8 → 16 → 16
The generated file	Sharpen.Filter.Model.cs
How to bake one in	As in How to bake one in

The script carrying the attribute

using Tsukimi;
using UnityEngine;

// The attribute argument is the path to the model file, seen from this script.
[OnnxGpu("filter.onnx")]
public partial class Sharpen : TsukimiBehaviour
{
    // The source image and the surface the result is shown on. Assign them in the Inspector.
    public Texture source;
    public Renderer display;

    void Start()
    {
        // Creates the buffers. Their sizes come from the model, so you do not give them.
        LoadFilter();

        // Puts the source image into the input buffer. A different size is scaled to fit.
        Gpu.Load(FilterInput, source);
    }

    void Update()
    {
        // Runs one pass. One kernel per layer runs in order.
        RunFilter();

        // The result only stays in FilterOutput, so it is shown on the surface as it is.
        Gpu.Show(FilterOutput, display);
    }
}

The generated file

This is an excerpt, with the layer bodies left out.

using Tsukimi;
using UnityEngine;

public partial class Sharpen
{
    /// Input. Write values from 0 to 1 (mapped to 0 to 1).
    public GpuBuffer2D FilterInput;
    /// Output. Read values from 0 to 1 (corresponding to -2.39134 to 1.3269).
    public GpuBuffer2D FilterOutput;

    // Buffers between layers. The names begin with the model's name.
    private GpuBuffer2D Filterb1;
    // …(one for every layer)…

    // The weight table for each layer.
    public Vector4[] Filterw0;
    // …(one for every layer)…

    /// Loads the weights and creates the buffers. Call once from <c>Start</c>.
    private void LoadFilter()
    {
        // The entry is 64×64. The channel count is 4 or fewer, so it does not extend sideways.
        FilterInput = Gpu.Buffer(64, 64);
        // The next layer has 8 channels — two bands of 4 components — so the width doubles.
        Filterb1 = Gpu.Buffer(128, 64);
        // …(one for every layer)…
        FilterOutput = Gpu.Buffer(64, 64);
    }

    // One kernel is generated per layer.
    [Kernel]
    static Color4 FilterLayer0(KernelId id, GpuBuffer2D prev, [Capacity(128)] Vector4[] t)
    {
        // …(the multiply-accumulates of the convolution follow)…
    }

    /// Runs the model once. Write <c>FilterInput</c> before calling, then read <c>FilterOutput</c>.
    private void RunFilter()
    {
        Gpu.Run(nameof(FilterLayer0), Filterb1, FilterInput, Filterw0);
        // …(one for every layer)…
    }
}

The generated names

Generated names	Description
FilterInput	The input buffer (the same type as in Running a GPU kernel). The range of values to assign to a cell is written in the generated file, next to the input field.
FilterOutput	The output buffer. For an 8-bit model, read its cells as values from 0 to 1; for a floating-point model, read the values as they are
LoadFilter()	Creates the buffers and loads the weights. Call it once from Start
RunFilter()	Runs one pass
FilterSteps	The number of steps in one pass
RunFilter(int first, int count)	Runs steps first to first + count - 1 only. Running all the steps in order gives the same result as RunFilter(), however many frames they are spread over. Write FilterInput after the last step, and read FilterOutput after the last step too. Each step is one draw, so this spreads the steps over frames; each step stays as heavy as before
FilterLayer0	The kernel for each layer. Gpu.Run launches it, so you never call it directly
FilterInput1 / FilterInput2 …	Generated only for a floating-point model with two or more inputs. The number follows the order of the model’s inputs
SetFilterTokens(int[] ids)	Generated only for a model that takes token ids. Instead of filling FilterInput, it writes one id per position
FilterStyles	Generated only when 2 or more paths are listed. The number of paths listed
UseFilter(int style)	Generated only when 2 or more paths are listed. Chooses which weights to use by their position in the list (counting from 0). Right after LoadFilter() it is 0

Value range	For an 8-bit model, assigning a value between 0 and 1 to a cell passes it through correctly, because the 8-bit steps line up exactly with the range the model accepts. For a model that takes its input as 16-bit floats, and for a floating-point model, assign the values as they are. Which one applies is decided by the model, and is written in the input field of the generated file.
What the output means	The range the output values actually stand for is written in the generated file.
8-bit format	Values between layers can be unsigned (uint8) or signed (int8). Weights are signed
Input of an 8-bit model	A picture, [N, C, H, W] (N is the batch size, from 1 to 16; see Models with a batch size above 1). An input whose number of axes is not 4 is an error that names the input (the inputs of a floating-point model are described under “Batch and shape” in Floating-point models)
Scales of an 8-bit convolution	A convolution whose input or output has a scale for each channel is an error. Weight scales are taken only for each output channel (an axis of 0 on the weight’s DequantizeLinear); scales for each input channel (an axis other than 0) are an error. The default axis is 1, so write axis=0 for scales for each output channel
Reading it back	The result lands in a buffer. To use it as numbers, take it out with Gpu.Texture and hand it to the runtime’s asynchronous readback
Swapping weights	Models that share one graph and differ only in their weights become one when their paths are listed, as in [OnnxGpu("a.onnx", "b.onnx")]. The buffers and kernels stay as one set, and only the weights are kept once per listed model

Floating-point models

A model with no quantization operators at all can also be baked in with the GPU attribute.

Number format	Computes in 16-bit floating point. The weights are also rounded to 16 bits when the file is generated
Input and output	Assign the values as they are and read them back as they are. The buffers are GpuFormat.Half
Layout of the values	One cell holds four channels. Which axis goes to which row and cell is written above the declarations of the inputs and the output in the generated file
Token ids	A model that looks up a vocabulary table by id (a language model) takes its ids through SetFilterTokens. The number of positions is fixed at export, and converting between text and ids is not included
Two or more inputs	Generated as FilterInput, FilterInput1, FilterInput2, in the order of the model’s inputs
Only the weights in 8 or 4 bits	When the weights are turned back by DequantizeLinear, they are kept at that size. Eight bits halve their size and four bits take about a quarter, plus a 16-bit scale for each block when the scales are per block. The values between layers stay 16-bit, and each layer can use its own number of bits
Weight integers and scales	The integers are int8, uint8, int4 or uint4. The scale is one for the whole weight, one per output channel, or one per block of block_size values (a multiple of 4). Blocks run along the input axis of a MatMul or Gemm and along the input channels of a Conv
MatMulNBits	MatMulNBits from the com.microsoft domain is read as the same product with 4-bit weights. accuracy_level is ignored, and the values it multiplies stay 16-bit. Float zero points, g_idx, a bias input and bits other than 4 are errors
Weights kept in 16 bits	Weights that a layer cannot use at their reduced size (blocks along the output channels, a Conv with group, a ConvTranspose, or a Gather table, for example) are turned back into real numbers and kept as 16-bit values, so they take as much memory as before. Compilation succeeds, and one TUKI0123 warning per layer names the layer and the reason
Batch and shape	Every size has to be fixed when the model is exported. Each input has 2 to 4 axes ([N, C, H, W], [N, T, D] or [N, D], where N is the batch size, from 1 to 16; see Models with a batch size above 1), and an input that a MatMul or Gemm reads only as its weight (its second input) is a matrix with one row. An input outside this is an error that gives its number of axes or its number of rows. A model with no shape annotations at all should be run through shape inference before it is saved. The layers of a model do not have to be listed in the order they run (like the original runtime, the conversion puts each layer after the layers whose outputs it reads). Layers that read each other’s outputs in a loop cannot be put in such an order, so the model is an error that names those layers (the original runtime refuses such a model too)
Inputs that give a shape, axes or a size	The shape of Reshape and Expand, the axes of Squeeze, Unsqueeze and the reductions, the repeats of Tile, the scales and sizes of Resize, and the starts, ends, axes and steps of Slice are constants. One given as an input of the model, or computed from one, is an error that names the layer and the input. The output shape that Reshape, Squeeze, Unsqueeze, Flatten, Expand, Tile, Slice, Split and a nearest Resize compute from these values has to match the shape the model records for that output (the original runtime goes by the values and drops a recorded shape that differs). A model where they differ is an error; running shape inference again rewrites the record
Inputs left unread	An input of the model that a layer reads, but that the converted model would leave unread, is an error that names the input
Element types	Tensors are float32 (integer token ids read by a Gather are the exception). A tensor of another type (an integer, float16, float64 and so on) is an error that names the type. The weights of a convolution are constants; weights given as an input of the model or computed by it are an error
The 16-bit limit	A 16-bit float holds magnitudes up to 65504, and a larger value becomes 65504 (or -65504 when negative), not infinity. A layer that can go past it (ReduceProd, sums of squares and variances, powers, functions that grow faster than their input such as Exp, products of two values whose sizes are not known, a layer that divides by a computed value, and a layer that goes past it even with a known input range) compiles, and a TUKI0135 warning names it. Layers with the same reason share one warning, which says how large the input has to get
Operators whose result jumps	The input and the values between layers are rounded to 16 bits when the GPU writes them (toward zero, as measured on one PC graphics card). A value just past a point where the result jumps can be rounded back across that point, and Ceil, Floor, Round, Sign, Mod, ThresholdedRelu and Shrink then give a different answer from the original model (the difference for each operator is in ONNX model operators (experimental)). “Just past” means within one 16-bit step, and the step is about 0.004 near 5, about 0.06 near 100, and 1 from 1024 to 2048
Weights in a separate file	A model whose weights were written to a separate file beside it (the default of recent PyTorch exports) is read when that file sits next to the model
NaN	NaN in the input is not carried through to the output as it is in the original model. A layer read from a table gives its value at 0 (0.5 for Sigmoid), and the shader compiler turns the comparisons in Relu, Max, Min and Clip into max and min instructions, which can give the other value or the bound instead of NaN. The generated code cannot change this, so keep NaN out of the input
Tensor constants	A constant of more than one value that a layer combines with a computed value (a weight tensor in an Add, Sub, Mul or Div, or a BatchNormalization turned into one) is held in 16 bits as well. A single value is kept as it is. A multiplier, or a constant divided by a computed value, that is above 65504 or so small that 16 bits hold 0 for it (below about 6e-8) while the result could still be held in 16 bits, is an error that names the layer and the value (a divisor is judged by its reciprocal). An added or subtracted constant above 65504 in size is an error
A mask before Softmax	A mask of -65504 or less added just before a Softmax is taken when every row the Softmax normalizes keeps at least one value 16 bits can hold (the masked positions come out 0 either way). A mask covering a whole row is an error (the original model makes that row even). Added or subtracted constants too small for 16 bits, and the constants of Max, Min and Where, are taken (the result changes by less than one 16-bit step)

The layout line (excerpt from the generated file)

    /// Input. Write the values as they are (not mapped to 0 to 1).
    /// Layout (shape [1, 4, 8, 8]): y is the index along axis 2, x the index along axis 3; channel c is the index along axis 1; the value is component c % 4 of cell (x + (c / 4) * 8, y).
    public GpuBuffer2D FilterInput;

Models with a batch size above 1

Accepted form	A model whose inputs all have the same batch size (1 to 16) on axis 0, and whose every layer computes each item of the batch on its own. This holds for floating-point models and for models reduced to eight bits
How it runs	The generated code runs all the steps once for each item, one item after another, and gives each item its own input and output buffers
Buffers for each item	FilterInput and FilterOutput hold item 0, FilterInputItem1 and FilterOutputItem1 hold item 1, and so on (FilterInput1Item1 for item 1 of a second input)
FilterItems	The batch size
Cost	A run takes the batch size times the steps and the GPU work of one item. FilterSteps counts all of them, so RunFilter(first, count) spreads them over frames in the same way
Layers that mix the items	An error that names the layer. These are a Softmax, LayerNormalization, RMSNormalization, Concat, Split, Gather, ArgMax or reduction along axis 0 (a reduction without axes reduces axis 0 too), a MeanVarianceNormalization with its default axes (0, 2 and 3), a Transpose that moves axis 0, a Reshape or Flatten that merges the batch with other axes, a Squeeze, Unsqueeze, Tile, Pad, Slice or Resize that changes axis 0, a MatMul or Gemm that multiplies along axis 0 or moves the batch off axis 0 by broadcasting, an Expand whose shape has more axes than its input (the added axes go in front, so the batch moves off axis 0), and a BatchNormalization in training mode
Other forms that are errors	A layer not known to compute each item on its own, a constant that differs from one item to another, inputs with different batch sizes, and token ids with a batch size above 1
How axes are counted	On the model as exported, with its batch size. An axis counted from the end (such as -1) of a value whose shape the model does not record is an error, since it cannot be told whether it is axis 0 (shape inference records the shapes). A Squeeze without axes is an error too, because with a batch size of 1 it would remove axis 0 as well (write its axes)
Exporting with a batch size of 1 and running each item	A model that is an error can still be exported with a batch size of 1 and run once for each item. That gives the same result as the original only when the original computes each item on its own; with a layer that mixes the items, the result differs

Models that are accepted

Forms that produce an error	What to do	Reason
Only partly quantized	Quantize the whole model, or export it with no quantization operators at all	A model with even one quantization operator is read as an 8-bit model, so values without a scale become errors. A floating-point model whose weights alone are in 8 or 4 bits is not this case
The scale is decided at run time	Export it in a form that is fixed at export time	To bake it in, the value steps have to be fixed at compile time
Not an 8-bit integer	Quantize to 8-bit integers	Even at 8 bits, a format that represents fractions is not accepted
Uses comparisons or outputs integer indices	Make the condition constant. For indices such as TopK, output the scores and pick on the Udon side	Comparisons between outputs of earlier operators, and operators that return integer indices, are not accepted. 16-bit values hold integers exactly only up to 2048. In a real-valued model, ArgMax and ArgMin are accepted when the axis they pick along has at most 2048 values and their output is used only as a model output
Dividing by a Clip whose lower bound 16 bits cannot hold	Raise the lower bound to a value 16 bits can hold	The smallest 16-bit value is about 6e-8, so a lower bound such as 1e-10 cannot be held, and the quotient comes out far off wherever the bound applies. A Div that reads it as the divisor directly, or through layers that only change the shape (Reshape, Transpose, Unsqueeze, Squeeze, Identity, Flatten), is an error. So is a Log or Reciprocal that reads it through those layers or Expand (a Log or Reciprocal directly after the Clip is computed together with it, so it is not an error)
Adding a value too small for 16 bits	Use a value 16 bits can hold	An Add or Sum of 1e-8, or a Max taking the larger of the input and that value, comes out 0 in 16 bits where it should be 1e-8, and dividing there gives infinity or NaN. It is an error when that result is read as the divisor of a Div, or by a Reciprocal, a Log or a Pow with a negative exponent (Pow(x, -0.5), for example). The same holds through layers that keep a 0 as 0 (layers that only change the shape and Expand, functions such as Sqrt, Abs, Neg and Relu, reductions such as ReduceSum and ReduceMean, and layers that combine it with other computed values by Add, Sub, Sum, Mean, Max, Min or a product), for example (x + 1e-8) + y and Sqrt(ReduceSum(Relu(x) + 1e-8)). A MatMul, Gemm or convolution on the way is not followed through, since its result is 0 only when every value it reads is. The same holds for a run of element-wise layers (such as an Add or Sqrt after a Clip with a small lower bound) whose result, for some input, is 0 in 16 bits or above 65504
An operator from a nonstandard domain	Export with standard ONNX operators	An operator from another domain is not accepted even if its name is the same. Only MatMulNBits from com.microsoft is read
The shape is dynamic	Fix the input shape before exporting	When an axis length is decided at run time, the texture size cannot be determined
A listed model differs from the first in more than its weights	Export it with the same graph, changing only the weights	Names the layer that differs. A model whose scales built into the computation differ is also an error here
An unsupported operator is present	Rewrite it into a form that does not use that operator	It names what was present. It never leaves one out and passes the model anyway

The error speaks in terms of the model, so fix the exporting side.

Limits

One cell holds 4 components	For a model whose channel count exceeds 4, the texture extends sideways
Where the weights are kept	Small layers are kept as constants and large ones as textures. Which one applies is decided per layer
Stride	strides other than 1 are handled for Conv and MaxPool
Size	Each layer’s output texture can be at most 16384 texels wide and 16384 texels tall. The length of a one-dimensional tensor goes on the height, so the length is also limited to 16384
Class declaration	partial is required. The generated file becomes the rest of that class