ops

type AddOp ¶

type AddOp struct{}

AddOp implements the ONNX Add operator. It computes the elementwise sum of two input tensors with broadcasting support.

func (AddOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type AveragePoolOp ¶

type AveragePoolOp struct{}

AveragePoolOp implements the ONNX AveragePool operator for 2D NCHW inputs. It supports configurable kernel_shape, strides, pads and the count_include_pad flag. ceil_mode is supported for both 0 (floor, the default) and 1 (ceil) output-size rounding; dilations are assumed 1 (the default — the F5DCTNet v1/v1s trunks only emit plain 2x2/s2 downsample pools).

Why this op exists in a "minimal CNN runtime": the v1 and v1s F5DCTNet trunks downsample with AveragePool (the v1ss trunk replaced it with strided convolution). Running the v44 v1/v1s ensemble members in pure Go therefore requires this operator.

func (AveragePoolOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type BatchNormalizationOp ¶

type BatchNormalizationOp struct{}

BatchNormalizationOp implements the ONNX BatchNormalization operator in inference mode (test mode):

y = (x - mean) / sqrt(var + epsilon) * scale + bias

Inputs (positional, all 1-D length C except x):

• x: [N, C, ...spatial...]
• scale: [C]
• bias: [C]
• mean: [C]
• var: [C]

The channel axis is fixed at index 1 (NCHW layout). For ranks > 2 the remaining trailing dims are treated as spatial. For rank-2 input [N, C] the operator degenerates to a per-element affine transform.

Only the data output (output 0) is produced — the optional auxiliary outputs (running mean / running var) used during training do not exist in inference graphs we care about.

func (BatchNormalizationOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type CastOp ¶

type CastOp struct{}

CastOp implements the ONNX Cast operator. The runtime tensor type is Float32 in every internal kernel, so Cast is functionally a no-op when the requested target type is also a real-valued type (Float32 / Float16 / Float64) and is otherwise a value-by-value coercion to the integer target type.

Because the operator's job is to bridge ONNX-declared graph types to the Float32 runtime, the implementation does not change the bit pattern of the underlying float32 buffer — it simply returns a tensor aliasing the input data with the same shape. Models that rely on Cast to perform integer truncation (e.g. Cast(float → int32) → downstream Reshape that uses the int32 buffer as an index list) are not supported here; that workload is rare in inference graphs and every observed F5 / Real-ESRGAN model uses Cast only to bridge an integer-typed graph input into the float32 compute domain.

Attribute contract: the ONNX `to` attribute carries the target DataType as an int. CastOp records the value purely so future ops that introspect the cast can reason about it; the Execute method itself is type-agnostic.

func (CastOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ClipOp ¶

type ClipOp struct{}

ClipOp implements the ONNX Clip operator (opset 11+). It clamps each element of the input tensor between optional min and max scalar bounds provided as tensor inputs.

func (op ClipOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ConcatFusedTailOp ¶

type ConcatFusedTailOp struct{}

ConcatFusedTailOp is the runtime counterpart of the graph.ApplyConcatWriteFusion pass. It implements a fused (Conv-variant → Concat) tail pattern where the producer Conv writes its output directly into the tail slice of the joint Concat output buffer, eliminating both the producer's standalone output materialisation and the corresponding memcpy that an unfused Concat would perform on that tail slice.

Input layout (set by the fusion pass):

[prefix_0, prefix_1, ..., prefix_{P-1}, x, w, bias_or_nil, slope?]

where P = "prefix_count" attribute, the prefix_i are the original Concat inputs other than the last, and the trailing inputs are the producer Conv's inputs in their canonical positions (slope is present only when producer_op == "ConvPRelu"). Bias may be nil for two-input Convs.

Attribute contract:

axis           int    — Concat axis; always 1 in the production graphs but
                        carried through so the op is self-describing.
prefix_count   int    — Number of prefix inputs (0 ≤ P ≤ len(inputs)-2).
producer_op    string — "Conv", "ConvLeakyRelu", or "ConvPRelu". Tells the
                        runtime which underlying kernel to dispatch.
strides, pads, dilations — Conv attributes, forwarded from the producer.
alpha          float  — LeakyReLU slope; required when producer_op =="ConvLeakyRelu".

The op allocates the joint output tensor of shape [N, sum(prefix_C)+producer_C, H, W] from the supplied allocator, copies each prefix input into its slice with the same block-copy strategy ConcatOp uses, then dispatches the producer Conv kernel via a custom allocator that returns a "view tensor" backed by the tail slice of the joint buffer. The Conv kernel's writes therefore land directly in the joint output, avoiding the separate output buffer and the Concat memcpy.

Batch fallback: when batch != 1, the per-batch tail slice of the joint NCHW buffer is non-contiguous (strides interleave the prefix and producer regions) and the view-tensor trick would corrupt the layout. The op falls back to running the producer to a fresh tensor and copying like a normal Concat. The production SR pipeline always uses batch == 1, so this fallback is a correctness safety net rather than a hot path.

func (ConcatFusedTailOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ConcatOp ¶

type ConcatOp struct{}

ConcatOp implements the ONNX Concat operator. It concatenates variadic input tensors along the specified axis using block-copy memory transfers.

func (ConcatOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ConvLeakyReluOp ¶

type ConvLeakyReluOp struct{}

ConvLeakyReluOp is the fused Conv → LeakyReLU operator. It performs the same im2col + GEMM as ConvOp, but the gemmTransposedB write fuses the per-output-channel bias add AND the LeakyReLU activation into one pass — eliminating the intermediate output materialisation that an unfused Conv → LeakyReLU pair would incur.

In Real-ESRGAN the Conv layer is followed almost universally by LeakyReLU (279 of the 351 Conv outputs feed directly into LeakyRelu), so the savings compound at the model level.

Attribute contract (read from the merged Conv + LeakyRelu attributes produced by the graph fusion pass):

strides:    int[2]   — Conv stride; defaults to [1, 1]
pads:       int[4]   — Conv pads (top, left, bottom, right); defaults to zeros
dilations:  int[2]   — Conv dilations; defaults to [1, 1]
alpha:      float32  — LeakyReLU slope for negative values; defaults to 0.01

Inputs are the same as Conv: [x, w, bias?].

func (ConvLeakyReluOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ConvOp ¶

type ConvOp struct{}

ConvOp implements the ONNX Conv operator using im2col matrix unrolling followed by GEMM. It supports 2D convolution with configurable strides, padding, and dilations for NCHW layout.

func (ConvOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ConvPReluOp ¶

type ConvPReluOp struct{}

ConvPReluOp is the fused Conv → PReLU operator. It performs the same im2col + GEMM as ConvOp, but the gemmTransposedBPReLU write fuses the per-output-channel bias add AND the per-output-channel PReLU activation into one pass — eliminating the intermediate output materialisation that an unfused Conv → PReLU pair would incur.

In realesr-general-x4v3, 33 of the 34 Conv layers feed directly into a per-channel PReLU; collapsing those pairs is the dominant v3-specific speed win because every elimination removes one full read+write pass over the activation tensor.

Attribute contract (read from the Conv attributes carried over by the graph fusion pass; the PReLU op carries no attributes of its own):

strides:    int[2]   — Conv stride; defaults to [1, 1]
pads:       int[4]   — Conv pads (top, left, bottom, right); defaults to zeros
dilations:  int[2]   — Conv dilations; defaults to [1, 1]

Inputs: [x, w, bias_or_nil, slope]. Bias is optional (3-input Conv passes a non-nil tensor; 2-input Conv passes nil at index 2). Slope is the PReLU slope tensor and must be either a scalar (1 element, broadcast across all channels) or per-channel-broadcastable using the same logic as PReluOp ([C], [C, 1, 1], [1, C, 1, 1], …).

func (ConvPReluOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type DepthToSpaceOp ¶

type DepthToSpaceOp struct{}

DepthToSpaceOp implements the ONNX DepthToSpace operator. It rearranges data from the depth (channel) dimension into spatial blocks, producing an output tensor of shape [N, C/(r*r), H*r, W*r] from an input of shape [N, C, H, W] and a blocksize attribute r.

Two modes are supported:

• "DCR" (depth-column-row): the legacy default for opset < 13. The channel axis is interpreted as (blocksize, blocksize, C_out).
• "CRD" (column-row-depth): the default for opset >= 13. The channel axis is interpreted as (C_out, blocksize, blocksize).

We default to "DCR" to match the ONNX spec defaults; callers running opset >= 13 models that omit the attribute (which is rare since the node is usually exported with mode set explicitly) should configure it at the model-import layer.

func (DepthToSpaceOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type DropoutOp ¶

type DropoutOp struct{}

DropoutOp implements the ONNX Dropout operator. In an inference-only runtime Dropout is the identity function: it returns a copy of the input tensor regardless of the ratio attribute / input or the training_mode flag.

ONNX changed Dropout's signature across opsets:

• opset < 12: 1 input (data), `ratio` is an attribute.
• opset >= 12: up to 3 inputs (data, ratio, training_mode) and a `seed` attribute.

Both forms produce one or two outputs (data, mask). Because we never materialise the boolean mask in inference, this implementation emits only the data output and ignores ratio / training_mode entirely.

func (DropoutOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type EqualOp ¶

type EqualOp struct{}

EqualOp implements the ONNX Equal operator. For each pair of broadcasted elements it emits 1.0 when the inputs match and 0.0 otherwise. ONNX defines the output as bool, but the runtime is float32-only so the boolean is encoded as 1.0/0.0 in a float32 tensor — downstream Cast and arithmetic ops treat the encoding transparently.

func (EqualOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type FlattenOp ¶

type FlattenOp struct{}

FlattenOp implements the ONNX Flatten operator. It reshapes the input tensor into a 2-D tensor whose first dimension is the product of all input dims with index < axis and whose second dimension is the product of the remaining dims.

The axis attribute defaults to 1. Negative values count from the end: axis = -1 is equivalent to rank - 1. The row-major data layout is preserved verbatim — Flatten is a metadata-only operation in row-major storage.

func (FlattenOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type GemmOp ¶

type GemmOp struct{}

GemmOp implements the ONNX Gemm operator (General Matrix Multiply):

Y = alpha * A^* @ B^* + beta * C

where A^* = A^T when transA == 1 (otherwise A), and likewise for B. C is an optional bias tensor that is broadcast over [M, N] following ONNX/NumPy rules.

Attributes (all optional):

• alpha float, default 1.0
• beta float, default 1.0
• transA int, default 0
• transB int, default 0

The matmul work is delegated to the package-internal gemm() helper in gemm.go which selects a SIMD or scalar kernel via build tags. This operator wrapper is responsible only for transposition (when requested), the alpha/beta scaling, and the broadcast bias add.

func (GemmOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type GlobalAveragePoolOp ¶

type GlobalAveragePoolOp struct{}

GlobalAveragePoolOp implements the ONNX GlobalAveragePool operator. Given an input of shape [N, C, D1, D2, ...] it computes the mean over the spatial dims (everything from axis 2 onward), producing an output of shape [N, C, 1, 1, ...] with the spatial axes preserved as size-1 dims (matching the ONNX spec, which keeps rank rather than squeezing).

func (GlobalAveragePoolOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type LeakyReluOp ¶

type LeakyReluOp struct{}

LeakyReluOp implements the ONNX LeakyRelu operator. For each element, it outputs x if x >= 0, otherwise alpha * x. The alpha attribute defaults to 0.01 when absent.

func (op LeakyReluOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type MulAddOp ¶

type MulAddOp struct{}

MulAddOp is the fused Mul + Add operator. It computes

y = (a * b) + c

in a single pass over the output, eliminating the intermediate elementwise tensor that an unfused Mul → Add pair would materialise.

Inputs are exactly three tensors [a, b, c] in this order. The graph fusion pass canonicalises the input ordering so the kernel can assume the operand arrangement above (Mul is commutative, Add is commutative — the fusion picks an ordering and this op trusts it).

Broadcasting follows ONNX/NumPy rules and matches the shape semantics of an unfused Mul + Add: the Mul intermediate has shape broadcast(a, b), and the final output has shape broadcast(broadcast(a, b), c). The fast-path covers the common Real-ESRGAN residual-scaling shape pattern where one of {a, b} is a scalar (shape [1]) or a per-channel vector ([1, C, 1, 1]) and c shares the same N-dimensional shape as the other Mul operand.

func (MulAddOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type MulOp ¶

type MulOp struct{}

MulOp implements the ONNX Mul operator. It computes the elementwise product of two input tensors with broadcasting support.

func (MulOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type Operator ¶

type Operator interface {
	Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)
}

Operator defines the execution interface for a single ONNX operator. Each implementation is a stateless struct that validates inputs, performs computation, and returns exactly one output tensor on success.

The alloc parameter controls how output tensors are allocated. Passing tensor.New uses heap allocation; passing a pool-backed allocator reuses buffers. A nil alloc falls back to defaultAlloc (tensor.New).

func Get(opType string) (Operator, error)

type PReluOp ¶

type PReluOp struct{}

PReluOp implements the ONNX PRelu operator. For each element, it outputs x when x >= 0, otherwise slope * x. The slope is supplied as a runtime input tensor (input[1]) and is broadcast against the data input across the channel dimension.

This implementation accepts any slope tensor that is unidirectionally broadcastable to a per-channel layout: a single scalar, a 1-D [C], or any rank-N shape whose product is C with all non-channel dimensions equal to 1 (e.g., [C, 1, 1] or [1, C, 1, 1] — both appear in realesr-general-x4v3 exports). Anything else is rejected with an error wrapping ErrIncompatibleShapes.

func (PReluOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ReduceMeanOp ¶

type ReduceMeanOp struct{}

ReduceMeanOp implements the ONNX ReduceMean operator. It computes the arithmetic mean of the input tensor across the named axes.

Axis source has shifted between opset versions:

• opset <= 17: `axes` is an int attribute list.
• opset >= 18: `axes` is supplied as the (optional) second input tensor, with float32 entries (the runtime is float32-only).

This implementation accepts both forms with the second-input form taking precedence when present, matching the pattern in unsqueeze.go.

Attributes:

• axes ints, optional (opset <= 17)
• keepdims int, default 1
• noop_with_empty_axes int, default 0 (only relevant when axes is empty)

When axes is absent (and not provided as an input), ONNX reduces over every axis, producing a scalar — unless `noop_with_empty_axes == 1` (opset >= 18), in which case the input is returned unchanged.

func (ReduceMeanOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ReluOp ¶

type ReluOp struct{}

ReluOp implements the ONNX Relu operator. For each element it computes y = max(0, x). This is a stateless, attribute-free, single-input elementwise activation.

func (ReluOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ReshapeOp ¶

type ReshapeOp struct{}

ReshapeOp implements the ONNX Reshape operator. It returns a tensor with a new shape but the same underlying float32 data, copied verbatim.

The shape input is a 1-D tensor (stored as float32 in the runtime; values are taken as their integer truncation). A value of -1 in the shape means "infer this dimension from the remaining elements". A value of 0 in the shape means "copy the corresponding dimension from the input" when the allowzero attribute is absent or 0; when allowzero is 1, a 0 is taken literally (which in practice yields an empty tensor and is not exercised by inference workloads in this codebase).

func (ReshapeOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type ResizeOp ¶

type ResizeOp struct{}

ResizeOp implements the ONNX Resize operator (opset 11+). It supports bilinear and nearest neighbor interpolation with configurable coordinate transformation modes (asymmetric, half_pixel, align_corners).

func (ResizeOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type SigmoidOp ¶

type SigmoidOp struct{}

SigmoidOp implements the ONNX Sigmoid operator. For each element it computes 1 / (1 + exp(-x)).

Performance: SPAN issues 18 Sigmoid ops per Run, each over a feature map of up to 256K elements. The naive math.Exp path costs ~30 ns per element on M-series. We replace it with a 1024-entry lookup table over [-8, 8] with linear interpolation, which is ~5× faster (~6 ns per element) and within 5e-5 max abs error vs the true sigmoid — well below the precision relevant to image SR (sigmoid output is element-wise multiplied with a feature map, so any error <1e-3 is imperceptible after pipeline-end u8 quantization).

func (SigmoidOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type SqueezeOp ¶

type SqueezeOp struct{}

SqueezeOp implements the ONNX Squeeze operator. It removes dimensions of size 1. When axes is supplied, only those positions are squeezed (and each must reference a size-1 dim). When axes is absent, every size-1 dimension is removed.

Like Unsqueeze, the axes specification moved from an attribute (opset <= 12) to a second input tensor (opset 13+). This implementation accepts either form: it checks input slot 1 first, then the `axes` attribute.

Underlying data is unchanged; only shape metadata shrinks.

func (SqueezeOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type SubOp ¶

type SubOp struct{}

SubOp implements the ONNX Sub operator. It computes the elementwise difference of two input tensors with NumPy-style broadcasting support.

func (SubOp) Execute(inputs []*tensor.Tensor, _ map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type TransposeOp ¶

type TransposeOp struct{}

TransposeOp implements the ONNX Transpose operator. It permutes the dimensions of the input tensor according to the perm attribute. When perm is absent or empty, the dimensions are reversed (matching the ONNX spec default).

Unlike Reshape, Transpose changes the row-major element order, so data must be physically copied to the new layout. The implementation walks every output element, derives its multi-dimensional index from the output strides, maps that index back through perm to the input coordinate, and reads from the input data at the computed offset.

func (TransposeOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type UnsqueezeOp ¶

type UnsqueezeOp struct{}

UnsqueezeOp implements the ONNX Unsqueeze operator. It inserts dimensions of size 1 at the positions named by `axes`. Negative axes are interpreted relative to the rank of the OUTPUT (post-unsqueeze) tensor.

The axes specification moved from an attribute (opset <= 12) to a second input tensor (opset 13+). This implementation accepts either form: when a second input is present its float32 data is interpreted as the axes list; otherwise the `axes` attribute is consulted.

The underlying data layout is unchanged — only the shape metadata grows.

func (UnsqueezeOp) Execute(inputs []*tensor.Tensor, attributes map[string]*onnx.Attribute, alloc tensor.Allocator) ([]*tensor.Tensor, error)

type WorkerPool ¶

type WorkerPool struct {
	// contains filtered or unexported fields
}

WorkerPool is a per-Engine-Run helper-goroutine budget. It is created at the start of Engine.Run, lives for the duration of a single inference call, and is shut down at the end of Run. Fanout sites (Winograd cells, GEMM M-split, Conv im2col fanout) submit work items to the pool; the pool's worker count caps how many helper goroutines may be in flight at once across the whole Run.

Design history: an earlier iteration kept a fixed roster of long-lived worker goroutines parked on a shared channel between work items, hoping to amortise goroutine spawn cost over a Run. On the M-series target it regressed end-to-end at every benchmarked SR shape — channel send/receive per Submit added ~0.5–1µs that swamped the cheap `go fn()` it replaced. The current design keeps the per-Run *budget* (a counted-token slice of the global concurrency limit, see ops/concurrency.go) but lets Submit ad-hoc-spawn a goroutine per work item. Worker spawn cost is paid once per fanout item, not once per Run, but the budget acquisition that used to happen per fanout-site now happens once per Run — eliminating the pthread_cond churn from the old fairHelperCount + tryAcquireWorkers pattern that profiles flagged as 91% of CPU time.

Reentrancy: each Run() creates its own pool, so concurrent Run() calls (different engines, or pipeline tile workers) do not share a pool. The package-level ActiveWorkerPool atomic.Pointer carries "the current pool" for fanout sites to discover. When two Runs race on the global, last-writer-wins is correctness-safe because pools are interchangeable from a work-execution standpoint — any live pool dispatches submitted work correctly.

Sizing: NewWorkerPool tries to acquire up to `requested` tokens from the global workerBudget (capped at GOMAXPROCS-1). The count actually acquired becomes the pool's worker count. This shares the budget across concurrent Runs: when 4 tile workers each call Run simultaneously, each gets ~NumCPU/4 helpers, keeping the OS-thread total bounded by NumCPU. Single-Run mode gets the full budget.

func NewWorkerPool(requested int) *WorkerPool

func (p *WorkerPool) Shutdown()

for i := 0; i < n-1; i++ {
    wg.Add(1)
    pool.Submit(func() { defer wg.Done(); ... })
}
// run last chunk in caller
doWork(...)
wg.Wait() // waits for the helpers' wg.Done

func (p *WorkerPool) Submit(fn func())

func (p *WorkerPool) Workers() int