mxnet
include
mxnet
storage.h
Go to the documentation of this file.
1
/*
2
* Licensed to the Apache Software Foundation (ASF) under one
3
* or more contributor license agreements. See the NOTICE file
4
* distributed with this work for additional information
5
* regarding copyright ownership. The ASF licenses this file
6
* to you under the Apache License, Version 2.0 (the
7
* "License"); you may not use this file except in compliance
8
* with the License. You may obtain a copy of the License at
9
*
10
* http://www.apache.org/licenses/LICENSE-2.0
11
*
12
* Unless required by applicable law or agreed to in writing,
13
* software distributed under the License is distributed on an
14
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
* KIND, either express or implied. See the License for the
16
* specific language governing permissions and limitations
17
* under the License.
18
*/
19
24
#ifndef MXNET_STORAGE_H_
25
#define MXNET_STORAGE_H_
26
27
#include <memory>
28
#include <string>
29
#include <vector>
30
#include "
./base.h
"
31
32
namespace
mxnet
{
33
34
#define MXNET_STORAGE_DEFAULT_PROFILER_SCOPE_CSTR "<unk>:"
35
#define MXNET_STORAGE_DEFAULT_NAME_CSTR "unknown"
36
40
class
Storage
{
41
public
:
45
struct
SyncObj
{
46
#if MXNET_USE_CUDA
47
50
std::vector<std::weak_ptr<cudaEvent_t>>
events
;
51
#endif
52
};
56
struct
Handle
{
60
void
*
dptr
{
nullptr
};
64
size_t
size
{0};
68
Context
ctx
;
72
int
shared_pid
{-1};
73
int
shared_id
{-1};
77
std::string
profiler_scope
{
MXNET_STORAGE_DEFAULT_PROFILER_SCOPE_CSTR
};
78
std::string
name
{
MXNET_STORAGE_DEFAULT_NAME_CSTR
};
83
SyncObj
sync_obj
;
84
};
92
Handle
Alloc
(
size_t
size,
Context
ctx,
bool
failsafe =
false
) {
93
Handle
hd;
94
hd.
size
= size;
95
hd.
ctx
= ctx;
96
this->
Alloc
(&hd, failsafe);
97
return
hd;
98
}
103
virtual
void
Alloc
(Handle* handle,
bool
failsafe =
false
) = 0;
108
virtual
void
SharedIncrementRefCount
(Handle handle) = 0;
113
virtual
void
Free
(Handle handle) = 0;
123
virtual
void
DirectFree
(Handle handle) = 0;
131
virtual
void
ReleaseAll
(
Context
ctx) = 0;
135
virtual
~Storage
() {}
139
std::mutex&
GetMutex
(
Context::DeviceType
dev) {
140
if
(dev ==
Context::kCPU
) {
141
return
cpu_mutex_;
142
}
else
{
143
return
gpu_mutex_;
144
}
145
}
149
static
Storage
*
Get
();
158
static
const
std::shared_ptr<Storage>&
_GetSharedRef
();
159
160
private
:
161
std::mutex cpu_mutex_;
162
std::mutex gpu_mutex_;
163
};
// class Storage
164
}
// namespace mxnet
165
#endif // MXNET_STORAGE_H_
mxnet
namespace of mxnet
Definition:
api_registry.h:33
MXNET_STORAGE_DEFAULT_PROFILER_SCOPE_CSTR
#define MXNET_STORAGE_DEFAULT_PROFILER_SCOPE_CSTR
Definition:
storage.h:34
mxnet::Context::kCPU
@ kCPU
Definition:
base.h:93
mxnet::Storage::SyncObj::events
std::vector< std::weak_ptr< cudaEvent_t > > events
All the events from the engine variable.
Definition:
storage.h:50
mxnet::Storage::Handle::dptr
void * dptr
Pointer to the data.
Definition:
storage.h:60
mxnet::Storage::GetMutex
std::mutex & GetMutex(Context::DeviceType dev)
Returns mutex used by storage manager.
Definition:
storage.h:139
mxnet::Storage::Get
static Storage * Get()
MXNET_STORAGE_DEFAULT_NAME_CSTR
#define MXNET_STORAGE_DEFAULT_NAME_CSTR
Definition:
storage.h:35
mxnet::Storage::ReleaseAll
virtual void ReleaseAll(Context ctx)=0
Release all memory from device if using a pooled storage manager.
mxnet::Storage::~Storage
virtual ~Storage()
Destructor.
Definition:
storage.h:135
mxnet::Storage::Handle::size
size_t size
Size of the storage.
Definition:
storage.h:64
mxnet::Storage::_GetSharedRef
static const std::shared_ptr< Storage > & _GetSharedRef()
Get shared pointer reference to storage singleton. Most user should not call this function....
mxnet::Storage::SyncObj
Storage sync object.
Definition:
storage.h:45
mxnet::Storage::Handle::sync_obj
SyncObj sync_obj
Used to pass events back and forth between the engine Var and the storage manager.
Definition:
storage.h:83
mxnet::Context::DeviceType
DeviceType
Type of device.
Definition:
base.h:92
mxnet::Storage::Handle
Storage handle.
Definition:
storage.h:56
mxnet::Storage::SharedIncrementRefCount
virtual void SharedIncrementRefCount(Handle handle)=0
Increase ref counter on shared memory.
mxnet::Storage::Handle::shared_pid
int shared_pid
Id for IPC shared memory.
Definition:
storage.h:72
mxnet::Context
Context information about the execution environment.
Definition:
base.h:90
mxnet::Storage::Handle::name
std::string name
Definition:
storage.h:78
mxnet::Storage::Free
virtual void Free(Handle handle)=0
Free storage.
mxnet::Storage::Handle::profiler_scope
std::string profiler_scope
Attributes for tracking storage allocations.
Definition:
storage.h:77
mxnet::Storage::Handle::ctx
Context ctx
Context information about device and ID.
Definition:
storage.h:68
mxnet::Storage::Alloc
Handle Alloc(size_t size, Context ctx, bool failsafe=false)
Allocate a new contiguous memory for a given size.
Definition:
storage.h:92
mxnet::Storage::Handle::shared_id
int shared_id
Definition:
storage.h:73
mxnet::Storage
Storage manager across multiple devices.
Definition:
storage.h:40
base.h
configuration of MXNet as well as basic data structure.
mxnet::Storage::DirectFree
virtual void DirectFree(Handle handle)=0
Free storage directly, without putting it into memory pool. This can synchronization of all previous ...
Generated on Thu Jan 5 2023 03:47:40 for mxnet by
1.8.17