Custom kernels for ReferenceEvaluator#

The C++ runtime in onnx_light/onnx_kernels dispatches every NodeProto against a static onnx::onnx_kernels::KernelDispatchTable(). Any operator that is not built in — typically an operator from a user-defined domain, an experimental op, or a stand-in for one not yet implemented — would otherwise fail with unsupported op_type.

The custom-kernel hook documented here lets callers extend (or override) the runtime without modifying the built-in dispatch table. It is exposed at three layers:

  • C++: onnx::onnx_kernels::RuntimeContext::RegisterCustomKernel(),

  • low-level Python: RuntimeContext.register_custom_kernel(domain, op_type, fn),

  • high-level Python: register_custom_kernel().

All three share the same model: a name-keyed onnx::onnx_kernels::CustomKernelMap stored on the active RuntimeContext. RunNode consults this map before the built-in dispatch table; model-local FunctionProto definitions and the dedicated control-flow paths (If, Loop, Scan, SequenceMap) still take precedence so existing graphs are unaffected.

Dispatch precedence#

For each NodeProto evaluated by RunNode:

  1. If / Loop / Scan / SequenceMap are routed to their dedicated subgraph runners.

  2. Model-local FunctionProto definitions found on the RuntimeContext are inlined.

  3. The custom-kernel map is consulted by the "<domain>:<op_type>" key (the empty default ONNX domain is normalised to "ai.onnx"). A registered entry overrides the static table.

  4. KernelDispatchTable resolves the built-in kernel.

  5. Otherwise the dispatcher throws unsupported op_type.

C++ usage#

A custom kernel is any callable compatible with onnx::onnx_kernels::CustomKernelFn, i.e. std::function<void(const NodeProto &, RuntimeContext &)>. The callback is in charge of reading its inputs from the context’s tensor map and writing its outputs back through RuntimeContext::Put() (or SetOutput()).

#include "onnx_kernels/run_nodes.h"
#include "onnx_kernels/runtime_context.h"

using namespace onnx::onnx_kernels;

RuntimeContext rt(KernelContext(/*opset=*/18));
rt.Set("x", Tensor::FromFloat("x", {3}, {1.0f, 2.0f, 3.0f}));

// Register a "my.domain.Scale" custom kernel that multiplies its
// single input by the "factor" attribute.
rt.RegisterCustomKernel(
    "my.domain", "Scale",
    [](const NodeProto &node, RuntimeContext &ctx) {
      float factor = 1.0f;
      for (int i = 0; i < node.attribute_size(); ++i) {
        if (node.attribute(i).name() == "factor") {
          factor = node.attribute(i).f();
        }
      }
      const Tensor &in = ctx.Get(node.input(0));
      std::vector<float> out(static_cast<size_t>(in.element_count()));
      const float *src = in.AsFloat();
      for (size_t i = 0; i < out.size(); ++i) {
        out[i] = src[i] * factor;
      }
      ctx.Put(node.output(0),
              Tensor::FromFloat(node.output(0), in.shape, out));
    });

// node has op_type="Scale", domain="my.domain", input "x", output "y".
RunNode(node, rt);

The empty domain is normalised to "ai.onnx", so registering ("", "Abs", fn) overrides the built-in Abs entry.

The same mechanism is the entry point for C++ extension modules: a shared library exposing additional kernels can register them on a RuntimeContext (either directly, or through the Python binding when loaded as a Python module) without rebuilding lib_onnx_kernels.

Low-level Python binding#

The RuntimeContext exposed by onnx_light.onnx_py._onnxpykernels mirrors the C++ API:

from onnx_light.onnx_py._onnxpykernels import runtime as rt

ctx = rt.RuntimeContext(rt.KernelContext(rt.default_opset(18)))
ctx.set("x", ...)

def scale(node, c):
    # Direct access to the raw NodeProto / RuntimeContext.
    x = c.get(str(node.input[0]))
    ...
    c.put(str(node.output[0]), ..., "output")

ctx.register_custom_kernel("my.domain", "Scale", scale)
rt.run_model(model, ctx)

This binding is intentionally close to the C++ API: the callback receives the raw NodeProto and RuntimeContext and is responsible for any tensor encoding/decoding. It is the appropriate entry point for C++ extension modules and for kernels that need to inspect or mutate the runtime context beyond their declared inputs and outputs (for example to read sequences or to participate in the event log).

High-level Python API#

For most users the register_custom_kernel() wrapper is the easiest entry point: it accepts numpy inputs and a numpy output (or a tuple/list of arrays for multi-output ops).

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

model = parser.parse_model(
    '<ir_version: 10, opset_import: ["" : 18, "my.domain" : 1]>'
    'agraph (float[3] x) => (float[3] y) {'
    '  y = my.domain.Scale<factor=3.0>(x)'
    '}'
)

def scale(node, x):
    factor = next(a for a in node.attribute if str(a.name) == "factor").f
    return x * float(factor)

sess = ReferenceEvaluator(model)
sess.register_custom_kernel("my.domain", "Scale", scale)
(y,) = sess.run(None, {"x": np.array([1.0, 2.0, 3.0], dtype=np.float32)})
# y == [3., 6., 9.]

Registrations are stored on the evaluator and reapplied to the fresh RuntimeContext built on every run() call, so the same evaluator can be reused across runs.

Symmetry with shape inference#

The custom-kernel hook is intentionally symmetric to onnx::onnx_optim::shapes::ShapesContext::SetCustomShapeInferenceFunction() in the shape-inference layer: both allow callers to plug a callable keyed by (domain, op_type) into a per-invocation context and both are consulted before the built-in domain dispatch.

See also#