Custom ops

onnx_light_ort/custom_ops/ builds onnx_light_ort::custom_ops, a shared library exporting only RegisterCustomOps. Its current profile registers onnx_light::LeakyRelu at custom opset 1, with one floating-point CPU input and output and an explicit float alpha attribute. It also registers onnx_light_cpu::Log at custom opset 1, with the same CPU input/output profile and no attributes. The numerical implementations are onnx-light-cpu’s LogFloat32, LogFloat64, LogFloat16 and LogBFloat16WithTuning, using the dependency’s default tuning and dispatch. Both operators support changing shapes, scalars and empty tensors up to rank 16.

Each operator exposes float16, bfloat16, float32 and float64. Four concrete typed descriptors share the same domain, name and custom opset; ORT selects the matching implementation from the model’s element types. Input and output types must match. There is no bridge-level conversion or fallback to another operator. LeakyRelu delegates all four types to the native onnx-light kernel; its alpha attribute remains an ONNX float.

This component owns descriptors, domain registration and callback wiring. It depends on Custom-op bridge: onnx-light called by ORT, which uses Shared support for onnx-light in ORT for native state, tensor validation and status conversion, never on the opposite bridge or the EP library. It requests the supplied API 18 table and has no link dependency on ORT. See Architecture for domain, kernel and borrowed-buffer lifetimes.

Tests

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --parallel "$(nproc)"
ctest --test-dir build -L custom_ops --output-on-failure

tests/test_custom_ops.cc calls the real shared-library export with an explicit minimal C API test host. It verifies stable domain reuse, failed registration, unsupported API negotiation, failed construction clearing its result, native execution through callbacks, exception-to-status translation, metadata cleanup and reuse after failure. No ORT runtime or mocking framework is needed. This fixture is not the production compatibility bridge.

With ONNX_LIGHT_ORT_TEST_WITH_ORT=ON, the same label also selects the real-ORT integration suite and runnable examples. The real-ORT and standalone suites exercise all four types for both operators, including changing shapes, scalars, empty tensors, non-finite values and float64 values that cannot be represented as float32. The ORT 1.18 Python API cannot construct BFLOAT16 OrtValues directly. Its tests use float32 graph inputs/outputs with explicit Cast nodes around the BFLOAT16 custom node, whose intermediate types are explicitly declared. The native custom kernel still receives and produces BFLOAT16, without changing either bridge or skipping this type. The standalone suite uses BFLOAT16 arrays directly.

BFLOAT16 with ORT’s Python API

To avoid NumPy feed/output dtype-mapping limitations for ml_dtypes.bfloat16, bind typed OrtValues over the raw 16-bit storage. The session’s model input/output types must be TensorProto.BFLOAT16:

This snippet requires an ORT Python build providing ortvalue_from_numpy_with_onnx_type; that convenience method is absent in ORT 1.18.1. With that older Python interface, use explicit Cast nodes at the graph boundaries as in the integration tests, or construct BFLOAT16 tensors through the C API’s CreateTensorWithDataAsOrtValue.

import numpy
import onnxruntime
from ml_dtypes import bfloat16
from onnx_light.onnx import TensorProto

values = numpy.array([-4, 2], dtype=bfloat16)
output = numpy.empty_like(values)
input_value = onnxruntime.OrtValue.ortvalue_from_numpy_with_onnx_type(
    values.view(numpy.uint16), TensorProto.BFLOAT16
)
output_value = onnxruntime.OrtValue.ortvalue_from_numpy_with_onnx_type(
    output.view(numpy.uint16), TensorProto.BFLOAT16
)
binding = session.io_binding()
binding.bind_ortvalue_input("X", input_value)
binding.bind_ortvalue_output("Y", output_value)
session.run_with_iobinding(binding)

output retains the bfloat16 dtype. The standalone CustomOpSession accepts these arrays directly. Keep bound arrays alive until inference completes.

Python example

From the repository root on Linux:

python -m pip install -r requirements/ort-tests.txt
cmake -S . -B build -DONNX_LIGHT_ORT_TEST_WITH_ORT=ON
cmake --build build --parallel "$(nproc)"
export PYTHONPATH="$PWD:$PWD/build/python:$PYTHONPATH"
python examples/custom_ops/custom_ops.py \
  --library build/onnx_light_ort/custom_ops/libonnx_light_ort_custom_ops.so \
  --model-output build/custom_ops_example.onnx

The example constructs and checks the custom-domain model with onnx_light.onnx (not the onnx package), loads the plugin, runs inference and checks the result [-0.5, -0.125, 0, 2]. The saved model can then be consumed by the C++ example.

"""Runs the onnx-light CPU custom op in a real ONNX Runtime session."""

import argparse
from pathlib import Path

import numpy
import onnxruntime
from onnx_light.onnx import TensorProto, checker, helper, save


def make_model():
    """Creates a model explicitly calling the onnx-light custom operator."""
    graph = helper.make_graph(
        [helper.make_node("LeakyRelu", ["X"], ["Y"], domain="onnx_light", alpha=0.125)],
        "custom-ops-example",
        [helper.make_tensor_value_info("X", TensorProto.FLOAT, [4])],
        [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4])],
    )
    return helper.make_model(
        graph,
        opset_imports=[
            helper.make_opsetid("", 18),
            helper.make_opsetid("onnx_light", 1),
        ],
        ir_version=9,
    )


def main():
    """Runs the native custom op in real ORT and optionally saves its model."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--library", type=Path, required=True)
    parser.add_argument("--model-output", type=Path)
    args = parser.parse_args()
    model = make_model()
    checker.check_model(model)
    if args.model_output is not None:
        save(model, args.model_output)
    options = onnxruntime.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    options.register_custom_ops_library(str(args.library.resolve(strict=True)))
    session = onnxruntime.InferenceSession(
        model.SerializeToString(), options, providers=["CPUExecutionProvider"]
    )
    values = numpy.array([-4, -1, 0, 2], dtype=numpy.float32)
    result = session.run(["Y"], {"X": values})[0]
    numpy.testing.assert_array_equal(result, numpy.array([-0.5, -0.125, 0, 2], numpy.float32))
    print(result)


if __name__ == "__main__":
    main()

C++ example

The native build compiles the example without installing or linking ORT. At execution, provide the absolute path to an existing ORT shared library, the custom-op library and the model saved above:

build/examples/custom_ops/onnx_light_ort_custom_ops_example \
  /absolute/path/to/libonnxruntime.so \
  build/onnx_light_ort/custom_ops/libonnx_light_ort_custom_ops.so \
  build/custom_ops_example.onnx

The example is a real ORT application, so it deliberately loads ORT and resolves OrtGetApiBase. This is not a fallback in either production bridge. It registers the custom-op library through RegisterCustomOpsLibrary_V2, creates a real session and verifies the same output. Scoped owners release values, sessions and options before the runtime library unloads.

Use the platform-specific library suffix and executable location on macOS or Windows. Multi-configuration generators additionally use Release/. Set ONNX_LIGHT_ORT_BUILD_EXAMPLES=OFF to disable the C++ example build.

To run both examples through CTest, enable real-host tests and pass -DONNX_LIGHT_ORT_RUNTIME_LIBRARY=/absolute/path/to/the/ORT/library. CTest first runs the Python example to generate the model, then the C++ example; fixture dependencies also apply when selecting only the C++ test.

// SPDX-License-Identifier: Apache-2.0

#include "onnxruntime_c_api.h"

#include <array>
#include <filesystem>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>

#ifdef _WIN32
#include <windows.h>
#else
#include <dlfcn.h>
#endif

namespace {

class RuntimeLibrary {
public:
  explicit RuntimeLibrary(const std::filesystem::path &path) {
#ifdef _WIN32
    handle_ = LoadLibraryW(path.c_str());
    if (handle_ == nullptr)
      throw std::runtime_error("Cannot load ORT: Windows error " + std::to_string(GetLastError()));
#else
    handle_ = dlopen(path.c_str(), RTLD_NOW | RTLD_LOCAL);
    if (handle_ == nullptr)
      throw std::runtime_error(dlerror());
#endif
  }
  RuntimeLibrary(const RuntimeLibrary &) = delete;
  RuntimeLibrary &operator=(const RuntimeLibrary &) = delete;
  ~RuntimeLibrary() {
#ifdef _WIN32
    FreeLibrary(handle_);
#else
    dlclose(handle_);
#endif
  }

  const OrtApi &Api() const {
    using GetApiBase = const OrtApiBase *(ORT_API_CALL *)();
#ifdef _WIN32
    const auto get_api_base =
        reinterpret_cast<GetApiBase>(GetProcAddress(handle_, "OrtGetApiBase"));
#else
    const auto get_api_base = reinterpret_cast<GetApiBase>(dlsym(handle_, "OrtGetApiBase"));
#endif
    if (get_api_base == nullptr)
      throw std::runtime_error("The runtime library does not export OrtGetApiBase.");
    const OrtApi *api = get_api_base()->GetApi(18);
    if (api == nullptr)
      throw std::runtime_error("The runtime does not support ORT API 18.");
    return *api;
  }

private:
#ifdef _WIN32
  HMODULE handle_ = nullptr;
#else
  void *handle_ = nullptr;
#endif
};

void Check(const OrtApi &api, OrtStatus *status) {
  if (status != nullptr) {
    const std::string message = api.GetErrorMessage(status);
    api.ReleaseStatus(status);
    throw std::runtime_error(message);
  }
}

template <typename T> using Owner = std::unique_ptr<T, void(ORT_API_CALL *)(T *)>;

} // namespace

int main(int argc, char **argv) {
  if (argc != 4)
    throw std::invalid_argument("Usage: custom_ops_cpp ORT_LIBRARY CUSTOM_OPS_LIBRARY MODEL");
  const RuntimeLibrary runtime(std::filesystem::absolute(argv[1]));
  const OrtApi &api = runtime.Api();
  OrtEnv *raw_env = nullptr;
  Check(api, api.CreateEnv(ORT_LOGGING_LEVEL_WARNING, "custom-ops-example", &raw_env));
  const Owner<OrtEnv> env(raw_env, api.ReleaseEnv);
  OrtSessionOptions *raw_options = nullptr;
  Check(api, api.CreateSessionOptions(&raw_options));
  const Owner<OrtSessionOptions> options(raw_options, api.ReleaseSessionOptions);
  Check(api, api.SetIntraOpNumThreads(options.get(), 1));
  Check(api, api.SetInterOpNumThreads(options.get(), 1));
  const auto plugin = std::filesystem::absolute(argv[2]);
  Check(api, api.RegisterCustomOpsLibrary_V2(options.get(), plugin.c_str()));
  OrtSession *raw_session = nullptr;
  const auto model = std::filesystem::absolute(argv[3]);
  Check(api, api.CreateSession(env.get(), model.c_str(), options.get(), &raw_session));
  const Owner<OrtSession> session(raw_session, api.ReleaseSession);
  OrtMemoryInfo *raw_memory = nullptr;
  Check(api, api.CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &raw_memory));
  const Owner<OrtMemoryInfo> memory(raw_memory, api.ReleaseMemoryInfo);
  std::array<float, 4> values{-4, -1, 0, 2};
  const std::array<int64_t, 1> shape{4};
  OrtValue *raw_input = nullptr;
  Check(api, api.CreateTensorWithDataAsOrtValue(memory.get(), values.data(), sizeof(values),
                                                shape.data(), shape.size(),
                                                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &raw_input));
  const Owner<OrtValue> input(raw_input, api.ReleaseValue);
  const OrtValue *inputs[]{input.get()};
  const char *input_names[]{"X"};
  const char *output_names[]{"Y"};
  OrtValue *raw_output = nullptr;
  Check(api, api.Run(session.get(), nullptr, input_names, inputs, 1, output_names, 1, &raw_output));
  const Owner<OrtValue> output(raw_output, api.ReleaseValue);
  OrtTensorTypeAndShapeInfo *raw_info = nullptr;
  Check(api, api.GetTensorTypeAndShape(output.get(), &raw_info));
  const Owner<OrtTensorTypeAndShapeInfo> info(raw_info, api.ReleaseTensorTypeAndShapeInfo);
  size_t rank = 0;
  Check(api, api.GetDimensionsCount(info.get(), &rank));
  if (rank != 1)
    throw std::runtime_error("The custom op did not return a vector.");
  int64_t extent = 0;
  Check(api, api.GetDimensions(info.get(), &extent, 1));
  ONNXTensorElementDataType type;
  Check(api, api.GetTensorElementType(info.get(), &type));
  if (extent != static_cast<int64_t>(values.size()) || type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)
    throw std::runtime_error("The custom op returned an unexpected shape or type.");
  void *data = nullptr;
  Check(api, api.GetTensorMutableData(output.get(), &data));
  const auto *result = static_cast<const float *>(data);
  const std::array<float, 4> expected{-0.5f, -0.125f, 0, 2};
  for (size_t i = 0; i < expected.size(); ++i) {
    if (result[i] != expected[i])
      throw std::runtime_error("The custom op returned an unexpected value.");
    std::cout << result[i] << (i + 1 == expected.size() ? '\n' : ' ');
  }
}

CPU Log in both custom-op hosts

The outgoing gallery example executes onnx_light_cpu::Log in real ORT. The incoming gallery example loads the same custom-op library through the standalone onnx-light host, without importing or loading ORT. Neither path uses an EP. Only onnx-light-cpu’s standalone ml_cpu_compute subproject is configured, not its Python bindings or another onnx-light runtime.

Reproduce the comparison from the repository root after building both bridges and the onnx-light model bindings in Release mode:

export PYTHONPATH="$PWD:$PWD/build/python:$PYTHONPATH"
python docs/examples/plot_benchmark_log_custom_ops.py \
  --library build/onnx_light_ort/custom_ops/libonnx_light_ort_custom_ops.so \
  --host build/onnx_light_ort/bridge_ops_ort_called_by_light/libonnx_light_ort_bridge_ops_ort_called_by_light.so \
  --output build/cpu-custom-ops-log-benchmark.json

The script compares steady-state session.run calls, including output allocation and Python interfaces, not isolated arithmetic. It uses identical float32 inputs from default_rng(42).uniform(0.01, 10) and validates complete outputs against NumPy float64 Log rounded to float32, before and after timing. Each backend runs in its own child process, with sequential, one-thread ORT sessions and no private CPU-kernel worker pool. Three trials rotate backend order; each size receives ten warmups followed by 31 timed batches of 20 calls. The report retains every per-call batch sample and records versions, affinity and measurement settings. The table reports the median of the 93 batch samples.

The following local diagnostic results used a Xeon Platinum 8480C, affinity restricted to CPU 0 with taskset -c 0, GCC 13.2.0, Release mode, runtime SIMD dispatch, Python 3.12.15 and source-built ORT 1.32.0. The onnx-light revision was 5c899bdd2c1011e2825728b4fc950dcc88fc6aeb and onnx-light-cpu was 054cdd984d9f85a7692dc4f678b2f13275e8ede1.

Steady-state Log inference, microseconds per call

Elements	Standard ORT Log	CPU custom op called by ORT	ORT / custom op	CPU custom op called by standalone onnx-light
1,024	5.076	4.700	1.08x	15.596
32,768	52.071	16.833	3.09x	42.551
65,536	100.364	30.519	3.29x	81.738
262,144	395.220	118.298	3.34x	437.091
1,048,576	1565.080	458.078	3.42x	1919.574

For 1,048,576 elements, the three trial medians ranged from 1558.594 to 1575.189 microseconds for standard ORT and from 444.595 to 458.625 microseconds for the custom op called by ORT. These machine-specific measurements are not portable thresholds or a multithreaded performance claim. Small tensors can be dominated by dispatch overhead. The standalone Python interface includes protective input/output copies and is slower than ORT for the larger sizes; its timings must not be presented as the outgoing bridge’s speedup.