//===- JIT.cpp - Target independent JIT infrastructure --------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
//===----------------------------------------------------------------------===//

#include "JIT.h"

#include "Shared/Debug.h"
#include "Shared/Utils.h"

#include "PluginInterface.h"
#include "omptarget.h"

#include "llvm/ADT/SmallVector.h"
#include "llvm/CodeGen/CommandFlags.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LLVMRemarkStreamer.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/InitializePasses.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Object/IRObjectFile.h"
#include "llvm/Passes/OptimizationLevel.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/TimeProfiler.h"
#include "llvm/Support/ToolOutputFile.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/TargetParser/SubtargetFeature.h"

#include <mutex>
#include <shared_mutex>
#include <system_error>

using namespace llvm;
using namespace llvm::object;
using namespace omp;
using namespace omp::target;

namespace {

bool isImageBitcode(const __tgt_device_image &Image) {
  StringRef Binary(reinterpret_cast<const char *>(Image.ImageStart),
                   utils::getPtrDiff(Image.ImageEnd, Image.ImageStart));

  return identify_magic(Binary) == file_magic::bitcode;
}

Expected<std::unique_ptr<Module>>
createModuleFromMemoryBuffer(std::unique_ptr<MemoryBuffer> &MB,
                             LLVMContext &Context) {
  SMDiagnostic Err;
  auto Mod = parseIR(*MB, Err, Context);
  if (!Mod)
    return error::createOffloadError(error::ErrorCode::UNKNOWN,
                                     "failed to create module");
  return std::move(Mod);
}
Expected<std::unique_ptr<Module>>
createModuleFromImage(const __tgt_device_image &Image, LLVMContext &Context) {
  StringRef Data((const char *)Image.ImageStart,
                 utils::getPtrDiff(Image.ImageEnd, Image.ImageStart));
  std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer(
      Data, /*BufferName=*/"", /*RequiresNullTerminator=*/false);
  return createModuleFromMemoryBuffer(MB, Context);
}

OptimizationLevel getOptLevel(unsigned OptLevel) {
  switch (OptLevel) {
  case 0:
    return OptimizationLevel::O0;
  case 1:
    return OptimizationLevel::O1;
  case 2:
    return OptimizationLevel::O2;
  case 3:
    return OptimizationLevel::O3;
  }
  llvm_unreachable("Invalid optimization level");
}

Expected<std::unique_ptr<TargetMachine>>
createTargetMachine(Module &M, std::string CPU, unsigned OptLevel) {
  Triple TT(M.getTargetTriple());
  std::optional<CodeGenOptLevel> CGOptLevelOrNone =
      CodeGenOpt::getLevel(OptLevel);
  assert(CGOptLevelOrNone && "Invalid optimization level");
  CodeGenOptLevel CGOptLevel = *CGOptLevelOrNone;

  std::string Msg;
  const Target *T = TargetRegistry::lookupTarget(M.getTargetTriple(), Msg);
  if (!T)
    return error::createOffloadError(error::ErrorCode::INVALID_BINARY,
                                     Msg.data());

  SubtargetFeatures Features;
  Features.getDefaultSubtargetFeatures(TT);

  std::optional<Reloc::Model> RelocModel;
  if (M.getModuleFlag("PIC Level"))
    RelocModel =
        M.getPICLevel() == PICLevel::NotPIC ? Reloc::Static : Reloc::PIC_;

  std::optional<CodeModel::Model> CodeModel = M.getCodeModel();

  TargetOptions Options = codegen::InitTargetOptionsFromCodeGenFlags(TT);

  std::unique_ptr<TargetMachine> TM(
      T->createTargetMachine(M.getTargetTriple(), CPU, Features.getString(),
                             Options, RelocModel, CodeModel, CGOptLevel));
  if (!TM)
    return error::createOffloadError(error::ErrorCode::INVALID_BINARY,
                                     "failed to create target machine");
  return std::move(TM);
}

} // namespace

JITEngine::JITEngine(Triple::ArchType TA) : TT(Triple::getArchTypeName(TA)) {
  codegen::RegisterCodeGenFlags();
#ifdef LIBOMPTARGET_JIT_NVPTX
  if (TT.isNVPTX()) {
    LLVMInitializeNVPTXTargetInfo();
    LLVMInitializeNVPTXTarget();
    LLVMInitializeNVPTXTargetMC();
    LLVMInitializeNVPTXAsmPrinter();
  }
#endif
#ifdef LIBOMPTARGET_JIT_AMDGPU
  if (TT.isAMDGPU()) {
    LLVMInitializeAMDGPUTargetInfo();
    LLVMInitializeAMDGPUTarget();
    LLVMInitializeAMDGPUTargetMC();
    LLVMInitializeAMDGPUAsmPrinter();
  }
#endif
}

void JITEngine::opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
                    unsigned OptLevel) {
  PipelineTuningOptions PTO;
  std::optional<PGOOptions> PGOOpt;

  LoopAnalysisManager LAM;
  FunctionAnalysisManager FAM;
  CGSCCAnalysisManager CGAM;
  ModuleAnalysisManager MAM;
  ModulePassManager MPM;

  PassBuilder PB(TM, PTO, PGOOpt, nullptr);

  FAM.registerPass([&] { return TargetLibraryAnalysis(*TLII); });

  // Register all the basic analyses with the managers.
  PB.registerModuleAnalyses(MAM);
  PB.registerCGSCCAnalyses(CGAM);
  PB.registerFunctionAnalyses(FAM);
  PB.registerLoopAnalyses(LAM);
  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);

  MPM.addPass(PB.buildPerModuleDefaultPipeline(getOptLevel(OptLevel)));
  MPM.run(M, MAM);
}

void JITEngine::codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII,
                        Module &M, raw_pwrite_stream &OS) {
  legacy::PassManager PM;
  PM.add(new TargetLibraryInfoWrapperPass(*TLII));
  MachineModuleInfoWrapperPass *MMIWP = new MachineModuleInfoWrapperPass(TM);
  TM->addPassesToEmitFile(PM, OS, nullptr,
                          TT.isNVPTX() ? CodeGenFileType::AssemblyFile
                                       : CodeGenFileType::ObjectFile,
                          /*DisableVerify=*/false, MMIWP);

  PM.run(M);
}

Expected<std::unique_ptr<MemoryBuffer>>
JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
                   unsigned OptLevel) {

  auto RemarksFileOrErr = setupLLVMOptimizationRemarks(
      M.getContext(), /*RemarksFilename=*/"", /*RemarksPasses=*/"",
      /*RemarksFormat=*/"", /*RemarksWithHotness=*/false);
  if (Error E = RemarksFileOrErr.takeError())
    return std::move(E);
  if (*RemarksFileOrErr)
    (*RemarksFileOrErr)->keep();

  auto TMOrErr = createTargetMachine(M, ComputeUnitKind, OptLevel);
  if (!TMOrErr)
    return TMOrErr.takeError();

  std::unique_ptr<TargetMachine> TM = std::move(*TMOrErr);
  TargetLibraryInfoImpl TLII(TT);

  if (PreOptIRModuleFileName.isPresent()) {
    std::error_code EC;
    raw_fd_stream FD(PreOptIRModuleFileName.get(), EC);
    if (EC)
      return createStringError(
          EC, "Could not open %s to write the pre-opt IR module\n",
          PreOptIRModuleFileName.get().c_str());
    M.print(FD, nullptr);
  }

  if (!JITSkipOpt)
    opt(TM.get(), &TLII, M, OptLevel);

  if (PostOptIRModuleFileName.isPresent()) {
    std::error_code EC;
    raw_fd_stream FD(PostOptIRModuleFileName.get(), EC);
    if (EC)
      return createStringError(
          error::ErrorCode::HOST_IO,
          "Could not open %s to write the post-opt IR module\n",
          PostOptIRModuleFileName.get().c_str());
    M.print(FD, nullptr);
  }

  // Prepare the output buffer and stream for codegen.
  SmallVector<char> CGOutputBuffer;
  raw_svector_ostream OS(CGOutputBuffer);

  codegen(TM.get(), &TLII, M, OS);

  return MemoryBuffer::getMemBufferCopy(OS.str());
}

Expected<std::unique_ptr<MemoryBuffer>>
JITEngine::getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
                              const std::string &ComputeUnitKind) {

  // Check if the user replaces the module at runtime with a finished object.
  if (ReplacementObjectFileName.isPresent()) {
    auto MBOrErr =
        MemoryBuffer::getFileOrSTDIN(ReplacementObjectFileName.get());
    if (!MBOrErr)
      return createStringError(MBOrErr.getError(),
                               "Could not read replacement obj from %s\n",
                               ReplacementModuleFileName.get().c_str());
    return std::move(*MBOrErr);
  }

  Module *Mod = nullptr;
  // Check if the user replaces the module at runtime or we read it from the
  // image.
  // TODO: Allow the user to specify images per device (Arch + ComputeUnitKind).
  if (!ReplacementModuleFileName.isPresent()) {
    auto ModOrErr = createModuleFromImage(Image, Ctx);
    if (!ModOrErr)
      return ModOrErr.takeError();
    Mod = ModOrErr->release();
  } else {
    auto MBOrErr =
        MemoryBuffer::getFileOrSTDIN(ReplacementModuleFileName.get());
    if (!MBOrErr)
      return createStringError(MBOrErr.getError(),
                               "Could not read replacement module from %s\n",
                               ReplacementModuleFileName.get().c_str());
    auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), Ctx);
    if (!ModOrErr)
      return ModOrErr.takeError();
    Mod = ModOrErr->release();
  }

  return backend(*Mod, ComputeUnitKind, JITOptLevel);
}

Expected<const __tgt_device_image *>
JITEngine::compile(const __tgt_device_image &Image,
                   const std::string &ComputeUnitKind,
                   PostProcessingFn PostProcessing) {
  std::lock_guard<std::mutex> Lock(ComputeUnitMapMutex);

  // Check if we JITed this image for the given compute unit kind before.
  ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];
  if (CUI.TgtImageMap.contains(&Image))
    return CUI.TgtImageMap[&Image].get();

  auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind);
  if (!ObjMBOrErr)
    return ObjMBOrErr.takeError();

  auto ImageMBOrErr = PostProcessing(std::move(*ObjMBOrErr));
  if (!ImageMBOrErr)
    return ImageMBOrErr.takeError();

  CUI.JITImages.insert({&Image, std::move(*ImageMBOrErr)});
  auto &ImageMB = CUI.JITImages[&Image];
  CUI.TgtImageMap.insert({&Image, std::make_unique<__tgt_device_image>()});
  auto &JITedImage = CUI.TgtImageMap[&Image];
  *JITedImage = Image;
  JITedImage->ImageStart = const_cast<char *>(ImageMB->getBufferStart());
  JITedImage->ImageEnd = const_cast<char *>(ImageMB->getBufferEnd());

  return JITedImage.get();
}

Expected<const __tgt_device_image *>
JITEngine::process(const __tgt_device_image &Image,
                   target::plugin::GenericDeviceTy &Device) {
  const std::string &ComputeUnitKind = Device.getComputeUnitKind();

  PostProcessingFn PostProcessing = [&Device](std::unique_ptr<MemoryBuffer> MB)
      -> Expected<std::unique_ptr<MemoryBuffer>> {
    return Device.doJITPostProcessing(std::move(MB));
  };

  if (isImageBitcode(Image))
    return compile(Image, ComputeUnitKind, PostProcessing);

  return &Image;
}

void JITEngine::erase(const __tgt_device_image &Image,
                      target::plugin::GenericDeviceTy &Device) {
  std::lock_guard<std::mutex> Lock(ComputeUnitMapMutex);
  const std::string &ComputeUnitKind = Device.getComputeUnitKind();
  ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];

  CUI.TgtImageMap.erase(&Image);
  CUI.JITImages.erase(&Image);
}