//==--------- global_handler.cpp --- Global objects handler ----------------==//
//
// 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
//
//===----------------------------------------------------------------------===//

#ifdef ENABLE_STACK_TRACE
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Signals.h"
#endif

#include <detail/adapter.hpp>
#include <detail/config.hpp>
#include <detail/global_handler.hpp>
#include <detail/platform_impl.hpp>
#include <detail/program_manager/program_manager.hpp>
#include <detail/scheduler/scheduler.hpp>
#include <detail/thread_pool.hpp>
#include <detail/ur.hpp>
#include <detail/xpti_registry.hpp>
#include <sycl/detail/device_filter.hpp>
#include <sycl/detail/spinlock.hpp>

#ifdef _WIN32
#include <windows.h>
#endif

#include <vector>

namespace sycl {
inline namespace _V1 {
namespace detail {

using LockGuard = std::lock_guard<SpinLock>;
SpinLock GlobalHandler::MSyclGlobalHandlerProtector{};

// forward decl
void shutdown_win(); // TODO: win variant will go away soon
void shutdown_early();
void shutdown_late();

// Utility class to track references on object.
// Used for GlobalHandler now and created as thread_local object on the first
// Scheduler usage. Origin idea is to track usage of Scheduler from main and
// other used threads - they increment MCounter; and to use but not add extra
// reference by our thread_pool threads. For this control MIncrementCounter
// class member is used.
class ObjectUsageCounter {
public:
  ObjectUsageCounter(bool ModifyCounter) : MModifyCounter(ModifyCounter) {
    if (MModifyCounter)
      MCounter++;
  }
  ~ObjectUsageCounter() {
    try {
      if (!MModifyCounter)
        return;

      LockGuard Guard(GlobalHandler::MSyclGlobalHandlerProtector);
      MCounter--;
      GlobalHandler *RTGlobalObjHandler = GlobalHandler::getInstancePtr();
      if (RTGlobalObjHandler) {
        RTGlobalObjHandler->prepareSchedulerToRelease(!MCounter);
      }
    } catch (std::exception &e) {
      __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in ~ObjectUsageCounter", e);
    }
  }

private:
  static std::atomic_uint MCounter;
  bool MModifyCounter;
};
std::atomic_uint ObjectUsageCounter::MCounter{0};

GlobalHandler::GlobalHandler() = default;
GlobalHandler::~GlobalHandler() = default;

void GlobalHandler::InitXPTI() {
#ifdef XPTI_ENABLE_INSTRUMENTATION
  // Let subscribers know a new stream is being initialized
  getXPTIRegistry().initializeStream(SYCL_STREAM_NAME, GMajVer, GMinVer,
                                     GVerStr);
  xpti::payload_t SYCLPayload("SYCL Runtime Exceptions");
  uint64_t SYCLInstanceNo;
  GSYCLCallEvent = xptiMakeEvent("SYCL Try-catch Exceptions", &SYCLPayload,
                                 xpti::trace_algorithm_event, xpti_at::active,
                                 &SYCLInstanceNo);
#endif
}

void GlobalHandler::TraceEventXPTI(const char *Message) {
  if (!Message)
    return;
#ifdef XPTI_ENABLE_INSTRUMENTATION
  static std::once_flag InitXPTIFlag;
  if (xptiTraceEnabled()) {
    std::call_once(InitXPTIFlag, [&]() { InitXPTI(); });

    // We have to handle the cases where: (1) we may have just the code location
    // set and not UID and (2) UID set
    detail::tls_code_loc_t Tls;
    auto CodeLocation = Tls.query();

    // Creating a tracepoint will convert a CodeLocation to UID, if not set
    xpti::framework::tracepoint_t TP(
        CodeLocation.fileName(), CodeLocation.functionName(),
        CodeLocation.lineNumber(), CodeLocation.columnNumber(), nullptr);

    // The call to notify will have the signature of:
    // (1) the stream defined in .stream()
    // (2) The trace type equal to what is set by .trace_type()
    // (3) Parent event set to NULL
    // (4) Current event set to one created from CodeLocation and UID
    // (5) An instance ID that records the number of times this code location
    // has been seen (6) The message generated by the exception handler
    TP.stream(SYCL_STREAM_NAME)
        .trace_type(xpti::trace_point_type_t::diagnostics)
        .notify(static_cast<const void *>(Message));
  }

#endif
}

GlobalHandler *&GlobalHandler::getInstancePtr() {
  static GlobalHandler *RTGlobalObjHandler = new GlobalHandler();
  return RTGlobalObjHandler;
}

GlobalHandler &GlobalHandler::instance() {
  GlobalHandler *RTGlobalObjHandler = GlobalHandler::getInstancePtr();
  assert(RTGlobalObjHandler && "Handler must not be deallocated earlier");
  return *RTGlobalObjHandler;
}

template <typename T, typename... Types>
T &GlobalHandler::getOrCreate(InstWithLock<T> &IWL, Types... Args) {
  const LockGuard Lock{IWL.Lock};

  if (!IWL.Inst)
    IWL.Inst = std::make_unique<T>(Args...);

  return *IWL.Inst;
}

void GlobalHandler::attachScheduler(Scheduler *Scheduler) {
  // The method is used in unit tests only. Do not protect with lock since
  // releaseResources will cause dead lock due to host queue release
  if (MScheduler.Inst)
    prepareSchedulerToRelease(true);
  MScheduler.Inst.reset(Scheduler);
}

static void enableOnCrashStackPrinting() {
#ifdef ENABLE_STACK_TRACE
  static std::once_flag PrintStackFlag;
  std::call_once(PrintStackFlag, []() {
    llvm::sys::PrintStackTraceOnErrorSignal(llvm::StringRef());
  });
#endif
}

Scheduler &GlobalHandler::getScheduler() {
  getOrCreate(MScheduler);
  registerSchedulerUsage();
  // On Windows the registration of the signal handler before main function
  // (e.g. from DLLMain or from constructors of program scope objects) doesn't
  // work. So, registering signal handler here because:
  // 1) getScheduler is likely to be called for any non-trivial application;
  // 2) first call to getScheduler is likely to be done after main starts.
  // The same is done in getAdapters.
  enableOnCrashStackPrinting();
  return *MScheduler.Inst;
}

bool GlobalHandler::isSchedulerAlive() const { return MScheduler.Inst.get(); }

void GlobalHandler::registerSchedulerUsage(bool ModifyCounter) {
  thread_local ObjectUsageCounter SchedulerCounter(ModifyCounter);
}

ProgramManager &GlobalHandler::getProgramManager() {
  return getOrCreate(MProgramManager);
}

std::unordered_map<PlatformImplPtr, ContextImplPtr> &
GlobalHandler::getPlatformToDefaultContextCache() {
  return getOrCreate(MPlatformToDefaultContextCache);
}

std::mutex &GlobalHandler::getPlatformToDefaultContextCacheMutex() {
  return getOrCreate(MPlatformToDefaultContextCacheMutex);
}

Sync &GlobalHandler::getSync() { return getOrCreate(MSync); }

std::vector<PlatformImplPtr> &GlobalHandler::getPlatformCache() {
  return getOrCreate(MPlatformCache);
}

std::mutex &GlobalHandler::getPlatformMapMutex() {
  return getOrCreate(MPlatformMapMutex);
}

std::mutex &GlobalHandler::getFilterMutex() {
  return getOrCreate(MFilterMutex);
}

std::vector<AdapterPtr> &GlobalHandler::getAdapters() {
  enableOnCrashStackPrinting();
  return getOrCreate(MAdapters);
}

ods_target_list &
GlobalHandler::getOneapiDeviceSelectorTargets(const std::string &InitValue) {
  return getOrCreate(MOneapiDeviceSelectorTargets, InitValue);
}

XPTIRegistry &GlobalHandler::getXPTIRegistry() {
  return getOrCreate(MXPTIRegistry);
}

ThreadPool &GlobalHandler::getHostTaskThreadPool() {
  int Size = SYCLConfig<SYCL_QUEUE_THREAD_POOL_SIZE>::get();
  ThreadPool &TP = getOrCreate(MHostTaskThreadPool, Size);

  return TP;
}

void GlobalHandler::releaseDefaultContexts() {
  // Release shared-pointers to SYCL objects.
  // Note that on Windows the destruction of the default context
  // races with the detaching of the DLL object that calls urLoaderTearDown.

  MPlatformToDefaultContextCache.Inst.reset(nullptr);
}

struct EarlyShutdownHandler {
  ~EarlyShutdownHandler() {
    try {
#ifdef _WIN32
      // on Windows we keep to the existing shutdown procedure
      GlobalHandler::instance().releaseDefaultContexts();
#else
      shutdown_early();
#endif
    } catch (std::exception &e) {
      __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in ~EarlyShutdownHandler",
                                        e);
    }
  }
};

void GlobalHandler::registerEarlyShutdownHandler() {
  static EarlyShutdownHandler handler{};
}

bool GlobalHandler::isOkToDefer() const { return OkToDefer; }

void GlobalHandler::endDeferredRelease() { OkToDefer = false; }

// Note: Split from shutdown so it is available to the unittests for ensuring
//       that the mock adapter is the lone adapter.
void GlobalHandler::unloadAdapters() {
  // Call to GlobalHandler::instance().getAdapters() initializes adapters. If
  // user application has loaded SYCL runtime, and never called any APIs,
  // there's no need to load and unload adapters.
  if (MAdapters.Inst) {
    for (const auto &Adapter : getAdapters()) {
      Adapter->release();
    }
  }

  UrFuncInfo<UrApiKind::urLoaderTearDown> loaderTearDownInfo;
  auto loaderTearDown =
      loaderTearDownInfo.getFuncPtrFromModule(ur::getURLoaderLibrary());
  loaderTearDown();
  // urLoaderTearDown();

  // Clear after unload to avoid uses after unload.
  getAdapters().clear();
}

void GlobalHandler::prepareSchedulerToRelease(bool Blocking) {
#ifndef _WIN32
  if (Blocking)
    drainThreadPool();
  if (MScheduler.Inst)
    MScheduler.Inst->releaseResources(Blocking ? BlockingT::BLOCKING
                                               : BlockingT::NON_BLOCKING);
#endif
}

void GlobalHandler::drainThreadPool() {
  if (MHostTaskThreadPool.Inst)
    MHostTaskThreadPool.Inst->drain();
}

#ifdef _WIN32
// because of something not-yet-understood on Windows
// threads may be shutdown once the end of main() is reached
// making an orderly shutdown difficult. Fortunately, Windows
// itself is very aggressive about reclaiming memory. Thus,
// we focus solely on unloading the adapters, so as to not
// accidentally retain device handles. etc
void shutdown_win() {
  GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
  Handler->unloadAdapters();
}
#else
void shutdown_early() {
  const LockGuard Lock{GlobalHandler::MSyclGlobalHandlerProtector};
  GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
  if (!Handler)
    return;

  // Now that we are shutting down, we will no longer defer MemObj releases.
  Handler->endDeferredRelease();

  // Ensure neither host task is working so that no default context is accessed
  // upon its release
  Handler->prepareSchedulerToRelease(true);

  if (Handler->MHostTaskThreadPool.Inst)
    Handler->MHostTaskThreadPool.Inst->finishAndWait();

  // This releases OUR reference to the default context, but
  // other may yet have refs
  Handler->releaseDefaultContexts();
}

void shutdown_late() {
  const LockGuard Lock{GlobalHandler::MSyclGlobalHandlerProtector};
  GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
  if (!Handler)
    return;

  // First, release resources, that may access adapters.
  Handler->MPlatformCache.Inst.reset(nullptr);
  Handler->MScheduler.Inst.reset(nullptr);
  Handler->MProgramManager.Inst.reset(nullptr);

  // Clear the adapters and reset the instance if it was there.
  Handler->unloadAdapters();
  if (Handler->MAdapters.Inst)
    Handler->MAdapters.Inst.reset(nullptr);

  Handler->MXPTIRegistry.Inst.reset(nullptr);

  // Release the rest of global resources.
  delete Handler;
  Handler = nullptr;
}
#endif

#ifdef _WIN32
extern "C" __SYCL_EXPORT BOOL WINAPI DllMain(HINSTANCE hinstDLL,
                                             DWORD fdwReason,
                                             LPVOID lpReserved) {
  bool PrintUrTrace = false;
  try {
    PrintUrTrace =
        sycl::detail::ur::trace(sycl::detail::ur::TraceLevel::TRACE_CALLS);
  } catch (std::exception &e) {
    __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in DllMain", e);
    return FALSE;
  }

  // Perform actions based on the reason for calling.
  switch (fdwReason) {
  case DLL_PROCESS_DETACH:
    if (PrintUrTrace)
      std::cout << "---> DLL_PROCESS_DETACH syclx.dll\n" << std::endl;

#ifdef XPTI_ENABLE_INSTRUMENTATION
    if (xptiTraceEnabled())
      return TRUE; // When doing xpti tracing, we can't safely call shutdown.
                   // TODO: figure out what XPTI is doing that prevents
                   // release.
#endif

    try {
      shutdown_win();
    } catch (std::exception &e) {
      __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in shutdown_win", e);
      return FALSE;
    }
    break;
  case DLL_PROCESS_ATTACH:
    if (PrintUrTrace)
      std::cout << "---> DLL_PROCESS_ATTACH syclx.dll\n" << std::endl;
    break;
  case DLL_THREAD_ATTACH:
    break;
  case DLL_THREAD_DETACH:
    break;
  }
  return TRUE; // Successful DLL_PROCESS_ATTACH.
}
#else
// Setting low priority on destructor ensures it runs after all other global
// destructors. Priorities 0-100 are reserved by the compiler. The priority
// value 110 allows SYCL users to run their destructors after runtime library
// deinitialization.
__attribute__((destructor(110))) static void syclUnload() { shutdown_late(); }
#endif
} // namespace detail
} // namespace _V1
} // namespace sycl