Contrib ops

Contents

The contrib ops domain contains ops that are built in to the runtime by default. Only selected operators are added as contrib ops to avoid increasing the binary size of the core runtime package. When possible, custom operators should be used.

Contrib Op List

The contrib operator schemas are documented in the ONNX Runtime repository.

Adding Contrib ops

The custom op’s schema and shape inference function should be added in contrib_defs.cc using ONNX_CONTRIB_OPERATOR_SCHEMA. Example: Inverse op

ONNX_CONTRIB_OPERATOR_SCHEMA(Inverse)
    .SetDomain(kMSDomain) // kMSDomain = "com.microsoft"
    .SinceVersion(1) // Same version used at op (symbolic) registration
    ...

A new operator should have complete reference implementation tests and shape inference tests.

Reference implementation python tests should be added in onnxruntime/test/python/contrib_ops. E.g., aten_op_tests.py

Shape inference C++ tests should be added in onnxruntime/test/contrib_ops. E.g., trilu_shape_inference_test.cc

The operator kernel should be implemented using Compute function under contrib namespace in onnxruntime/contrib_ops/cpu/ for CPU and onnxruntime/contrib_ops/cuda/ for CUDA.

namespace onnxruntime {
namespace contrib {

class Inverse final : public OpKernel {
 public:
  explicit Inverse(const OpKernelInfo& info) : OpKernel(info) {}
  Status Compute(OpKernelContext* ctx) const override;

 private:
 ...
};

ONNX_OPERATOR_KERNEL_EX(
    Inverse,
    kMSDomain,
    1,
    kCpuExecutionProvider,
    KernelDefBuilder()
        .TypeConstraint("T", BuildKernelDefConstraints<float, double, MLFloat16>()),
    Inverse);

Status Inverse::Compute(OpKernelContext* ctx) const {
... // kernel implementation
}

}  // namespace contrib
}  // namespace onnxruntime

The kernel should be registered in cpu_contrib_kernels.cc for CPU and cuda_contrib_kernels.cc for CUDA.

Now you should be able to build and install ONNX Runtime to start using your custom op.

Contrib Op Tests

Tests should be added in onnxruntime/test/contrib_ops/. For example:

namespace onnxruntime {
namespace test {

// Add a comprehensive set of unit tests for custom op kernel implementation

TEST(InverseContribOpTest, two_by_two_float) {
  OpTester test("Inverse", 1, kMSDomain); // custom opset version and domain
  test.AddInput<float>("X", {2, 2}, {4, 7, 2, 6});
  test.AddOutput<float>("Y", {2, 2}, {0.6f, -0.7f, -0.2f, 0.4f});
  test.Run();
}

...

}  // namespace test
}  // namespace onnxruntime