From dc53ddef7a2fbd242a3e10b38a45463557060d33 Mon Sep 17 00:00:00 2001 From: cao lei Date: Sun, 23 Apr 2023 21:54:35 -0700 Subject: [PATCH] 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 --- .../onnxruntime/core/framework/op_kernel_context.h | 6 ++++++ .../onnxruntime/core/session/onnxruntime_c_api.h | 12 ++++++++++++ .../onnxruntime/core/session/onnxruntime_cxx_api.h | 1 + .../core/session/onnxruntime_cxx_inline.h | 6 ++++++ onnxruntime/core/framework/op_kernel.cc | 9 ++++++--- onnxruntime/core/session/custom_ops.cc | 12 ++++++++++++ onnxruntime/core/session/onnxruntime_c_api.cc | 3 ++- onnxruntime/core/session/ort_apis.h | 2 ++ onnxruntime/test/shared_lib/custom_op_utils.cc | 13 +++++++++++++ 9 files changed, 60 insertions(+), 4 deletions(-) diff --git a/include/onnxruntime/core/framework/op_kernel_context.h b/include/onnxruntime/core/framework/op_kernel_context.h index 58e2891dff..5f9eee9d23 100644 --- a/include/onnxruntime/core/framework/op_kernel_context.h +++ b/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); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 76c88b9c42..3f74f24156 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/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); }; /* diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 68ba0a59a7..0c8170b7d6 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -1702,6 +1702,7 @@ struct KernelContext { UnownedValue GetOutput(size_t index, const std::vector& dims) const; void* GetGPUComputeStream() const; Logger GetLogger() const; + OrtAllocator* GetAllocator(const OrtMemoryInfo& memory_info) const; private: OrtKernelContext* ctx_; diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index e47b678b58..25dd033369 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/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)); diff --git a/onnxruntime/core/framework/op_kernel.cc b/onnxruntime/core/framework/op_kernel.cc index e4f4b9b5d3..5590a1de4b 100644 --- a/onnxruntime/core/framework/op_kernel.cc +++ b/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()) { diff --git a/onnxruntime/core/session/custom_ops.cc b/onnxruntime/core/session/custom_ops.cc index 2ab9fdb338..ad2a0d8bb6 100644 --- a/onnxruntime/core/session/custom_ops.cc +++ b/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(context)->GetAllocator(*mem_info); + if (!allocator) { + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "No requested allocator available"); + } + std::unique_ptr p = std::make_unique(std::move(allocator)); + *out = p.release(); + return nullptr; + API_IMPL_END +}; + #ifdef _WIN32 #pragma warning(pop) #endif diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index f4ea9536a7..2b6c9a4d48 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/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' diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index efd21e5da6..f51e807374 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/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 diff --git a/onnxruntime/test/shared_lib/custom_op_utils.cc b/onnxruntime/test/shared_lib/custom_op_utils.cc index eb012190c0..bf7efacdbb 100644 --- a/onnxruntime/test/shared_lib/custom_op_utils.cc +++ b/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 @@ -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