//===- OmptTracingBuffer.h - Target independent OpenMP target RTL -- C++ *-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Interface to be used for generating and flushing OMPT device trace records.
//
//===----------------------------------------------------------------------===//

#ifndef OPENMP_LIBOMPTARGET_OMPTTRACINGBUFFER_H
#define OPENMP_LIBOMPTARGET_OMPTTRACINGBUFFER_H

#ifdef OMPT_SUPPORT

#include <atomic>
#include <cassert>
#include <condition_variable>
#include <cstdint>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <omp-tools.h>

#include "Shared/EnvironmentVar.h"

// Maximum number of devices supported in device tracing. No device tracing
// will be performed for any device-id larger than 1023.
#define MAX_NUM_DEVICES 1024

// TODO Start with 1 helper thread and add dynamically if required
// Number of helper threads must not execeed 32 since the
// thread-wait-tracker is 32 bits in length.
#define OMPT_NUM_HELPER_THREADS 1

/*
 * Buffer manager for trace records generated by OpenMP master and
 * worker threads. During device init, a tool may register a
 * buffer-request and a buffer-completion callback. The buffer-request
 * callback should be used to allocate new buffers as required. The
 * buffer-complete callback should be used to return trace records to
 * the tool.
 *
 * In addition to trace records, this class manages the helper threads
 * for dispatching a range of trace records to the tool.
 */
class OmptTracingBufferMgr {
public:
  /*
   * A trace record (TR) holds the trace data. Its type
   * can be ompt or native. Currently, only ompt type is implemented.
   */

  /*
   * A TR can be in the following states:
   * TR_init: initial state
   * TR_ready: An OpenMP thread marks a TR ready when it is done
   * populating the TR
   * TR_released: A helper thread marks a TR released after it has
   * completed returning the TR to the tool
   */
  enum TRStatus { TR_init, TR_ready, TR_released };

  /*
   * Metadata capturing the state of a buffer of trace records. Once a
   * buffer is allocated by an OpenMP worker thread, trace records are
   * carved out from that buffer by that same OpenMP thread alone. Thus
   * the allocated buffer is thread-specific from the allocation/population
   * standpoint. But it may be manipulated by helper threads.
   *
   * Id, DeviceId, Start, and TotalBytes are not changed once set.
   * RemainingBytes could be written multiple times but only by the same
   * thread. But Cursor and IsFull may be read/written by an OpenMP worker
   * thread and read by helper threads. Hence, accesses of
   * this 2nd set of locations need to be atomic or synchronized.
   */
  struct Buffer {
    uint64_t Id;                // Unique identifier of the buffer
    int64_t DeviceId;           // Device for which this buffer is allocated
    void *Start;                // Start of allocated space for trace records
    size_t TotalBytes;          // Total number of bytes in the allocated space
    size_t RemainingBytes;      // Total number of unused bytes
                                // corresponding to Cursor
    std::atomic<void *> Cursor; // Address of the last trace record carved out
    std::atomic<bool> IsFull;   // true if no more trace records can be
                                // accomodated, otherwise false
    Buffer(uint64_t BufId, int64_t DevId, void *S, size_t Bytes, size_t Rem,
           void *C, bool F)
        : Id(BufId), DeviceId(DevId), Start(S), TotalBytes(Bytes),
          RemainingBytes(Rem), Cursor(C), IsFull(F) {}
    Buffer() = delete;
    Buffer(const Buffer &) = delete;
    Buffer &operator=(const Buffer &) = delete;
  };
  using BufPtr = std::shared_ptr<Buffer>;

private:
  /// Envar to control whether a buffer should be flushed when it gets full.
  BoolEnvar OMPX_FlushOnBufferFull;

  /// Envar to control whether all buffers should be flushed during shutdown.
  BoolEnvar OMPX_FlushOnShutdown;

  // Internal variable for tracking threads to wait for flush
  uint32_t ThreadFlushTracker;

  // Internal variable for tracking threads shutting down
  uint32_t ThreadShutdownTracker;

  using MapId2Buf = std::map<uint64_t, BufPtr>;

  // Map from id to corresponding buffer. The ids are assigned in
  // increasing order of creation.
  MapId2Buf Id2BufferMap;

  // Trace record. We currently support OMPT data type only. The state
  // (TRStatus type) is maintained inline in the trace record. The
  // tool is expected to access only the OMPT record.
  struct TraceRecord {
    ompt_record_ompt_t TR;
    std::atomic<TRStatus> TRState;
  };

  // Thread-specific array of pointers to a buffer. The buffer pointed to
  // is the last one allocated by this thread for a given device. The ith
  // element points to the buffer for the ith device. At most MAX_NUM_DEVICES
  // devices are supported.
  static thread_local BufPtr ArrayOfBufPtr[MAX_NUM_DEVICES];

  /*
   * A buffer is flushed when it fills up or when the tool invokes
   * flush_trace. So it's possible that the same buffer may be flushed
   * more than once. When a buffer is flushed the first time, a unique
   * id (flush-id) is generated and assigned to that buffer. Even if
   * it is flushed again, the previously assigned id is maintained for
   * that buffer. This id is loosely used to determine the order in
   * which the buffers are processed and the corresponding trace
   * records released to the tool.
   */

  struct FlushInfo {
    uint64_t FlushId;
    void *FlushCursor;
    BufPtr FlushBuf;
    FlushInfo() = default;
    FlushInfo(uint64_t Id, void *CR, BufPtr Buf)
        : FlushId{Id}, FlushCursor{CR}, FlushBuf{Buf} {}
  };

  /*
   * A buffer may be in the following states:
   * Flush_waiting: when a buffer is flushed, either because it is
   * full or because the tool invokes ompt_flush_trace
   * Flush_processing: when a helper thread claims the waiting buffer
   * and is in the process of dispatching buffer-completion callbacks
   * on an associated range of trace records. If all trace records are
   * not released, the state may be reset to Flush_waiting after the
   * buffer-completion callbacks return
   */
  enum BufferFlushStatus { Flush_waiting, Flush_processing };
  struct FlushMd {
    void *FlushCursor;
    BufPtr FlushBuf;
    BufferFlushStatus FlushStatus;
    FlushMd(void *CR, BufPtr Buf, BufferFlushStatus Status)
        : FlushCursor{CR}, FlushBuf{Buf}, FlushStatus{Status} {}
    FlushMd() = delete;
  };

  using MapId2Md = std::map<uint64_t, FlushMd>;

  /*
   * A map from a flush-id to metadata containing the current
   * cursor. the corresponding buffer, and its flushed status. If a
   * buffer is flushed multiple times, the cursor is updated to the
   * furthest one
   */
  MapId2Md Id2FlushMdMap;

  using UMapBufPtr2Id = std::unordered_map<BufPtr, uint64_t>;

  // A hash map from a buffer address to the corresponding flush-id
  UMapBufPtr2Id FlushBufPtr2IdMap;

  using USetCursor = std::unordered_set<void *>;

  USetCursor LastCursors;

  using UMapThd2Id = std::unordered_map<std::thread::id, uint32_t>;

  // A hash map from a helper thread id to an integer
  UMapThd2Id HelperThreadIdMap;

  // Mutex to protect Id2BufferMap and Cursor2BufMdMap
  std::mutex BufferMgrMutex;

  // Mutex to protect FlushBufPtr2IdMap and Id2FlushMdMap
  std::mutex FlushMutex;

  // Mutex to protect metadata tracking last cursors of buffer-completion
  // callbacks
  std::mutex LastCursorMutex;

  // Condition variable used by helper thread to signal that flush is requested
  std::condition_variable FlushCv;

  // Condition variable used while waiting for flushing to complete
  std::condition_variable ThreadFlushCv;

  // Condition variable used while waiting for threads to shutdown
  std::condition_variable ThreadShutdownCv;

  // TODO Separate out the helper thread into its own class
  std::vector<std::thread> CompletionThreads;

  /// Called when a buffer \p Buf may be flushed with \p Cursor as the
  /// last allocated trace record in the buffer.
  /// triggerFlushOnBufferFull should be called without holding any lock.
  void triggerFlushOnBufferFull(void *Cursor, BufPtr Buf);

  // Called to dispatch buffer-completion callbacks for the trace records in
  // this buffer
  void flushBuffer(FlushInfo);

  // Dispatch a buffer-completion callback with a range of trace records
  void dispatchCallback(int64_t DeviceId, void *Buffer, void *FirstCursor,
                        void *LastCursor);

  // Add a last cursor
  void addLastCursor(void *Cursor) {
    std::unique_lock<std::mutex> Lock(LastCursorMutex);
    LastCursors.emplace(Cursor);
  }

  // Remove a last cursor
  void removeLastCursor(void *Cursor) {
    std::unique_lock<std::mutex> Lock(LastCursorMutex);
    assert(LastCursors.find(Cursor) != LastCursors.end());
    LastCursors.erase(Cursor);
  }

  // Given a trace record pointer, initialize its metadata
  void initTraceRecordMetaData(void *Rec);

  // Given a device-id, get/set a pointer to the last allocated buffer metadata.
  BufPtr getDeviceSpecificBuffer(int64_t DevId);
  void setDeviceSpecificBuffer(int64_t DevId, BufPtr Buf);

  // Reserve a candidate buffer for flushing, preventing other helper threads
  // from accessing it
  FlushInfo findAndReserveFlushedBuf(uint64_t FlushId);

  // Unreserve a buffer so that other helper threads can process it
  void unreserveFlushedBuf(const FlushInfo &);

  // All done with this buffer, so the buffer and its metadata can be removed
  void destroyFlushedBuf(const FlushInfo &);

  // Add a new buffer by an OpenMP thread so that a helper thread can process it
  uint64_t addNewFlushEntry(BufPtr Buf, void *Cursor);

  // Get the next trace record
  void *getNextTR(void *TR);

  // Given a buffer, return the latest cursor
  void *getBufferCursor(BufPtr);

  // Is no more space remaining for trace records in this buffer?
  bool isBufferFull(const FlushInfo &);

  // Have all trace records in this buffer been returned to the tool?
  bool isBufferOwned(const FlushInfo &);

  // Dispatch a buffer-completion callback and indicate that the buffer can be
  // deallocated
  void dispatchBufferOwnedCallback(const FlushInfo &);

  // Main entry point for a helper thread
  void driveCompletion();

  // Examine the flushed buffers and dispatch buffer-completion callbacks
  void invokeCallbacks();

  // The caller does not hold a lock while calling this method
  void waitForFlushCompletion();

  // Given a thread number, set the corresponding bit in the flush
  // tracker. The caller must hold the flush lock.
  void setThreadFlush(uint32_t ThreadNum) {
    ThreadFlushTracker |= (1 << ThreadNum);
  }

  // Reset this thread's flush bit. The caller must hold the flush lock
  void resetThisThreadFlush() {
    std::thread::id ID = std::this_thread::get_id();
    assert(HelperThreadIdMap.find(ID) != HelperThreadIdMap.end());
    ThreadFlushTracker &= ~(1 << HelperThreadIdMap[ID]);
  }

  // Given a thread number, set the corresponding bit in the shutdown
  // tracker. The caller must hold the flush lock.
  void setThreadShutdown(uint32_t ThreadNum) {
    ThreadShutdownTracker |= (1 << ThreadNum);
  }

  // Reset this thread's shutdown bit. The caller must hold the flush
  // lock
  void resetThisThreadShutdown() {
    std::thread::id ID = std::this_thread::get_id();
    assert(HelperThreadIdMap.find(ID) != HelperThreadIdMap.end());
    ThreadShutdownTracker &= ~(1 << HelperThreadIdMap[ID]);
  }

  // Return true if this thread's flush bit is set. The caller must
  // hold the flush lock
  bool isThisThreadFlushWaitedUpon() {
    std::thread::id ID = std::this_thread::get_id();
    assert(HelperThreadIdMap.find(ID) != HelperThreadIdMap.end());
    return (ThreadFlushTracker & (1 << HelperThreadIdMap[ID])) != 0;
  }

  // Return true if this thread's shutdown bit is set. The caller must
  // hold the flush lock
  bool isThisThreadShutdownWaitedUpon() {
    std::thread::id ID = std::this_thread::get_id();
    assert(HelperThreadIdMap.find(ID) != HelperThreadIdMap.end());
    return (ThreadShutdownTracker & (1 << HelperThreadIdMap[ID])) != 0;
  }

  // The caller must not hold the flush lock
  bool amIHelperThread() {
    std::unique_lock<std::mutex> flush_lock(FlushMutex);
    if (HelperThreadIdMap.find(std::this_thread::get_id()) !=
        HelperThreadIdMap.end())
      return true;
    return false;
  }

  // The caller must not hold the flush lock
  bool areHelperThreadsAvailable();

  // The caller must hold the appropriate lock
  void init();

  // The caller must hold the flush lock
  void createHelperThreads();

  // The caller must hold the flush lock
  void destroyHelperThreads();

public:
  OmptTracingBufferMgr()
      : OMPX_FlushOnBufferFull("LIBOMPTARGET_OMPT_FLUSH_ON_BUFFER_FULL", true),
        OMPX_FlushOnShutdown("LIBOMPTARGET_OMPT_FLUSH_ON_SHUTDOWN", true) {
    // no need to hold locks for init() since object is getting constructed
    // here.
    init();
  }

  OmptTracingBufferMgr(const OmptTracingBufferMgr &) = delete;
  OmptTracingBufferMgr &operator=(const OmptTracingBufferMgr &) = delete;

  // The caller must not hold the flush lock
  void startHelperThreads();

  // The caller must not hold the flush lock. The helper threads are shut down
  // without flushing any outstanding trace records.
  void shutdownHelperThreads();

  // The caller must not hold the flush lock. The helper threads are shut down
  // after flushing all outstanding trace records for all devices.
  void flushAndShutdownHelperThreads();

  // Assign a cursor for a new trace record. This will assign a trace record
  // for the provided device-id, allocating a new buffer if required.
  void *assignCursor(ompt_callbacks_t Type, int64_t DeviceId);

  // Get the size of a trace record
  size_t getTRSize() { return sizeof(TraceRecord); }

  // Get the status of a trace record. This function does not acquire
  // a lock. If locking is required, the caller must hold a lock.
  TRStatus getTRStatus(void *Rec);

  // Set the status of a trace record. This function does not acquire
  // a lock. If locking is required, the caller must hold a lock.
  void setTRStatus(void *Rec, TRStatus);

  // Is this a last cursor of a buffer completion callback?
  bool isLastCursor(void *Cursor) {
    std::unique_lock<std::mutex> Lock(LastCursorMutex);
    return LastCursors.find(Cursor) != LastCursors.end();
  }

  // Called for flushing outstanding buffers for the provided device-id.
  int flushAllBuffers(int DeviceId);
};

#else
class OmptTracingBufferMgr {};
#endif // OMPT_SUPPORT

#endif // OPENMP_LIBOMPTARGET_OMPTTRACINGBUFFER_H