nvtx_profiler.h

// NIST-developed software is provided by NIST as a public service. You may use, copy and distribute copies of the
// software in any medium, provided that you keep intact this entire notice. You may improve, modify and create
// derivative works of the software or any portion of the software, and you may copy and distribute such modifications
// or works. Modified works should carry a notice stating that you changed the software and should note the date and
// nature of any such change. Please explicitly acknowledge the National Institute of Standards and Technology as the
// source of the software. NIST-developed software is expressly provided "AS IS." NIST MAKES NO WARRANTY OF ANY KIND,
// EXPRESS, IMPLIED, IN FACT OR ARISING BY OPERATION OF LAW, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTY OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT AND DATA ACCURACY. NIST NEITHER REPRESENTS NOR
// WARRANTS THAT THE OPERATION OF THE SOFTWARE WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT ANY DEFECTS WILL BE
// CORRECTED. NIST DOES NOT WARRANT OR MAKE ANY REPRESENTATIONS REGARDING THE USE OF THE SOFTWARE OR THE RESULTS
// THEREOF, INCLUDING BUT NOT LIMITED TO THE CORRECTNESS, ACCURACY, RELIABILITY, OR USEFULNESS OF THE SOFTWARE. You
// are solely responsible for determining the appropriateness of using and distributing the software and you assume
// all risks associated with its use, including but not limited to the risks and costs of program errors, compliance
// with applicable laws, damage to or loss of data, programs or equipment, and the unavailability or interruption of 
// operation. This software is not intended to be used in any situation where a failure could cause risk of injury or
// damage to property. The software developed by NIST employees is not subject to copyright protection within the
// United States.


#ifndef HEDGEHOG_NVTX_PROFILER_H
#define HEDGEHOG_NVTX_PROFILER_H

#ifdef HH_USE_NVTX
#include <nvtx3/nvToolsExt.h>
#endif

#include <string_view>
#include <cstring>
#include <iostream>

namespace hh {
#define NVTX_COLOR_INITIALIZING    0xFF123456
#define NVTX_COLOR_EXECUTING       0xFF72ff68
#define NVTX_COLOR_WAITING         0xFFff7f83
#define NVTX_COLOR_WAITING_FOR_MEM 0xFFffc86a
#define NVTX_COLOR_RELEASE_MEM     0xFF7fbdff
#define NVTX_COLOR_SHUTTING_DOWN   0xFF654321

class NvtxProfiler {
 private:
#ifdef HH_USE_NVTX
  std::string_view initializeName_{}; 
  std::string_view executeName_{}; 
  std::string_view waitName_{}; 
  std::string_view waitForMemName_{}; 
  std::string_view releaseMemName_{}; 
  std::string_view shutdownName_{}; 

  nvtxDomainHandle_t taskDomain_; 

  nvtxStringHandle_t initializeString_{}; 
  nvtxStringHandle_t executeString_{}; 
  nvtxStringHandle_t waitString_{}; 
  nvtxStringHandle_t waitForMemString_{}; 
  nvtxStringHandle_t releaseMemString_{}; 
  nvtxStringHandle_t shutdownString_{}; 

  nvtxEventAttributes_t *initializeAttrib_; 
  nvtxEventAttributes_t *executeAttrib_; 
  nvtxEventAttributes_t *waitAttrib_; 
  nvtxEventAttributes_t *waitForMemAttrib_; 
  nvtxEventAttributes_t *releaseMemAttrib_; 
  nvtxEventAttributes_t *shutdownAttrib_; 

  nvtxRangeId_t initializeRangeId_ = 0; 
  nvtxRangeId_t executeRangeId_ = 0; 
  nvtxRangeId_t waitRangeId_ = 0; 
  nvtxRangeId_t waitForMemRangeId_ = 0; 
  nvtxRangeId_t shutdownRangeId_ = 0; 


  nvtxEventAttributes_t *createEventAttribute(uint32_t color) {
    nvtxEventAttributes_t *event = new nvtxEventAttributes_t;
    bzero(event, NVTX_EVENT_ATTRIB_STRUCT_SIZE);
    event->version = NVTX_VERSION;
    event->size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
    event->colorType = NVTX_COLOR_ARGB;
    event->color = color;
    return event;
  }
#endif

 public:

  NvtxProfiler() = delete;


#ifdef HH_USE_NVTX
  explicit NvtxProfiler(std::string_view const & taskName) {
    taskDomain_ = nvtxDomainCreateA(taskName.data());

    initializeAttrib_ = createEventAttribute(NVTX_COLOR_INITIALIZING);
    executeAttrib_ = createEventAttribute(NVTX_COLOR_EXECUTING);

    waitAttrib_ = createEventAttribute(NVTX_COLOR_WAITING);
    waitAttrib_->payloadType = NVTX_PAYLOAD_TYPE_UNSIGNED_INT64;

    waitAttrib_->payload.ullValue = 0;

    waitForMemAttrib_ = createEventAttribute(NVTX_COLOR_WAITING_FOR_MEM);
    releaseMemAttrib_ = createEventAttribute(NVTX_COLOR_RELEASE_MEM);
    shutdownAttrib_ = createEventAttribute(NVTX_COLOR_SHUTTING_DOWN);
  }
#else //HH_USE_NVTX
  explicit NvtxProfiler(std::string_view const &) {}
#endif //HH_USE_NVTX

#ifdef HH_USE_NVTX
  ~NvtxProfiler() {
    delete initializeAttrib_;
    delete executeAttrib_;
    delete waitAttrib_;
    delete waitForMemAttrib_;
    delete releaseMemAttrib_;
    delete shutdownAttrib_;
    nvtxDomainDestroy(taskDomain_);
  }
#else //HH_USE_NVTX
  ~NvtxProfiler() = default;
#endif //HH_USE_NVTX

  void initialize([[maybe_unused]]int threadId) {
#ifdef HH_USE_NVTX
    std::string prefixName(std::to_string(threadId));
    initializeName_ = prefixName + ":Initializing";
    executeName_ = prefixName + ":Executing";
    waitName_ = prefixName + ":Waiting";
    waitForMemName_ = prefixName + ":MemWait";
    releaseMemName_ = prefixName + ":Release";
    shutdownName_ = prefixName + ":Shutdown";

    initializeString_ = nvtxDomainRegisterStringA(taskDomain_, initializeName_.data());
    executeString_ = nvtxDomainRegisterStringA(taskDomain_, executeName_.data());
    waitString_ = nvtxDomainRegisterStringA(taskDomain_, waitName_.data());
    waitForMemString_ = nvtxDomainRegisterStringA(taskDomain_, waitForMemName_.data());
    releaseMemString_ = nvtxDomainRegisterStringA(taskDomain_, releaseMemName_.data());
    shutdownString_ = nvtxDomainRegisterStringA(taskDomain_, shutdownName_.data());

    initializeAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    initializeAttrib_->message.registered = initializeString_;

    executeAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    executeAttrib_->message.registered = executeString_;

    waitAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    waitAttrib_->message.registered = waitString_;

    waitForMemAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    waitForMemAttrib_->message.registered = waitForMemString_;

    releaseMemAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    releaseMemAttrib_->message.registered = releaseMemString_;

    shutdownAttrib_->messageType = NVTX_MESSAGE_TYPE_REGISTERED;
    shutdownAttrib_->message.registered = shutdownString_;
#endif
  }

  void addReleaseMarker() {
#ifdef HH_USE_NVTX
    nvtxDomainMarkEx(taskDomain_, releaseMemAttrib_);
#endif
  }

  void startRangeInitializing() {
#ifdef HH_USE_NVTX
    initializeRangeId_ = nvtxDomainRangeStartEx(taskDomain_, initializeAttrib_);
#endif
  }

  void startRangeExecuting() {
#ifdef HH_USE_NVTX
    executeRangeId_ = nvtxDomainRangeStartEx(taskDomain_, executeAttrib_);
#endif
  }

  void startRangeWaiting([[maybe_unused]]uint64_t const &queueSize) {
#ifdef HH_USE_NVTX
    waitAttrib_->payload.ullValue = queueSize;
    waitRangeId_ = nvtxDomainRangeStartEx(taskDomain_, waitAttrib_);
#endif
  }

  void startRangeWaitingForMemory() {
#ifdef HH_USE_NVTX
    waitForMemRangeId_ = nvtxDomainRangeStartEx(taskDomain_, waitForMemAttrib_);
#endif
  }

  void startRangeShuttingDown() {
#ifdef HH_USE_NVTX
    shutdownRangeId_ = nvtxDomainRangeStartEx(taskDomain_, shutdownAttrib_);
#endif
  }

  void endRangeInitializing() {
#ifdef HH_USE_NVTX
    nvtxDomainRangeEnd(taskDomain_, initializeRangeId_);
#endif
  }

  void endRangeExecuting() {
#ifdef HH_USE_NVTX
    nvtxDomainRangeEnd(taskDomain_, executeRangeId_);
#endif
  }

  void endRangeWaiting() {
#ifdef HH_USE_NVTX
    nvtxDomainRangeEnd(taskDomain_, waitRangeId_);
#endif
  }

  void endRangeWaitingForMem() {
#ifdef HH_USE_NVTX
    nvtxDomainRangeEnd(taskDomain_, waitForMemRangeId_);
#endif
  }

  void endRangeShuttingDown() {
#ifdef HH_USE_NVTX
    nvtxDomainRangeEnd(taskDomain_, shutdownRangeId_);
#endif
  }
#ifdef HH_USE_NVTX


  nvtxDomainHandle_t taskDomain() const {
    return taskDomain_;
  }
#endif

};
}
#endif //HEDGEHOG_NVTX_PROFILER_H