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 "
./base.h
"
29
30
namespace
mxnet
{
31
35
class
Storage
{
36
public
:
40
struct
Handle
{
44
void
*
dptr
{
nullptr
};
48
size_t
size
{0};
52
Context
ctx
;
56
int
shared_pid
{-1};
57
int
shared_id
{-1};
58
};
65
Handle
Alloc
(
size_t
size
,
Context
ctx
) {
66
Handle
hd;
67
hd.
size
=
size
;
68
hd.
ctx
=
ctx
;
69
this->
Alloc
(&hd);
70
return
hd;
71
}
76
virtual
void
Alloc
(
Handle
* handle) = 0;
81
virtual
void
SharedIncrementRefCount
(
Handle
handle) = 0;
86
virtual
void
Free
(
Handle
handle) = 0;
96
virtual
void
DirectFree
(
Handle
handle) = 0;
104
virtual
void
ReleaseAll
(
Context
ctx
) = 0;
108
virtual
~Storage
() {}
112
std::mutex&
GetMutex
(
Context::DeviceType
dev) {
113
if
(dev ==
Context::kCPU
) {
114
return
cpu_mutex_;
115
}
else
{
116
return
gpu_mutex_;
117
}
118
}
122
static
Storage
*
Get
();
131
static
std::shared_ptr<Storage>
_GetSharedRef
();
132
133
private
:
134
std::mutex cpu_mutex_;
135
std::mutex gpu_mutex_;
136
};
// class Storage
137
}
// namespace mxnet
138
#endif // MXNET_STORAGE_H_
mxnet::Storage::DirectFree
virtual void DirectFree(Handle handle)=0
Free storage directly, without putting it into memory pool. This can synchronization of all previous ...
mxnet::Storage::_GetSharedRef
static std::shared_ptr< Storage > _GetSharedRef()
Get shared pointer reference to storage singleton. Most user should not call this function...
mxnet
namespace of mxnet
Definition:
api_registry.h:33
mxnet::Storage
Storage manager across multiple devices.
Definition:
storage.h:35
mxnet::Storage::~Storage
virtual ~Storage()
Destructor.
Definition:
storage.h:108
mxnet::Storage::Free
virtual void Free(Handle handle)=0
Free storage.
mxnet::Storage::Handle::dptr
void * dptr
Pointer to the data.
Definition:
storage.h:44
mxnet::Storage::Handle::ctx
Context ctx
Context information about device and ID.
Definition:
storage.h:52
mxnet::Storage::GetMutex
std::mutex & GetMutex(Context::DeviceType dev)
Returns mutex used by storage manager.
Definition:
storage.h:112
mxnet::Storage::Handle::shared_id
int shared_id
Definition:
storage.h:57
mxnet::Context::kCPU
Definition:
base.h:104
mxnet::Storage::Get
static Storage * Get()
mxnet::Context::DeviceType
DeviceType
Type of device.
Definition:
base.h:103
mxnet::Storage::Handle
Storage handle.
Definition:
storage.h:40
mxnet::Storage::Alloc
Handle Alloc(size_t size, Context ctx)
Allocate a new contiguous memory for a given size.
Definition:
storage.h:65
mxnet::Storage::SharedIncrementRefCount
virtual void SharedIncrementRefCount(Handle handle)=0
Increase ref counter on shared memory.
base.h
mxnet::Storage::Handle::size
size_t size
Size of the storage.
Definition:
storage.h:48
mxnet::Context
Context information about the execution environment.
Definition:
base.h:101
mxnet::Storage::Handle::shared_pid
int shared_pid
Id for IPC shared memory.
Definition:
storage.h:56
mxnet::Storage::ReleaseAll
virtual void ReleaseAll(Context ctx)=0
Release all memory from device if using a pooled storage manager.
Generated on Thu Jan 5 2023 00:58:42 for mxnet by
1.8.13