[L0 v2] make device allocation resident

on all devices in the context to match the default
behavior of legacy adapter.

Author: igchor

source/adapters/level_zero/v2/context.cpp

@@ -13,12 +13,41 @@
 #include "context.hpp"
 #include "event_provider_normal.hpp"
 
+static std::vector<ur_device_handle_t>
+filterP2PDevices(ur_device_handle_t hSourceDevice,
+                 const std::vector<ur_device_handle_t> &devices) {
+  std::vector<ur_device_handle_t> p2pDevices;
+  for (auto &device : devices) {
+    if (device == hSourceDevice) {
+      continue;
+    }
+
+    ze_bool_t p2p;
+    ZE2UR_CALL_THROWS(zeDeviceCanAccessPeer,
+                      (hSourceDevice->ZeDevice, device->ZeDevice, &p2p));
+
+    if (p2p) {
+      p2pDevices.push_back(device);
+    }
+  }
+  return p2pDevices;
+}
+
+static std::vector<std::vector<ur_device_handle_t>>
+populateP2PDevices(size_t maxDevices,
+                   const std::vector<ur_device_handle_t> &devices) {
+  std::vector<std::vector<ur_device_handle_t>> p2pDevices(maxDevices);
+  for (auto &device : devices) {
+    p2pDevices[device->Id.value()] = filterP2PDevices(device, devices);
+  }
+  return p2pDevices;
+}
+
 ur_context_handle_t_::ur_context_handle_t_(ze_context_handle_t hContext,
                                            uint32_t numDevices,
                                            const ur_device_handle_t *phDevices,
                                            bool ownZeContext)
-    : hContext(hContext, ownZeContext),
-      hDevices(phDevices, phDevices + numDevices), commandListCache(hContext),
+    : commandListCache(hContext),
       eventPoolCache(phDevices[0]->Platform->getNumDevices(),
                      [context = this,
                       platform = phDevices[0]->Platform](DeviceId deviceId) {
@@ -28,6 +57,10 @@ ur_context_handle_t_::ur_context_handle_t_(ze_context_handle_t hContext,
                            context, device, v2::EVENT_COUNTER,
                            v2::QUEUE_IMMEDIATE);
                      }),
+      hContext(hContext, ownZeContext),
+      hDevices(phDevices, phDevices + numDevices),
+      p2pAccessDevices(populateP2PDevices(
+          phDevices[0]->Platform->getNumDevices(), this->hDevices)),
       defaultUSMPool(this, nullptr) {}
 
 ur_result_t ur_context_handle_t_::retain() {
@@ -65,6 +98,11 @@ ur_usm_pool_handle_t ur_context_handle_t_::getDefaultUSMPool() {
   return &defaultUSMPool;
 }
 
+const std::vector<ur_device_handle_t> &
+ur_context_handle_t_::getP2PDevices(ur_device_handle_t hDevice) const {
+  return p2pAccessDevices[hDevice->Id.value()];
+}
+
 namespace ur::level_zero {
 ur_result_t urContextCreate(uint32_t deviceCount,
                             const ur_device_handle_t *phDevices,

source/adapters/level_zero/v2/context.hpp

@@ -28,14 +28,22 @@ struct ur_context_handle_t_ : _ur_object {
   ur_platform_handle_t getPlatform() const;
   const std::vector<ur_device_handle_t> &getDevices() const;
   ur_usm_pool_handle_t getDefaultUSMPool();
+  const std::vector<ur_device_handle_t> &
+  getP2PDevices(ur_device_handle_t hDevice) const;
 
   // Checks if Device is covered by this context.
   // For that the Device or its root devices need to be in the context.
   bool isValidDevice(ur_device_handle_t Device) const;
 
-  const v2::raii::ze_context_handle_t hContext;
-  const std::vector<ur_device_handle_t> hDevices;
   v2::command_list_cache_t commandListCache;
   v2::event_pool_cache eventPoolCache;
+
+private:
+  const v2::raii::ze_context_handle_t hContext;
+  const std::vector<ur_device_handle_t> hDevices;
+
+  // P2P devices for each device in the context, indexed by device id.
+  const std::vector<std::vector<ur_device_handle_t>> p2pAccessDevices;
+
   ur_usm_pool_handle_t_ defaultUSMPool;
 };

source/adapters/level_zero/v2/usm.cpp

@@ -85,8 +85,22 @@ makePool(umf_disjoint_pool_params_t *poolParams,
   params.level_zero_device_handle =
       poolDescriptor.hDevice ? poolDescriptor.hDevice->ZeDevice : nullptr;
   params.memory_type = urToUmfMemoryType(poolDescriptor.type);
-  // TODO: handle memory residency:
-  // set resident_device_handles and resident_device_count
+
+  std::vector<ze_device_handle_t> residentZeHandles;
+
+  if (poolDescriptor.type == UR_USM_TYPE_DEVICE) {
+    assert(params.level_zero_device_handle);
+    auto residentHandles =
+        poolDescriptor.hContext->getP2PDevices(poolDescriptor.hDevice);
+    residentZeHandles.push_back(params.level_zero_device_handle);
+    for (auto &device : residentHandles) {
+      residentZeHandles.push_back(device->ZeDevice);
+    }
+
+    params.resident_device_handles = residentZeHandles.data();
+    params.resident_device_count = residentZeHandles.size();
+  }
+
   auto [ret, provider] =
       umf::providerMakeUniqueFromOps(umfLevelZeroMemoryProviderOps(), &params);
   if (ret != UMF_RESULT_SUCCESS) {

test/adapters/level_zero/v2/CMakeLists.txt

@@ -45,3 +45,12 @@ add_unittest(level_zero_event_pool
         ${PROJECT_SOURCE_DIR}/source/adapters/level_zero/v2/event_provider_counter.cpp
         ${PROJECT_SOURCE_DIR}/source/adapters/level_zero/v2/event.cpp
 )
+
+add_adapter_test(level_zero_memory_residency
+    FIXTURE DEVICES
+    SOURCES
+        memory_residency.cpp
+    ENVIRONMENT
+        "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_level_zero_v2>\""
+        "ZES_ENABLE_SYSMAN=1"
+)

test/adapters/level_zero/v2/memory_residency.cpp

@@ -0,0 +1,48 @@
+// Copyright (C) 2024 Intel Corporation
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+// See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "ur_print.hpp"
+#include "uur/fixtures.h"
+#include "uur/raii.h"
+#include "uur/utils.h"
+
+#include <map>
+#include <string>
+
+using urMemoryResidencyTest = uur::urMultiDeviceContextTestTemplate<1>;
+
+TEST_F(urMemoryResidencyTest, allocatingDeviceMemoryWillResultInOOM) {
+    static constexpr size_t allocSize = 1024 * 1024;
+
+    if (!uur::isPVC(uur::DevicesEnvironment::instance->devices[0])) {
+        GTEST_SKIP() << "Test requires a PVC device";
+    }
+
+    size_t initialMemFree = 0;
+    ASSERT_SUCCESS(
+        urDeviceGetInfo(uur::DevicesEnvironment::instance->devices[0],
+                        UR_DEVICE_INFO_GLOBAL_MEM_FREE, sizeof(size_t),
+                        &initialMemFree, nullptr));
+
+    if (initialMemFree < allocSize) {
+        GTEST_SKIP() << "Not enough device memory available";
+    }
+
+    void *ptr = nullptr;
+    ASSERT_SUCCESS(
+        urUSMDeviceAlloc(context, uur::DevicesEnvironment::instance->devices[0],
+                         nullptr, nullptr, allocSize, &ptr));
+
+    size_t currentMemFree = 0;
+    ASSERT_SUCCESS(
+        urDeviceGetInfo(uur::DevicesEnvironment::instance->devices[0],
+                        UR_DEVICE_INFO_GLOBAL_MEM_FREE, sizeof(size_t),
+                        &currentMemFree, nullptr));
+
+    // amount of free memory should decrease after making a memory allocation resident
+    ASSERT_LE(currentMemFree, initialMemFree);
+
+    ASSERT_SUCCESS(urUSMFree(context, ptr));
+}

test/conformance/testing/include/uur/utils.h

@@ -483,6 +483,18 @@ getDriverVersion(ur_device_handle_t hDevice) {
         }                                                                      \
     } while (0)
 
+// Is this a Data Center GPU Max series (aka PVC)?
+// TODO: change to use
+// https://spec.oneapi.io/level-zero/latest/core/api.html#ze-device-ip-version-ext-t
+// when that is stable.
+static inline bool isPVC(ur_device_handle_t hDevice) {
+    uint32_t deviceId;
+    EXPECT_EQ(urDeviceGetInfo(hDevice, UR_DEVICE_INFO_DEVICE_ID,
+                              sizeof(uint32_t), &deviceId, nullptr),
+              UR_RESULT_SUCCESS);
+    return (deviceId & 0xff0) == 0xbd0 || (deviceId & 0xff0) == 0xb60;
+}
+
 } // namespace uur
 
 #endif // UR_CONFORMANCE_INCLUDE_UTILS_H_INCLUDED