API Reference

C++ API

The public C++ API is declared in onnx_light_cpu/cpu_kernels.h. Every kernel dispatches at runtime to the best available SIMD path.

namespace onnx_light_cpu {

enum class SimdLevel : int {
  kNone = 0,   // Scalar fallback (no SIMD).
  kSSE2 = 1,   // SSE2 (128-bit).
  kAVX = 2,    // AVX (256-bit).
  kAVX2 = 3,   // AVX2 (256-bit with FMA, integer ops).
  kAVX512 = 4, // AVX-512F (512-bit).
};

SimdLevel DetectSimdLevel();

void AbsFloat32(const float *input, float *output, std::size_t count);
void AbsFloat64(const double *input, double *output, std::size_t count);
void AbsInt32(const std::int32_t *input, std::int32_t *output, std::size_t count);
void AbsInt64(const std::int64_t *input, std::int64_t *output, std::size_t count);

} // namespace onnx_light_cpu

onnx-light kernel class

When onnx-light-cpu is built with -DONNX_LIGHT_CPU_WITH_ONNX_LIGHT=ON (which requires the onnx-light C++ package), an additional library lib_onnx_light_cpu_kernels is produced. It declares onnx_light_cpu/onnx_light/abs_kernel.h:

namespace onnx_light_cpu {

// Derives from onnx_light::core::runtime::KernelBase and delegates to the
// SIMD Abs* routines above.
class AbsKernel : public onnx_light::core::runtime::KernelBase { ... };

// Registers AbsKernel into onnx-light's shared KernelDispatchTable for the
// default ONNX domain / CPU device, overriding the built-in Abs.
void RegisterKernels();

} // namespace onnx_light_cpu

AbsKernel is a full KernelBase subclass, so once RegisterKernels() has run every Abs node executed by onnx-light’s runtime (and therefore any model run through ReferenceEvaluator) resolves to the SIMD-accelerated kernel.

Python API

onnx_light_cpu.onnx_py._cpukernels.detect_simd_level() → int

Returns the detected SIMD level: 0=None, 1=SSE2, 2=AVX, 3=AVX2, 4=AVX512.

onnx_light_cpu.onnx_py._cpukernels.abs(input)

Computes the elementwise absolute value of a 1-D array using the optimized SIMD dispatch. input must be a contiguous CPU array with dtype float32, float64, int32 or int64; the function dispatches on the dtype and returns a new array of the same dtype, like numpy.abs().

onnx_light_cpu.onnx_py._cpukernels.has_cpu_kernels() → bool

Returns True when the CPU kernel extension is available.

Registering kernels with onnx-light

onnx_light_cpu.register_kernels(sess, domain='')

Registers the onnx-light-cpu kernels on an onnx-light ReferenceEvaluator (any object exposing a compatible register_custom_kernel(domain, op_type, fn) method). After this call, every Abs node evaluated by sess dispatches to the SIMD-accelerated onnx-light-cpu kernel instead of the built-in one, so any ONNX model using Abs benefits from the optimized kernel. Returns sess so calls can be chained.

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

sess = ReferenceEvaluator(model)
register_kernels(sess)
(y,) = sess.run(None, {"x": np.array([-1.0, 2.0], dtype=np.float32)})