//===--- Target RTLs Implementation ---------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// RTL for SPIR-V/Xe machine.
//
//===----------------------------------------------------------------------===//

#include <level_zero/zes_api.h>

#include "L0Device.h"
#include "L0Interop.h"
#include "L0Kernel.h"
#include "L0Plugin.h"
#include "L0Trace.h"

#include "llvm/Object/OffloadBinary.h"

namespace llvm::omp::target::plugin {

using namespace llvm::omp::target;
using namespace error;

Expected<int32_t> LevelZeroPluginTy::findDevices() {
  ze_result_t RC;
  CALL_ZE(RC, zeInit, ZE_INIT_FLAG_GPU_ONLY);
  if (RC != ZE_RESULT_SUCCESS) {
    ODBG(OLDT_Init) << "Cannot initialize Level Zero library. Error: "
                    << getZeErrorName(RC);
    return 0;
  }

  uint32_t NumDrivers = 0;
  CALL_ZE(RC, zeDriverGet, &NumDrivers, nullptr);
  if (RC != ZE_RESULT_SUCCESS) {
    ODBG(OLDT_Init) << "Error getting number of drivers. Error: "
                    << getZeErrorName(RC);
    return 0;
  }
  if (NumDrivers == 0) {
    ODBG(OLDT_Init) << "Cannot find any drivers.";
    return 0;
  }

  // We expect multiple drivers on Windows to support different device types,
  // so we need to maintain multiple drivers and contexts in general.
  llvm::SmallVector<ze_driver_handle_t> FoundDrivers(NumDrivers);
  CALL_ZE_RET_ERROR(zeDriverGet, &NumDrivers, FoundDrivers.data());

  struct RootInfoTy {
    uint32_t OrderId;
    ze_device_handle_t ZeDevice;
    L0ContextTy *Driver;
    bool IsDiscrete;
  };
  llvm::SmallVector<RootInfoTy> RootDevices;

  uint32_t OrderId = 0;
  for (uint32_t DriverId = 0; DriverId < NumDrivers; DriverId++) {
    const auto &Driver = FoundDrivers[DriverId];
    uint32_t DeviceCount = 0;
    ze_result_t RC;
    CALL_ZE(RC, zeDeviceGet, Driver, &DeviceCount, nullptr);
    if (RC != ZE_RESULT_SUCCESS || DeviceCount == 0) {
      ODBG(OLDT_Init) << "Cannot find any devices from driver " << Driver
                      << ".";
      continue;
    }

    // We have a driver that supports at least one device.
    ContextList.emplace_back(*this, Driver, DriverId);
    auto &DrvInfo = ContextList.back();
    if (auto Err = DrvInfo.init()) {
      // Remove the partially initialized context from the list
      ContextList.pop_back();
      return std::move(Err);
    }
    llvm::SmallVector<ze_device_handle_t> FoundDevices(DeviceCount);
    CALL_ZE_RET_ERROR(zeDeviceGet, Driver, &DeviceCount, FoundDevices.data());

    for (auto &zeDevice : FoundDevices)
      RootDevices.push_back(
          {OrderId++, zeDevice, &DrvInfo, L0DeviceTy::isDiscrete(zeDevice)});
  }

  // Move discrete devices to the front.
  std::sort(RootDevices.begin(), RootDevices.end(),
            [](const RootInfoTy &A, const RootInfoTy &B) {
              // If both are discrete, order by OrderId.
              // If both are not discrete, order by OrderId.
              // Otherwise, discrete goes first.

              if (A.IsDiscrete && B.IsDiscrete)
                return A.OrderId < B.OrderId;
              if (!A.IsDiscrete && !B.IsDiscrete)
                return A.OrderId < B.OrderId;
              return A.IsDiscrete;
            });

  for (size_t RootId = 0; RootId < RootDevices.size(); RootId++) {
    const auto ZeDevice = RootDevices[RootId].ZeDevice;
    auto *RootDriver = RootDevices[RootId].Driver;
    DetectedDevices.push_back(DeviceInfoTy{
        {ZeDevice, static_cast<int32_t>(RootId), -1, -1}, RootDriver});
  }
  int32_t NumDevices = DetectedDevices.size();

  ODBG_OS(OLDT_Init, [&](llvm::raw_ostream &O) {
    O << "Found " << NumDevices << " devices.\n"
      << "List of devices (DeviceID[.SubID[.CCSID]])\n";
    for (auto &DeviceInfo : DetectedDevices)
      O << "-- Device " << DeviceInfo.Id.RootId
        << (DeviceInfo.Id.SubId < 0
                ? ""
                : ("." + std::to_string(DeviceInfo.Id.SubId)))
        << (DeviceInfo.Id.CCSId < 0
                ? ""
                : ("." + std::to_string(DeviceInfo.Id.CCSId)))
        << "\n";
  });
  return NumDevices;
}

Expected<int32_t> LevelZeroPluginTy::initImpl() {
  ODBG(OLDT_Init) << "Level0 NG plugin initialization";
  // Process options before anything else.
  Options.init();
  return findDevices();
}

Error LevelZeroPluginTy::deinitImpl() {
  ODBG(OLDT_Deinit) << "Deinit Level0 plugin!";
  if (auto Err = ContextTLSTable.deinit())
    return Err;
  for (auto &Context : ContextList)
    if (auto Err = Context.deinit())
      return Err;
  ContextList.clear();
  ODBG(OLDT_Deinit) << "Level0 plugin deinitialized successfully";
  return Plugin::success();
}

GenericDeviceTy *LevelZeroPluginTy::createDevice(GenericPluginTy &Plugin,
                                                 int32_t DeviceId,
                                                 int32_t NumDevices) {
  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));

  return new L0DeviceTy(static_cast<LevelZeroPluginTy &>(Plugin), DeviceId,
                        NumDevices, zeDevice, *zeDriver, std::move(IdStr),
                        CCSId < 0 ? 0 : CCSId /* ComputeIndex */);
}

GenericGlobalHandlerTy *LevelZeroPluginTy::createGlobalHandler() {
  return new L0GlobalHandlerTy();
}

Error LevelZeroPluginTy::flushQueueImpl(omp_interop_val_t *Interop) {
  return Plugin::success();
}

Expected<bool> LevelZeroPluginTy::isImageCompatible(StringRef Image) const {
  switch (identify_magic(Image)) {
  case file_magic::spirv_object:
    // Handle SPIRV objects directly
    return true;
  case file_magic::offload_binary: {
    // Handle OffloadBinary format
    MemoryBufferRef Buffer(Image, "offload_binary");
    auto BinariesOrErr = OffloadBinary::create(Buffer);
    if (!BinariesOrErr)
      return BinariesOrErr.takeError();

    auto &Binaries = *BinariesOrErr;
    if (Binaries.size() != 1)
      return false;

    const OffloadBinary *InnerBinary = Binaries[0].get();
    ImageKind ImageKind = InnerBinary->getImageKind();
    llvm::Triple Triple(InnerBinary->getTriple());

    if (Triple.getArch() != getTripleArch())
      return false;

    if (ImageKind != llvm::object::IMG_SPIRV &&
        ImageKind != llvm::object::IMG_Object)
      return false;

    return true;
  }
  default:
    // Unknown format
    return false;
  }
}

Error LevelZeroPluginTy::syncBarrierImpl(omp_interop_val_t *Interop) {
  if (!Interop) {
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "Invalid/inconsistent OpenMP interop " DPxMOD "\n",
                         DPxPTR(Interop));
  }
  if (!Interop->async_info || !Interop->async_info->Queue)
    return Plugin::success();

  const auto L0 = static_cast<L0Interop::Property *>(Interop->rtl_property);

  ODBG(OLDT_Sync) << "LevelZeroPluginTy::sync_barrier: Synchronizing "
                  << Interop << " with ImmCmdList barrier";
  auto ImmCmdList = L0->ImmCmdList;

  CALL_ZE_RET_ERROR(zeCommandListHostSynchronize, ImmCmdList, L0DefaultTimeout);

  return Plugin::success();
}

Error LevelZeroPluginTy::asyncBarrierImpl(omp_interop_val_t *Interop) {
  if (!Interop) {
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "Invalid/inconsistent OpenMP interop " DPxMOD "\n",
                         DPxPTR(Interop));
  }
  if (!Interop->async_info || !Interop->async_info->Queue)
    return Plugin::success();

  const auto L0 = static_cast<L0Interop::Property *>(Interop->rtl_property);
  if (Interop->attrs.inorder)
    return Plugin::success();

  ODBG(OLDT_Sync) << "LevelZeroPluginTy::async_barrier: Appending ImmCmdList "
                  << "barrier to " << Interop;
  auto ImmCmdList = L0->ImmCmdList;
  CALL_ZE_RET_ERROR(zeCommandListAppendBarrier, ImmCmdList, nullptr, 0,
                    nullptr);

  return Plugin::success();
}

Error LevelZeroPluginContextTy::initAsyncInfoImpl(
    GenericDeviceTy &Device, AsyncInfoWrapperTy &AsyncInfoWrapper) {
  auto &L0Device = static_cast<L0DeviceTy &>(Device);
  auto QueueOrErr = L0Device.getOrCreateQueue(AsyncInfoWrapper, this);
  return QueueOrErr ? Plugin::success() : QueueOrErr.takeError();
}

Error LevelZeroPluginContextTy::deinit() {
  if (auto Err = QueueCache.deinit())
    return Err;
  // Tear down allocators before their ze_context.
  for (auto &KV : DeviceAllocators)
    if (auto Err = KV.second->deinit())
      return Err;
  DeviceAllocators.clear();
  if (HostAllocator) {
    if (auto Err = HostAllocator->deinit())
      return Err;
    HostAllocator.reset();
  }
  if (OwnsZeContext && ZeContext) {
    CALL_ZE_RET_ERROR(zeContextDestroy, ZeContext);
    ZeContext = nullptr;
    OwnsZeContext = false;
  }
  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 (const auto &[Device, Allocator] : DeviceAllocators) {
    if (auto *Info = Allocator->getAllocInfo(Raw))
      return PluginAllocInfoTy{Device, static_cast<TargetAllocTy>(Info->Kind),
                               Info->Base, Info->ReqSize};
  }
  if (HostAllocator) {
    if (auto *Info = HostAllocator->getAllocInfo(Raw))
      return PluginAllocInfoTy{nullptr, static_cast<TargetAllocTy>(Info->Kind),
                               Info->Base, Info->ReqSize};
  }

  return Plugin::error(error::ErrorCode::NOT_FOUND,
                       "pointer is not a known allocation in this context");
}

Expected<std::unique_ptr<PluginContextTy>>
LevelZeroPluginTy::createPluginContext(
    llvm::ArrayRef<GenericDeviceTy *> Devices) {
  if (Devices.empty())
    return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                         "createPluginContext called with no devices");

  // All devices must share the same L0 driver context, since a ze_context is
  // scoped to a single driver.
  auto &First = static_cast<L0DeviceTy &>(*Devices[0]);
  L0ContextTy &DriverCtx = First.getL0Context();
  ze_driver_handle_t Driver = DriverCtx.getZeDriver();
  for (auto *D : Devices.drop_front()) {
    auto &L0D = static_cast<L0DeviceTy &>(*D);
    if (&L0D.getL0Context() != &DriverCtx)
      return Plugin::error(
          ErrorCode::INVALID_DEVICE,
          "all devices in an L0 context must share the same driver");
  }

  // When the user requests every device on the driver, share the driver's
  // default ze_context so we interop cleanly with other L0 clients on the
  // same driver. Partial device sets get their own fresh context.
  size_t DriverDeviceCount = 0;
  for (int32_t I = 0, N = getNumDevices(); I != N; ++I) {
    auto &L0D = static_cast<const L0DeviceTy &>(getDevice(I));
    if (&L0D.getL0Context() == &DriverCtx)
      ++DriverDeviceCount;
  }
  const bool IsFullDriver = (Devices.size() == DriverDeviceCount);

  ze_context_handle_t ZeContext = nullptr;
  bool OwnsZeContext = false;
  if (IsFullDriver && DriverCtx.DriverGetDefaultContext.available())
    ZeContext = DriverCtx.DriverGetDefaultContext(Driver);
  if (!ZeContext) {
    ze_context_desc_t Desc{ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
    CALL_ZE_RET_ERROR(zeContextCreate, Driver, &Desc, &ZeContext);
    OwnsZeContext = true;
  }

  auto Ctx = std::make_unique<LevelZeroPluginContextTy>(
      *this, Devices, Driver, ZeContext, OwnsZeContext);
  if (auto Err = Ctx->initAllocators())
    return std::move(Err);
  return Ctx;
}

} // namespace llvm::omp::target::plugin

extern "C" {
llvm::omp::target::plugin::GenericPluginTy *createPlugin_level_zero() {
  return new llvm::omp::target::plugin::LevelZeroPluginTy();
}
}