この記事はもともと別場所にて公開していました。

cudaStreamBeginCaptureなどを使えば CUDA Graph は比較的簡単に作ることができるが、複数のストリームを使って並列に動くようなプログラムを capture することができないので、そのような場合は CUDA Graphs の API を叩きつつ自前でグラフを作る必要がある。これは capture が使える場合と比べて大変面倒なので、それを回避するために C++でクラスを書いた。内部では適当にcudaStreamBeginCaptureしつつ、できたグラフをcudaGraphAddChildGraphNodeで一つにまとめている。

cuFHE で使うために書いたので cuFHE 用になっている（stream をcufhe::Streamで扱っているなど）が、ちゃんと一般的に書けば一般的に使えるはず。一般的にしたコードは……そのうち……。

#include <cufhe.h>
#include <cufhe_gpu.cuh>

class CUDAGraphBuilder {
  friend class CUDAGraph;

public:
  using HostFunc = std::function<void()>;

private:
  cudaGraph_t graph_;
  std::vector<std::shared_ptr<HostFunc>> hostFuncInsts_;

private:
  static void handlerHostNode(void *userData) {
    HostFunc *body = static_cast<HostFunc *>(userData);
    (*body)();
  }

public:
  CUDAGraphBuilder() { cudaGraphCreate(&graph_, 0); }

  void addDep(cudaGraphNode_t from, cudaGraphNode_t to) {
    cudaGraphAddDependencies(graph_, &from, &to, 1);
  }

  template <class Fun>
  cudaGraphNode_t addKernel(const std::vector<cudaGraphNode_t> &deps, Fun fun) {
    cufhe::Stream stream;
    stream.Create();

    cudaGraph_t emb_graph;
    cudaStreamBeginCapture(stream.st(), cudaStreamCaptureModeGlobal);
    fun(stream);
    cudaStreamEndCapture(stream.st(), &emb_graph);
    cudaGraphNode_t node;
    cudaGraphAddChildGraphNode(&node, graph_, deps.data(), deps.size(),
                               emb_graph);

    stream.Destroy();
    return node;
  }

  template <class Fun> cudaGraphNode_t addKernel(Fun &&fun) {
    return addKernel({}, std::forward<Fun>(fun));
  }

  cudaGraphNode_t addHost(const std::vector<cudaGraphNode_t> &deps,
                          const std::function<void()> &fun) {
    cudaGraphNode_t node;
    auto funptr = std::make_shared<HostFunc>(fun);
    hostFuncInsts_.push_back(funptr);
    const cudaHostNodeParams param = {handlerHostNode,
                                      static_cast<void *>(funptr.get())};
    cudaGraphAddHostNode(&node, graph_, deps.data(), deps.size(), &param);
    return node;
  }

  template <class... Args> cudaGraphNode_t addHost(Args &&... args) {
    return addHost({}, std::forward<Args>(args)...);
  }

  cudaGraphExec_t instantiate() {
    cudaGraphExec_t instance;
    cudaGraphInstantiate(&instance, graph_, NULL, NULL, 0);
    return instance;
  }
};

class CUDAGraph {
private:
  cudaGraphExec_t instance_;
  std::vector<std::shared_ptr<CUDAGraphBuilder::HostFunc>> hostFuncInsts_;

public:
  CUDAGraph(CUDAGraphBuilder &&b)
      : hostFuncInsts_(std::move(b.hostFuncInsts_)) {
    cudaGraphInstantiate(&instance_, b.graph_, NULL, NULL, 0);
  }

  void runSync(cudaStream_t st) {
    cudaGraphLaunch(instance_, st);
    cudaStreamSynchronize(st);
  }
};

次のように使うと、aとbが並列に走ってくれるのが嬉しい。

CUDAGraphBuilder bld;

// グラフを作る
auto a = bld.addKernel([&](auto &&s) { /* sを使ってkernelをlaunchする */ });
auto b = bld.addKernel([&](auto &&s) { /* sを使ってkernelをlaunchする */ });
auto c = bld.addHost([&] { /* hostで行う処理を書く */ });
bld.addDep(a, c); // aが終わってからcが走る
bld.addDep(b, c); // bが終わってからcが走る

// instantiateする
CUDAGraph graph{std::move(bld)};
// streamを指定して使う
cufhe::Stream stream;
stream.Create();
graph.runSync(stream);