File tree Expand file tree Collapse file tree 2 files changed +23
-9
lines changed
source/adapters/level_zero Expand file tree Collapse file tree 2 files changed +23
-9
lines changed Original file line number Diff line number Diff line change @@ -100,6 +100,9 @@ struct ur_context_handle_t_ : _ur_object {
100
100
l0_command_list_cache_info>>>
101
101
ZeCopyCommandListCache;
102
102
103
+ std::unordered_map<ur_device_handle_t , std::list<ur_device_handle_t >>
104
+ P2PDeviceCache;
105
+
103
106
// Store USM pool for USM shared and device allocations. There is 1 memory
104
107
// pool per each pair of (context, device) per each memory type.
105
108
std::unordered_map<ze_device_handle_t , umf::pool_unique_handle_t >
Original file line number Diff line number Diff line change @@ -154,15 +154,26 @@ static ur_result_t USMAllocationMakeResident(
154
154
} else {
155
155
Devices.push_back (Device);
156
156
if (ForceResidency == USMAllocationForceResidencyType::P2PDevices) {
157
- ze_bool_t P2P;
158
- for (const auto &D : Context->Devices ) {
159
- if (D == Device)
160
- continue ;
161
- // TODO: Cache P2P devices for a context
162
- ZE2UR_CALL (zeDeviceCanAccessPeer,
163
- (D->ZeDevice , Device->ZeDevice , &P2P));
164
- if (P2P)
165
- Devices.push_back (D);
157
+ // Check if the P2P devices are already cached
158
+ auto it = Context->P2PDeviceCache .find (Device);
159
+ if (it != Context->P2PDeviceCache .end ()) {
160
+ // Use cached P2P devices
161
+ Devices.insert (Devices.end (), it->second .begin (), it->second .end ());
162
+ } else {
163
+ // Query for P2P devices and update the cache
164
+ std::list<ur_device_handle_t > P2PDevices;
165
+ ze_bool_t P2P;
166
+ for (const auto &D : Context->Devices ) {
167
+ if (D == Device)
168
+ continue ;
169
+ ZE2UR_CALL (zeDeviceCanAccessPeer,
170
+ (D->ZeDevice , Device->ZeDevice , &P2P));
171
+ if (P2P)
172
+ P2PDevices.push_back (D);
173
+ }
174
+ // Update the cache
175
+ Context->P2PDeviceCache [Device] = P2PDevices;
176
+ Devices.insert (Devices.end (), P2PDevices.begin (), P2PDevices.end ());
166
177
}
167
178
}
168
179
}
You can’t perform that action at this time.
0 commit comments