Move DeviceRange to PluginInterface and remove internal Device list

Author: adurang

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -1266,7 +1266,7 @@ struct GenericPluginTy {
   virtual GenericGlobalHandlerTy *createGlobalHandler() = 0;
 
   /// Get the reference to the device with a certain device id.
-  GenericDeviceTy &getDevice(int32_t DeviceId) {
+  GenericDeviceTy &getDevice(int32_t DeviceId) const {
     assert(isValidDeviceId(DeviceId) && "Invalid device id");
     assert(Devices[DeviceId] && "Device is uninitialized");
 
@@ -1527,6 +1527,23 @@ struct GenericPluginTy {
   /// object and return immediately.
   int32_t async_barrier(omp_interop_val_t *Interop);
 
+  struct DevicesRangeTy {
+    using iterator = llvm::SmallVector<GenericDeviceTy *>::iterator;
+
+    iterator BeginIt;
+    iterator EndIt;
+
+    DevicesRangeTy(iterator BeginIt, iterator EndIt)
+        : BeginIt(BeginIt), EndIt(EndIt) {}
+
+    auto &begin() { return BeginIt; }
+    auto &end() { return EndIt; }
+  };
+
+  DevicesRangeTy getDevicesRange() {
+    return DevicesRangeTy(Devices.begin(), Devices.end());
+  }
+
 private:
   /// Indicates if the platform runtime has been fully initialized.
   bool Initialized = false;

offload/plugins-nextgen/level_zero/include/L0Plugin.h

@@ -26,16 +26,16 @@ namespace llvm::omp::target::plugin {
 /// Class implementing the LevelZero specific functionalities of the plugin.
 class LevelZeroPluginTy final : public GenericPluginTy {
 private:
-  /// Number of devices available including subdevices
-  uint32_t NumDevices = 0;
+  struct DeviceInfoTy {
+    L0DeviceIdTy Id;
+    L0ContextTy *Driver;
+    bool isRoot() const { return Id.SubId < 0 && Id.CCSId < 0; }
+  };
+  llvm::SmallVector<DeviceInfoTy> DetectedDevices;
 
   /// Context (and Driver) specific data
   std::list<L0ContextTy> ContextList;
 
-  /// L0 device used by each OpenMP device
-  using DeviceContainerTy = llvm::SmallVector<L0DeviceTy *>;
-  DeviceContainerTy L0Devices;
-
   // Table containing per-thread information using TLS
   L0ThreadTblTy ThreadTLSTable;
   // Table containing per-thread information for each device using TLS
@@ -51,6 +51,10 @@ class LevelZeroPluginTy final : public GenericPluginTy {
 
   auto &getTLS() { return ThreadTLSTable.get(); }
 
+  /// Find L0 devices and initialize device properties.
+  /// Returns number of devices reported to omptarget.
+  Expected<int32_t> findDevices();
+
 public:
   LevelZeroPluginTy() : GenericPluginTy(getTripleArch()) {}
   virtual ~LevelZeroPluginTy() {}
@@ -62,35 +66,10 @@ class LevelZeroPluginTy final : public GenericPluginTy {
 
   static const auto &getOptions() { return Options; }
 
-  struct DevicesRangeTy {
-    using iterator = DeviceContainerTy::iterator;
-
-    iterator BeginIt;
-    iterator EndIt;
-
-    DevicesRangeTy(iterator BeginIt, iterator EndIt)
-        : BeginIt(BeginIt), EndIt(EndIt) {}
-
-    auto &begin() { return BeginIt; }
-    auto &end() { return EndIt; }
-  };
-
-  auto getDevicesRange() {
-    return DevicesRangeTy(L0Devices.begin(), L0Devices.end());
-  }
-
-  /// Find L0 devices and initialize device properties.
-  /// Returns number of devices reported to omptarget.
-  Expected<int32_t> findDevices();
-
   L0DeviceTy &getDeviceFromId(int32_t DeviceId) const {
-    assert("Invalid device ID" && DeviceId >= 0 &&
-           DeviceId < static_cast<int32_t>(L0Devices.size()));
-    return *L0Devices[DeviceId];
+    return static_cast<L0DeviceTy &>(getDevice(DeviceId));
   }
 
-  uint32_t getNumRootDevices() const { return NumDevices; }
-
   AsyncQueueTy *getAsyncQueue() {
     auto *Queue = getTLS().getAsyncQueue();
     if (!Queue)

offload/plugins-nextgen/level_zero/src/L0Plugin.cpp

@@ -86,57 +86,29 @@ Expected<int32_t> LevelZeroPluginTy::findDevices() {
               return A.IsDiscrete;
             });
 
-  struct DeviceInfoTy {
-    L0DeviceIdTy Id;
-    L0ContextTy *Driver;
-    bool isRoot() const { return Id.SubId < 0 && Id.CCSId < 0; }
-  };
-
-  llvm::SmallVector<DeviceInfoTy> DevicesToAdd;
-
   for (size_t RootId = 0; RootId < RootDevices.size(); RootId++) {
     const auto zeDevice = RootDevices[RootId].zeDevice;
     auto *RootDriver = RootDevices[RootId].Driver;
-    DevicesToAdd.push_back(
-        {{zeDevice, static_cast<int32_t>(RootId), -1, -1}, RootDriver});
-  }
-  NumDevices = DevicesToAdd.size();
-  auto DeviceId = 0;
-  for (auto &DeviceInfo : DevicesToAdd) {
-    auto RootId = DeviceInfo.Id.RootId;
-    auto SubId = DeviceInfo.Id.SubId;
-    auto CCSId = DeviceInfo.Id.CCSId;
-    auto zeDevice = DeviceInfo.Id.zeId;
-    auto *Driver = DeviceInfo.Driver;
-
-    std::string IdStr = std::to_string(RootId) +
-                        (SubId < 0 ? "" : "." + std::to_string(SubId)) +
-                        (CCSId < 0 ? "" : "." + std::to_string(CCSId));
-
-    L0Devices.push_back(new L0DeviceTy(*this, DeviceId, getNumRootDevices(),
-                                       zeDevice, *Driver, std::move(IdStr),
-                                       CCSId < 0 ? 0 : CCSId /* ComputeIndex */
-                                       ));
-    DeviceId++;
+    DetectedDevices.push_back(DeviceInfoTy{
+        {zeDevice, static_cast<int32_t>(RootId), -1, -1}, RootDriver});
   }
+  int32_t NumDevices = DetectedDevices.size();
 
-  DP("Found %" PRIu32 " root devices, %" PRIu32 " total devices.\n",
-     getNumRootDevices(), NumDevices);
+  DP("Found %" PRIu32 " devices.\n", NumDevices);
   DP("List of devices (DeviceID[.SubID[.CCSID]])\n");
-  for (auto &l0Device : L0Devices) {
-    DP("-- %s\n", l0Device->getZeIdCStr());
-    (void)l0Device; // silence warning
+  for (auto &DeviceInfo : DetectedDevices) {
+    (void)DeviceInfo; // to avoid unused variable warning in non-debug builds
+    DP("-- Device %" PRIu32 "%s%s (zeDevice=%p) from Driver %p\n",
+       DeviceInfo.Id.RootId,
+       (DeviceInfo.Id.SubId < 0
+            ? ""
+            : ("." + std::to_string(DeviceInfo.Id.SubId)).c_str()),
+       (DeviceInfo.Id.CCSId < 0
+            ? ""
+            : ("." + std::to_string(DeviceInfo.Id.CCSId)).c_str()),
+       DPxPTR(DeviceInfo.Id.zeId), DPxPTR(DeviceInfo.Id.Driver));
   }
-
-  if (getDebugLevel() > 0) {
-    DP("Root Device Information\n");
-    for (uint32_t I = 0; I < getNumRootDevices(); I++) {
-      auto &l0Device = getDeviceFromId(I);
-      l0Device.reportDeviceInfo();
-    }
-  }
-
-  return getNumRootDevices();
+  return NumDevices;
 }
 
 Expected<int32_t> LevelZeroPluginTy::initImpl() {
@@ -163,7 +135,25 @@ Error LevelZeroPluginTy::deinitImpl() {
 GenericDeviceTy *LevelZeroPluginTy::createDevice(GenericPluginTy &Plugin,
                                                  int32_t DeviceId,
                                                  int32_t NumDevices) {
-  return &getDeviceFromId(DeviceId);
+  auto &DeviceInfo = DetectedDevices[DeviceId];
+  auto RootId = DeviceInfo.Id.RootId;
+  auto SubId = DeviceInfo.Id.SubId;
+  auto CCSId = DeviceInfo.Id.CCSId;
+  auto zeDevice = DeviceInfo.Id.zeId;
+  auto *zeDriver = DeviceInfo.Driver;
+
+  std::string IdStr = std::to_string(RootId) +
+                      (SubId < 0 ? "" : "." + std::to_string(SubId)) +
+                      (CCSId < 0 ? "" : "." + std::to_string(CCSId));
+
+  auto *NewDevice = new L0DeviceTy(
+      static_cast<LevelZeroPluginTy &>(Plugin), DeviceId, NumDevices, zeDevice,
+      *zeDriver, std::move(IdStr), CCSId < 0 ? 0 : CCSId /* ComputeIndex */);
+  if (NewDevice && getDebugLevel() > 0) {
+    DP("Device %" PRIi32 " information\n", DeviceId);
+    NewDevice->reportDeviceInfo();
+  }
+  return NewDevice;
 }
 
 GenericGlobalHandlerTy *LevelZeroPluginTy::createGlobalHandler() {