Add KernelBuilder aliasing API functions. Add a relu kernel to test it.

Author: adrianlizarraga

cmake/onnxruntime_unittests.cmake

@@ -2105,6 +2105,8 @@ if (onnxruntime_BUILD_SHARED_LIB AND
           "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/data_types.cc"
           "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/squeeze.h"
           "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/squeeze.cc"
+          "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/relu.h"
+          "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/relu.cc"
           "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/mul.h"
           "${TEST_SRC_DIR}/autoep/library/example_kernel_plugin_ep/kernels/mul.cc")
   onnxruntime_add_shared_library_module(example_kernel_plugin_ep ${onnxruntime_autoep_test_example_kernel_plugin_ep_src})

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -3331,6 +3331,8 @@ struct KernelDef : detail::ConstKernelDefImpl<OrtKernelDef> {
  * Used by plugin EPs to build a kernel definition.
  */
 struct KernelDefBuilder : detail::Base<OrtKernelDefBuilder> {
+  using InOutAliasPair = std::pair<int, int>;
+
   KernelDefBuilder();                           ///< Wraps OrtEpApi::CreateKernelDefBuilder
   explicit KernelDefBuilder(std::nullptr_t) {}  ///< Create an empty object, must be assigned a valid one to be used
   explicit KernelDefBuilder(OrtKernelDefBuilder* ort_kernel_def_builder);
@@ -3343,6 +3345,10 @@ struct KernelDefBuilder : detail::Base<OrtKernelDefBuilder> {
   KernelDefBuilder& SetOutputMemType(size_t output_index, OrtMemType mem_type);
   KernelDefBuilder& AddTypeConstraint(const char* arg_name, const OrtMLDataType* data_type);
   KernelDefBuilder& AddTypeConstraint(const char* arg_name, const std::vector<const OrtMLDataType*>& data_types);
+  KernelDefBuilder& AddInputOutputAlias(int input_index, int output_index);
+  KernelDefBuilder& AddInputOutputAliases(const std::vector<InOutAliasPair>& aliases);
+  KernelDefBuilder& AddInputOutputMutableAlias(int input_index, int output_index);
+  KernelDefBuilder& AddInputOutputMutableAliases(const std::vector<InOutAliasPair>& aliases);
 
   KernelDef Build();
 };

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -3644,6 +3644,48 @@ inline KernelDefBuilder& KernelDefBuilder::AddTypeConstraint(const char* arg_nam
   return *this;
 }
 
+inline KernelDefBuilder& KernelDefBuilder::AddInputOutputAlias(int input_index, int output_index) {
+  ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputAliases(p_, &input_index, &output_index, 1));
+  return *this;
+}
+
+inline KernelDefBuilder& KernelDefBuilder::AddInputOutputAliases(const std::vector<InOutAliasPair>& aliases) {
+  std::vector<int> input_indices;
+  std::vector<int> output_indices;
+
+  input_indices.reserve(aliases.size());
+  output_indices.reserve(aliases.size());
+  for (const std::pair<int, int>& alias : aliases) {
+    input_indices.push_back(alias.first);
+    output_indices.push_back(alias.second);
+  }
+
+  ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputAliases(p_, input_indices.data(), output_indices.data(),
+                                                                 input_indices.size()));
+  return *this;
+}
+
+inline KernelDefBuilder& KernelDefBuilder::AddInputOutputMutableAlias(int input_index, int output_index) {
+  ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputMutableAliases(p_, &input_index, &output_index, 1));
+  return *this;
+}
+
+inline KernelDefBuilder& KernelDefBuilder::AddInputOutputMutableAliases(const std::vector<InOutAliasPair>& aliases) {
+  std::vector<int> input_indices;
+  std::vector<int> output_indices;
+
+  input_indices.reserve(aliases.size());
+  output_indices.reserve(aliases.size());
+  for (const std::pair<int, int>& alias : aliases) {
+    input_indices.push_back(alias.first);
+    output_indices.push_back(alias.second);
+  }
+
+  ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputMutableAliases(p_, input_indices.data(), output_indices.data(),
+                                                                        input_indices.size()));
+  return *this;
+}
+
 inline KernelDef KernelDefBuilder::Build() {
   OrtKernelDef* kernel_def = nullptr;
   ThrowOnError(GetEpApi().KernelDefBuilder_Build(p_, &kernel_def));

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -683,7 +683,7 @@ struct OrtEpApi {
   ORT_API2_STATUS(KernelDefBuilder_SetOutputMemType, _In_ OrtKernelDefBuilder* kernel_def_builder,
                   _In_ size_t output_index, _In_ OrtMemType mem_type);
 
-  /** \brief Sets type constraints for a kernel argument represented as a string (e.g., "T").
+  /** \brief Adds type constraints for a kernel argument represented as a string (e.g., "T").
    *
    * \param[in] kernel_def_builder The OrtKernelDefBuilder instance.
    * \param[in] arg_name A null-terminated string representing the argument to constrain (e.g., "T").
@@ -699,6 +699,45 @@ struct OrtEpApi {
                   _In_ const char* arg_name, _In_reads_(num_types) const OrtMLDataType* const* types,
                   _In_ size_t num_types);
 
+  /** \brief Adds aliases for the given input and output pairs.
+   *
+   * \note Used for operators like Identity and Reshape to allow ORT to reuse the input buffer for the output
+   *       without modification.
+   *
+   * \param[in] kernel_def_builder The OrtKernelDefBuilder instance.
+   * \param[in] input_indices Array of input indices. Array must contain `num_io_indices` elements.
+   * \param[in] output_indices Array of output indices. Each output index is aliased with a corresponding
+   *                           input index in `input_indices`. Array must contain `num_io_indices` elements.
+   * \param[in] num_io_indices The number of input/output index pairs to alias.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(KernelDefBuilder_AddInputOutputAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                  _In_reads_(num_io_indices) int const* input_indices,
+                  _In_reads_(num_io_indices) int const* output_indices,
+                  _In_ size_t num_io_indices);
+
+  /** \brief Adds mutable aliases for the given input and output pairs.
+   *
+   * \note Allows ORT to reuse and *modify* an input buffer (in-place) for the output buffer.
+   *
+   * \param[in] kernel_def_builder The OrtKernelDefBuilder instance.
+   * \param[in] input_indices Array of input indices. Array must contain `num_io_indices` elements.
+   * \param[in] output_indices Array of output indices. Each output index is aliased with a corresponding
+   *                           input index in `input_indices`. Array must contain `num_io_indices` elements.
+   * \param[in] num_io_indices The number of input/output index pairs to alias.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(KernelDefBuilder_AddInputOutputMutableAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                  _In_reads_(num_io_indices) int const* input_indices,
+                  _In_reads_(num_io_indices) int const* output_indices,
+                  _In_ size_t num_io_indices);
+
   /** \brief Creates a OrtKernelDef instance from the given kernel definition builder.
    *
    * \param[in] kernel_def_builder The OrtKernelDefBuilder instance.

onnxruntime/core/session/plugin_ep/ep_api.cc

@@ -364,6 +364,76 @@ ORT_API_STATUS_IMPL(KernelDefBuilder_AddTypeConstraint, _In_ OrtKernelDefBuilder
   API_IMPL_END
 }
 
+ORT_API_STATUS_IMPL(KernelDefBuilder_AddInputOutputAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                    _In_reads_(num_io_indices) int const* input_indices,
+                    _In_reads_(num_io_indices) int const* output_indices,
+                    _In_ size_t num_io_indices) {
+  API_IMPL_BEGIN
+  if (num_io_indices == 0) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Must specify at least one input/output alias");
+  }
+
+  if (input_indices == nullptr) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Must specify a valid array of input indices to alias");
+  }
+
+  if (output_indices == nullptr) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Must specify a valid array of output indices to alias");
+  }
+
+  if (num_io_indices == 1) {
+    kernel_def_builder->Alias(input_indices[0], output_indices[0]);
+  } else {
+    std::vector<std::pair<int, int>> pairs;
+    pairs.reserve(num_io_indices);
+
+    for (size_t i = 0; i < num_io_indices; ++i) {
+      pairs.push_back({input_indices[i], output_indices[i]});
+    }
+
+    kernel_def_builder->Alias(pairs);
+  }
+
+  return nullptr;
+  API_IMPL_END
+}
+
+ORT_API_STATUS_IMPL(KernelDefBuilder_AddInputOutputMutableAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                    _In_reads_(num_io_indices) int const* input_indices,
+                    _In_reads_(num_io_indices) int const* output_indices,
+                    _In_ size_t num_io_indices) {
+  API_IMPL_BEGIN
+  if (num_io_indices == 0) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Must specify at least one input/output alias (mutable)");
+  }
+
+  if (input_indices == nullptr) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
+                                 "Must specify a valid array of input indices to alias (mutable)");
+  }
+
+  if (output_indices == nullptr) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
+                                 "Must specify a valid array of output indices to alias (mutable)");
+  }
+
+  if (num_io_indices == 1) {
+    kernel_def_builder->MayInplace(input_indices[0], output_indices[0]);
+  } else {
+    std::vector<std::pair<int, int>> pairs;
+    pairs.reserve(num_io_indices);
+
+    for (size_t i = 0; i < num_io_indices; ++i) {
+      pairs.push_back({input_indices[i], output_indices[i]});
+    }
+
+    kernel_def_builder->MayInplace(pairs);
+  }
+
+  return nullptr;
+  API_IMPL_END
+}
+
 ORT_API_STATUS_IMPL(KernelDefBuilder_Build, _In_ OrtKernelDefBuilder* kernel_def_builder,
                     _Outptr_ OrtKernelDef** kernel_def_out) {
   API_IMPL_BEGIN
@@ -550,6 +620,8 @@ static constexpr OrtEpApi ort_ep_api = {
     &OrtExecutionProviderApi::KernelDefBuilder_SetInputMemType,
     &OrtExecutionProviderApi::KernelDefBuilder_SetOutputMemType,
     &OrtExecutionProviderApi::KernelDefBuilder_AddTypeConstraint,
+    &OrtExecutionProviderApi::KernelDefBuilder_AddInputOutputAliases,
+    &OrtExecutionProviderApi::KernelDefBuilder_AddInputOutputMutableAliases,
     &OrtExecutionProviderApi::KernelDefBuilder_Build,
     &OrtExecutionProviderApi::ReleaseKernelDef,
     &OrtExecutionProviderApi::KernelDef_GetOperatorType,

onnxruntime/core/session/plugin_ep/ep_api.h

@@ -73,6 +73,14 @@ ORT_API_STATUS_IMPL(KernelDefBuilder_SetOutputMemType, _In_ OrtKernelDefBuilder*
 ORT_API_STATUS_IMPL(KernelDefBuilder_AddTypeConstraint, _In_ OrtKernelDefBuilder* kernel_def_builder,
                     _In_ const char* arg_name, _In_reads_(num_types) const OrtMLDataType* const* types,
                     _In_ size_t num_types);
+ORT_API_STATUS_IMPL(KernelDefBuilder_AddInputOutputAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                    _In_reads_(num_io_indices) int const* input_indices,
+                    _In_reads_(num_io_indices) int const* output_indices,
+                    _In_ size_t num_io_indices);
+ORT_API_STATUS_IMPL(KernelDefBuilder_AddInputOutputMutableAliases, _In_ OrtKernelDefBuilder* kernel_def_builder,
+                    _In_reads_(num_io_indices) int const* input_indices,
+                    _In_reads_(num_io_indices) int const* output_indices,
+                    _In_ size_t num_io_indices);
 ORT_API_STATUS_IMPL(KernelDefBuilder_Build, _In_ OrtKernelDefBuilder* kernel_def_builder,
                     _Outptr_ OrtKernelDef** kernel_def_out);
 

onnxruntime/test/autoep/library/example_kernel_plugin_ep/ep_kernel_registration.cc

@@ -8,11 +8,13 @@
 
 // Include kernels:
 #include "kernels/mul.h"
+#include "kernels/relu.h"
 #include "kernels/squeeze.h"
 
 // Table of BuildKernelCreateInfo functions for each operator
 static const BuildKernelCreateInfoFn build_kernel_create_info_funcs[] = {
     BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kOnnxDomain, 7, 24, Mul)>,
+    BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kOnnxDomain, 14, 24, Relu)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kOnnxDomain, 13, 24, Squeeze)>,
 };
 

onnxruntime/test/autoep/library/example_kernel_plugin_ep/kernels/relu.cc

@@ -0,0 +1,83 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "relu.h"
+
+#include <gsl/span>
+#include <algorithm>
+#include <cassert>
+
+#include "utils.h"
+
+ONNX_OPERATOR_VERSIONED_KERNEL_EX(
+    Relu,
+    kOnnxDomain,
+    14, 24,
+    (Ort::KernelDefBuilder()
+         .AddTypeConstraint("T", MLDataTypes::GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT))
+         .AddInputOutputMutableAlias(0, 0)),
+    Relu)
+
+Relu::Relu(const OrtKernelInfo* info, void* state, PrivateTag)
+    : info_(info),
+      state_(state) {
+  ort_version_supported = ORT_API_VERSION;
+  Compute = ComputeImpl;
+  Release = ReleaseImpl;
+}
+
+/*static*/
+OrtStatus* Relu::Create(const OrtKernelInfo* info, void* state, /*out*/ std::unique_ptr<Relu>& kernel) {
+  Ort::ConstKernelInfo kernel_info(info);
+
+  try {
+    kernel = std::make_unique<Relu>(info, state, PrivateTag{});
+  } catch (const Ort::Exception& ex) {
+    Ort::Status status(ex);
+    return status.release();
+  } catch (const std::exception& ex) {
+    Ort::Status status(ex.what(), ORT_EP_FAIL);
+    return status.release();
+  }
+
+  return nullptr;
+}
+
+/*static*/
+OrtStatus* ORT_API_CALL Relu::ComputeImpl(OrtKernelImpl* this_ptr, OrtKernelContext* kernel_ctx) noexcept {
+  Relu* relu = static_cast<Relu*>(this_ptr);
+  return relu->DoCompute(kernel_ctx);
+}
+
+/*static*/
+void ORT_API_CALL Relu::ReleaseImpl(OrtKernelImpl* this_ptr) noexcept {
+  delete static_cast<Relu*>(this_ptr);
+}
+
+OrtStatus* Relu::DoCompute(OrtKernelContext* kernel_ctx) noexcept {
+  // const OrtEpApi& ep_api = Ort::GetEpApi();
+  Ort::KernelContext kernel_context(kernel_ctx);
+  (void)this->state_;  // NOTE: Unused in this example.
+  (void)this->info_;   // NOTE: Unused in this example.
+
+  try {
+    gsl::span<const float> input0;
+    std::vector<int64_t> shape0;
+    RETURN_IF_ERROR(GetKernelInputDataAndShape<float>(kernel_context, 0, input0, shape0));
+
+    Ort::UnownedValue output = kernel_context.GetOutput(0, shape0);
+    float* output_data = output.GetTensorMutableData<float>();
+
+    for (size_t i = 0; i < input0.size(); ++i) {
+      output_data[i] = std::max(0.0f, input0[i]);
+    }
+  } catch (const Ort::Exception& ex) {
+    Ort::Status status(ex);
+    return status.release();
+  } catch (const std::exception& ex) {
+    Ort::Status status(ex.what(), ORT_EP_FAIL);
+    return status.release();
+  }
+
+  return nullptr;
+}

onnxruntime/test/autoep/library/example_kernel_plugin_ep/kernels/relu.h

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "utils.h"
+#include "../../plugin_ep_utils.h"
+
+// Forward declarations of kernel classes used as template args for BuildKernelCreateInfo
+class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kOnnxDomain, 14, 24, Relu);
+
+struct Relu : public OrtKernelImpl {
+ private:
+  struct PrivateTag {};
+
+ public:
+  static OrtStatus* Create(const OrtKernelInfo* info, void* state, /*out*/ std::unique_ptr<Relu>& kernel);
+
+  Relu(const OrtKernelInfo* info, void* state, PrivateTag);
+
+  static OrtStatus* ORT_API_CALL ComputeImpl(OrtKernelImpl* this_ptr, OrtKernelContext* kernel_ctx) noexcept;
+  static void ORT_API_CALL ReleaseImpl(OrtKernelImpl* this_ptr) noexcept;
+
+  OrtStatus* DoCompute(OrtKernelContext* kernel_ctx) noexcept;
+
+ private:
+  const OrtKernelInfo* info_;
+  void* state_{nullptr};  // Custom state passed from OrtEp
+};

onnxruntime/test/autoep/library/example_kernel_plugin_ep/kernels/squeeze.cc

@@ -14,7 +14,8 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
     13, 24,
     (Ort::KernelDefBuilder()
          .AddTypeConstraint("T", MLDataTypes::GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT))
-         .AddTypeConstraint("axes", MLDataTypes::GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64))),
+         .AddTypeConstraint("axes", MLDataTypes::GetTensorType(ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64))
+         .AddInputOutputAlias(0, 0)),
     Squeeze)
 
 Squeeze::Squeeze(const OrtKernelInfo* info, void* state, PrivateTag)