//===--- Level Zero Target RTL 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
//
//===----------------------------------------------------------------------===//
//
//  Level Zero Context abstraction.
//
//===----------------------------------------------------------------------===//

#include "L0Context.h"
#include "L0Plugin.h"

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

L0ContextTy::L0ContextTy(LevelZeroPluginTy &Plugin, ze_driver_handle_t zeDriver,
                         int32_t DriverId)
    : Plugin(Plugin), zeDriver(zeDriver) {}

L0ContextTy::~L0ContextTy() = default;

Error L0ContextTy::init() {
  auto CleanupOnError = [&]() {
    if (zeContext) {
      zeContextDestroy(zeContext);
      zeContext = nullptr;
    }
  };
  CALL_ZE_RET_ERROR(zeDriverGetApiVersion, zeDriver, &APIVersion);
  ODBG(OLDT_Init) << "Driver API version is "
                  << llvm::format(PRIx32, APIVersion);

  ze_context_desc_t Desc{ZE_STRUCTURE_TYPE_CONTEXT_DESC, nullptr, 0};
  CALL_ZE_RET_ERROR(zeContextCreate, zeDriver, &Desc, &zeContext);

  const auto &Options = Plugin.getOptions();
  bool UseCounterBasedEvents = Options.CommandMode == CommandModeTy::InOrder ||
                               Options.CommandMode == CommandModeTy::Sync;
  if (UseCounterBasedEvents)
    ODBG(OLDT_Init) << "Using counter-based events for "
                    << (Options.CommandMode == CommandModeTy::InOrder
                            ? "InOrder"
                            : "Sync")
                    << " command mode";

  if (auto Err = EventPool.init(zeContext, UseCounterBasedEvents,
                                /* Flags */ 0)) {
    CleanupOnError();
    return Err;
  }
  if (auto Err = HostMemAllocator.initHostPool(*this, Plugin.getOptions())) {
    if (auto DeinitErr = EventPool.deinit())
      Err = joinErrors(std::move(Err), std::move(DeinitErr));
    CleanupOnError();
    return Err;
  }

  ze_result_t RC;
  CALL_ZE(RC, zeDriverGetExtensionFunctionAddress, zeDriver,
          "zexKernelGetArgumentSize", (void **)&zexKernelGetArgumentSize);
  if (RC != ZE_RESULT_SUCCESS)
    zexKernelGetArgumentSize = nullptr;

  CALL_ZE(RC, zeDriverGetExtensionFunctionAddress, zeDriver,
          "zeCommandListAppendHostFunction",
          (void **)&zeCommandListAppendHostFunction);
  if (RC != ZE_RESULT_SUCCESS)
    zeCommandListAppendHostFunction = nullptr;

  CALL_ZE(RC, zeDriverGetExtensionFunctionAddress, zeDriver,
          "zeDriverGetDefaultContext", (void **)&zeDriverGetDefaultContext);
  if (RC != ZE_RESULT_SUCCESS)
    zeDriverGetDefaultContext = nullptr;

  DefaultUserCtx = std::make_unique<LevelZeroPluginContextTy>(
      Plugin, /*Devices=*/llvm::ArrayRef<GenericDeviceTy *>{}, zeDriver,
      zeContext, /*OwnsZeContext=*/false);

  return Plugin::success();
}

Error L0ContextTy::deinit() {
  // Release the default context (drains its queue cache) before zeContext.
  if (DefaultUserCtx) {
    if (auto Err = DefaultUserCtx->deinit())
      return Err;
    DefaultUserCtx.reset();
  }
  if (auto Err = EventPool.deinit())
    return Err;
  if (auto Err = HostMemAllocator.deinit())
    return Err;
  if (zeContext)
    CALL_ZE_RET_ERROR(zeContextDestroy, zeContext);
  return Plugin::success();
}

StagingBufferTy &L0ContextTy::getStagingBuffer() {
  auto &TLS = Plugin.getContextTLS(getZeContext());
  auto &Buffer = TLS.getStagingBuffer();
  const auto &Options = Plugin.getOptions();
  if (!Buffer.initialized())
    Buffer.init(getZeContext(), Options.StagingBufferSize,
                Options.StagingBufferCount);
  return Buffer;
}

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