cuda_task.hpp

Go to the documentation of this file.

#pragma once
 
#include "cuda_graph.hpp"
 
namespace tf {
 
// ----------------------------------------------------------------------------
// cudaTask Types
// ----------------------------------------------------------------------------
 
enum class cudaTaskType : int {
  EMPTY = 0,
  HOST,
  MEMSET,
  MEMCPY,
  KERNEL,
  SUBFLOW,
  CAPTURE,
  UNDEFINED
};
 
constexpr const char* to_string(cudaTaskType type) {
  switch(type) {
    case cudaTaskType::EMPTY:   return "empty";
    case cudaTaskType::HOST:    return "host";
    case cudaTaskType::MEMSET:  return "memset";
    case cudaTaskType::MEMCPY:  return "memcpy";
    case cudaTaskType::KERNEL:  return "kernel";
    case cudaTaskType::SUBFLOW: return "subflow";
    case cudaTaskType::CAPTURE: return "capture";
    default:                    return "undefined";
  }
}
 
// ----------------------------------------------------------------------------
// cudaTask
// ----------------------------------------------------------------------------
 
class cudaTask {
 
  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);
 
    cudaTask& name(const std::string& name);
 
    const std::string& name() const;
 
    size_t num_successors() const;
 
    size_t num_dependents() const;
 
    bool empty() const;
 
    cudaTaskType type() const;
 
    template <typename T>
    void dump(T& ostream) const;
 
    template <typename V>
    void for_each_successor(V&& visitor) const;
 
    template <typename V>
    void for_each_dependent(V&& visitor) const;
 
  private:
 
    cudaTask(cudaNode*);
 
    cudaNode* _node {nullptr};
};
 
// Constructor
inline cudaTask::cudaTask(cudaNode* node) : _node {node} {
}
 
// Function: precede
template <typename... Ts>
cudaTask& cudaTask::precede(Ts&&... tasks) {
  (_node->_precede(tasks._node), ...);
  return *this;
}
 
// Function: succeed
template <typename... Ts>
cudaTask& cudaTask::succeed(Ts&&... tasks) {
  (tasks._node->_precede(_node), ...);
  return *this;
}
 
// Function: empty
inline bool cudaTask::empty() const {
  return _node == nullptr;
}
 
// Function: name
inline cudaTask& cudaTask::name(const std::string& name) {
  _node->_name = name;
  return *this;
}
 
// Function: name
inline const std::string& cudaTask::name() const {
  return _node->_name;
}
 
// Function: num_successors
inline size_t cudaTask::num_successors() const {
  return _node->_successors.size();
}
 
// Function: num_dependents
inline size_t cudaTask::num_dependents() const {
  return _node->_dependents.size();
}
 
// Function: type
inline cudaTaskType cudaTask::type() const {
  switch(_node->_handle.index()) {
    case cudaNode::EMPTY:   return cudaTaskType::EMPTY;
    case cudaNode::HOST:    return cudaTaskType::HOST;
    case cudaNode::MEMSET:  return cudaTaskType::MEMSET;
    case cudaNode::MEMCPY:  return cudaTaskType::MEMCPY;
    case cudaNode::KERNEL:  return cudaTaskType::KERNEL;
    case cudaNode::SUBFLOW: return cudaTaskType::SUBFLOW;
    case cudaNode::CAPTURE: return cudaTaskType::CAPTURE;
    default:                return cudaTaskType::UNDEFINED;
  }
}
 
// Procedure: dump
template <typename T>
void cudaTask::dump(T& os) const {
  os << "cudaTask ";
  if(_node->_name.empty()) os << _node;
  else os << _node->_name;
  os << " [type=" << to_string(type()) << ']';
}
 
// Function: for_each_successor
template <typename V>
void cudaTask::for_each_successor(V&& visitor) const {
  for(size_t i=0; i<_node->_successors.size(); ++i) {
    visitor(cudaTask(_node->_successors[i]));
  }
}
 
// Function: for_each_dependent
template <typename V>
void cudaTask::for_each_dependent(V&& visitor) const {
  for(size_t i=0; i<_node->_dependents.size(); ++i) {
    visitor(cudaTask(_node->_dependents[i]));
  }
}
 
// ----------------------------------------------------------------------------
// global ostream
// ----------------------------------------------------------------------------
 
inline std::ostream& operator << (std::ostream& os, const cudaTask& ct) {
  ct.dump(os);
  return os;
}
 
}  // end of namespace tf -----------------------------------------------------