forked from apache/tvm
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
10 changed files
with
331 additions
and
40 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| /*! | ||
| * Copyright (c) 2019 by Contributors | ||
| * \file tvm/runtime/memory_manager.h | ||
| * \brief Abstract device memory management API | ||
| */ | ||
| #ifndef TVM_RUNTIME_MEMORY_MANAGER_H_ | ||
| #define TVM_RUNTIME_MEMORY_MANAGER_H_ | ||
|
|
||
| #include <memory> | ||
| #include <mutex> | ||
| #include <vector> | ||
| #include <unordered_map> | ||
| #include "c_runtime_api.h" | ||
|
|
||
| namespace std { | ||
| template<> | ||
| struct hash<TVMContext> { | ||
| std::size_t operator()(const TVMContext& ctx) const { | ||
| return ((ctx.device_id << 8) | ctx.device_type); | ||
| } | ||
| }; | ||
|
|
||
| template<> | ||
| struct equal_to<TVMContext> { | ||
| bool operator()(const TVMContext& lhs, const TVMContext& rhs) const { | ||
| return (lhs.device_type == rhs.device_type && | ||
| lhs.device_id == rhs.device_id); | ||
| } | ||
| }; | ||
|
|
||
| } // namespace std | ||
|
|
||
| namespace tvm { | ||
| namespace runtime { | ||
|
|
||
| struct Buffer { | ||
| // data pointer | ||
| void* data{nullptr}; | ||
| // Buffer size in bytes | ||
| size_t size{0}; | ||
| // TVM Context | ||
| TVMContext ctx; | ||
| }; | ||
|
|
||
| class Allocator { | ||
| public: | ||
| Allocator(TVMContext ctx) : ctx_(ctx) {} | ||
|
|
||
| virtual Buffer Alloc(size_t nbytes, size_t alignment, TVMType type_hint) = 0; | ||
| virtual void Free(const Buffer& buffer) = 0; | ||
| virtual size_t UsedMemory() = 0; | ||
| virtual ~Allocator() = default; | ||
|
|
||
| protected: | ||
| TVMContext ctx_; | ||
| }; | ||
|
|
||
| class MemoryManager { | ||
| public: | ||
| static MemoryManager* Global(); | ||
|
|
||
| Allocator* GetAllocator(TVMContext ctx); | ||
|
|
||
| private: | ||
| MemoryManager() {} | ||
|
|
||
| private: | ||
| std::mutex mu_; | ||
| std::unordered_map<TVMContext, std::unique_ptr<Allocator>> allocators_; | ||
| }; | ||
|
|
||
| } // namespace runtime | ||
| } // namespace tvm | ||
|
|
||
| #endif // TVM_RUNTIME_MEMORY_MANAGER_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,26 @@ | ||
| #include <tvm/runtime/memory_manager.h> | ||
| #include "naive_allocator.h" | ||
| #include "pooled_allocator.h" | ||
|
|
||
| namespace tvm { | ||
| namespace runtime { | ||
|
|
||
| MemoryManager* MemoryManager::Global() { | ||
| static MemoryManager memory_manager; | ||
| return &memory_manager; | ||
| } | ||
|
|
||
| Allocator* MemoryManager::GetAllocator(TVMContext ctx) { | ||
| std::lock_guard<std::mutex> lock(mu_); | ||
| if (allocators_.find(ctx) == allocators_.end()) { | ||
| LOG(INFO) << "New allocator for " << DeviceName(ctx.device_type) << "(" | ||
| << ctx.device_id << ")"; | ||
| std::unique_ptr<Allocator> alloc(new NaiveAllocator(ctx)); | ||
| //std::unique_ptr<Allocator> alloc(new PooledAllocator(ctx, 128)); | ||
| allocators_.emplace(ctx, std::move(alloc)); | ||
| } | ||
| return allocators_.at(ctx).get(); | ||
| } | ||
|
|
||
| } // namespace runtime | ||
| } // namespace tvm |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| /*! | ||
| * Copyright (c) 2019 by Contributors | ||
| * \file runtime/naive_allocator.h | ||
| */ | ||
| #ifndef TVM_RUNTIME_NAIVE_ALLOCATOR_H_ | ||
| #define TVM_RUNTIME_NAIVE_ALLOCATOR_H_ | ||
|
|
||
| #include <atomic> | ||
| #include <tvm/runtime/device_api.h> | ||
| #include <tvm/runtime/memory_manager.h> | ||
|
|
||
| namespace tvm { | ||
| namespace runtime { | ||
|
|
||
| class NaiveAllocator final : public Allocator { | ||
| public: | ||
| NaiveAllocator(TVMContext ctx) : Allocator(ctx), used_memory_(0) {} | ||
|
|
||
| Buffer Alloc(size_t nbytes, size_t alignment, TVMType type_hint) override { | ||
| Buffer buf; | ||
| buf.ctx = ctx_; | ||
| buf.size = nbytes; | ||
| buf.data = DeviceAPI::Get(ctx_)->AllocDataSpace( | ||
| ctx_, nbytes, alignment, type_hint); | ||
| used_memory_.fetch_add(nbytes, std::memory_order_relaxed); | ||
| LOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B"; | ||
| return buf; | ||
| } | ||
|
|
||
| void Free(const Buffer& buffer) override { | ||
| DeviceAPI::Get(ctx_)->FreeDataSpace(buffer.ctx, buffer.data); | ||
| used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed); | ||
| LOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_ << " B"; | ||
| } | ||
|
|
||
| size_t UsedMemory() override { | ||
| return used_memory_.load(std::memory_order_relaxed); | ||
| } | ||
|
|
||
| private: | ||
| std::atomic<size_t> used_memory_; | ||
| }; | ||
|
|
||
| } // namespace runtime | ||
| } // namespace tvm | ||
|
|
||
| #endif // TVM_RUNTIME_NAIVE_ALLOCATOR_H_ |
Oops, something went wrong.