What's the right way to construct custom ops with the same name but different output types?

[yuanyao-nv]

I'm following the examples here to construct custom ops. In my particular use case I want to use an attribute to determine the output type of the custom op, similar to how the standard Cast op chooses its output type depending on the "to" attribute. However, if I register several custom ops with the same name but different output types, anything other than the first registered version will give a type inference error.

Here's a sketch of my code:

template <typename out_t>
struct KernelCast {
  int32_t output_dtype_attr;

  KernelCast(const OrtApi* ort_api, const OrtKernelInfo* info) {
    int64_t output_dtype = 1;
    if(ort_api->KernelInfoGetAttribute_int64(info, "output_dtype", &output_dtype) != nullptr)
    {
      std::cout << "Read output_dtype attr error, using default value of 1" << std::endl;
    }
    output_dtype_attr = output_dtype;
  }

  void Compute(const Ort::Custom::Tensor<float_t> &x_,
                Ort::Custom::Tensor<out_t> &y_)
  {
    // Do something
  }
};

void RegisterOps(Ort::CustomOpDomain &domain)
{
  static const std::unique_ptr<OrtLiteCustomOp> c_CustomOpCastFloat16{Ort::Custom::CreateLiteCustomOp<KernelCast<Ort::Float16_t>>("CustomCast", "CPUExecutionProvider")};
  domain.Add(c_CustomOpCastFloat16.get());
  static const std::unique_ptr<OrtLiteCustomOp> c_CustomOpCastFloat{Ort::Custom::CreateLiteCustomOp<KernelCast<float>>("CustomCast", "CPUExecutionProvider")};
  domain.Add(c_CustomOpCastFloat.get());
}

And here's the error: onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Type Error: Type (tensor(float)) of output arg (y) of node () does not match expected type (tensor(float16)).

What's the right way to let it dispatch to different output type kernels depending on the attribute? Any examples available to reference?

[skottmckay]

What happens if you have a template type for both the input and output types? Possibly the parsing to infer type constraints requires that. If it doesn't infer the output type constraint correctly it won't match the kernel correctly.

[yuanyao-nv]

@skottmckay If I also template on the input type and there are two versions with the same input type but different output types, then I run into the issue again where only the first registered version works. However, if each input type corresponds to a different output type then it works.
In other words, when deciding which kernel to dispatch to, ORT only seems to look at the input type.

Btw, I see this example https://github.com/maminus/runtime_custom_op/blob/main/fma_use_t/fma_use_t.cpp#L65 where the logic for kernel dispatch is delayed till the Compute() method. But in this example it's based on input type. Does something like this work on output type?

[skottmckay]

It might. You'd have to define the custom kernel using Ort::CustomOpBase (vs using CreateLiteCustomOp) so you can return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED from GetOutputType as that represents 'any type'. That may let you do it in one kernel.

https://github.com/maminus/runtime_custom_op/blob/fde922b2b3f12a50073ecc819f9cad999e0fffab/fma_use_t/fma_use_t.cpp#L109C10-L109C49

Side note: prefer setting WithStatus to true (3rd template param of CustomOpBase) so errors are handled more gracefully. Otherwise the kernel has to throw which is not ideal in some scenarios (e.g. build with exceptions disables means that's an abort()).

[skottmckay]

Working example overriding the ONNX Cast op. I don't think you need to do it by overriding the ONNX operator though, and I'm not sure if we officially support that.

I believe the key pieces are to use ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED for the output type and implement type inferencing. I got a similar error to yours until I added that. See CustomCast::InferOutputShape below.

Tested using an ONNX Cast test case. Model is here.

Was able to load and execute with it hitting the 'not implemented' error.

struct CustomCastKernel {
  CustomCastKernel(const OrtApi& ort_api, const OrtKernelInfo* /*info*/)
      : ort_(ort_api) {
  }

  OrtStatusPtr ComputeV2(OrtKernelContext* context) {
    return Ort::Status("CustomCastKernel::ComputeV2 is not implemented", ORT_NOT_IMPLEMENTED).release();
  }

 private:
  const OrtApi& ort_;
};

struct CustomCast : Ort::CustomOpBase<CustomCast, CustomCastKernel, true> {
  explicit CustomCast(const char* provider) : provider_(provider) {
    start_ver_ = 7;  // set as needed. should match minimum ONNX schema you implement if overriding ONNX op
    // end_ver_ = ...; should match maximum ONNX schema you implement or unset for unlimited.
  }

  static OrtStatusPtr InferOutputShape(Ort::ShapeInferContext& context) {
    // output type should be based on 'to'
    auto to = context.GetAttrInt("to");
    if (to != ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) {
      return Ort::Status("Unexpected type", ORT_INVALID_ARGUMENT).release();
    }

    auto shape = context.GetInputShape(0);
    // set shape and type
    context.SetOutputShape(0, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE);
    return nullptr;
  }

  OrtStatusPtr CreateKernelV2(const OrtApi& api, const OrtKernelInfo* info, void** op_kernel) const {
    Ort::ConstKernelInfo ki(info);
    *op_kernel = new CustomCastKernel(api, info);
    return nullptr;
  };

  const char* GetName() const { return "Cast"; };
  const char* GetExecutionProviderType() const { return provider_; };

  size_t GetInputTypeCount() const { return 1; };
  ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
  };

  size_t GetOutputTypeCount() const { return 1; };
  ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; };

 private:
  const char* provider_{"CPUExecutionProvider"};
};

Usage to add custom op in ONNX domain to override the ONNX op. Use a custom domain name to avoid that.

  CustomCast custom_op{onnxruntime::kCpuExecutionProvider};

  Ort::CustomOpDomain custom_op_domain("");  // onnx domain is empty string
  custom_op_domain.Add(&custom_op);

[yuanyao-nv]

@skottmckay Thank you for your example, and apologies I was pulled into something else and didn't get a chance to respond.

I tried the example above and ran into a segfault before even reaching the ComputeV2 function. In detail, I have two files in my directory custom_op_library.h and custom_op_library.cc which follows the structure given here.

custom_op_library.h:

#pragma once
#include "core/session/onnxruntime_c_api.h"

#ifdef __cplusplus
extern "C" {
#endif

ORT_EXPORT OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api);

#ifdef __cplusplus
}
#endif

custom_op_library.cc (kernel definition copied over from your code above):

#include "custom_op_library.h"

#define ORT_API_MANUAL_INIT
#include "core/session/onnxruntime_cxx_api.h"
#undef ORT_API_MANUAL_INIT

#include <iostream>


struct CustomCastKernel {
  CustomCastKernel(const OrtApi& ort_api, const OrtKernelInfo* /*info*/)
      : ort_(ort_api) {
  }

  OrtStatusPtr ComputeV2(OrtKernelContext* context) {
    return Ort::Status("CustomCastKernel::ComputeV2 is not implemented", ORT_NOT_IMPLEMENTED).release();
  }

 private:
  const OrtApi& ort_;
};

struct CustomCast : Ort::CustomOpBase<CustomCast, CustomCastKernel, true> {
  explicit CustomCast(const char* provider) : provider_(provider) {
    start_ver_ = 7;  // set as needed. should match minimum ONNX schema you implement if overriding ONNX op
    // end_ver_ = ...; should match maximum ONNX schema you implement or unset for unlimited.
  }

  static OrtStatusPtr InferOutputShape(Ort::ShapeInferContext& context) {
    // output type should be based on 'to'
    auto to = context.GetAttrInt("to");
    if (to != ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) {
      return Ort::Status("Unexpected type", ORT_INVALID_ARGUMENT).release();
    }

    auto shape = context.GetInputShape(0);
    // set shape and type
    context.SetOutputShape(0, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE);
    return nullptr;
  }

  OrtStatusPtr CreateKernelV2(const OrtApi& api, const OrtKernelInfo* info, void** op_kernel) const {
    Ort::ConstKernelInfo ki(info);
    *op_kernel = new CustomCastKernel(api, info);
    return nullptr;
  };

  const char* GetName() const { return "Cast"; };
  const char* GetExecutionProviderType() const { return provider_; };

  size_t GetInputTypeCount() const { return 1; };
  ONNXTensorElementDataType GetInputType(size_t /*index*/) const {
    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
  };

  size_t GetOutputTypeCount() const { return 1; };
  ONNXTensorElementDataType GetOutputType(size_t /*index*/) const { return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; };

 private:
  const char* provider_{"CPUExecutionProvider"};
};

OrtStatus *ORT_API_CALL RegisterCustomOps(OrtSessionOptions *options, const OrtApiBase *api)
{
  Ort::Global<void>::api_ = api->GetApi(ORT_API_VERSION);
  OrtStatus *result = nullptr;

  Ort::CustomOpDomain custom_op_domain("");  // onnx domain is empty string
  CustomCast custom_op{"CPUExecutionProvider"};
  custom_op_domain.Add(&custom_op);

  Ort::UnownedSessionOptions session_options(options);
  session_options.Add(custom_op_domain);

  std::cout << "Registered custom operators";
  return result;
}

The build command is

g++ -shared -fPIC -Wall -Wextra -std=c++17 -g -o custom_op_library.so custom_op_library.cc -Ionnxruntime

where the onnxruntime directory contains the header files from here.

Then to run the test model, I use

model = onnx.load("model.onnx")
named_inputs = {"input": np.random.rand(3,4).astype(np.float32)}
so = ort.SessionOptions()
so.register_custom_ops_library(os.path.realpath("custom_op_library.so"))
so.log_severity_level = 0 # verbose
sess = ort.InferenceSession(model.SerializeToString(), sess_options=so, providers=["CPUExecutionProvider"])

I get:

2025-03-09 13:33:12.278718796 [I:onnxruntime:, inference_session.cc:592 TraceSessionOptions] Session Options {  execution_mode:0 execution_order:DEFAULT enable_profiling:0 optimized_model_filepath:"" enable_mem_pattern:1 enable_mem_reuse:1 enable_cpu_mem_arena:1 profile_file_prefix:onnxruntime_profile_ session_logid: session_log_severity_level:0 session_log_verbosity_level:0 max_num_graph_transformation_steps:10 graph_optimization_level:3 intra_op_param:OrtThreadPoolParams { thread_pool_size: 0 auto_set_affinity: 0 allow_spinning: 1 dynamic_block_base_: 0 stack_size: 0 affinity_str:  set_denormal_as_zero: 0 } inter_op_param:OrtThreadPoolParams { thread_pool_size: 0 auto_set_affinity: 0 allow_spinning: 1 dynamic_block_base_: 0 stack_size: 0 affinity_str:  set_denormal_as_zero: 0 } use_per_session_threads:1 thread_pool_allow_spinning:1 use_deterministic_compute:0 config_options: {  } }
2025-03-09 13:33:12.278746906 [I:onnxruntime:, inference_session.cc:412 operator()] Flush-to-zero and denormal-as-zero are off
2025-03-09 13:33:12.278757567 [I:onnxruntime:, inference_session.cc:420 ConstructorCommon] Creating and using per session threadpools since use_per_session_threads_ is true
2025-03-09 13:33:12.278770120 [I:onnxruntime:, inference_session.cc:438 ConstructorCommon] Dynamic block base set to 0
Segmentation fault (core dumped)

Do you spot anything wrong?

[yuanyao-nv]

Here's the backtrace from a debug build:

Thread 1 "python3" received signal SIGSEGV, Segmentation fault.
0x00007bb7aefe5780 in std::_Hash_bytes(void const*, unsigned long, unsigned long) () from /usr/lib/x86_64-linux-gnu/libstdc++.so.6
(gdb) bt
#0 0x00007bb7aefe5780 in std::_Hash_bytes(void const*, unsigned long, unsigned long) () from /usr/lib/x86_64-linux-gnu/libstdc++.so.6
#1 0x00007bb77816d83a in std::_Hash_impl::hash (__seed=3339675911, __clength=, __ptr=) at /usr/include/c++/13/bits/functional_hash.h:206
#2 std::__str_hash_base<char, std::allocator, std::__cxx11::basic_string<char, std::char_traits, std::allocator > >::operator() (__s=...,
this=0x7bb7790e9700 onnx::OpSchemaRegistry::DomainToVersionRange::Instance()::domain_to_version_range) at /usr/include/c++/13/bits/basic_string.h:4408
#3 std::__detail::_Hash_code_base<std::__cxx11::basic_string<char, std::char_traits, std::allocator >, std::pair<std::__cxx11::basic_string<char, std::char_traits, std::allocator > const, std::pair<int, int> >, std::__detail::_Select1st, std::hash<std::__cxx11::basic_string<char, std::char_traits, std::allocator > >, std::__detail::_Mod_range_hashing, std::__detail::_Default_ranged_hash, true>::_M_hash_code (__k=...,
this=0x7bb7790e9700 onnx::OpSchemaRegistry::DomainToVersionRange::Instance()::domain_to_version_range) at /usr/include/c++/13/bits/hashtable_policy.h:1310
#4 std::__detail::_Map_base<std::__cxx11::basic_string<char, std::char_traits, std::allocator >, std::pair<std::__cxx11::basic_string<char, std::char_traits, std::allocator > const, std::pair<int, int> >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits, std::allocator > const, std::pair<int, int> > >, std::__detail::_Select1st, std::equal_to<std::__cxx11::basic_string<char, std::char_traits, std::allocator > >, std::hash<std::__cxx11::basic_string<char, std::char_traits, std::allocator > >, std::__detail::_Mod_range_hashing, std::__detail::_Default_ranged_hash, std::__detail::_Prime_rehash_policy, std::__detail::_Hashtable_traits<true, false, true>, true>::operator[] (this=this@entry=0x7bb7790e9700 onnx::OpSchemaRegistry::DomainToVersionRange::Instance()::domain_to_version_range,
__k="\220p~\023\000\000\000\000\002\000\000\000\000\000\000\000\377\377\377\377\001\000\000\0001", '\000' <repeats 15 times>, "\b\200z<\351W\2659\270^\264\000\000\000\000\000\270^\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "C\325\001\230t\215\241\237\350X\264\000\000\000\000\000\350X\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "zXZ\370\313Z^\034\210X\264\000\000\000\000\000\210X\264\000\000\000\000\0000\000\000\000\000\000\000\0001\000\000\000\000\000\000\000\260\313m\023\000\000\000\000\n\000\000\000\000\000\000\000\v\272\337x\267{\000\000"...) at /usr/include/c++/13/bits/hashtable_policy.h:813
#5 0x00007bb77816e36b in std::unordered_map<std::__cxx11::basic_string<char, std::char_traits, std::allocator >, std::pair<int, int>, std::hash<std::__cxx11::basic_string<char, std::char_traits, std::allocator > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits, std::allocator > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits, std::allocator > const, std::pair<int, int> > > >::operator[] (
__k="\220p~\023\000\000\000\000\002\000\000\000\000\000\000\000\377\377\377\377\001\000\000\0001", '\000' <repeats 15 times>, "\b\200z<\351W\2659\270^\264\000\000\000\000\000\270^\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "C\325\001\230t\215\241\237\350X\264\000\000\000\000\000\350X\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "zXZ\370\313Z^\034\210X\264\000\000\000\000\000\210X\264\000\000\000\000\0000\000\000\000\000\000\000\0001\000\000\000\000\000\000\000\260\313m\023\000\000\000\000\n\000\000\000\000\000\000\000\v\272\337x\267{\000\000"..., this=0x7bb7790e9700 onnx::OpSchemaRegistry::DomainToVersionRange::Instance()::domain_to_version_range)
at /usr/include/c++/13/bits/unordered_map.h:986
#6 onnx::OpSchemaRegistry::DomainToVersionRange::AddDomainToVersion (this=0x7bb7790e9700 onnx::OpSchemaRegistry::DomainToVersionRange::Instance()::domain_to_version_range,
domain="\220p~\023\000\000\000\000\002\000\000\000\000\000\000\000\377\377\377\377\001\000\000\0001", '\000' <repeats 15 times>, "\b\200z<\351W\2659\270^\264\000\000\000\000\000\270^\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "C\325\001\230t\215\241\237\350X\264\000\000\000\000\000\350X\264", '\000' <repeats 13 times>, "1", '\000' <repeats 15 times>, "zXZ\370\313Z^\034\210X\264\000\000\000\000\000\210X\264\000\000\000\000\0000\000\000\000\000\000\000\0001\000\000\000\000\000\000\000\260\313m\023\000\000\000\000\n\000\000\000\000\000\000\000\v\272\337x\267{\000\000"..., min_version=1, max_version=1000, last_release_version=-1)
at /home/yuanyao/onnxruntime/build/Linux/RelWithDebInfo/_deps/onnx-src/onnx/defs/schema.h:1259
#7 0x00007bb77820a76e in onnxruntime::CreateCustomRegistry (op_domains=..., output=std::shared_ptronnxruntime::CustomRegistry (use count 1, weak count 0) = {...})
at /home/yuanyao/onnxruntime/onnxruntime/core/session/custom_ops.cc:1232
#8 0x00007bb778191e48 in onnxruntime::InferenceSession::AddCustomOpDomains (this=0x13d549a0, op_domains=...)
at /home/yuanyao/onnxruntime/onnxruntime/core/session/inference_session.cc:866
#9 0x00007bb77812ad67 in onnxruntime::python::RegisterCustomOpDomains (sess=0x13cf5370, so=...) at /home/yuanyao/onnxruntime/onnxruntime/python/onnxruntime_pybind_state.cc:1338
#10 0x00007bb77812be2a in operator() (so=...,
arg="\b\v\022\fbackend-test:p\n \n\005input\022\006output"\004Cast*\t\n\002to\030\v\240\001\002\022\031test_cast_FLOAT_to_DOUBLEZ\027\n\005input\022\016\n\f\b\001\022\b\n\002\b\003\n\002\b\004b\030\n\006output\022\016\n\f\b\v\022\b\n\002\b\003\n\002\b\004B\004\n\000\020\027", is_arg_file_name=is_arg_file_name@entry=false,
load_config_from_model=load_config_from_model@entry=false, __closure=) at /usr/include/c++/13/bits/unique_ptr.h:199
#11 0x00007bb778149119 in operator() (__closure=, args#3=, args#2=, args#1="", args#0=..., v_h=...)

[yuanyao-nv]

@skottmckay Any thoughts on the above?

Another question is in the real use case supporting multiple output types, should the logic for dispatching to different kernels be in the ComputeV2() function? If so then there should be a way to access the attribute value from the ComputeV2() function. I'm not sure what would be the right way to do that.