Custom Kernels for onnxruntime#

onnxruntime implements a C API which allows the user to add custom implementation for any new operator. This mechanism is described on onnxruntime documentation Custom operators. This packages implements a couple of custom operators for CPU and GPU (NVIDIA). The first steps is to register an assembly to let onnxruntime use them.

from onnxruntime import InferenceSession, SessionOptions
from onnx_extended.ortops.optim.cpu import get_ort_ext_libs

opts = SessionOptions()
opts.register_custom_ops_library(get_ort_ext_libs()[0])

sess = InferenceSession(
    "<model_name_or_bytes>", opts, providers=[..., "CPUExecutionProvider"]
)

It supports any onnxruntime C API greater than version:

<<<

from onnx_extended.ortcy.wrap.ortinf import get_ort_c_api_supported_version

print(get_ort_c_api_supported_version())

>>>

    16

Next section introduces the list of operators and assemblies this package implements.

onnx_extended.ortops.tutorial.cpu#

<<<

from onnx_extended.ortops.tutorial.cpu import get_ort_ext_libs

print(get_ort_ext_libs())

>>>

    ['/home/xadupre/github/onnx-extended/onnx_extended/ortops/tutorial/cpu/libortops_tutorial_cpu.so']

onnx_extented.ortops.tutorial.cpu.DynamicQuantizeLinear#

Implements DynamicQuantizeLinear opset 20.

Provider

CPUExecutionProvider

Attributes

  • to: quantized type

Inputs

  • X (T1): tensor of type T

Outputs

  • Y (T2): quantized X

  • scale (TS): scale

  • Y (T2): zero point

Constraints

  • T1: float, float 16

  • TS: float

  • T2: int8, uint8, float8e4m3fn, float8e4m3fnuz, float8e5m2, float8e5m2fnuz

onnx_extented.ortops.tutorial.cpu.MyCustomOp#

It does the sum of two tensors.

Provider

CPUExecutionProvider

Inputs

  • X (T): tensor of type T

  • Y (T): tensor of type T

Outputs

  • Z (T): addition of X, Y

Constraints

  • T: float

onnx_extented.ortops.tutorial.cpu.MyCustomOpWithAttributes#

It does the sum of two tensors + a constant equal to cst = att_float + att_int64 + att_string[0] + att_tensot[0].

Provider

CPUExecutionProvider

Attributes

  • att_float: a float

  • att_int64: an integer

  • att_tensor: a tensor of any type and shape

  • att_string: a string

Inputs

  • X (T): tensor of type T

  • Y (T): tensor of type T

Outputs

  • Z (T): addition of X, Y + cst

Constraints

  • T: float

onnx_extended.ortops.tutorial.cuda#

<<<

from onnx_extended.ortops.tutorial.cuda import get_ort_ext_libs

try:
    print(get_ort_ext_libs())
except RuntimeError as e:
    print(f"CUDA is not enabled: {e}")

>>>

    ['/home/xadupre/github/onnx-extended/onnx_extended/ortops/tutorial/cuda/libortops_tutorial_cuda.so']

onnx_extented.ortops.tutorial.cuda.CustomGemm#

It calls CUDA library for Gemm \alpha A B + \beta C.

Provider

CUDAExecutionProvider

Inputs

  • A (T): tensor of type T

  • B (T): tensor of type T

  • C (T): tensor of type T

  • D (T): tensor of type T

  • E (T): tensor of type T

Outputs

  • Z (T): \alpha A B + \beta C

Constraints

  • T: float, float16, bfloat16

onnx_extended.ortops.optim.cpu#

<<<

from onnx_extended.ortops.optim.cpu import get_ort_ext_libs

print(get_ort_ext_libs())

>>>

    ['/home/xadupre/github/onnx-extended/onnx_extended/ortops/optim/cpu/libortops_optim_cpu.so']

onnx_extented.ortops.option.cpu.TreeEnsembleClassifier#

It does the sum of two tensors.

Provider

CPUExecutionProvider

Inputs

  • X (T1): tensor of type T1

Outputs

  • label (T3): labels of type T3

  • Y (T2): probabilities of type T2

Constraints

  • T1: float, double

  • T2: float, double

  • T3: int64

Attributes

onnx_extented.ortops.option.cpu.TreeEnsembleRegressor#

It does the sum of two tensors.

Provider

CPUExecutionProvider

Inputs

  • X (T1): tensor of type T1

Outputs

  • Y (T2): prediction of type T2

Constraints

  • T1: float, double

  • T2: float, double

Attributes