| 119 | } |
| 120 | |
| 121 | void EndCapture() |
| 122 | { |
| 123 | auto err = cudaStreamEndCapture(stream, &graph); |
| 124 | if (err != cudaSuccess) { |
| 125 | ngs_cuda_stream = prev_stream; |
| 126 | throw ngstd::Exception(std::string("[CudaGraph] cudaStreamEndCapture FAILED: ") |
| 127 | + cudaGetErrorString(err)); |
| 128 | } |
| 129 | |
| 130 | size_t numnodes = 0; |
| 131 | cudaGraphGetNodes(graph, nullptr, &numnodes); |
| 132 | std::cerr << "[CudaGraph] captured nodes: " << numnodes << std::endl; |
| 133 | if (numnodes == 0) |
| 134 | std::cerr << "[CudaGraph] WARNING: 0 nodes — ops may not be on capture stream!" << std::endl; |
| 135 | |
| 136 | err = cudaGraphInstantiate(&instance, graph, NULL, NULL, 0); |
| 137 | if (err != cudaSuccess) { |
| 138 | ngs_cuda_stream = prev_stream; |
| 139 | throw ngstd::Exception(std::string("[CudaGraph] cudaGraphInstantiate FAILED: ") |
| 140 | + cudaGetErrorString(err)); |
| 141 | } |
| 142 | |
| 143 | capture_ok = true; |
| 144 | ngs_cuda_stream = prev_stream; |
| 145 | if (stream_change_callback) { |
| 146 | stream_change_callback(ngs_cuda_stream); |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | void Launch() |
| 151 | { |