[llvm] [offload] add context to olMemAlloc* (PR #224930)

[Kazu Hirata via llvm-commits]

=?utf-8?q?Łukasz?= Plewa <lukasz.plewa at intel.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/224930 at github.com>


================
@@ -273,6 +283,91 @@ Error LevelZeroPluginContextTy::deinit() {
   return Plugin::success();
 }
 
+Error LevelZeroPluginContextTy::initAllocators() {
+  const auto &Options = static_cast<LevelZeroPluginTy &>(Plugin).getOptions();
+  for (auto *D : Devices) {
+    auto &L0Device = static_cast<L0DeviceTy &>(*D);
+    auto Alloc = std::make_unique<MemAllocatorTy>();
+    if (auto Err = Alloc->initDevicePools(L0Device, Options, ZeContext))
+      return Err;
+    DeviceAllocators.try_emplace(&L0Device, std::move(Alloc));
+  }
+  if (Devices.empty())
+    return Plugin::success();
+  auto &First = static_cast<L0DeviceTy &>(*Devices.front());
+  HostAllocator = std::make_unique<MemAllocatorTy>();
+  if (auto Err =
+          HostAllocator->initHostPool(First.getL0Context(), Options, ZeContext))
+    return Err;
+  // Host MaxAllocSize = min over devices, matching L0ContextTy's driver pool.
+  for (auto *D : Devices)
+    HostAllocator->updateMaxAllocSize(static_cast<L0DeviceTy &>(*D));
+  return Plugin::success();
+}
+
+Expected<void *> LevelZeroPluginContextTy::allocate(GenericDeviceTy &Device,
+                                                    int64_t Size,
+                                                    void * /*HostPtr*/,
+                                                    TargetAllocTy Kind,
+                                                    size_t Alignment) {
+  MemAllocatorTy *Allocator = nullptr;
+  int32_t ResolvedKind = Kind;
+  if (Kind == TARGET_ALLOC_HOST) {
+    if (!HostAllocator)
+      return Plugin::error(ErrorCode::INVALID_ARGUMENT,
+                           "host allocator not initialized");
+    Allocator = HostAllocator.get();
+  } else {
+    if (ResolvedKind == TARGET_ALLOC_DEFAULT)
+      ResolvedKind = TARGET_ALLOC_DEVICE;
+    auto &L0Device = static_cast<L0DeviceTy &>(Device);
+    auto It = DeviceAllocators.find(&L0Device);
+    if (It == DeviceAllocators.end())
+      return Plugin::error(ErrorCode::INVALID_DEVICE,
+                           "device is not part of this context");
+    Allocator = It->second.get();
+  }
+  return Allocator->alloc(Size, Alignment, ResolvedKind, /*Offset=*/0,
+                          /*UserAlloc=*/true, /*DevMalloc=*/false,
+                          /*MemAdvice=*/
+                          std::numeric_limits<uint32_t>::max(),
+                          AllocOptionTy::ALLOC_OPT_NONE);
+}
+
+Error LevelZeroPluginContextTy::deallocate(GenericDeviceTy &Device, void *Ptr,
+                                           TargetAllocTy Kind) {
+  if (Kind == TARGET_ALLOC_HOST) {
+    if (!HostAllocator)
+      return Plugin::error(ErrorCode::NOT_FOUND,
+                           "no host allocation tracked in this context");
+    return HostAllocator->dealloc(Ptr);
+  }
+  auto &L0Device = static_cast<L0DeviceTy &>(Device);
+  auto It = DeviceAllocators.find(&L0Device);
+  if (It == DeviceAllocators.end())
+    return Plugin::error(ErrorCode::INVALID_DEVICE,
+                         "device is not part of this context");
+  return It->second->dealloc(Ptr);
+}
+
+Expected<PluginAllocInfoTy>
+LevelZeroPluginContextTy::getAllocInfo(const void *Ptr) {
+  void *Raw = const_cast<void *>(Ptr);
+  for (auto &KV : DeviceAllocators) {
+    if (auto *Info = KV.second->getAllocInfo(Raw))
+      return PluginAllocInfoTy{KV.first, static_cast<TargetAllocTy>(Info->Kind),
----------------
kazutakahirata wrote:

May I suggest structured bindings here?
```suggestion
  for (const auto &[Device, Allocator] : DeviceAllocators) {
    if (auto *Info = Allocator->getAllocInfo(Raw))
      return PluginAllocInfoTy{Device, static_cast<TargetAllocTy>(Info->Kind),
```

https://github.com/llvm/llvm-project/pull/224930