From 8faa4b8d05ff929d274eb713908df964ceeac751 Mon Sep 17 00:00:00 2001
From: Sergey V Maslov <sergey.v.maslov@intel.com>
Date: Wed, 11 Aug 2021 14:17:12 -0700
Subject: [PATCH] [SYCL] create memory allocators for contexts with single
 root-device

Signed-off-by: Sergey V Maslov <sergey.v.maslov@intel.com>
---
 sycl/plugins/level_zero/pi_level_zero.hpp | 54 ++++++++++++++---------
 1 file changed, 34 insertions(+), 20 deletions(-)

diff --git a/sycl/plugins/level_zero/pi_level_zero.hpp b/sycl/plugins/level_zero/pi_level_zero.hpp
index 2de84223529ae..b341ca382ec36 100644
--- a/sycl/plugins/level_zero/pi_level_zero.hpp
+++ b/sycl/plugins/level_zero/pi_level_zero.hpp
@@ -384,6 +384,9 @@ struct _pi_context : _pi_object {
         OwnZeContext{OwnZeContext}, Devices{Devs, Devs + NumDevices},
         ZeCommandListInit{nullptr}, ZeEventPool{nullptr},
         NumEventsAvailableInEventPool{}, NumEventsUnreleasedInEventPool{} {
+    // NOTE: one must additionally call initialize() to complete
+    // PI context creation.
+
     // Create USM allocator context for each pair (device, context).
     for (uint32_t I = 0; I < NumDevices; I++) {
       pi_device Device = Devs[I];
@@ -395,8 +398,6 @@ struct _pi_context : _pi_object {
           std::piecewise_construct, std::make_tuple(Device),
           std::make_tuple(std::unique_ptr<SystemMemory>(
               new USMDeviceMemoryAlloc(this, Device))));
-      // NOTE: one must additionally call initialize() to complete
-      // PI context creation.
     }
     // Create USM allocator context for host. Device and Shared USM allocations
     // are device-specific. Host allocations are not device-dependent therefore
@@ -406,27 +407,40 @@ struct _pi_context : _pi_object {
 
     if (NumDevices == 1) {
       SingleRootDevice = Devices[0];
-      return;
-    }
-
-    // Check if we have context with subdevices of the same device (context may
-    // include root device itself as well)
-    SingleRootDevice =
-        Devices[0]->RootDevice ? Devices[0]->RootDevice : Devices[0];
+    } else {
+
+      // Check if we have context with subdevices of the same device (context
+      // may include root device itself as well)
+      SingleRootDevice =
+          Devices[0]->RootDevice ? Devices[0]->RootDevice : Devices[0];
+
+      // For context with sub subdevices, the SingleRootDevice might still
+      // not be the root device.
+      // Check whether the SingleRootDevice is the subdevice or root device.
+      if (SingleRootDevice->isSubDevice()) {
+        SingleRootDevice = SingleRootDevice->RootDevice;
+      }
 
-    // For context with sub subdevices, the SingleRootDevice might still
-    // not be the root device.
-    // Check whether the SingleRootDevice is the subdevice or root device.
-    if (SingleRootDevice->isSubDevice()) {
-      SingleRootDevice = SingleRootDevice->RootDevice;
+      for (auto &Device : Devices) {
+        if ((!Device->RootDevice && Device != SingleRootDevice) ||
+            (Device->RootDevice && Device->RootDevice != SingleRootDevice)) {
+          SingleRootDevice = nullptr;
+          break;
+        }
+      }
     }
 
-    for (auto &Device : Devices) {
-      if ((!Device->RootDevice && Device != SingleRootDevice) ||
-          (Device->RootDevice && Device->RootDevice != SingleRootDevice)) {
-        SingleRootDevice = nullptr;
-        break;
-      }
+    // We may allocate memory to this root device so create allocators.
+    if (SingleRootDevice && DeviceMemAllocContexts.find(SingleRootDevice) ==
+                                DeviceMemAllocContexts.end()) {
+      SharedMemAllocContexts.emplace(
+          std::piecewise_construct, std::make_tuple(SingleRootDevice),
+          std::make_tuple(std::unique_ptr<SystemMemory>(
+              new USMSharedMemoryAlloc(this, SingleRootDevice))));
+      DeviceMemAllocContexts.emplace(
+          std::piecewise_construct, std::make_tuple(SingleRootDevice),
+          std::make_tuple(std::unique_ptr<SystemMemory>(
+              new USMDeviceMemoryAlloc(this, SingleRootDevice))));
     }
   }