saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Expose additional shared_provider APIs (#8478)

Browse source
This commit is contained in:
Ankur Verma 2021-07-26 18:03:12 -07:00 committed by GitHub
parent 534c22d769
commit 91936864ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/core/providers/shared_library/provider_interfaces.h
View file

@ -113,6 +113,8 @@ struct Node__EdgeIterator {
// calls the virtual function (which will lead to infinite recursion in the bridge). There is no known way to get the non virtual member
// function pointer implementation in this case.
struct ProviderHost {
virtual const OrtApiBase* OrtGetApiBase() = 0;
virtual void* HeapAllocate(size_t size) = 0;
virtual void HeapFree(void*) = 0;
@ -726,14 +728,14 @@ struct ProviderHost {
virtual void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) = 0;
virtual AllocatorPtr AllocatorManager__GetAllocator(const AllocatorManager* p, int id, OrtMemType mem_type) = 0;
#ifdef USE_CUDA
virtual PhiloxGenerator& PhiloxGenerator__Default() = 0;
#if defined(ENABLE_TRAINING) && defined(ORT_USE_NCCL)
virtual training::DistributedRunContext& GetDistributedRunContextInstance() = 0;
#endif
#ifdef USE_CUDA
virtual PhiloxGenerator& PhiloxGenerator__Default() = 0;
#ifdef ENABLE_TRAINING_TORCH_INTEROP
virtual void contrib__PythonOpBase__Init(contrib::PythonOpBase* p, const OpKernelInfo& info) = 0;
virtual void contrib__PythonOpBase__Clear(contrib::PythonOpBase* p) = 0;

12
onnxruntime/core/session/provider_bridge_ort.cc
View file

@ -33,7 +33,7 @@
#include "orttraining/core/framework/torch/refcount_tracker.h"
#endif
#endif
#if defined(USE_CUDA) && defined(ORT_USE_NCCL)
#if defined(ORT_USE_NCCL)
#include "orttraining/training_ops/cuda/communication/nccl_service.h"
#include "orttraining/core/framework/distributed_run_context.h"
#endif
@ -135,6 +135,8 @@ struct Node__EdgeIterator_Impl : Node__EdgeIterator {
// wrapped = The internal object is exposed as an opaque pointer, so we wrap it in a class that forwards every call to the real calls. No members are ever directly accessed
// direct = Same implementation is used for shared providers & core code, but some of the methods need to be routed through here to make the linker happy
struct ProviderHostImpl : ProviderHost {
const OrtApiBase* OrtGetApiBase() override { return ::OrtGetApiBase(); }
void* HeapAllocate(size_t size) override { return new uint8_t[size]; }
void HeapFree(void* p) override { delete[] reinterpret_cast<uint8_t*>(p); }
@ -809,14 +811,14 @@ struct ProviderHostImpl : ProviderHost {
void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) override { p->AllocatorManager::InsertAllocator(allocator); }
AllocatorPtr AllocatorManager__GetAllocator(const AllocatorManager* p, int id, OrtMemType mem_type) override { return p->AllocatorManager::GetAllocator(id, mem_type); };
#ifdef USE_CUDA
PhiloxGenerator& PhiloxGenerator__Default() override { return PhiloxGenerator::Default(); }
#if defined(ENABLE_TRAINING) && defined(ORT_USE_NCCL)
training::DistributedRunContext& GetDistributedRunContextInstance() override { return training::DistributedRunContext::GetInstance(); }
#endif
#ifdef USE_CUDA
PhiloxGenerator& PhiloxGenerator__Default() override { return PhiloxGenerator::Default(); }
#ifdef ENABLE_TRAINING_TORCH_INTEROP
void contrib__PythonOpBase__Init(contrib::PythonOpBase* p, const OpKernelInfo& info) override { p->PythonOpBase::Init(info); }
void contrib__PythonOpBase__Clear(contrib::PythonOpBase* p) override { p->PythonOpBase::Clear(); }