Create a new C API KernelContext_GetAllocator() for Custom Op scenario (#15591)

### Description
Create a new C API KernelContext_GetAllocator() for Custom Op scenario



### Motivation and Context
Create a new C API KernelContext_GetAllocator() for Custom Op scenario

include/onnxruntime/core/framework/op_kernel_context.h

@@ -180,6 +180,12 @@ class OpKernelContext {
     return true;
   }
 
+  /**
+  Returns Allocator from a specific OrtMemoryInfo object.
+  TODO(leca): Replace GetTempSpaceAllocator() and GetTempSpaceCPUAllocator() with this API in the future
+  */
+  AllocatorPtr GetAllocator(const OrtMemoryInfo& memory_info) const;
+
  protected:
   OpKernelContext(concurrency::ThreadPool* threadpool, const logging::Logger& logger, Stream* stream);
 

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -4162,6 +4162,18 @@ struct OrtApi {
    * \snippet{doc} snippets.dox OrtStatus Return Value
    */
   ORT_API2_STATUS(GetResizedStringTensorElementBuffer, _Inout_ OrtValue* value, _In_ size_t index, _In_ size_t length_in_bytes, _Inout_ char** buffer);
+
+  /** \brief Get Allocator from KernelContext for a specific memoryInfo.
+   *
+   * \param[in] context OrtKernelContext instance
+   * \param[in] mem_info OrtMemoryInfo instance
+   * \param[out] out A pointer to OrtAllocator.
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value. Always returns nullptr.
+   *
+   * \since Version 1.15.
+   */
+  ORT_API2_STATUS(KernelContext_GetAllocator, _In_ const OrtKernelContext* context, _In_ const OrtMemoryInfo* mem_info, _Outptr_ OrtAllocator** out);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -1702,6 +1702,7 @@ struct KernelContext {
   UnownedValue GetOutput(size_t index, const std::vector<int64_t>& dims) const;
   void* GetGPUComputeStream() const;
   Logger GetLogger() const;
+  OrtAllocator* GetAllocator(const OrtMemoryInfo& memory_info) const;
 
  private:
   OrtKernelContext* ctx_;

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -1563,6 +1563,12 @@ inline void* KernelContext::GetGPUComputeStream() const {
   return out;
 }
 
+inline OrtAllocator* KernelContext::GetAllocator(const OrtMemoryInfo& memory_info) const {
+  OrtAllocator* out = nullptr;
+  Ort::ThrowOnError(GetApi().KernelContext_GetAllocator(ctx_, &memory_info, &out));
+  return out;
+}
+
 inline Logger KernelContext::GetLogger() const {
   const OrtLogger* out = nullptr;
   ThrowOnError(GetApi().KernelContext_GetLogger(this->ctx_, &out));

onnxruntime/core/framework/op_kernel.cc

@@ -93,7 +93,7 @@ int OpKernelContext::NumVariadicInputs(size_t arg_num) const {
 }
 
 Status OpKernelContext::GetTempSpaceAllocator(AllocatorPtr* output) const {
-  *output = execution_frame_->GetAllocator(kernel_->Allocator(OrtMemTypeDefault));
+  *output = GetAllocator(kernel_->Allocator(OrtMemTypeDefault));
   if (!*output)
     return Status(common::ONNXRUNTIME, common::FAIL, "TempSpace allocator not found");
   return Status::OK();
@@ -104,8 +104,7 @@ Status OpKernelContext::GetTempSpaceCPUAllocator(AllocatorPtr* output) const {
   // (which is called via ExecutionFrame), the allocator lookup
   // logic doesn't key on OrtAllocatorType, so any OrtAllocatorType
   // is good here.
-  *output = execution_frame_->GetAllocator(
-      OrtMemoryInfo(CPU, OrtAllocatorType::OrtArenaAllocator));
+  *output = GetAllocator(OrtMemoryInfo(CPU, OrtAllocatorType::OrtArenaAllocator));
   if (!*output)
     return Status(common::ONNXRUNTIME, common::FAIL, "CPU allocator not found");
   return Status::OK();
@@ -183,6 +182,10 @@ OrtValue* OpKernelContext::GetOutputMLValue(int index) {
   return execution_frame_->GetMutableNodeInputOrOutputMLValue(output_arg_index);
 }
 
+AllocatorPtr OpKernelContext::GetAllocator(const OrtMemoryInfo& memory_info) const {
+  return execution_frame_->GetAllocator(memory_info);
+}
+
 #ifdef ENABLE_ATEN
 Status OpKernelContext::SetOutputMLValue(int index, const OrtValue& ort_value) {
   if (index < 0 || index >= OutputCount()) {

onnxruntime/core/session/custom_ops.cc

@@ -106,6 +106,18 @@ ORT_API_STATUS_IMPL(OrtApis::KernelContext_GetGPUComputeStream, _In_ const OrtKe
   API_IMPL_END
 };
 
+ORT_API_STATUS_IMPL(OrtApis::KernelContext_GetAllocator, _In_ const OrtKernelContext* context, _In_ const OrtMemoryInfo* mem_info, _Outptr_ OrtAllocator** out) {
+  API_IMPL_BEGIN
+  onnxruntime::AllocatorPtr allocator = reinterpret_cast<const onnxruntime::OpKernelContext*>(context)->GetAllocator(*mem_info);
+  if (!allocator) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "No requested allocator available");
+  }
+  std::unique_ptr<onnxruntime::OrtAllocatorImplWrappingIAllocator> p = std::make_unique<onnxruntime::OrtAllocatorImplWrappingIAllocator>(std::move(allocator));
+  *out = p.release();
+  return nullptr;
+  API_IMPL_END
+};
+
 #ifdef _WIN32
 #pragma warning(pop)
 #endif

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -2724,7 +2724,8 @@ static constexpr OrtApi ort_api_1_to_15 = {
     &OrtApis::KernelInfoGetConstantInput_tensor,
     &OrtApis::CastTypeInfoToOptionalTypeInfo,
     &OrtApis::GetOptionalContainedTypeInfo,
-    &OrtApis::GetResizedStringTensorElementBuffer};
+    &OrtApis::GetResizedStringTensorElementBuffer,
+    &OrtApis::KernelContext_GetAllocator};
 
 // Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other)
 // If any of these asserts hit, read the above 'Rules on how to add a new Ort API version'

onnxruntime/core/session/ort_apis.h

@@ -462,4 +462,6 @@ ORT_API_STATUS_IMPL(GetOptionalContainedTypeInfo, _In_ const OrtOptionalTypeInfo
 
 ORT_API_STATUS_IMPL(GetResizedStringTensorElementBuffer, _Inout_ OrtValue* value,
                     _In_ size_t index, _In_ size_t length_in_bytes, _Inout_ char**);
+
+ORT_API_STATUS_IMPL(KernelContext_GetAllocator, _In_ const OrtKernelContext* context, _In_ const OrtMemoryInfo* mem_info, _Outptr_ OrtAllocator** out);
 }  // namespace OrtApis

onnxruntime/test/shared_lib/custom_op_utils.cc

@@ -5,6 +5,8 @@
 
 #include "custom_op_utils.h"
 #include "core/common/common.h"
+#include "core/framework/ortdevice.h"
+#include "core/framework/ortmemoryinfo.h"
 
 #ifdef USE_CUDA
 #include <cuda_runtime.h>
@@ -31,6 +33,17 @@ void MyCustomKernel::Compute(OrtKernelContext* context) {
   auto output_info = output.GetTensorTypeAndShapeInfo();
   int64_t size = output_info.GetElementCount();
 
+#ifdef USE_CUDA
+  OrtMemoryInfo mem_info("", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, 0));
+#else
+  OrtMemoryInfo mem_info("", OrtAllocatorType::OrtArenaAllocator, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::DEFAULT, 0));
+#endif
+  OrtAllocator* allocator;
+  Ort::ThrowOnError(ort_.KernelContext_GetAllocator(context, &mem_info, &allocator));
+  void* allocated = allocator->Alloc(allocator, 2);
+  EXPECT_NE(allocated, nullptr) << "KernelContext_GetAllocator() can successfully allocate some memory";
+  allocator->Free(allocator, allocated);
+
   // Do computation
 #ifdef USE_CUDA
   // Launch on stream 0 or user provided stream