OmniSciDB  a5dc49c757
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
DeviceMemoryAllocationMap.h
Go to the documentation of this file.
1 /*
2  * Copyright 2023 HEAVY.AI, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <functional>
20 #include <map>
21 #include <memory>
22 #include <unordered_map>
23 
24 #include "Shared/uuid.h"
25 
26 namespace CudaMgr_Namespace {
27 
29  public:
30  using DevicePtr = uint64_t;
31  struct Allocation {
32  uint64_t size = 0u;
33  uint64_t handle = 0u;
35  int device_num = -1;
36  bool is_slab = false;
37  };
38  using Map = std::map<DevicePtr, Allocation>;
39  using MapChangedCBID = uint32_t;
40  using MapChangedCB = std::function<void(const heavyai::UUID, const bool)>;
41 
43  ~DeviceMemoryAllocationMap() = default;
44 
45  const Map& getMap() const;
46  const bool mapEmpty() const;
47 
48  void addAllocation(const DevicePtr device_ptr,
49  const uint64_t size,
50  const uint64_t handle,
51  const heavyai::UUID uuid,
52  const int device_num,
53  const bool is_slab);
54  Allocation removeAllocation(const DevicePtr device_ptr);
55  std::pair<DevicePtr, Allocation> getAllocation(const DevicePtr device_ptr);
56 
58  void unregisterMapChangedCB(const MapChangedCBID cbid);
59  void notifyMapChanged(const heavyai::UUID device_uuid, const bool is_slab) const;
60 
61  private:
64  std::unordered_map<MapChangedCBID, MapChangedCB> map_changed_cbs_;
65 };
66 
67 using DeviceMemoryAllocationMapUqPtr = std::unique_ptr<DeviceMemoryAllocationMap>;
68 
69 } // namespace CudaMgr_Namespace
std::unordered_map< MapChangedCBID, MapChangedCB > map_changed_cbs_
void notifyMapChanged(const heavyai::UUID device_uuid, const bool is_slab) const
Allocation removeAllocation(const DevicePtr device_ptr)
No-frills UUID type class to allow easy containerization and comparison of device UUIDs from differen...
const MapChangedCBID registerMapChangedCB(MapChangedCB cb)
std::function< void(const heavyai::UUID, const bool)> MapChangedCB
void addAllocation(const DevicePtr device_ptr, const uint64_t size, const uint64_t handle, const heavyai::UUID uuid, const int device_num, const bool is_slab)
std::pair< DevicePtr, Allocation > getAllocation(const DevicePtr device_ptr)
std::unique_ptr< DeviceMemoryAllocationMap > DeviceMemoryAllocationMapUqPtr