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;
101  virtual ~Storage() {}
105  std::mutex& GetMutex(Context::DeviceType dev) {
106  if (dev == Context::kCPU) {
107  return cpu_mutex_;
108  } else {
109  return gpu_mutex_;
110  }
111  }
115  static Storage* Get();
124  static std::shared_ptr<Storage> _GetSharedRef();
125 
126  private:
127  std::mutex cpu_mutex_;
128  std::mutex gpu_mutex_;
129 }; // class Storage
130 } // namespace mxnet
131 #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:127
Storage manager across multiple devices.
Definition: storage.h:36
virtual ~Storage()
Destructor.
Definition: storage.h:101
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:105
int shared_id
Definition: storage.h:58
Definition: base.h:145
static Storage * Get()
DeviceType
Type of device.
Definition: base.h:144
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:142
int shared_pid
Id for IPC shared memory.
Definition: storage.h:57