#include "Profiler/Cupti/CuptiProfiler.h"
#include "Context/Context.h"
#include "Data/Metric.h"
#include "Device.h"
#include "Driver/GPU/CudaApi.h"
#include "Driver/GPU/CuptiApi.h"
#include "Driver/GPU/NvtxApi.h"
#include "Profiler/Cupti/CuptiPCSampling.h"
#include "Profiler/Graph.h"
#include "Runtime/CudaRuntime.h"
#include "Utility/Env.h"
#include "Utility/Errors.h"
#include "Utility/Map.h"
#include "Utility/String.h"
#include "Utility/Vector.h"

#include <algorithm>
#include <array>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <unordered_map>
#include <unordered_set>

namespace proton {

template <>
thread_local GPUProfiler<CuptiProfiler>::ThreadState
    GPUProfiler<CuptiProfiler>::threadState(CuptiProfiler::instance());

namespace {

std::unique_ptr<Metric>
convertKernelActivityToMetric(CUpti_Activity *activity,
                              bool isMetricKernel = false) {
  std::unique_ptr<Metric> metric;
  auto *kernel = reinterpret_cast<CUpti_ActivityKernel5 *>(activity);
  if (kernel->start < kernel->end) {
    metric =
        std::make_unique<KernelMetric>(static_cast<uint64_t>(kernel->start),
                                       static_cast<uint64_t>(kernel->end), 1,
                                       static_cast<uint64_t>(kernel->deviceId),
                                       static_cast<uint64_t>(DeviceType::CUDA),
                                       static_cast<uint64_t>(kernel->streamId),
                                       static_cast<uint64_t>(isMetricKernel));
  } // else: not a valid kernel activity
  return metric;
}

uint32_t processActivityKernel(
    CuptiProfiler::CorrIdToExternIdMap &corrIdToExternId,
    CuptiProfiler::ExternIdToStateMap &externIdToState,
    std::map<uint64_t, std::reference_wrapper<CuptiProfiler::ExternIdState>>
        &externIdToStateCache,
    DataPhases &dataPhases, CUpti_Activity *activity) {
  // Support CUDA >= 11.0
  auto *kernel = reinterpret_cast<CUpti_ActivityKernel5 *>(activity);
  auto correlationId = kernel->correlationId;
  size_t externId = 0;
  if (!/*not valid*/ corrIdToExternId.withRead(
          correlationId, [&externId](size_t value) { externId = value; })) {
    corrIdToExternId.erase(correlationId);
    return correlationId;
  }
  if (kernel->graphId == 0) { // XXX: This is a misnomer confirmed by NVIDIA,
                              // actually it refers to graphExecId
    // Non-graph kernels
    bool isMissingName = false;
    DataToEntryMap dataToEntry;
    externIdToState.withRead(externId,
                             [&](const CuptiProfiler::ExternIdState &state) {
                               isMissingName = state.isMissingName;
                               dataToEntry = state.dataToEntry;
                             });
    if (!isMissingName) {
      for (auto &[data, entry] : dataToEntry) {
        if (auto kernelMetric = convertKernelActivityToMetric(activity)) {
          entry.upsertMetric(std::move(kernelMetric));
          detail::updateDataPhases(dataPhases, data, entry.phase);
        }
      }
    } else {
      for (auto &[data, entry] : dataToEntry) {
        if (auto kernelMetric = convertKernelActivityToMetric(activity)) {
          auto childEntry =
              data->addOp(entry.phase, entry.id, {Context(kernel->name)});
          childEntry.upsertMetric(std::move(kernelMetric));
          detail::updateDataPhases(dataPhases, data, entry.phase);
        }
      }
    }
    externIdToState.erase(externId);
    corrIdToExternId.erase(correlationId);
  } else {
    // Graph kernels
    // A single graph launch can trigger multiple kernels.
    // Our solution is to construct the following maps:
    // --- Application threads ---
    // If graph creation has been captured:
    // - parentId, nodeId -> launch context + capture context
    // Otherwise:
    // - parentId -> launch context
    // --- CUPTI thread ---
    // - corrId -> numNodes
    auto iter = externIdToStateCache.find(externId);
    CuptiProfiler::ExternIdState *state = nullptr;
    if (iter != externIdToStateCache.end()) {
      state = &iter->second.get();
    } else {
      // Cache miss, fetch from the main map
      auto ref = externIdToState.find(externId);
      // Update the cache
      externIdToStateCache.emplace(externId, ref.value());
      state = &ref.value().get();
    }
    auto &externState = *state;
    // We have a graph creation captured
    auto *nodeIdToState = externState.nodeIdToState;
    if (nodeIdToState) {
      const GraphState::NodeState &nodeState = nodeIdToState->at(
          kernel->graphNodeId); // nodeIdToState must have the nodeId
      if (nodeState.status.isMissingName()) {
        throw makeLogicError("Kernel name is missing for a graph node.");
      }
      const bool isMetricKernel = nodeState.status.isMetricNode();
      for (auto &[data, entry] : externState.dataToGraphEntry) {
        auto targetEntryIdIter = nodeState.dataToEntryId.find(data);
        if (targetEntryIdIter != nodeState.dataToEntryId.end()) {
          auto targetEntryId = targetEntryIdIter->second;
          if (auto kernelMetric =
                  convertKernelActivityToMetric(activity, isMetricKernel)) {
            entry.upsertLinkedMetric(std::move(kernelMetric), targetEntryId);
            detail::updateDataPhases(dataPhases, data, entry.phase);
          }
        }
      } // else the profiling session has been deactivated during graph
        // capture when encountering this node
    } else {
      // This can happen when graph creation is not captured.
      // Since we don't have per-node info, we just attach the kernel metric to
      // the graph launch entry without creating a child entry for the node.
      for (auto &[data, entry] : externState.dataToEntry) {
        if (auto kernelMetric = convertKernelActivityToMetric(activity)) {
          auto childEntry =
              data->addOp(entry.phase, entry.id, {Context(kernel->name)});
          childEntry.upsertMetric(std::move(kernelMetric));
          detail::updateDataPhases(dataPhases, data, childEntry.phase);
        }
      }
    }
    // Decrease the expected kernel count
    if (externState.numNodes > 0) {
      externState.numNodes--;
    }
    // If all kernels have been processed, clean up
    if (externState.numNodes == 0) {
      externIdToState.erase(externId);
      corrIdToExternId.erase(correlationId);
    }
  }
  return correlationId;
}

uint32_t processActivity(
    CuptiProfiler::CorrIdToExternIdMap &corrIdToExternId,
    CuptiProfiler::ExternIdToStateMap &externIdToState,
    std::map<uint64_t, std::reference_wrapper<CuptiProfiler::ExternIdState>>
        &externIdToStateCache,
    DataPhases &dataPhases, CUpti_Activity *activity) {
  auto correlationId = 0;
  switch (activity->kind) {
  case CUPTI_ACTIVITY_KIND_KERNEL:
  case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: {
    correlationId =
        processActivityKernel(corrIdToExternId, externIdToState,
                              externIdToStateCache, dataPhases, activity);
    break;
  }
  default:
    break;
  }
  return correlationId;
}

CuptiProfiler::ExternIdState *
buildGraphNodeEntries(const DataToEntryMap &dataToEntry, GraphState &graphState,
                      CuptiProfiler::ExternIdToStateMap &externIdToState,
                      size_t externId) {
  if (dataToEntry.empty())
    return nullptr;

  auto &externIdState = externIdToState[externId];
  for (auto &[data, entry] : dataToEntry) {
    auto nodeStateIt = graphState.dataToEntryIdToNodeStates.find(data);
    if (nodeStateIt == graphState.dataToEntryIdToNodeStates.end())
      // This is a new data which was not enabled during graph capture
      continue;
    externIdState.dataToGraphEntry.insert({data, entry});
  }
  externIdState.nodeIdToState = &graphState.nodeIdToState;
  return &externIdState;
}

void queueGraphMetrics(PendingGraphPool *pendingGraphPool,
                       const CUpti_CallbackData *callbackData,
                       const GraphState &graphState,
                       const CuptiProfiler::ExternIdState *externIdState) {
  if (graphState.metricSeqIdToNodeId.empty()) {
    return;
  }
  PendingGraphQueue::SeqIdToStateMap seqIdToState;
  size_t phase = Data::kNoCompletePhase;
  for (const auto &[seqId, nodeId] : graphState.metricSeqIdToNodeId) {
    const auto &metricNodeState = graphState.metricNodeIdToState.at(nodeId);
    if (externIdState) {
      for (const auto &[data, graphEntry] : externIdState->dataToGraphEntry) {
        phase = graphEntry.phase;
        auto &pendingMetricNode =
            seqIdToState
                .emplace(seqId,
                         PendingGraphQueue::MetricNodeState{
                             metricNodeState.metricId, {}})
                .first->second;
        auto entryId = Scope::DummyScopeId;
        if (auto entryIdIter = metricNodeState.dataToEntryId.find(data);
            entryIdIter != metricNodeState.dataToEntryId.end()) {
          entryId = entryIdIter->second;
        }
        // Otherwise, a DummyScopeId entry indicates that we'll call
        // upsertFlexibleMetric instead of upsertLinkedFlexibleMetric in
        // queueGraphMetrics, so that the flexible metric can be attached to the
        // graph launch entry.
        pendingMetricNode.dataToEntry.emplace(
            data, DataEntry(entryId, phase, graphEntry.metricSet.get()));
      }
    }
  }
  // Whether a data is active or not, the GPU will write the metric data into
  // the metric buffer if there is a metric node. So we need the complete number
  // of metrics of a graph.
  if (callbackData->context != nullptr)
    pendingGraphPool->flushIfNeeded(graphState.numMetricWords);
  pendingGraphPool->push(phase, graphState.numMetricWords,
                         std::move(seqIdToState));
}

constexpr std::array<CUpti_CallbackId, 11> kGraphCallbacks = {
    CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch,
    CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch_ptsz,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_ptsz,
    CUPTI_DRIVER_TRACE_CBID_cuStreamEndCapture,
    CUPTI_DRIVER_TRACE_CBID_cuStreamEndCapture_ptsz,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_v2,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_v2_ptsz,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCaptureToGraph,
    CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCaptureToGraph_ptsz,
    CUPTI_DRIVER_TRACE_CBID_cuStreamEndCapture};

#define PROTON_KERNEL_CALLBACK_LIST(X)                                         \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunch)                                          \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchGrid)                                      \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchGridAsync)                                 \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchKernel)                                    \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchKernel_ptsz)                               \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchKernelEx)                                  \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchKernelEx_ptsz)                             \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchCooperativeKernel)                         \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchCooperativeKernel_ptsz)                    \
  X(CUPTI_DRIVER_TRACE_CBID_cuLaunchCooperativeKernelMultiDevice)

#define PROTON_KERNEL_CB_AS_ID(cbId) cbId,
constexpr std::array<CUpti_CallbackId, 10> kKernelCallbacks = {
    PROTON_KERNEL_CALLBACK_LIST(PROTON_KERNEL_CB_AS_ID)};
#undef PROTON_KERNEL_CB_AS_ID

constexpr std::array<CUpti_CallbackId, 2> kGraphResourceCallbacks = {
    CUPTI_CBID_RESOURCE_GRAPHNODE_CREATED,
    CUPTI_CBID_RESOURCE_GRAPHNODE_CLONED};

constexpr std::array<CUpti_CallbackId, 4> kResourceCallbacks = {
    CUPTI_CBID_RESOURCE_MODULE_LOADED,
    CUPTI_CBID_RESOURCE_MODULE_UNLOAD_STARTING,
    CUPTI_CBID_RESOURCE_CONTEXT_CREATED,
    CUPTI_CBID_RESOURCE_CONTEXT_DESTROY_STARTING,
};

constexpr std::array<CUpti_CallbackId, 2> kNvtxCallbacks = {
    CUPTI_CBID_NVTX_nvtxRangePushA,
    CUPTI_CBID_NVTX_nvtxRangePop,
};

void setLaunchCallbacks(CUpti_SubscriberHandle subscriber, bool enable) {
  for (auto cbId : kKernelCallbacks) {
    cupti::enableCallback<true>(static_cast<uint32_t>(enable), subscriber,
                                CUPTI_CB_DOMAIN_DRIVER_API, cbId);
  }
}

void setGraphCallbacks(CUpti_SubscriberHandle subscriber, bool enable) {
  for (auto cbId : kGraphCallbacks) {
    cupti::enableCallback<true>(static_cast<uint32_t>(enable), subscriber,
                                CUPTI_CB_DOMAIN_DRIVER_API, cbId);
  }
  for (auto cbId : kGraphResourceCallbacks) {
    cupti::enableCallback<true>(static_cast<uint32_t>(enable), subscriber,
                                CUPTI_CB_DOMAIN_RESOURCE, cbId);
  }
}

void setResourceCallbacks(CUpti_SubscriberHandle subscriber, bool enable) {
  for (auto cbId : kResourceCallbacks) {
    cupti::enableCallback<true>(static_cast<uint32_t>(enable), subscriber,
                                CUPTI_CB_DOMAIN_RESOURCE, cbId);
  }
}

void setNvtxCallbacks(CUpti_SubscriberHandle subscriber, bool enable) {
  for (auto cbId : kNvtxCallbacks) {
    cupti::enableCallback<true>(static_cast<uint32_t>(enable), subscriber,
                                CUPTI_CB_DOMAIN_NVTX, cbId);
  }
}

bool isKernel(CUpti_CallbackId cbId) {
  switch (cbId) {
#define PROTON_KERNEL_CB_AS_CASE(cbId)                                         \
  case cbId:                                                                   \
    return true;
    PROTON_KERNEL_CALLBACK_LIST(PROTON_KERNEL_CB_AS_CASE)
#undef PROTON_KERNEL_CB_AS_CASE
  default:
    return false;
  }
}

bool isGraphLaunch(CUpti_CallbackId cbId) {
  return cbId == CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuGraphLaunch_ptsz;
}

bool isLaunch(CUpti_CallbackId cbId) {
  return isKernel(cbId) || isGraphLaunch(cbId);
}

bool isStreamCaptureBegin(CUpti_CallbackId cbId) {
  return cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_ptsz ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_v2 ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCapture_v2_ptsz ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCaptureToGraph ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamBeginCaptureToGraph_ptsz;
}

bool isStreamCaptureEnd(CUpti_CallbackId cbId) {
  return cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamEndCapture ||
         cbId == CUPTI_DRIVER_TRACE_CBID_cuStreamEndCapture_ptsz;
}

#undef PROTON_KERNEL_CALLBACK_LIST

} // namespace

struct CuptiProfiler::CuptiProfilerPimpl
    : public GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface {
  CuptiProfilerPimpl(CuptiProfiler &profiler)
      : GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface(profiler) {
    auto runtime = &CudaRuntime::instance();
    profiler.metricBuffer = std::make_unique<MetricBuffer>(
        getIntEnv("TRITON_PROFILE_METRIC_BUFFER_SIZE", 64 * 1024 * 1024),
        runtime, /*mapped=*/true);
    profiler.pendingGraphPool =
        std::make_unique<PendingGraphPool>(profiler.metricBuffer.get());
  }
  virtual ~CuptiProfilerPimpl() = default;

  void doStart() override;
  void doFlush() override;
  void doStop() override;

  static void allocBuffer(uint8_t **buffer, size_t *bufferSize,
                          size_t *maxNumRecords);
  static void completeBuffer(CUcontext context, uint32_t streamId,
                             uint8_t *buffer, size_t size, size_t validSize);
  static void callbackFn(void *userData, CUpti_CallbackDomain domain,
                         CUpti_CallbackId cbId, const void *cbData);

  static constexpr size_t AlignSize = 8;
  static constexpr size_t AttributeSize = sizeof(size_t);
  static constexpr const char *CaptureTag = "<captured_at>";

  CUpti_SubscriberHandle subscriber{};
  CuptiPCSampling pcSampling;

  ThreadSafeMap<uint32_t, GraphState> graphStates;

private:
  void handleGraphResourceCallbacks(CuptiProfiler &profiler,
                                    CUpti_CallbackId cbId,
                                    CUpti_GraphData *graphData);
  void handleResourceCallbacks(CuptiProfiler &profiler, CUpti_CallbackId cbId,
                               const void *cbData);
  void handleNvtxCallbacks(CUpti_CallbackId cbId, const void *cbData);

  void handleStreamCaptureCallbacks(CUpti_CallbackId cbId,
                                    const CUpti_CallbackData *callbackData);
  void handleApiEnterLaunchCallbacks(CuptiProfiler &profiler,
                                     CUpti_CallbackId cbId,
                                     const CUpti_CallbackData *callbackData);
  void handleApiExitLaunchCallbacks(CuptiProfiler &profiler,
                                    CUpti_CallbackId cbId,
                                    const CUpti_CallbackData *callbackData);
  void handleApiCallbacks(CuptiProfiler &profiler, CUpti_CallbackId cbId,
                          const void *cbData);
};

void CuptiProfiler::CuptiProfilerPimpl::allocBuffer(uint8_t **buffer,
                                                    size_t *bufferSize,
                                                    size_t *maxNumRecords) {
  const auto envBufferSize =
      getIntEnv("TRITON_PROFILE_BUFFER_SIZE", 64 * 1024 * 1024);
  *buffer = static_cast<uint8_t *>(aligned_alloc(AlignSize, envBufferSize));
  if (*buffer == nullptr) {
    throw makeRuntimeError("aligned_alloc failed");
  }
  *bufferSize = envBufferSize;
  *maxNumRecords = 0;
}

void CuptiProfiler::CuptiProfilerPimpl::completeBuffer(CUcontext ctx,
                                                       uint32_t streamId,
                                                       uint8_t *buffer,
                                                       size_t size,
                                                       size_t validSize) {
  CuptiProfiler &profiler = threadState.profiler;
  uint32_t maxCorrelationId = 0;
  DataPhases dataPhases;
  CUptiResult status;
  CUpti_Activity *activity = nullptr;
  std::map<uint64_t, std::reference_wrapper<CuptiProfiler::ExternIdState>>
      externIdToStateCache;
  do {
    status = cupti::activityGetNextRecord<false>(buffer, validSize, &activity);
    if (status == CUPTI_SUCCESS) {
      auto correlationId =
          processActivity(profiler.correlation.corrIdToExternId,
                          profiler.correlation.externIdToState,
                          externIdToStateCache, dataPhases, activity);
      maxCorrelationId = std::max(maxCorrelationId, correlationId);
    } else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) {
      break;
    } else {
      throw makeRuntimeError("cupti::activityGetNextRecord failed");
    }
  } while (true);

  std::free(buffer);

  profiler.correlation.complete(maxCorrelationId);
  profiler.flushDataPhases(dataPhases, profiler.pendingGraphPool.get());
}

void CuptiProfiler::CuptiProfilerPimpl::handleGraphResourceCallbacks(
    CuptiProfiler &profiler, CUpti_CallbackId cbId,
    CUpti_GraphData *graphData) {
  uint32_t graphId = 0;
  uint32_t graphExecId = 0;
  if (graphData->graph)
    cupti::getGraphId<true>(graphData->graph, &graphId);
  if (graphData->graphExec)
    cupti::getGraphExecId<true>(graphData->graphExec, &graphExecId);
  if (cbId == CUPTI_CBID_RESOURCE_GRAPHNODE_CREATED ||
      cbId == CUPTI_CBID_RESOURCE_GRAPHNODE_CLONED) {
    if (graphData->nodeType != CU_GRAPH_NODE_TYPE_KERNEL) {
      // We only care about kernel nodes
      return;
    }
    uint64_t nodeId = 0;
    cupti::getGraphNodeId<true>(graphData->node, &nodeId);
    if (cbId == CUPTI_CBID_RESOURCE_GRAPHNODE_CREATED) {
      if (!profiler.isOpInProgress()) {
        // else no op in progress; creation triggered by graph
        // clone/instantiate, we don't increase the numNodes because the
        // original graph has already accounted for it
        return;
      }
      auto &graphState = graphStates[graphId];
      auto &nodeState = graphState.nodeIdToState[nodeId];
      nodeState.nodeId = nodeId;
      const auto &name = threadState.scopeStack.back().name;
      if (name.empty())
        nodeState.status.setMissingName();
      const bool isMetricKernelNode = threadState.isMetricKernelLaunching;
      if (isMetricKernelNode) {
        nodeState.status.setMetricNode();
        auto metricKernelLaunchInfo =
            threadState.metricKernelLaunchInfoQueue.front();
        threadState.metricKernelLaunchInfoQueue.pop_front();
        graphState.metricNodeIdToState.insert_or_assign(
            nodeId,
            GraphState::MetricNodeState{metricKernelLaunchInfo.seqId,
                                        metricKernelLaunchInfo.metricId,
                                        metricKernelLaunchInfo.numWords});
        graphState.metricSeqIdToNodeId.insert_or_assign(
            metricKernelLaunchInfo.seqId, nodeId);
        graphState.numMetricWords += metricKernelLaunchInfo.numWords;
      }
      for (auto *data : profiler.dataSet) {
        auto currentContexts = data->getContexts();
        std::vector<Context> contexts;
        contexts.emplace_back(GraphState::captureTag);
        for (const auto &context : currentContexts) {
          contexts.push_back(context);
        }
        if (isMetricKernelNode) {
          auto flexibleMetricContexts = data->getContexts(false);
          std::vector<Context> flexibleMetricEntryContexts;
          flexibleMetricEntryContexts.emplace_back(GraphState::captureTag);
          for (const auto &context : flexibleMetricContexts) {
            flexibleMetricEntryContexts.push_back(context);
          }
          if (!threadState.isApiExternOp) { // Triton ops
            flexibleMetricEntryContexts.emplace_back(name);
          }
          contexts.emplace_back(GraphState::metricTag);
          flexibleMetricEntryContexts.emplace_back(GraphState::metricTag);

          // For metrics nodes, timing info is attributed to a frame under
          // a metadata state.
          // Flexible metrics are attributed to the current GPU operation
          auto staticEntry =
              data->addOp(Data::kVirtualPhase, Data::kRootEntryId, contexts);
          nodeState.dataToEntryId.insert_or_assign(data, staticEntry.id);
          graphState.dataToEntryIdToNodeStates[data][staticEntry.id].insert(
              &nodeState);
          auto flexibleMetricEntry =
              data->addOp(Data::kVirtualPhase, Data::kRootEntryId,
                          flexibleMetricEntryContexts);
          graphState.metricNodeIdToState.at(nodeId)
              .dataToEntryId.insert_or_assign(data, flexibleMetricEntry.id);
        } else {
          contexts.emplace_back(name);
          auto staticEntry =
              data->addOp(Data::kVirtualPhase, Data::kRootEntryId, contexts);
          nodeState.dataToEntryId.insert_or_assign(data, staticEntry.id);
          graphState.dataToEntryIdToNodeStates[data][staticEntry.id].insert(
              &nodeState);
        }
      }
    } else { // CUPTI_CBID_RESOURCE_GRAPHNODE_CLONED
      // When a graph is cloned under the stream capture mode, graphId is the
      // same as the graphExecId to be created
      uint32_t originalGraphId = 0;
      uint64_t originalNodeId = 0;
      cupti::getGraphId<true>(graphData->originalGraph, &originalGraphId);
      cupti::getGraphNodeId<true>(graphData->originalNode, &originalNodeId);
      auto &originalGraphState = graphStates[originalGraphId];
      auto &graphState = graphStates[graphId];
      graphState.nodeIdToState[nodeId] =
          originalGraphState.nodeIdToState[originalNodeId];
      auto &nodeState = graphState.nodeIdToState[nodeId];
      nodeState.nodeId = nodeId;
      for (const auto &[data, entryId] : nodeState.dataToEntryId) {
        graphState.dataToEntryIdToNodeStates[data][entryId].insert(&nodeState);
      }
      auto originalMetricNodeIt =
          originalGraphState.metricNodeIdToState.find(originalNodeId);
      if (originalMetricNodeIt !=
          originalGraphState.metricNodeIdToState.end()) {
        auto metricNodeState = originalMetricNodeIt->second;
        graphState.metricNodeIdToState.insert_or_assign(nodeId,
                                                        metricNodeState);
        graphState.metricSeqIdToNodeId.insert_or_assign(metricNodeState.seqId,
                                                        nodeId);
        graphState.numMetricWords += metricNodeState.numWords;
      }
    }
  }
}

void CuptiProfiler::CuptiProfilerPimpl::handleResourceCallbacks(
    CuptiProfiler &profiler, CUpti_CallbackId cbId, const void *cbData) {
  auto *resourceData =
      static_cast<CUpti_ResourceData *>(const_cast<void *>(cbData));
  if (cbId == CUPTI_CBID_RESOURCE_MODULE_LOADED) {
    auto *moduleResource = static_cast<CUpti_ModuleResourceData *>(
        resourceData->resourceDescriptor);
    if (profiler.pcSamplingEnabled)
      pcSampling.loadModule(moduleResource->pCubin, moduleResource->cubinSize);
  } else if (cbId == CUPTI_CBID_RESOURCE_MODULE_UNLOAD_STARTING) {
    auto *moduleResource = static_cast<CUpti_ModuleResourceData *>(
        resourceData->resourceDescriptor);
    if (profiler.pcSamplingEnabled)
      pcSampling.unloadModule(moduleResource->pCubin,
                              moduleResource->cubinSize);
  } else if (cbId == CUPTI_CBID_RESOURCE_CONTEXT_CREATED) {
    if (profiler.pcSamplingEnabled)
      pcSampling.initialize(resourceData->context);
  } else if (cbId == CUPTI_CBID_RESOURCE_CONTEXT_DESTROY_STARTING) {
    if (profiler.pcSamplingEnabled)
      pcSampling.finalize(resourceData->context);
  } else {
    auto *graphData =
        static_cast<CUpti_GraphData *>(resourceData->resourceDescriptor);
    handleGraphResourceCallbacks(profiler, cbId, graphData);
  }
}

void CuptiProfiler::CuptiProfilerPimpl::handleNvtxCallbacks(
    CUpti_CallbackId cbId, const void *cbData) {
  auto *nvtxData = static_cast<const CUpti_NvtxData *>(cbData);
  if (cbId == CUPTI_CBID_NVTX_nvtxRangePushA) {
    auto message = nvtx::getMessageFromRangePushA(nvtxData->functionParams);
    threadState.enterScope(message);
  } else if (cbId == CUPTI_CBID_NVTX_nvtxRangePop) {
    threadState.exitScope();
  } // TODO: else handle other NVTX range functions
}

void CuptiProfiler::CuptiProfilerPimpl::handleStreamCaptureCallbacks(
    CUpti_CallbackId cbId, const CUpti_CallbackData *callbackData) {
  if (callbackData->callbackSite != CUPTI_API_ENTER)
    return;

  if (isStreamCaptureBegin(cbId)) {
    threadState.isStreamCapturing = true;
    profiler.metricBuffer->reserve();
  } else if (isStreamCaptureEnd(cbId)) {
    threadState.isStreamCapturing = false;
  }
}

void CuptiProfiler::CuptiProfilerPimpl::handleApiEnterLaunchCallbacks(
    CuptiProfiler &profiler, CUpti_CallbackId cbId,
    const CUpti_CallbackData *callbackData) {
  size_t numNodes = 1;
  if (isGraphLaunch(cbId)) {
    threadState.enterOp(Scope(""));
  } else {
    // Symbol name is only available for kernel launch APIs.
    const auto symbolName = callbackData->context && callbackData->symbolName
                                ? std::string(callbackData->symbolName)
                                : "";
    threadState.enterOp(Scope(std::move(symbolName)));
  }

  auto &dataToEntry = threadState.dataToEntry;
  if (threadState.isStreamCapturing) // Do not correlate stream captured kernels
    return;

  const auto &scope = threadState.scopeStack.back();
  if (isGraphLaunch(cbId)) {
    auto graphExec =
        static_cast<const cuGraphLaunch_params *>(callbackData->functionParams)
            ->hGraph;
    uint32_t graphExecId = 0;
    cupti::getGraphExecId<true>(graphExec, &graphExecId);
    numNodes = std::numeric_limits<size_t>::max();
    auto findGraph = false;
    if (graphStates.contain(graphExecId)) {
      if (!graphStates[graphExecId].captureStatusChecked)
        numNodes = graphStates[graphExecId].nodeIdToState.size();
      findGraph = true;
    }
    if (!findGraph && !graphStates[graphExecId].captureStatusChecked) {
      graphStates[graphExecId].captureStatusChecked = true;
      std::cerr << "[PROTON] Cannot find graph for graphExecId: " << graphExecId
                << ", and it may cause memory leak. To avoid this problem, "
                   "please start profiling before the graph is created."
                << std::endl;
    } else if (findGraph && !graphStates[graphExecId].captureStatusChecked) {
      auto &graphState = graphStates[graphExecId];
      static const bool timingEnabled =
          getBoolEnv("PROTON_GRAPH_LAUNCH_TIMING", false);
      using Clock = std::chrono::steady_clock;
      auto t0 = decltype(Clock::now()){};
      if (timingEnabled)
        t0 = Clock::now();

      auto *externIdState = buildGraphNodeEntries(
          dataToEntry, graphState, profiler.correlation.externIdToState,
          scope.scopeId);

      if (timingEnabled) {
        auto t1 = Clock::now();
        auto elapsed =
            std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0)
                .count();
        std::cerr << "[PROTON] Graph launch call path time: " << elapsed
                  << " us for graphExecId: " << graphExecId << std::endl;
        t0 = Clock::now();
      }

      queueGraphMetrics(profiler.pendingGraphPool.get(), callbackData,
                        graphState, externIdState);

      if (timingEnabled) {
        auto t1 = Clock::now();
        auto elapsed =
            std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0)
                .count();
        std::cerr << "[PROTON] Graph launch metric time: " << elapsed
                  << " us for graphExecId: " << graphExecId << std::endl;
      }
    }
  }

  if (dataToEntry.empty()) // Profiler is deactivated
    return;

  profiler.correlation.correlate(callbackData->correlationId, scope.scopeId,
                                 numNodes, scope.name.empty(), dataToEntry);
  if (profiler.pcSamplingEnabled)
    pcSampling.start(callbackData->context);
}

void CuptiProfiler::CuptiProfilerPimpl::handleApiExitLaunchCallbacks(
    CuptiProfiler &profiler, CUpti_CallbackId cbId,
    const CUpti_CallbackData *callbackData) {
  auto &dataToEntry = threadState.dataToEntry;
  bool deactivated = dataToEntry.empty();

  if (profiler.pcSamplingEnabled) {
    // XXX: Conservatively stop every GPU kernel for now.
    pcSampling.stop(callbackData->context, dataToEntry);
  }

  threadState.exitOp();

  if (threadState
          .isStreamCapturing) // Do not correlate for stream captured kernels
    return;
  if (deactivated) // Profiler is deactivated
    return;
  profiler.correlation.submit(callbackData->correlationId);
}

void CuptiProfiler::CuptiProfilerPimpl::handleApiCallbacks(
    CuptiProfiler &profiler, CUpti_CallbackId cbId, const void *cbData) {
  const CUpti_CallbackData *callbackData =
      static_cast<const CUpti_CallbackData *>(cbData);
  handleStreamCaptureCallbacks(cbId, callbackData);
  if (isLaunch(cbId)) {
    if (callbackData->callbackSite == CUPTI_API_ENTER) {
      handleApiEnterLaunchCallbacks(profiler, cbId, callbackData);
    } else if (callbackData->callbackSite == CUPTI_API_EXIT) {
      handleApiExitLaunchCallbacks(profiler, cbId, callbackData);
    }
  }
}

void CuptiProfiler::CuptiProfilerPimpl::callbackFn(void *userData,
                                                   CUpti_CallbackDomain domain,
                                                   CUpti_CallbackId cbId,
                                                   const void *cbData) {
  CuptiProfiler &profiler = threadState.profiler;
  auto *pImpl = dynamic_cast<CuptiProfilerPimpl *>(profiler.pImpl.get());
  if (domain == CUPTI_CB_DOMAIN_RESOURCE) {
    pImpl->handleResourceCallbacks(profiler, cbId, cbData);
  } else if (domain == CUPTI_CB_DOMAIN_NVTX) {
    pImpl->handleNvtxCallbacks(cbId, cbData);
  } else {
    pImpl->handleApiCallbacks(profiler, cbId, cbData);
  }
}

void CuptiProfiler::CuptiProfilerPimpl::doStart() {
  cupti::subscribe<true>(&subscriber, callbackFn, nullptr);
  if (profiler.pcSamplingEnabled) {
    setResourceCallbacks(subscriber, /*enable=*/true);
    // Continuous PC sampling is not compatible with concurrent kernel profiling
    cupti::activityEnable<true>(CUPTI_ACTIVITY_KIND_KERNEL);
  } else {
    cupti::activityEnable<true>(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL);
    if (getBoolEnv("TRITON_ENABLE_HW_TRACE", false))
      cupti::activityEnableHWTrace<true>(/*enable=*/1);
  }
  cupti::activityRegisterCallbacks<true>(allocBuffer, completeBuffer);
  setGraphCallbacks(subscriber, /*enable=*/true);
  setLaunchCallbacks(subscriber, /*enable=*/true);
  if (getBoolEnv("TRITON_ENABLE_NVTX", true)) {
    nvtx::enable();
    setNvtxCallbacks(subscriber, /*enable=*/true);
  }

  if (!profiler.isTimestampCalibrated) {
    profiler.timestampOffsetNs =
        detail::computeTimestampOffsetNs(cupti::getTimestamp<true>);
    profiler.isTimestampCalibrated = true;
  }
}

void CuptiProfiler::CuptiProfilerPimpl::doFlush() {
  // cuptiActivityFlushAll returns the activity records associated with all
  // contexts/streams.
  // This is a blocking call but it doesn’t issue any CUDA synchronization calls
  // implicitly thus it’s not guaranteed that all activities are completed on
  // the underlying devices.
  // We do an "opportunistic" synchronization here to try to ensure that all
  // activities are completed on the current context.
  // If the current context is not set, we don't do any synchronization.
  CUcontext cuContext = nullptr;
  cuda::ctxGetCurrent<false>(&cuContext);
  if (cuContext) {
    cuda::ctxSynchronize<true>();
  }
  profiler.correlation.flush(
      /*maxRetries=*/100, /*sleepUs=*/10,
      /*flush=*/[]() {
        cupti::activityFlushAll<true>(
            /*flag=*/0);
      });
  // CUPTI_ACTIVITY_FLAG_FLUSH_FORCED is used to ensure that even incomplete
  // activities are flushed so that the next profiling session can start with
  // new activities.
  cupti::activityFlushAll<true>(/*flag=*/CUPTI_ACTIVITY_FLAG_FLUSH_FORCED);
  // Flush the tensor metric buffer
  profiler.pendingGraphPool->flushAll();
}

void CuptiProfiler::CuptiProfilerPimpl::doStop() {
  if (profiler.pcSamplingEnabled) {
    profiler.pcSamplingEnabled = false;
    CUcontext cuContext = nullptr;
    cuda::ctxGetCurrent<false>(&cuContext);
    if (cuContext)
      pcSampling.finalize(cuContext);
    setResourceCallbacks(subscriber, /*enable=*/false);
    cupti::activityDisable<true>(CUPTI_ACTIVITY_KIND_KERNEL);
  } else {
    cupti::activityDisable<true>(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL);
    if (getBoolEnv("TRITON_ENABLE_HW_TRACE", false))
      cupti::activityEnableHWTrace<true>(/*enable=*/0);
  }
  profiler.periodicFlushingEnabled = false;
  profiler.periodicFlushingFormat.clear();
  // We have to clear the correlation maps before unsubscribing because CUPTI
  // will reset correlation ID after unsubscribing
  profiler.correlation.clear();
  setGraphCallbacks(subscriber, /*enable=*/false);
  setLaunchCallbacks(subscriber, /*enable=*/false);
  nvtx::disable();
  setNvtxCallbacks(subscriber, /*enable=*/false);
  cupti::unsubscribe<true>(subscriber);
  cupti::finalize<true>();
}

CuptiProfiler::CuptiProfiler() {
  pImpl = std::make_unique<CuptiProfilerPimpl>(*this);
}

CuptiProfiler::~CuptiProfiler() = default;

void CuptiProfiler::doSetMode(const std::vector<std::string> &modeAndOptions) {
  auto mode = modeAndOptions[0];
  if (proton::toLower(mode) == "pcsampling") {
    pcSamplingEnabled = true;
  } else if (proton::toLower(mode) == "periodic_flushing") {
    detail::setPeriodicFlushingMode(periodicFlushingEnabled,
                                    periodicFlushingFormat, modeAndOptions,
                                    "CuptiProfiler");
  } else if (!mode.empty()) {
    throw makeInvalidArgument("CuptiProfiler: unsupported mode: " + mode);
  }
}

} // namespace proton
