Registering kernels#

onnx-light-cpu does not run models on its own: it installs its kernels into onnx-light’s shared C++ KernelDispatchTable or a session’s RuntimeContext. For supported operators, newly resolved nodes use these implementations instead of the built-in ones, subject to the registration scope and precedence below.

This page explains how to register the shipped kernels, how to add a brand new kernel, and — most importantly — why a registration can appear to be ignored and how to avoid that.

Register the shipped kernels#

The shipped kernels are installed with a single call.

From Python:

import numpy as np
from onnx_light.onnx.reference import ReferenceEvaluator

from onnx_light_cpu import (
    MicrosoftKernelImplementation,
    register_kernel_for_session,
    register_kernel_global,
    register_kernels_global,
)

# Install the SIMD kernels into onnx-light's dispatch table. Import
# onnx-light *before* calling this so its built-in kernels are already
# registered; the onnx-light-cpu kernels then override them.
register_kernels_global()
sess = ReferenceEvaluator(model)  # any model containing an Abs node
(y,) = sess.run(None, {"x": np.array([-1.0, 2.0, -3.0], dtype=np.float32)})

register_kernels is only available in builds compiled with the onnx-light integration (-DONNX_LIGHT_CPU_WITH_ONNX_LIGHT=ON); it wraps the compiled internal registration binding.

register_kernels_global() installs the kernels process-wide: it populates onnx-light’s shared dispatch table, so every ReferenceEvaluator created in the process afterwards uses them. Existing sessions that have already resolved their nodes retain the cached factories. Register only one global kernel with register_kernel_global(domain, op_type).

To install the same native compiled implementation for a single session only, without touching the shared table, use register_kernel_for_session:

sess = ReferenceEvaluator(model)
register_kernel_for_session(sess, "", "Abs")  # this session only

The evaluator’s runtime context owns that registration and releases it with the evaluator. Other evaluators are unchanged. register_kernels_for_session provides the equivalent all-kernels operation. Session-local registrations take precedence over global and built-in kernels for that session.

Session-local registrations expose native factories to RuntimeSession. The runtime constructs and retains kernel instances, resolves shared tuning profiles, and applies preparation hooks before execution, just as for global registrations. Changing only tensor values reuses the configured instance. Replacing a registration affects future resolution; already-prepared sessions retain their previous instances.

The default installs the optimized com.microsoft family. Tests and diagnostic callers can explicitly install the complete scalar reference family instead:

register_kernels_global(
    microsoft_implementation=MicrosoftKernelImplementation.NAIVE
)

From C++:

#include <onnx_light_cpu/kernels/register_kernels.h>

onnx_light_cpu::RegisterKernelGlobal("", "Abs");
onnx_light_cpu::RegisterAllKernelsGlobal();

onnx_light::core::runtime::RuntimeContext session(kernel_context);
onnx_light_cpu::RegisterKernelForSession(session, "", "Abs");
onnx_light_cpu::RegisterAllKernelsForSession(session);

Global registration replaces an existing entry by default and is owned by the process-wide dispatch table. Session registration replaces an entry on that RuntimeContext by default and must occur before a runtime session resolves its nodes. Pass replace=false to either scope to keep an existing entry; the function then reports no installation. Repeated calls are therefore idempotent in state, whether they replace the same implementation or retain it.

RegisterAllKernels() is the C++ convenience entry point for process-wide registration with replacement enabled; register_kernels() is the Python compatibility entry point. RegisterAllKernelsGlobal and RegisterAllKernelsForSession install the same operator set, but provide explicit scope and replacement control and return the number of installations. For session-local calls, replace=false checks only existing local entries: it does not prevent overriding a global or built-in kernel.

RegisterMicrosoftKernels installs exactly one complete implementation family for the com.microsoft domain. Mixing families in one registration pass is intentionally unsupported, so one operator cannot silently replace a different variant registered earlier in that pass.

The registration overrides the built-in entry for the default ONNX domain, and it wins regardless of the order in which the built-in and the onnx-light-cpu kernels are registered: onnx-light’s bulk built-in registration never replaces an entry that is already present, and an explicit registration always replaces the built-in one.

The registered kernels are plain KernelBase subclasses. They do not install an onnx-light-cpu-specific scheduler bridge: when onnx-light runs a session it already places that session’s CpuExecutor on the calling thread, and the low-level SIMD helpers pick it up directly. Calling a kernel class directly outside a session therefore stays serial unless the caller installs an onnx-light CpuExecutorScope first, matching native onnx-light kernels.

Complete RegisterAllKernels inventory#

Every family below is installed for the CPU device. ai.onnx denotes the default ONNX domain (the empty string is equivalent); ai.onnx.ml and com.microsoft are separate domains. In particular, the two LinearAttention entries are distinct operators.

Kernel families installed by RegisterAllKernels#
Family registrar	Domain	Operators
RegisterAbsKernel	ai.onnx	Abs
RegisterAttentionKernel	ai.onnx	Attention
RegisterLinearAttentionKernel	ai.onnx	LinearAttention
RegisterRotaryEmbeddingKernels	ai.onnx	RotaryEmbedding
RegisterRotaryEmbeddingKernels	com.microsoft	RotaryEmbedding
RegisterBinaryKernels	ai.onnx	Add, Sub, Mul, Div, Mod, Pow, Equal, Greater, GreaterOrEqual, Less, LessOrEqual, And, Or, Xor, BitwiseAnd, BitwiseOr, BitwiseXor, BitShift, PRelu
RegisterExpKernel	ai.onnx	Exp
RegisterLogKernel	ai.onnx	Log
RegisterGemmKernel	ai.onnx	Gemm
RegisterCastKernel	ai.onnx	Cast
RegisterConcatKernel	ai.onnx	Concat
RegisterGatherKernel	ai.onnx	Gather
RegisterNonZeroKernel	ai.onnx	NonZero
RegisterReduceSumKernel	ai.onnx	ReduceSum
RegisterScatterNDKernel	ai.onnx	ScatterND
RegisterSliceKernel	ai.onnx	Slice
RegisterSplitKernel	ai.onnx	Split
RegisterMatMulKernel	ai.onnx	MatMul
RegisterIntegerMatMulKernels	ai.onnx	MatMulInteger, QLinearMatMul
RegisterNotKernel	ai.onnx	Not
RegisterNormalizationKernels	ai.onnx	BatchNormalization, GroupNormalization, InstanceNormalization, LayerNormalization, LpNormalization, MeanVarianceNormalization
RegisterRmsNormalizationKernel	ai.onnx	RMSNormalization
RegisterSigmoidKernel	ai.onnx	Sigmoid
RegisterTanhKernel	ai.onnx	Tanh
RegisterSimplifiedLayerNormalizationKernel	ai.onnx	SimplifiedLayerNormalization
RegisterSoftmaxKernel	ai.onnx	Softmax
RegisterSwiGLUKernel	ai.onnx	SwiGLU
RegisterTreeEnsembleKernel	ai.onnx.ml	TreeEnsemble
RegisterVariadicElementwiseKernels	ai.onnx	Sum, Mean, Min, Max
RegisterMicrosoftKernels	com.microsoft	BiasGelu, CDist, CausalConvWithState, GroupQueryAttention, MultiHeadAttention, LinearAttention, SkipLayerNormalization, SkipSimplifiedLayerNormalization
RegisterMatMulNBitsKernel	com.microsoft	MatMulNBits
RegisterGatherBlockQuantizedKernel	com.microsoft	GatherBlockQuantized
RegisterMoEKernel	com.microsoft	MoE
RegisterQMoEKernel	com.microsoft	QMoE

The table describes shipped registrations, not all operators supported by onnx-light, nor a guarantee of every type, opset, or shape for a listed operator. The Kernel inventory and usage API and the generated ByOp catalogue expose the registration metadata, including supported element types, opset bounds, and implementation names.

Microsoft implementation policy#

MicrosoftKernelImplementation.OPTIMIZED is the default for all global and session-local entry points. It selects the production implementations of all six operators registered by RegisterMicrosoftKernels above. NAIVE selects their reference implementations instead, using the RegisterNaive*Kernel registrars and Naive* kernel names (NaiveMicrosoftLinearAttention for com.microsoft::LinearAttention). These are correctness oracles for testing and diagnostics, not a performance option. The naive LinearAttention adapter delegates to onnx-light’s concrete reference kernel, as described below.

The policy changes only those six Microsoft implementations, not standalone Microsoft registrations such as MatMulNBits, MoE, and QMoE, RotaryEmbedding, nor any ai.onnx / ai.onnx.ml registration. The C++ policy overload is RegisterAllKernels(MicrosoftKernelImplementation::NAIVE); the explicit global and session APIs also accept the policy.

OPTIMIZED does not promise SIMD for every invocation. Production kernels use tuning, parallel execution, and SIMD where implemented and applicable, with portable/scalar paths depending on element type, shape, build configuration, and CPU capabilities. Selecting NAIVE is distinct from disabling SIMD in production kernels and does not turn the entire registered kernel set into reference implementations.

Keeping the inventory current#

registered_kernels(microsoft_implementation=...) calls CollectRegisteredKernels in C++, which executes the same RegisterAllKernels registration path that populates the dispatch table, but collects metadata without installing or executing kernels. It describes the shipped set for that policy, not the current contents of the global table or a particular session.

unittests/python/test_kernels_doc_e2e.py compares this table’s domain/operator/device set with that live inventory for both OPTIMIZED and NAIVE. Adding, removing, or renaming a registration therefore requires updating this table. Multiple opset records for one operator are collapsed here; the generated per-kernel catalogue retains their detailed metadata.

Checking which kernels are used#

Because the onnx-light-cpu kernels are drop-in replacements, a model produces the same numbers whether it runs them or onnx-light’s built-in kernels. To tell them apart, every onnx-light-cpu kernel carries a unique, library-qualified name (for example "onnx_light_cpu::Abs") that it can record when it runs. Recording is disabled by default and must be explicitly enabled for diagnostics on an individual session. Each session’s RuntimeContext retains the first 1024 invocations after each clear; later records are dropped until cleared. Copies of that context, including subgraph and function execution contexts, share the recording state; independent sessions do not. Disabled calls only load an atomic flag, without locking or allocating. The names can be inspected from Python:

from onnx_light.onnx.reference import ReferenceEvaluator
from onnx_light_cpu import (
    clear_used_kernel_names,
    register_kernels,
    registered_kernel_names,
    registered_kernels,
    set_kernel_usage_recording,
    used_kernel_names,
)

register_kernels()
registered_kernel_names()  # {'Abs': 'onnx_light_cpu::Abs', 'Exp': ...}
registered_kernels()       # (RegisteredKernel(domain='ai.onnx', op_type='Abs', ...), ...)

sess = ReferenceEvaluator(model)
set_kernel_usage_recording(sess, True)
clear_used_kernel_names(sess)
sess.run(None, feeds)  # run a model containing e.g. an Abs node
set_kernel_usage_recording(sess, False)
used_kernel_names(sess)   # ['onnx_light_cpu::Abs', ...]

Retrieval returns an independent snapshot without consuming the log. Recording, retrieval, and reset share a mutex; concurrent records are ordered by mutex acquisition and may fall before or after a reset. Reset does not toggle recording. Disabling waits for active appends and preserves existing entries.

Each recording function requires an explicit ReferenceEvaluator session; missing or invalid sessions raise TypeError. Enabling, disabling, or clearing one session’s log does not affect any independent session. Retrieval includes all backend kernel names recorded by that context, without filtering to onnx_light_cpu names.

The previous process-wide API without a session argument was intentionally removed; there is no global compatibility fallback. An updated onnx-light containing PR #4942 is required. That upstream change has not been released yet, so no minimum released version is specified.

If used_kernel_names(sess) does not include the expected onnx_light_cpu names after an explicitly recorded run whose operators onnx-light-cpu overrides, the registration did not take effect — see Troubleshooting: the registration is ignored below. The same names are available in C++ as the static AbsKernel::kName (etc.) members and through onnx_light_cpu::RegisteredKernelNames() in onnx_light_cpu/kernels/kernel_usage.h. C++ recording and retrieval use the session’s RuntimeContext, including rt.RecordKernelUsage(kName) and rt.GetKernelUsage() rather than process-wide recording functions. The recording bound is defined by RuntimeContext::kKernelUsageLimit.

rt.set_kernel_usage_enabled(true);
rt.ClearKernelUsage();
// Run the session using rt.
rt.set_kernel_usage_enabled(false);
const auto names = rt.GetKernelUsage();

Add a new kernel#

Kernel implementations must follow the borrowed executor contract in execution.h: borrow the current executor only for the invocation, keep range callbacks non-throwing, and submit parallel work through the shared range helpers. The contract defines lifetime, failure propagation, nested execution, participant limits, and which scheduling state must not be retained. See also Parallel kernel execution for the execution flow.

A brand new kernel is a subclass of onnx-light’s onnx_light::core::runtime::KernelBase that is registered into the dispatch table with RegisterKernelFn. The shipped kernels are the template to copy; onnx_light_cpu/kernels/math/abs_kernel.cc is the smallest complete example. The three steps are:

Derive from KernelBase and override Run(RuntimeContext &) to read the node’s inputs and write its outputs (see AbsKernel::Run). Give the class a unique static constexpr const char *kName (for example "onnx_light_cpu::MyOp") and call rt.RecordKernelUsage(kName) at the top of Run so the kernel can be recognised in rt.GetKernelUsage(); add its {op_type, kName} pair to RegisteredKernelNames() in onnx_light_cpu/kernels/kernel_usage.cc. Install a RuntimeExecutionScope from onnx_light_cpu/kernels/execution.h at the start of Run and each directly callable operator() so numerical calls borrow the session’s executor without coupling the computation library to ONNX.
Wrap the class in a NodeKernelFn factory that constructs the kernel and calls set_node.
Call RegisterKernelFn(domain, op_type, symbolic::Device::kCPU, factory). An empty domain is normalised to the default ONNX domain, so it overrides the corresponding built-in operator.

void RegisterMyKernel() {
  onnx_light::core::runtime::NodeKernelFn factory =
      [](const onnx_light::NodeProto &node,
         onnx_light::core::runtime::RuntimeContext &rt)
      -> std::unique_ptr<onnx_light::core::runtime::KernelBase> {
    auto kernel = std::make_unique<MyKernel>(rt.kernel_ctx());
    kernel->set_node(node);
    return kernel;
  };
  onnx_light::core::runtime::RegisterKernelFn(
      "", "MyOp", onnx_light::core::symbolic::Device::kCPU, std::move(factory));
}

Then add RegisterMyKernel() to RegisterAllKernels in onnx_light_cpu/kernels/register_kernels.cc so it is installed together with the other kernels.

Every new com.microsoft operator must instead provide two independent KernelBase adapters: a readable scalar Naive*Kernel correctness oracle and a production optimized kernel. Add both registrars to the corresponding branch of RegisterMicrosoftKernels and add differential tests covering their complete shared contract. Backend TEST cases generate expectations with the naive implementation; BENCHMARK cases use the optimized one.

For com.microsoft::CausalConvWithState, the two adapters share only input validation, output planning, and node attribute parsing. The naive adapter indexes a virtual concatenation of past state and input; the production adapter uses lag-based convolution over independently scheduled batch/channel tasks. Scheduling cost includes all convolution taps and state copies, so prefill and empty-sequence state handoffs are not mistaken for cheap decode work. Differential tests cover all three types, both layouts (including higher-rank channels-last inputs), dilation, optional bias and past state, activation modes, and empty or short sequences. The ordinary FLOAT/FLOAT16 contract is also checked against onnx-light’s concrete reference kernel.

For com.microsoft::LinearAttention, the test-only NaiveMicrosoftLinearAttentionKernel calls the concrete onnx-light onnx_kernels::kernel::LinearAttention reference kernel, not the optimized dispatcher. It expands shared query/key heads to the ONNX contract and preserves the Microsoft defaults, optional state, decay and beta layouts, and both outputs. No attention recurrence is duplicated in the adapter. FLOAT parity uses the backend tolerances (rtol=1e-3, atol=1e-7); default and global registration continue to select MicrosoftLinearAttentionKernel.

For a kernel that only needs to run from Python — for example a quick, model-specific override written in NumPy — use onnx-light’s per-evaluator ReferenceEvaluator.register_custom_kernel(domain, op_type, fn) hook instead; it does not require rebuilding this package.

Troubleshooting: the registration is ignored#

If register_kernels() (or RegisterAllKernels()) runs without error but the model still uses onnx-light’s built-in kernels, the most common cause is that the registration and the runtime end up using two different copies of onnx-light’s dispatch table.

The dispatch table is a singleton that lives inside onnx-light’s lib_onnx_core. It only works as a single shared registry when every extension links the same lib_onnx_core at run time. Building the integration from an onnx-light source tree while a separately installed onnx-light Python package is what actually runs the model breaks that assumption:

python setup.py build_ext --inplace --onnx-light-source (equivalently -DONNX_LIGHT_CPU_ONNX_LIGHT_SOURCE_DIR=...) compiles onnx-light from source with its Python build disabled, which produces a static lib_onnx_core that is embedded privately into _cpuregister.
The installed onnx_light Python package ships its own shared lib_onnx_core that its ReferenceEvaluator uses.

register_kernels() then populates the private copy inside _cpuregister, but the evaluator reads the copy inside onnx_light — so the SIMD kernels are never seen and the registration silently appears to be ignored.

To make the registration take effect, build the integration so it links the same shared lib_onnx_core that the running onnx_light uses. Install onnx-light so that its onnx_lightConfig.cmake is available and build with find_package instead of building from source:

python setup.py build_ext --inplace --onnx-light

Use --onnx-light-source only for a self-contained C++ integration where onnx-light-cpu owns the whole runtime (there is no separately installed onnx-light Python package to share state with).