mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-23 19:32:23 +00:00
29dcfb24ab
* Add support for sharing allocators * Incremental update * Address some PR comments, add unit tests, add documentation. * Address PR comments, add tests and some documentation. * Fix build and test issues * Remove RegisterAllocator API restoring the OrtAllocator interface changes. Changed docs to reflect this. Also fixed the orttraining segfault. The segfault was because in the case of training session, the CPU exec prov is not available at the time the transformers are applied. Changed it to create a new one.
25 lines
651 B
C++
25 lines
651 B
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#pragma once
|
|
#include "core/session/onnxruntime_c_api.h"
|
|
#include "core/framework/allocator.h"
|
|
#include "core/framework/arena.h"
|
|
#include "core/session/device_allocator.h"
|
|
|
|
namespace onnxruntime {
|
|
class AllocatorWrapper : public IAllocator {
|
|
public:
|
|
AllocatorWrapper(OrtAllocator* impl) : IAllocator(*impl->Info(impl)), impl_(impl) {}
|
|
void* Alloc(size_t size) override {
|
|
return impl_->Alloc(impl_, size);
|
|
}
|
|
void Free(void* p) override {
|
|
return impl_->Free(impl_, p);
|
|
}
|
|
|
|
private:
|
|
OrtAllocator* impl_;
|
|
};
|
|
|
|
} // namespace onnxruntime
|