[EP ABI] Add weight pre-packing support to kernel-based plugin EPs

cmake/onnxruntime_unittests.cmake

@@ -2104,6 +2104,9 @@ if (onnxruntime_BUILD_SHARED_LIB AND
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_factory.cc"
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep.h"
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep.cc"
+          "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_allocator.h"
+          "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_data_transfer.h"
+          "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_data_transfer.cc"
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_kernel_registration.h"
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_kernel_registration.cc"
           "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/utils.h"

include/onnxruntime/core/session/onnxruntime_ep_c_api.h

@@ -308,6 +308,32 @@ struct OrtKernelImpl {
    * \since Version 1.24.
    */
   ORT_API_T(void, Release, _In_ OrtKernelImpl* this_ptr);
+
+  /** \brief Optional function to pack an initialized constant tensor to the kernel's preferred data layout.
+   *
+   * For example, a Conv kernel can define this function to pack input W to the channel-last data layout.
+   *
+   * \note The kernel is responsible for storing/owning the packed data and related metadata if it sets `is_packed`
+   *       to true. In this case, ORT may release the original constant tensor, which must not be accessed in
+   *       the OrtKernelImpl::Compute() function.
+   *
+   * \param[in] this_ptr The OrtKernelImpl instance.
+   * \param[in] tensor The OrtValue instance representing the constant tensor. Do not cache in the kernel.
+   * \param[in] input_index The input index of the tensor in this kernel.
+   * \param[in] alloc Implementation should use this allocator for allocating the packed data. It will either be
+   *                  an allocator provided by the EP (default or read-only) or a shared allocator overridden by
+   *                  the application via OrtApi::CreateSharedAllocator().
+   * \param[out] is_packed Output parameter that the implementation sets to true if the kernel packed the tensor data.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \note Implementation of this function is optional. If not implemented (set to NULL), ORT assumes the kernel
+   *       does not pre-pack constant tensors.
+   *
+   * \since Version 1.24.
+   */
+  ORT_API2_STATUS(PrePackConstantTensor, _In_ OrtKernelImpl* this_ptr, _In_ const OrtValue* tensor,
+                  _In_ int input_index, _Inout_ OrtAllocator* alloc, _Out_ bool* is_packed);
 };
 
 /** \brief Type definition for a function that creates an OrtKernelImpl instance for an operator kernel.

onnxruntime/core/session/plugin_ep/ep_kernel_registration.cc

@@ -4,10 +4,13 @@
 #include "core/session/plugin_ep/ep_kernel_registration.h"
 
 #include <memory>
+#include <unordered_map>
 #include <utility>
 
 #include "core/framework/error_code_helper.h"
 #include "core/framework/kernel_registry.h"
+#include "core/framework/tensor.h"
+#include "core/session/allocator_adapters.h"
 #include "core/session/plugin_ep/ep_api.h"
 
 namespace onnxruntime {
@@ -37,8 +40,60 @@ class PluginEpOpKernel final : public OpKernel {
     return ToStatusAndRelease(kernel_impl_->Compute(kernel_impl_, reinterpret_cast<OrtKernelContext*>(ctx)));
   }
 
+  Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
+                 /*out*/ bool& is_packed, /*out*/ PrePackedWeights* /*prepacked_weights*/) override {
+    // Note: The `prepacked_weights` parameter is not used because sharing of prepacked weights is only
+    // really supported for ORT's CPU EP.
+    assert(kernel_impl_ != nullptr);  // Should be ensured by PluginEpOpKernel::Create().
+
+    if (kernel_impl_->PrePackConstantTensor == nullptr) {
+      // OrtKernelImpl did not define a PrePack implementation.
+      is_packed = false;
+      return Status::OK();
+    }
+
+    auto empty_tensor_deleter = [](void* /*data*/) -> void {
+      // Do not delete Tensor (as we do not own it).
+    };
+
+    // Create a non-owning OrtValue that wraps the const Tensor& with an empty deleter.
+    // This is passed to OrtKernelImpl::PrePackConstantTensor() as a const OrtValue*.
+    // The above reasons make the const_cast relatively "safe".
+    // Note: Documentation for OrtKernelImpl::PrePackConstantTensor disallows caching the OrtValue pointer.
+    const OrtValue ort_value(const_cast<Tensor*>(&tensor), DataTypeImpl::GetType<Tensor>(), empty_tensor_deleter);
+
+    ORT_RETURN_IF_ERROR(ToStatusAndRelease(
+        kernel_impl_->PrePackConstantTensor(kernel_impl_, &ort_value, input_idx,
+                                            GetCachedOrtAllocator(std::move(alloc)), &is_packed)));
+
+    return Status::OK();
+  }
+
  private:
+  /// <summary>
+  /// Gets the cached OrtAllocator for the given AllocatorPtr.
+  /// </summary>
+  /// <param name="alloc"></param>
+  /// <returns></returns>
+  OrtAllocator* GetCachedOrtAllocator(AllocatorPtr&& alloc) {
+    OrtAllocator* result = nullptr;
+    IAllocator* internal_alloc = alloc.get();
+
+    auto iter = prepack_allocs_.find(internal_alloc);
+    if (iter != prepack_allocs_.end()) {
+      result = iter->second.get();
+    } else {
+      auto ort_allocator = std::make_unique<OrtAllocatorImplWrappingIAllocator>(std::move(alloc));
+      result = ort_allocator.get();
+
+      prepack_allocs_.emplace(internal_alloc, std::move(ort_allocator));
+    }
+
+    return result;
+  }
+
   OrtKernelImpl* kernel_impl_ = nullptr;
+  std::unordered_map<IAllocator*, std::unique_ptr<OrtAllocatorImplWrappingIAllocator>> prepack_allocs_;
 };
 
 /*static*/

onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ep_allocator.h

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "../plugin_ep_utils.h"
+
+#include <memory>
+
+// `OrtAllocator` is a C API struct. `BaseAllocator` is a minimal C++ struct which inherits from `OrtAllocator`.
+// Notably, `BaseAllocator` has a virtual destructor to enable a derived class to be deleted through a `BaseAllocator`
+// pointer. Allocators which need to be deleted through a base class pointer should inherit from `BaseAllocator`.
+struct BaseAllocator : OrtAllocator {
+  virtual ~BaseAllocator() = default;
+};
+
+using AllocatorUniquePtr = std::unique_ptr<BaseAllocator>;
+
+struct CustomAllocator : BaseAllocator {
+  CustomAllocator(const OrtMemoryInfo* mem_info) : memory_info{mem_info} {
+    version = ORT_API_VERSION;
+    Alloc = AllocImpl;
+    Free = FreeImpl;
+    Info = InfoImpl;
+    Reserve = AllocImpl;  // no special reserve logic and most likely unnecessary unless you have your own arena
+    GetStats = nullptr;
+    AllocOnStream = nullptr;
+  }
+
+  static void* ORT_API_CALL AllocImpl(struct OrtAllocator* /*this_*/, size_t size) {
+    return malloc(size);
+  }
+
+  /// Free a block of memory previously allocated with OrtAllocator::Alloc
+  static void ORT_API_CALL FreeImpl(struct OrtAllocator* /*this_*/, void* p) {
+    return free(p);
+  }
+
+  /// Return a pointer to an ::OrtMemoryInfo that describes this allocator
+  static const struct OrtMemoryInfo* ORT_API_CALL InfoImpl(const struct OrtAllocator* this_) {
+    const CustomAllocator& impl = *static_cast<const CustomAllocator*>(this_);
+    return impl.memory_info;
+  }
+
+ private:
+  const OrtMemoryInfo* memory_info;
+};

onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ep_data_transfer.cc

@@ -0,0 +1,115 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "ep_data_transfer.h"
+
+#include <cassert>
+#include <gsl/span>
+
+/*static*/
+bool ORT_API_CALL ExampleDataTransfer::CanCopyImpl(const OrtDataTransferImpl* this_ptr,
+                                                   const OrtMemoryDevice* src_memory_device,
+                                                   const OrtMemoryDevice* dst_memory_device) noexcept {
+  const auto& impl = *static_cast<const ExampleDataTransfer*>(this_ptr);
+  bool src_is_our_device = impl.ep_api_.MemoryDevice_AreEqual(src_memory_device, impl.device_mem_info);
+  bool dst_is_our_device = impl.ep_api_.MemoryDevice_AreEqual(dst_memory_device, impl.device_mem_info);
+
+  if (src_is_our_device && dst_is_our_device) {
+    return true;
+  }
+
+  // implementation should check if the copy is possible, which may require checking the device type, the memory type
+  // and the vendor and device IDs as needed.
+  OrtMemoryInfoDeviceType src_device_type = impl.ep_api_.MemoryDevice_GetDeviceType(src_memory_device);
+  OrtMemoryInfoDeviceType dst_device_type = impl.ep_api_.MemoryDevice_GetDeviceType(dst_memory_device);
+  OrtDeviceMemoryType src_mem_type = impl.ep_api_.MemoryDevice_GetMemoryType(src_memory_device);
+  OrtDeviceMemoryType dst_mem_type = impl.ep_api_.MemoryDevice_GetMemoryType(dst_memory_device);
+
+  // we can copy to/from CPU or CPU accessible memory
+  if (src_is_our_device) {
+    return (dst_device_type == OrtMemoryInfoDeviceType_CPU || dst_mem_type == OrtDeviceMemoryType_HOST_ACCESSIBLE);
+  }
+
+  if (dst_is_our_device) {
+    return (src_device_type == OrtMemoryInfoDeviceType_CPU || src_mem_type == OrtDeviceMemoryType_HOST_ACCESSIBLE);
+  }
+
+  return false;
+}
+
+namespace {
+void CopyImpl(const void* src_data, void* dst_data, size_t bytes, OrtSyncStream* stream) {
+  // in our example setup this is really CPU to CPU
+
+  if (stream) {
+    // EP can do an async copy using the stream. e.g. an NVIDIA EP would provide the stream to cudaMemcpyAsync
+  }
+
+  if (src_data != dst_data) {
+    memcpy(dst_data, src_data, bytes);
+  }
+}
+}  // namespace
+
+// function to copy one or more tensors.
+// implementation can optionally use async copy if a stream is available for the input.
+/*static*/
+OrtStatus* ORT_API_CALL ExampleDataTransfer::CopyTensorsImpl(OrtDataTransferImpl* this_ptr,
+                                                             const OrtValue** src_tensors_ptr,
+                                                             OrtValue** dst_tensors_ptr,
+                                                             OrtSyncStream** streams_ptr,
+                                                             size_t num_tensors) noexcept {
+  auto& impl = *static_cast<ExampleDataTransfer*>(this_ptr);
+
+  auto src_tensors = gsl::make_span<const OrtValue*>(src_tensors_ptr, num_tensors);
+  auto dst_tensors = gsl::make_span<OrtValue*>(dst_tensors_ptr, num_tensors);
+
+  for (size_t i = 0; i < num_tensors; ++i) {
+    // the implementation for a 'real' EP would be something along these lines.
+    // See CudaDataTransferImpl in onnxruntime\core\providers\cuda\cuda_provider_factory.cc
+    const OrtMemoryDevice* src_device = impl.ep_api_.Value_GetMemoryDevice(src_tensors[i]);
+    const OrtMemoryDevice* dst_device = impl.ep_api_.Value_GetMemoryDevice(dst_tensors[i]);
+
+    OrtMemoryInfoDeviceType src_device_type = impl.ep_api_.MemoryDevice_GetDeviceType(src_device);
+    OrtMemoryInfoDeviceType dst_device_type = impl.ep_api_.MemoryDevice_GetDeviceType(dst_device);
+
+    //  OrtDeviceMemoryType src_mem_type = impl.ep_api.MemoryDevice_GetMemoryType(src_device);
+    //  OrtDeviceMemoryType dst_mem_type = impl.ep_api.MemoryDevice_GetMemoryType(dst_device);
+    //  bool copy_involves_host_accessible_memory = src_mem_type == OrtDeviceMemoryType_HOST_ACCESSIBLE ||
+    //                                              dst_mem_type == OrtDeviceMemoryType_HOST_ACCESSIBLE;
+
+    const void* src_data = nullptr;
+    void* dst_data = nullptr;
+    size_t bytes;
+
+    RETURN_IF_ERROR(impl.ort_api_.GetTensorData(src_tensors[i], &src_data));
+    RETURN_IF_ERROR(impl.ort_api_.GetTensorMutableData(dst_tensors[i], &dst_data));
+    RETURN_IF_ERROR(impl.ort_api_.GetTensorSizeInBytes(src_tensors[i], &bytes));
+
+    if (dst_device_type == OrtMemoryInfoDeviceType_GPU) {
+      if (src_device_type == OrtMemoryInfoDeviceType_GPU) {
+        // GPU -> GPU
+      } else {
+        // CPU -> GPU
+      }
+    } else if (src_device_type == OrtMemoryInfoDeviceType_GPU) {
+      // GPU -> CPU
+    } else {
+      // CPU -> CPU. may involve copy a to/from host accessible memory and a synchronize may be required first
+    }
+
+    // but in our example EP it's simpler as it's really a (fake) CPU to CPU copy
+    CopyImpl(src_data, dst_data, bytes, streams_ptr ? streams_ptr[i] : nullptr);
+  }
+
+  return nullptr;
+}
+
+/*static*/
+void ORT_API_CALL ExampleDataTransfer::ReleaseImpl(OrtDataTransferImpl* /*this_ptr*/) noexcept {
+  // In our setup the factory owns a shared ExampleDataTransfer instance so it will do the cleanup, and we ignore
+  // the call to Release from the plugin_ep::DataTransfer dtor (see /onnxruntime/core/framework/plugin_data_transfer.h)
+  //
+  // If you create a new instance on each call to OrtEpFactory::CreateDataTransfer you call `delete` here
+  // delete static_cast<ExampleDataTransfer*>(this_ptr);
+}

onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ep_data_transfer.h

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "../plugin_ep_utils.h"
+
+struct ExampleDataTransfer : OrtDataTransferImpl {
+  ExampleDataTransfer(const OrtApi& ort_api, const OrtEpApi& ep_api,
+                      const OrtMemoryDevice* device_mem_info_)
+      : ort_api_(ort_api), ep_api_(ep_api), device_mem_info{device_mem_info_} {
+    CanCopy = CanCopyImpl;
+    CopyTensors = CopyTensorsImpl;
+    Release = ReleaseImpl;
+  }
+
+  static bool ORT_API_CALL CanCopyImpl(const OrtDataTransferImpl* this_ptr,
+                                       const OrtMemoryDevice* src_memory_device,
+                                       const OrtMemoryDevice* dst_memory_device) noexcept;
+
+  // function to copy one or more tensors.
+  // implementation can optionally use async copy if a stream is available for the input.
+  static OrtStatus* ORT_API_CALL CopyTensorsImpl(OrtDataTransferImpl* this_ptr,
+                                                 const OrtValue** src_tensors_ptr,
+                                                 OrtValue** dst_tensors_ptr,
+                                                 OrtSyncStream** streams_ptr,
+                                                 size_t num_tensors) noexcept;
+  static void ORT_API_CALL ReleaseImpl(OrtDataTransferImpl* this_ptr) noexcept;
+
+ private:
+  const OrtApi& ort_api_;
+  const OrtEpApi& ep_api_;
+  const OrtMemoryDevice* device_mem_info;  // device our EP runs on
+};