cuda_graph.hpp

#pragma once
 
#include <filesystem>
 
#include "cuda_memory.hpp"
#include "cuda_stream.hpp"
#include "cuda_meta.hpp"
 
#include "../utility/traits.hpp"
 
namespace tf {
 
// ----------------------------------------------------------------------------
// cudaGraph_t routines
// ----------------------------------------------------------------------------

template <typename T,
  std::enable_if_t<!std::is_same_v<T, void>, void>* = nullptr
>
cudaMemcpy3DParms cuda_get_copy_parms(T* tgt, const T* src, size_t num) {
 
  using U = std::decay_t<T>;
 
  cudaMemcpy3DParms p;
 
  p.srcArray = nullptr;
  p.srcPos = ::make_cudaPos(0, 0, 0);
  p.srcPtr = ::make_cudaPitchedPtr(const_cast<T*>(src), num*sizeof(U), num, 1);
  p.dstArray = nullptr;
  p.dstPos = ::make_cudaPos(0, 0, 0);
  p.dstPtr = ::make_cudaPitchedPtr(tgt, num*sizeof(U), num, 1);
  p.extent = ::make_cudaExtent(num*sizeof(U), 1, 1);
  p.kind = cudaMemcpyDefault;
 
  return p;
}

inline cudaMemcpy3DParms cuda_get_memcpy_parms(
  void* tgt, const void* src, size_t bytes
)  {
 
  // Parameters in cudaPitchedPtr
  // d   - Pointer to allocated memory
  // p   - Pitch of allocated memory in bytes
  // xsz - Logical width of allocation in elements
  // ysz - Logical height of allocation in elements
  cudaMemcpy3DParms p;
  p.srcArray = nullptr;
  p.srcPos = ::make_cudaPos(0, 0, 0);
  p.srcPtr = ::make_cudaPitchedPtr(const_cast<void*>(src), bytes, bytes, 1);
  p.dstArray = nullptr;
  p.dstPos = ::make_cudaPos(0, 0, 0);
  p.dstPtr = ::make_cudaPitchedPtr(tgt, bytes, bytes, 1);
  p.extent = ::make_cudaExtent(bytes, 1, 1);
  p.kind = cudaMemcpyDefault;
 
  return p;
}

inline cudaMemsetParams cuda_get_memset_parms(void* dst, int ch, size_t count) {
 
  cudaMemsetParams p;
  p.dst = dst;
  p.value = ch;
  p.pitch = 0;
  //p.elementSize = (count & 1) == 0 ? ((count & 3) == 0 ? 4 : 2) : 1;
  //p.width = (count & 1) == 0 ? ((count & 3) == 0 ? count >> 2 : count >> 1) : count;
  p.elementSize = 1;  // either 1, 2, or 4
  p.width = count;
  p.height = 1;
 
  return p;
}

template <typename T, std::enable_if_t<
  is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>* = nullptr
>
cudaMemsetParams cuda_get_fill_parms(T* dst, T value, size_t count) {
 
  cudaMemsetParams p;
  p.dst = dst;
 
  // perform bit-wise copy
  p.value = 0;  // crucial
  static_assert(sizeof(T) <= sizeof(p.value), "internal error");
  std::memcpy(&p.value, &value, sizeof(T));
 
  p.pitch = 0;
  p.elementSize = sizeof(T);  // either 1, 2, or 4
  p.width = count;
  p.height = 1;
 
  return p;
}

template <typename T, std::enable_if_t<
  is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>* = nullptr
>
cudaMemsetParams cuda_get_zero_parms(T* dst, size_t count) {
 
  cudaMemsetParams p;
  p.dst = dst;
  p.value = 0;
  p.pitch = 0;
  p.elementSize = sizeof(T);  // either 1, 2, or 4
  p.width = count;
  p.height = 1;
 
  return p;
}

inline size_t cuda_graph_get_num_root_nodes(cudaGraph_t graph) {
  size_t num_nodes;
  TF_CHECK_CUDA(
    cudaGraphGetRootNodes(graph, nullptr, &num_nodes),
    "failed to get native graph root nodes"
  );
  return num_nodes;
}

inline size_t cuda_graph_get_num_nodes(cudaGraph_t graph) {
  size_t num_nodes;
  TF_CHECK_CUDA(
    cudaGraphGetNodes(graph, nullptr, &num_nodes),
    "failed to get native graph nodes"
  );
  return num_nodes;
}

inline size_t cuda_graph_get_num_edges(cudaGraph_t graph, cudaGraphNode_t* from, cudaGraphNode_t* to) {
  size_t num_edges;
  TF_CHECK_CUDA(
      TF_CUDA_PRE13(cudaGraphGetEdges(graph, from, to, &num_edges))
      TF_CUDA_POST13(cudaGraphGetEdges(graph, from, to, nullptr, &num_edges)),
      "failed to get native graph edges"
  );
  return num_edges;
}

inline size_t cuda_graph_node_get_dependencies(cudaGraphNode_t node, cudaGraphNode_t* dependencies) {
  size_t num_predecessors;
  TF_CHECK_CUDA(
      TF_CUDA_PRE13(cudaGraphNodeGetDependencies(node, dependencies, &num_predecessors))
      TF_CUDA_POST13(cudaGraphNodeGetDependencies(node, dependencies, nullptr, &num_predecessors)),
  "Failed to get number of dependencies");
  return num_predecessors;
}

inline size_t cuda_graph_node_get_dependent_nodes(cudaGraphNode_t node, cudaGraphNode_t *dependent_nodes) {
  size_t num_successors;
  TF_CHECK_CUDA(
      TF_CUDA_PRE13(cudaGraphNodeGetDependentNodes(node, dependent_nodes, &num_successors))
      TF_CUDA_POST13(cudaGraphNodeGetDependentNodes(node, dependent_nodes, nullptr, &num_successors)),
      "Failed to get CUDA dependent nodes");
  return num_successors;
}

inline void cuda_graph_add_dependencies(cudaGraph_t graph, const cudaGraphNode_t *from, const cudaGraphNode_t *to, size_t numDependencies) {
  TF_CHECK_CUDA(
      TF_CUDA_PRE13(cudaGraphAddDependencies(graph, from, to, numDependencies))
      TF_CUDA_POST13(cudaGraphAddDependencies(graph, from, to, nullptr, numDependencies)),
      "Failed to add CUDA graph node dependencies"
      );
}

inline size_t cuda_graph_get_num_edges(cudaGraph_t graph) {
  return cuda_graph_get_num_edges(graph, nullptr, nullptr);
}
 
 

inline std::vector<cudaGraphNode_t> cuda_graph_get_nodes(cudaGraph_t graph) {
  size_t num_nodes = cuda_graph_get_num_nodes(graph);
  std::vector<cudaGraphNode_t> nodes(num_nodes);
  TF_CHECK_CUDA(
    cudaGraphGetNodes(graph, nodes.data(), &num_nodes),
    "failed to get native graph nodes"
  );
  return nodes;
}

inline std::vector<cudaGraphNode_t> cuda_graph_get_root_nodes(cudaGraph_t graph) {
  size_t num_nodes = cuda_graph_get_num_root_nodes(graph);
  std::vector<cudaGraphNode_t> nodes(num_nodes);
  TF_CHECK_CUDA(
    cudaGraphGetRootNodes(graph, nodes.data(), &num_nodes),
    "failed to get native graph nodes"
  );
  return nodes;
}

inline std::vector<std::pair<cudaGraphNode_t, cudaGraphNode_t>>
cuda_graph_get_edges(cudaGraph_t graph) {
  size_t num_edges = cuda_graph_get_num_edges(graph);
  std::vector<cudaGraphNode_t> froms(num_edges), tos(num_edges);
  num_edges = cuda_graph_get_num_edges(graph, froms.data(), tos.data());
  std::vector<std::pair<cudaGraphNode_t, cudaGraphNode_t>> edges(num_edges);
  for(size_t i=0; i<num_edges; i++) {
    edges[i] = std::make_pair(froms[i], tos[i]);
  }
  return edges;
}

inline cudaGraphNodeType cuda_get_graph_node_type(cudaGraphNode_t node) {
  cudaGraphNodeType type;
  TF_CHECK_CUDA(
    cudaGraphNodeGetType(node, &type), "failed to get native graph node type"
  );
  return type;
}
 
// ----------------------------------------------------------------------------
// cudaTask Types
// ----------------------------------------------------------------------------

constexpr const char* to_string(cudaGraphNodeType type) {
  switch (type) {
    case cudaGraphNodeTypeKernel:             return "Kernel";
    case cudaGraphNodeTypeMemcpy:             return "Memcpy";
    case cudaGraphNodeTypeMemset:             return "Memset";
    case cudaGraphNodeTypeHost:               return "Host";
    case cudaGraphNodeTypeGraph:              return "Graph";
    case cudaGraphNodeTypeEmpty:              return "Empty";
    case cudaGraphNodeTypeWaitEvent:          return "WaitEvent";
    case cudaGraphNodeTypeEventRecord:        return "EventRecord";
    case cudaGraphNodeTypeExtSemaphoreSignal: return "ExtSemaphoreSignal";
    case cudaGraphNodeTypeExtSemaphoreWait:   return "ExtSemaphoreWait";
    case cudaGraphNodeTypeMemAlloc:           return "MemAlloc";
    case cudaGraphNodeTypeMemFree:            return "MemFree";
    case cudaGraphNodeTypeConditional:        return "Conditional";
    default:                                  return "undefined";
  }
}
 
// ----------------------------------------------------------------------------
// cudaTask
// ----------------------------------------------------------------------------

class cudaTask {
 
  template <typename Creator, typename Deleter>
  friend class cudaGraphBase;
  
  template <typename Creator, typename Deleter>
  friend class cudaGraphExecBase;
 
  friend class cudaFlow;
  friend class cudaFlowCapturer;
  friend class cudaFlowCapturerBase;
 
  friend std::ostream& operator << (std::ostream&, const cudaTask&);
 
  public:

    cudaTask() = default;

    cudaTask(const cudaTask&) = default;

    cudaTask& operator = (const cudaTask&) = default;

    template <typename... Ts>
    cudaTask& precede(Ts&&... tasks);

    template <typename... Ts>
    cudaTask& succeed(Ts&&... tasks);

    size_t num_successors() const;

    size_t num_predecessors() const;

    auto type() const;

    void dump(std::ostream& os) const;
 
  private:
 
    cudaTask(cudaGraph_t, cudaGraphNode_t);
    
    cudaGraph_t _native_graph {nullptr};
    cudaGraphNode_t _native_node {nullptr};
};
 
// Constructor
inline cudaTask::cudaTask(cudaGraph_t native_graph, cudaGraphNode_t native_node) : 
  _native_graph {native_graph}, _native_node  {native_node} {
}
  
// Function: precede
template <typename... Ts>
cudaTask& cudaTask::precede(Ts&&... tasks) {
  (
    cuda_graph_add_dependencies(
      _native_graph, &_native_node, &(tasks._native_node), 1
    ), ...
  );
  return *this;
}
 
// Function: succeed
template <typename... Ts>
cudaTask& cudaTask::succeed(Ts&&... tasks) {
  (tasks.precede(*this), ...);
  return *this;
}
 
// Function: num_predecessors
inline size_t cudaTask::num_predecessors() const {
  return cuda_graph_node_get_dependencies(_native_node, nullptr);
}
 
// Function: num_successors
inline size_t cudaTask::num_successors() const {
  return cuda_graph_node_get_dependent_nodes(_native_node, nullptr);
}
 
// Function: type
inline auto cudaTask::type() const {
  cudaGraphNodeType type;
  cudaGraphNodeGetType(_native_node, &type);
  return type;
}
 
// Function: dump
inline void cudaTask::dump(std::ostream& os) const {
  os << "cudaTask [type=" << to_string(type()) << ']';
}

inline std::ostream& operator << (std::ostream& os, const cudaTask& ct) {
  ct.dump(os);
  return os;
}
 
// ----------------------------------------------------------------------------
// cudaGraph
// ----------------------------------------------------------------------------

class cudaGraphCreator {
 
  public:

  cudaGraph_t operator () () const {
    cudaGraph_t g;
    TF_CHECK_CUDA(cudaGraphCreate(&g, 0), "failed to create a CUDA native graph");
    return g;
  }
  
  cudaGraph_t operator () (cudaGraph_t graph) const {
    return graph;
  }
 
};

class cudaGraphDeleter {
 
  public:
 
  void operator () (cudaGraph_t g) const {
    cudaGraphDestroy(g);
  }
};
  

template <typename Creator, typename Deleter>
class cudaGraphBase : public std::unique_ptr<std::remove_pointer_t<cudaGraph_t>, cudaGraphDeleter> {
  
  static_assert(std::is_pointer_v<cudaGraph_t>, "cudaGraph_t is not a pointer type");
 
  public:
  
  using base_type = std::unique_ptr<std::remove_pointer_t<cudaGraph_t>, Deleter>;

  template <typename... ArgsT>
  explicit cudaGraphBase(ArgsT&& ... args) : base_type(
    Creator{}(std::forward<ArgsT>(args)...), Deleter()
  ) {
  }  
  
  cudaGraphBase(cudaGraphBase&&) = default;

  cudaGraphBase& operator = (cudaGraphBase&&) = default;

  size_t num_nodes() const;
  
  size_t num_edges() const;

  bool empty() const;

  void dump(std::ostream& os);
 
  // ------------------------------------------------------------------------
  // Graph building routines
  // ------------------------------------------------------------------------

  cudaTask noop();

  template <typename C>
  cudaTask host(C&& callable, void* user_data);

  template <typename F, typename... ArgsT>
  cudaTask kernel(dim3 g, dim3 b, size_t s, F f, ArgsT... args);

  cudaTask memset(void* dst, int v, size_t count);

  cudaTask memcpy(void* tgt, const void* src, size_t bytes);

  template <typename T, std::enable_if_t<
    is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>* = nullptr
  >
  cudaTask zero(T* dst, size_t count);

  template <typename T, std::enable_if_t<
    is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>* = nullptr
  >
  cudaTask fill(T* dst, T value, size_t count);

  template <typename T,
    std::enable_if_t<!std::is_same_v<T, void>, void>* = nullptr
  >
  cudaTask copy(T* tgt, const T* src, size_t num);
  
  // ------------------------------------------------------------------------
  // generic algorithms
  // ------------------------------------------------------------------------

  template <typename C>
  cudaTask single_task(C c);
  
  template <typename I, typename C, typename E = cudaDefaultExecutionPolicy>
  cudaTask for_each(I first, I last, C callable);
  
  template <typename I, typename C, typename E = cudaDefaultExecutionPolicy>
  cudaTask for_each_index(I first, I last, I step, C callable);
  
  template <typename I, typename O, typename C, typename E = cudaDefaultExecutionPolicy>
  cudaTask transform(I first, I last, O output, C op);
  
  template <typename I1, typename I2, typename O, typename C, typename E = cudaDefaultExecutionPolicy>
  cudaTask transform(I1 first1, I1 last1, I2 first2, O output, C op);
 
  private:
 
  cudaGraphBase(const cudaGraphBase&) = delete;
  cudaGraphBase& operator = (const cudaGraphBase&) = delete;
};
 
// query the number of nodes
template <typename Creator, typename Deleter>
size_t cudaGraphBase<Creator, Deleter>::num_nodes() const {
  size_t n;
  TF_CHECK_CUDA(
    cudaGraphGetNodes(this->get(), nullptr, &n),
    "failed to get native graph nodes"
  );
  return n;
}
 
// query the emptiness
template <typename Creator, typename Deleter>
bool cudaGraphBase<Creator, Deleter>::empty() const {
  return num_nodes() == 0;
}
 
// query the number of edges
template <typename Creator, typename Deleter>
size_t cudaGraphBase<Creator, Deleter>::num_edges() const {
  return cuda_graph_get_num_edges(this->get());
}

//inline void cudaGraph::dump(std::ostream& os) {
//  
//  // acquire the native handle
//  auto g = this->get();
//
//  os << "digraph cudaGraph {\n";
//
//  std::stack<std::tuple<cudaGraph_t, cudaGraphNode_t, int>> stack;
//  stack.push(std::make_tuple(g, nullptr, 1));
//
//  int pl = 0;
//
//  while(stack.empty() == false) {
//
//    auto [graph, parent, l] = stack.top();
//    stack.pop();
//
//    for(int i=0; i<pl-l+1; i++) {
//      os << "}\n";
//    }
//
//    os << "subgraph cluster_p" << graph << " {\n"
//       << "label=\"cudaGraph-L" << l << "\";\n"
//       << "color=\"purple\";\n";
//
//    auto nodes = cuda_graph_get_nodes(graph);
//    auto edges = cuda_graph_get_edges(graph);
//
//    for(auto& [from, to] : edges) {
//      os << 'p' << from << " -> " << 'p' << to << ";\n";
//    }
//
//    for(auto& node : nodes) {
//      auto type = cuda_get_graph_node_type(node);
//      if(type == cudaGraphNodeTypeGraph) {
//
//        cudaGraph_t child_graph;
//        TF_CHECK_CUDA(cudaGraphChildGraphNodeGetGraph(node, &child_graph), "");
//        stack.push(std::make_tuple(child_graph, node, l+1));
//
//        os << 'p' << node << "["
//           << "shape=folder, style=filled, fontcolor=white, fillcolor=purple, "
//           << "label=\"cudaGraph-L" << l+1
//           << "\"];\n";
//      }
//      else {
//        os << 'p' << node << "[label=\""
//           << to_string(type)
//           << "\"];\n";
//      }
//    }
//
//    // precede to parent
//    if(parent != nullptr) {
//      std::unordered_set<cudaGraphNode_t> successors;
//      for(const auto& p : edges) {
//        successors.insert(p.first);
//      }
//      for(auto node : nodes) {
//        if(successors.find(node) == successors.end()) {
//          os << 'p' << node << " -> " << 'p' << parent << ";\n";
//        }
//      }
//    }
//
//    // set the previous level
//    pl = l;
//  }
//
//  for(int i=0; i<=pl; i++) {
//    os << "}\n";
//  }
//}
 
// dump the graph
template <typename Creator, typename Deleter>
void cudaGraphBase<Creator, Deleter>::dump(std::ostream& os) {
 
  // Generate a unique temporary filename in the system's temp directory using filesystem
  auto temp_path = std::filesystem::temp_directory_path() / "graph_";
  std::random_device rd;
  std::uniform_int_distribution<int> dist(100000, 999999); // Generates a random number
  temp_path += std::to_string(dist(rd)) + ".dot";
 
  // Call the original function with the temporary file
  TF_CHECK_CUDA(cudaGraphDebugDotPrint(this->get(), temp_path.string().c_str(), 0), "");
 
  // Read the file and write to the output stream
  std::ifstream file(temp_path);
  if (file) {
    os << file.rdbuf();  // Copy file contents to the stream
    file.close();
    std::filesystem::remove(temp_path);  // Clean up the temporary file
  } else {
    TF_THROW("failed to open ", temp_path, " for dumping the CUDA graph");
  }
}
 
// Function: noop
template <typename Creator, typename Deleter>
cudaTask cudaGraphBase<Creator, Deleter>::noop() {
 
  cudaGraphNode_t node;
 
  TF_CHECK_CUDA(
    cudaGraphAddEmptyNode(&node, this->get(), nullptr, 0),
    "failed to create a no-operation (empty) node"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: host
template <typename Creator, typename Deleter>
template <typename C>
cudaTask cudaGraphBase<Creator, Deleter>::host(C&& callable, void* user_data) {
 
  cudaGraphNode_t node;
  cudaHostNodeParams p {callable, user_data};
 
  TF_CHECK_CUDA(
    cudaGraphAddHostNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a host node"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: kernel
template <typename Creator, typename Deleter>
template <typename F, typename... ArgsT>
cudaTask cudaGraphBase<Creator, Deleter>::kernel(
  dim3 g, dim3 b, size_t s, F f, ArgsT... args
) {
 
  cudaGraphNode_t node;
  cudaKernelNodeParams p;
 
  void* arguments[sizeof...(ArgsT)] = { (void*)(&args)... };
 
  p.func = (void*)f;
  p.gridDim = g;
  p.blockDim = b;
  p.sharedMemBytes = s;
  p.kernelParams = arguments;
  p.extra = nullptr;
 
  TF_CHECK_CUDA(
    cudaGraphAddKernelNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a kernel task"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: zero
template <typename Creator, typename Deleter>
template <typename T, std::enable_if_t<
  is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>*
>
cudaTask cudaGraphBase<Creator, Deleter>::zero(T* dst, size_t count) {
 
  cudaGraphNode_t node;
  auto p = cuda_get_zero_parms(dst, count);
 
  TF_CHECK_CUDA(
    cudaGraphAddMemsetNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a memset (zero) task"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: fill
template <typename Creator, typename Deleter>
template <typename T, std::enable_if_t<
  is_pod_v<T> && (sizeof(T)==1 || sizeof(T)==2 || sizeof(T)==4), void>*
>
cudaTask cudaGraphBase<Creator, Deleter>::fill(T* dst, T value, size_t count) {
 
  cudaGraphNode_t node;
  auto p = cuda_get_fill_parms(dst, value, count);
  TF_CHECK_CUDA(
    cudaGraphAddMemsetNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a memset (fill) task"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: copy
template <typename Creator, typename Deleter>
template <
  typename T,
  std::enable_if_t<!std::is_same_v<T, void>, void>*
>
cudaTask cudaGraphBase<Creator, Deleter>::copy(T* tgt, const T* src, size_t num) {
 
  cudaGraphNode_t node;
  auto p = cuda_get_copy_parms(tgt, src, num);
 
  TF_CHECK_CUDA(
    cudaGraphAddMemcpyNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a memcpy (copy) task"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: memset
template <typename Creator, typename Deleter>
cudaTask cudaGraphBase<Creator, Deleter>::memset(void* dst, int ch, size_t count) {
 
  cudaGraphNode_t node;
  auto p = cuda_get_memset_parms(dst, ch, count);
 
  TF_CHECK_CUDA(
    cudaGraphAddMemsetNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a memset task"
  );
 
  return cudaTask(this->get(), node);
}
 
// Function: memcpy
template <typename Creator, typename Deleter>
cudaTask cudaGraphBase<Creator, Deleter>::memcpy(void* tgt, const void* src, size_t bytes) {
 
  cudaGraphNode_t node;
  auto p = cuda_get_memcpy_parms(tgt, src, bytes);
 
  TF_CHECK_CUDA(
    cudaGraphAddMemcpyNode(&node, this->get(), nullptr, 0, &p),
    "failed to create a memcpy task"
  );
 
  return cudaTask(this->get(), node);
}
 
 
 
 
 
}  // end of namespace tf -----------------------------------------------------