/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include "device/pal/palubercapturemgr.hpp"
#include "device/pal/paldevice.hpp"
#include "device/pal/palvirtual.hpp"
#include "device/pal/palprogram.hpp"
#include "device/pal/palkernel.hpp"
#include "device/pal/palblit.hpp"

#include "palPlatform.h"
#include "palTraceSession.h"
#include "palRenderOpTraceController.h"
#include "palCodeObjectTraceSource.h"
#include "palQueueTimingsTraceSource.h"

#include "devDriverServer.h"
#include "protocols/driverControlServer.h"
#include "util/ddStructuredReader.h"

namespace amd::pal {

// ================================================================================================
// Returns true if the given device properties/settings support tracing.
static inline bool GpuSupportsTracing(const Pal::DeviceProperties& props,
                                      const Settings& settings) {
  return props.gfxipProperties.flags.supportRgpTraces && !settings.rgpSqttForceDisable_;
}

// ================================================================================================
// Creates the GPU Open Developer Mode manager class.
UberTraceCaptureMgr* UberTraceCaptureMgr::Create(Pal::IPlatform* platform, const Device& device) {
  UberTraceCaptureMgr* mgr = new UberTraceCaptureMgr(platform, device);

  if (mgr != nullptr && !mgr->Init(platform)) {
    delete mgr;
    mgr = nullptr;
  }

  return mgr;
}

static void PAL_STDCALL UberTraceStateChangeCallback(const GpuUtil::TraceSession& pTraceSession,
                                                     GpuUtil::TraceSessionState newState,
                                                     void* pPrivateData)
{
    UberTraceCaptureMgr* mgr = static_cast<UberTraceCaptureMgr*>(pPrivateData);

    switch (newState)
    {
        // boundary for prepare-phase dispatches
        case GpuUtil::TraceSessionState::Preparing:
        // boundary for detailed capture
        case GpuUtil::TraceSessionState::Running:
        // boundary for end of detailed trace
#if (PAL_CLIENT_INTERFACE_MAJOR_VERSION >= 939)
        case GpuUtil::TraceSessionState::Postamble:
#else
        // boundary for end of trace
        case GpuUtil::TraceSessionState::Waiting:
#endif
        {
            VirtualGPU* current_gpu = mgr->GetCurrentGPU();
            if (current_gpu != nullptr) {
                bool flush_success = current_gpu->queue(MainEngine).flush();
                assert(flush_success);
            }
            break;
        }
        default:
            break;
    }
}

// ================================================================================================
UberTraceCaptureMgr::UberTraceCaptureMgr(Pal::IPlatform* platform, const Device& device)
    : device_(device),
      dev_driver_server_(platform->GetDevDriverServer()),
      global_disp_count_(1),  // Must start from 1 according to RGP spec
      user_event_(nullptr),
      current_event_id_(0),
      trace_session_(platform->GetTraceSession()),
      trace_controller_(nullptr),
      code_object_trace_source_(nullptr),
      queue_timings_trace_source_(nullptr),
      registered_trace_state_callback_(false) {}

// ================================================================================================
UberTraceCaptureMgr::~UberTraceCaptureMgr() { DestroyUberTraceResources(); }

// ================================================================================================
bool UberTraceCaptureMgr::CreateUberTraceResources(Pal::IPlatform* platform) {
  bool success = false;

  do {
    // Create the user event RGP marker
    user_event_ = new RgpSqttMarkerUserEventWithString;
    if (user_event_ == nullptr) {
      break;
    }

    // Initialize the renderop trace controller
    trace_controller_ = new GpuUtil::RenderOpTraceController(platform, device_.iDev());
    if (trace_controller_ == nullptr) {
      break;
    }

    Pal::Result result = trace_session_->RegisterController(trace_controller_);
    if (result != Pal::Result::Success) {
      break;
    }

    // Initialize the code object trace source
    code_object_trace_source_ = new GpuUtil::CodeObjectTraceSource(platform);
    if (code_object_trace_source_ == nullptr) {
      break;
    }

    result = trace_session_->RegisterSource(code_object_trace_source_);
    if (result != Pal::Result::Success) {
      break;
    }

    // Initialize the queue timings trace source
    queue_timings_trace_source_ = new GpuUtil::QueueTimingsTraceSource(platform);
    if (queue_timings_trace_source_ == nullptr) {
      break;
    }

    result = trace_session_->RegisterSource(queue_timings_trace_source_);
    if (result != Pal::Result::Success) {
      break;
    }

    result = trace_session_->RegisterTraceStateChangeCallback(UberTraceStateChangeCallback, this);
    if (result != Pal::Result::Success) {
      break;
    }

    registered_trace_state_callback_ = true;

    success = true;
  } while (false);

  return success;
}

// ================================================================================================
void UberTraceCaptureMgr::DestroyUberTraceResources() {
  // RGP user event marker
  if (user_event_ != nullptr) {
    delete user_event_;
    user_event_ = nullptr;
  }

  // RenderOp TraceController
  if (trace_controller_ != nullptr) {
    trace_session_->UnregisterController(trace_controller_);
    delete trace_controller_;
    trace_controller_ = nullptr;
  }

  // CodeObjects TraceSource
  if (code_object_trace_source_ != nullptr) {
    trace_session_->UnregisterSource(code_object_trace_source_);
    delete code_object_trace_source_;
    code_object_trace_source_ = nullptr;
  }

  // QueueTimings TraceSource
  if (queue_timings_trace_source_ != nullptr) {
    trace_session_->UnregisterSource(queue_timings_trace_source_);
    delete queue_timings_trace_source_;
    queue_timings_trace_source_ = nullptr;
  }

  if (registered_trace_state_callback_) {
    trace_session_->UnregisterTraceStateChangeCallback(UberTraceStateChangeCallback, this);
    registered_trace_state_callback_ = false;
  }
}

// ================================================================================================
bool UberTraceCaptureMgr::Init(Pal::IPlatform* platform) {
  // Finalize the devmode manager
  if (dev_driver_server_ == nullptr) {
    return false;
  }
  dev_driver_server_->Finalize();

  // Initialize the trace sources & controllers owned by the compute driver
  const bool success = CreateUberTraceResources(platform);

  if (!success) {
    DestroyUberTraceResources();
    return false;
  }

  return true;
}

// ================================================================================================
void UberTraceCaptureMgr::PreDispatch(VirtualGPU* gpu, const pal::Kernel& kernel, size_t x,
                                      size_t y, size_t z) {
  // Wait for the driver to be resumed in case it's been paused.
  WaitForDriverResume();

  // Increment dispatch count in RenderOp trace controller
  Pal::IQueue* pQueue = gpu->queue(MainEngine).iQueue_;
  GpuUtil::RenderOpCounts opCounts = {
      .dispatchCount = 1u,
  };
  current_gpu_ = gpu;
  trace_controller_->RecordRenderOps(pQueue, opCounts);
  current_gpu_ = nullptr;

  if (trace_session_->GetTraceSessionState() == GpuUtil::TraceSessionState::Running) {
    RgpSqttMarkerEventType apiEvent = RgpSqttMarkerEventType::CmdNDRangeKernel;

    if (kernel.prog().isInternal()) {
      constexpr RgpSqttMarkerEventType ApiEvents[KernelBlitManager::BlitTotal] = {
          RgpSqttMarkerEventType::CmdCopyImage,
          RgpSqttMarkerEventType::CmdCopyImage,
          RgpSqttMarkerEventType::CmdCopyImageToBuffer,
          RgpSqttMarkerEventType::CmdCopyBufferToImage,
          RgpSqttMarkerEventType::CmdCopyBuffer,
          RgpSqttMarkerEventType::CmdCopyBuffer,
          RgpSqttMarkerEventType::CmdCopyBuffer,
          RgpSqttMarkerEventType::CmdCopyBuffer,
          RgpSqttMarkerEventType::CmdFillBuffer,
          RgpSqttMarkerEventType::CmdFillImage,
          RgpSqttMarkerEventType::CmdScheduler};

      for (uint i = 0; i < KernelBlitManager::BlitTotal; ++i) {
        if (kernel.name().compare(BlitName[i]) == 0) {
          apiEvent = ApiEvents[i];
          break;
        }
      }
    }

    // Write the hash value
    WriteComputeBindMarker(gpu, kernel.prog().ApiHash());

    // Write dispatch marker
    WriteEventWithDimsMarker(gpu, apiEvent, static_cast<uint32_t>(x), static_cast<uint32_t>(y),
                             static_cast<uint32_t>(z));
  }

  // Increment the global dispatch counter
  global_disp_count_++;
}

// ================================================================================================
void UberTraceCaptureMgr::PostDispatch(VirtualGPU* gpu) {}

// ================================================================================================
// Waits for the driver to be resumed if it's currently paused.
void UberTraceCaptureMgr::WaitForDriverResume() {
  auto* pDriverControlServer = dev_driver_server_->GetDriverControlServer();

  assert(pDriverControlServer != nullptr);
  pDriverControlServer->DriverTick();
}

// ================================================================================================
void UberTraceCaptureMgr::PreDeviceDestroy() {
  if (trace_session_->GetTraceSessionState() == GpuUtil::TraceSessionState::Ready) {
    DestroyUberTraceResources();
  }
}

// ================================================================================================
void UberTraceCaptureMgr::FinishRGPTrace(VirtualGPU* gpu, bool aborted) {
  // Nothing to be done
}

// ================================================================================================
bool UberTraceCaptureMgr::IsQueueTimingActive() const {
  return ((queue_timings_trace_source_ != nullptr) &&
          (queue_timings_trace_source_->IsTimingInProgress()));
}

// ================================================================================================
bool UberTraceCaptureMgr::RegisterTimedQueue(uint32_t queue_id, Pal::IQueue* iQueue,
                                             bool* debug_vmid) const {
  // Get the OS context handle for this queue (this is a thing that RGP needs on DX clients;
  // it may be optional for Vulkan, but we provide it anyway if available).
  Pal::KernelContextInfo kernelContextInfo = {};
  Pal::Result result = iQueue->QueryKernelContextInfo(&kernelContextInfo);

  // QueryKernelContextInfo may fail.
  // If so, just use a context identifier of 0.
  uint64_t queueContext =
      (result == Pal::Result::Success) ? kernelContextInfo.contextIdentifier : 0;

  // Register the queue with the GPA session class for timed queue operation support.
  result = queue_timings_trace_source_->RegisterTimedQueue(iQueue, queue_id, queueContext);

  return (result == Pal::Result::Success);
}

// ================================================================================================
Pal::Result UberTraceCaptureMgr::TimedQueueSubmit(Pal::IQueue* queue, uint64_t cmdId,
                                                  const Pal::SubmitInfo& submitInfo) const {
  // Fill in extra meta-data information to associate the API command buffer data with
  // the generated timing information.
  GpuUtil::TimedSubmitInfo timedSubmitInfo = {};
  Pal::uint64 apiCmdBufIds = cmdId;
  Pal::uint32 sqttCmdBufIds = 0;

  timedSubmitInfo.pApiCmdBufIds = &apiCmdBufIds;
  timedSubmitInfo.pSqttCmdBufIds = &sqttCmdBufIds;
  timedSubmitInfo.frameIndex = 0;

  Pal::Result result = Pal::Result::Success;
  if (IsQueueTimingActive()) {
    result = queue_timings_trace_source_->TimedSubmit(queue, submitInfo, timedSubmitInfo);
    if (result != Pal::Result::Success) {
      result = queue->Submit(submitInfo);
    }
  } else {
    result = queue->Submit(submitInfo);
  }

  return result;
}

// ================================================================================================
bool UberTraceCaptureMgr::Update(Pal::IPlatform* platform) {
  Pal::Result result = queue_timings_trace_source_->Init(device_.iDev());
  return (result == Pal::Result::Success);
}

// ================================================================================================
uint64_t UberTraceCaptureMgr::AddElfBinary(const void* exe_binary, size_t exe_binary_size,
                                           const void* elf_binary, size_t elf_binary_size,
                                           Pal::IGpuMemory* pGpuMemory, size_t offset) {
  GpuUtil::ElfBinaryInfo elfBinaryInfo = {};
  elfBinaryInfo.pBinary = exe_binary;
  elfBinaryInfo.binarySize = exe_binary_size;  ///< FAT Elf binary size.
  elfBinaryInfo.pGpuMemory = pGpuMemory;       ///< GPU Memory where the compiled ISA resides.
  elfBinaryInfo.offset = static_cast<Pal::gpusize>(offset);

  elfBinaryInfo.originalHash = DevDriver::MetroHash::MetroHash64(
      reinterpret_cast<const DevDriver::uint8*>(elf_binary), elf_binary_size);

  elfBinaryInfo.compiledHash = DevDriver::MetroHash::MetroHash64(
      reinterpret_cast<const DevDriver::uint8*>(exe_binary), exe_binary_size);

  assert(code_object_trace_source_ != nullptr);

  code_object_trace_source_->RegisterElfBinary(elfBinaryInfo);

  return elfBinaryInfo.originalHash;
}

// ================================================================================================
void UberTraceCaptureMgr::WriteMarker(const VirtualGPU* gpu, const void* data,
                                      size_t data_size) const {
  assert((data_size % sizeof(uint32_t)) == 0);
  assert((data_size / sizeof(uint32_t)) > 0);

  Pal::RgpMarkerSubQueueFlags subQueueFlags = {};
  subQueueFlags.includeMainSubQueue = 1;

  gpu->queue(MainEngine)
      .iCmd()
      ->CmdInsertRgpTraceMarker(subQueueFlags, static_cast<uint32_t>(data_size / sizeof(uint32_t)),
                                data);
}

// ================================================================================================
// Inserts a compute bind marker
void UberTraceCaptureMgr::WriteComputeBindMarker(const VirtualGPU* gpu, uint64_t api_hash) const {
  RgpSqttMarkerPipelineBind marker = {};
  marker.identifier = RgpSqttMarkerIdentifierBindPipeline;
  marker.cbID = gpu->queue(MainEngine).cmdBufId();
  marker.bindPoint = 1;

  memcpy(marker.apiPsoHash, &api_hash, sizeof(api_hash));
  WriteMarker(gpu, &marker, sizeof(marker));
}

// ================================================================================================
// Inserts an RGP pre-dispatch marker
void UberTraceCaptureMgr::WriteEventWithDimsMarker(const VirtualGPU* gpu,
                                                   RgpSqttMarkerEventType apiType, uint32_t x,
                                                   uint32_t y, uint32_t z) const {
  assert(apiType != RgpSqttMarkerEventType::Invalid);

  RgpSqttMarkerEvent event = {};
  event.identifier = RgpSqttMarkerIdentifierEvent;
  event.apiType = static_cast<uint32_t>(apiType);
  event.cmdID = current_event_id_++;
  event.cbID = gpu->queue(MainEngine).cmdBufId();

  RgpSqttMarkerEventWithDims eventWithDims = {};
  eventWithDims.event = event;
  eventWithDims.event.hasThreadDims = 1;
  eventWithDims.threadX = x;
  eventWithDims.threadY = y;
  eventWithDims.threadZ = z;

  WriteMarker(gpu, &eventWithDims, sizeof(eventWithDims));
}

// ================================================================================================
void UberTraceCaptureMgr::WriteBarrierStartMarker(const VirtualGPU* gpu,
                                                  const Pal::Developer::BarrierData& data) const {
  if (trace_session_->GetTraceSessionState() == GpuUtil::TraceSessionState::Running) {
    amd::ScopedLock traceLock(&trace_mutex_);

    RgpSqttMarkerBarrierStart marker = {};
    marker.cbId = gpu->queue(MainEngine).cmdBufId();
    marker.identifier = RgpSqttMarkerIdentifierBarrierStart;
    marker.internal = true;
    marker.dword02 = data.reason;

    WriteMarker(gpu, &marker, sizeof(marker));
  }
}

// ================================================================================================
void UberTraceCaptureMgr::WriteBarrierEndMarker(const VirtualGPU* gpu,
                                                const Pal::Developer::BarrierData& data) const {
  if (trace_session_->GetTraceSessionState() == GpuUtil::TraceSessionState::Running) {
    amd::ScopedLock traceLock(&trace_mutex_);

    // Copy the operations part and include the same data from previous markers
    // within the same barrier sequence to create a full picture of all cache
    // syncs and pipeline stalls.
    Pal::Developer::BarrierOperations operations = data.operations;
    operations.pipelineStalls.u16All |= 0;
    operations.caches.u16All |= 0;

    RgpSqttMarkerBarrierEnd marker = {};
    marker.identifier = RgpSqttMarkerIdentifierBarrierEnd;
    marker.cbId = gpu->queue(MainEngine).cmdBufId();
    marker.numLayoutTransitions = 0;
    marker.waitOnEopTs = operations.pipelineStalls.eopTsBottomOfPipe;
    marker.vsPartialFlush = operations.pipelineStalls.vsPartialFlush;
    marker.psPartialFlush = operations.pipelineStalls.psPartialFlush;
    marker.csPartialFlush = operations.pipelineStalls.csPartialFlush;
    marker.pfpSyncMe = operations.pipelineStalls.pfpSyncMe;
    marker.syncCpDma = operations.pipelineStalls.syncCpDma;
    marker.invalTcp = operations.caches.invalTcp;
    marker.invalSqI = operations.caches.invalSqI$;
    marker.invalSqK = operations.caches.invalSqK$;
    marker.flushTcc = operations.caches.flushTcc;
    marker.invalTcc = operations.caches.invalTcc;
    marker.flushCb = operations.caches.flushCb;
    marker.invalCb = operations.caches.invalCb;
    marker.flushDb = operations.caches.flushDb;
    marker.invalDb = operations.caches.invalDb;

    WriteMarker(gpu, &marker, sizeof(marker));
  }
}

}  // namespace amd::pal