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 
25 #ifndef MXNET_STORAGE_H_
26 #define MXNET_STORAGE_H_
27 
28 #include <memory>
29 #include "./base.h"
30 
31 namespace mxnet {
32 
36 class Storage {
37  public:
41  struct Handle {
45  void* dptr{nullptr};
49  size_t size{0};
57  int shared_pid{-1};
58  int shared_id{-1};
59  };
67  Handle hd;
68  hd.size = size;
69  hd.ctx = ctx;
70  this->Alloc(&hd);
71  return hd;
72  }
77  virtual void Alloc(Handle* handle) = 0;
82  virtual void SharedIncrementRefCount(Handle handle) = 0;
87  virtual void Free(Handle handle) = 0;
97  virtual void DirectFree(Handle handle) = 0;
105  virtual void ReleaseAll(Context ctx) = 0;
109  virtual ~Storage() {}
113  std::mutex& GetMutex(Context::DeviceType dev) {
114  if (dev == Context::kCPU) {
115  return cpu_mutex_;
116  } else {
117  return gpu_mutex_;
118  }
119  }
123  static Storage* Get();
132  static std::shared_ptr<Storage> _GetSharedRef();
133 
134  private:
135  std::mutex cpu_mutex_;
136  std::mutex gpu_mutex_;
137 }; // class Storage
138 } // namespace mxnet
139 #endif // MXNET_STORAGE_H_
virtual void DirectFree(Handle handle)=0
Free storage directly, without putting it into memory pool. This can synchronization of all previous ...
static std::shared_ptr< Storage > _GetSharedRef()
Get shared pointer reference to storage singleton. Most user should not call this function...
namespace of mxnet
Definition: base.h:89
Storage manager across multiple devices.
Definition: storage.h:36
virtual ~Storage()
Destructor.
Definition: storage.h:109
virtual void Free(Handle handle)=0
Free storage.
void * dptr
Pointer to the data.
Definition: storage.h:45
Context ctx
Context information about device and ID.
Definition: storage.h:53
std::mutex & GetMutex(Context::DeviceType dev)
Returns mutex used by storage manager.
Definition: storage.h:113
int shared_id
Definition: storage.h:58
Definition: base.h:105
static Storage * Get()
DeviceType
Type of device.
Definition: base.h:104
Storage handle.
Definition: storage.h:41
Handle Alloc(size_t size, Context ctx)
Allocate a new contiguous memory for a given size.
Definition: storage.h:66
virtual void SharedIncrementRefCount(Handle handle)=0
Increase ref counter on shared memory.
size_t size
Size of the storage.
Definition: storage.h:49
Context information about the execution environment.
Definition: base.h:102
int shared_pid
Id for IPC shared memory.
Definition: storage.h:57
virtual void ReleaseAll(Context ctx)=0
Release all memory from device if using a pooled storage manager.