GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 0.0% 0 / 0 / 79
Functions: 0.0% 0 / 0 / 25
Branches: 0.0% 0 / 0 / 104
src/graph.cpp
Line Branch Exec Source
1 // Device-neutral graph capture (public API: include/brotensor/cuda_graph.h).
2 // CudaGraphCapture picks a device and opens a recorder from the factory its
3 // backend registered (detail/graph_backend.h); CudaGraph owns what the
4 // recorder's finish() returned. All backend mechanics stay in the backends.
5
6 #include <brotensor/cuda_graph.h>
7 #include <brotensor/detail/dispatch.h>
8 #include <brotensor/detail/graph_backend.h>
9 #include <brotensor/runtime.h>
10
11 #include <array>
12 #include <stdexcept>
13 #include <string>
14
15 namespace brotensor {
16
17 namespace detail {
18
19 namespace {
20
21 // Written by the backends' probes (single-threaded, inside init()), read by
22 // capture scopes; plain function pointers, zero-initialised before any static
23 // constructor runs.
24 std::array<GraphRecorderFactory, 5> g_factories{};
25
26 ✗ int slot(DeviceType dt) { return static_cast<int>(dt); }
27
28 } // namespace
29
30 ✗ void register_graph_backend(DeviceType dt, GraphRecorderFactory factory) {
31 ✗ const int i = slot(dt);
32 ✗ if (i < 0 || i >= static_cast<int>(g_factories.size())) {
33 ✗ throw std::runtime_error("brotensor: register_graph_backend: bad device type");
34 }
35 ✗ g_factories[static_cast<std::size_t>(i)] = factory;
36 ✗ }
37
38 ✗ GraphRecorderFactory graph_backend(DeviceType dt) {
39 ✗ const int i = slot(dt);
40 ✗ if (i < 0 || i >= static_cast<int>(g_factories.size())) return nullptr;
41 ✗ return g_factories[static_cast<std::size_t>(i)];
42 ✗ }
43
44 } // namespace detail
45
46 namespace {
47
48 ✗ bool capturable(Device d) {
49 ✗ return d.is_gpu() && detail::graph_backend(d.type) != nullptr && detail::is_registered(d);
50 }
51
52 // The device a default-constructed capture records on, as (type, index) with
53 // index -1 for "the backend's current device". Throws when there is none.
54 ✗ Device pick_capture_device() {
55 ✗ const Device d = detail::resolve_device_alias(default_device());
56 ✗ if (capturable(d)) return d;
57 // A CPU / Metal default device: the capture used to be CUDA's current
58 // stream whatever the default device was; keep that, then Vulkan.
59 ✗ if (capturable(Device::cuda(0))) return Device(DeviceType::CUDA, -1);
60 ✗ if (capturable(Device::vulkan(0))) return Device::vulkan(0);
61 ✗ throw std::runtime_error(
62 "brotensor: CudaGraphCapture: no GPU with graph capture (CUDA or Vulkan) is "
63 ✗ "registered; default device is " + to_string(d));
64 ✗ }
65
66 } // namespace
67
68 ✗ bool graph_capture_available(Device d) { return capturable(detail::resolve_device_alias(d)); }
69
70 ✗ bool graph_capture_available() {
71 try {
72 ✗ (void)pick_capture_device();
73 ✗ return true;
74 ✗ } catch (const std::exception&) {
75 ✗ return false;
76 ✗ }
77 ✗ }
78
79 // ─── CudaGraph ──────────────────────────────────────────────────────────────
80
81 ✗ struct CudaGraph::Impl {
82 ✗ Device device = Device::cpu();
83 std::unique_ptr<detail::GraphExec> exec;
84 };
85
86 ✗ CudaGraph::CudaGraph() = default;
87 ✗ CudaGraph::~CudaGraph() = default;
88 ✗ CudaGraph::CudaGraph(CudaGraph&&) noexcept = default;
89 ✗ CudaGraph& CudaGraph::operator=(CudaGraph&&) noexcept = default;
90
91 ✗ bool CudaGraph::valid() const { return impl_ && impl_->exec != nullptr; }
92
93 ✗ void CudaGraph::reset() { impl_.reset(); }
94
95 ✗ Device CudaGraph::device() const { return impl_ ? impl_->device : Device::cpu(); }
96
97 ✗ void CudaGraph::launch() {
98 ✗ if (!valid()) {
99 ✗ throw std::runtime_error("brotensor: CudaGraph::launch: no captured graph");
100 }
101 ✗ impl_->exec->launch();
102 ✗ }
103
104 // ─── CudaGraphCapture ───────────────────────────────────────────────────────
105
106 ✗ struct CudaGraphCapture::Impl {
107 ✗ Device device = Device::cpu();
108 std::unique_ptr<detail::GraphRecorder> recorder;
109 };
110
111 namespace {
112
113 ✗ std::unique_ptr<detail::GraphRecorder> open_recorder(Device d) {
114 ✗ detail::GraphRecorderFactory f = detail::graph_backend(d.type);
115 ✗ if (!f) {
116 ✗ throw std::runtime_error("brotensor: CudaGraphCapture: " + to_string(d) +
117 " has no graph capture");
118 }
119 ✗ auto r = f(d.index);
120 ✗ if (!r) throw std::runtime_error("brotensor: CudaGraphCapture: the backend refused to capture");
121 ✗ return r;
122 ✗ }
123
124 } // namespace
125
126 ✗ CudaGraphCapture::CudaGraphCapture() : impl_(new Impl) {
127 ✗ const Device d = pick_capture_device();
128 ✗ impl_->recorder = open_recorder(d);
129 ✗ impl_->device = Device(d.type, impl_->recorder->device_index());
130 ✗ }
131
132 ✗ CudaGraphCapture::CudaGraphCapture(Device d) : impl_(new Impl) {
133 ✗ d = detail::resolve_device_alias(d);
134 ✗ if (!capturable(d)) {
135 ✗ throw std::runtime_error("brotensor: CudaGraphCapture: " + to_string(d) +
136 " is not a registered GPU with graph capture");
137 }
138 ✗ impl_->recorder = open_recorder(d);
139 ✗ impl_->device = Device(d.type, impl_->recorder->device_index());
140 ✗ }
141
142 ✗ CudaGraphCapture::~CudaGraphCapture() = default; // recorder's dtor aborts
143
144 ✗ Device CudaGraphCapture::device() const {
145 ✗ return impl_ ? impl_->device : Device::cpu();
146 }
147
148 ✗ CudaGraph CudaGraphCapture::finish() {
149 ✗ if (!impl_ || !impl_->recorder) {
150 ✗ throw std::runtime_error("brotensor: CudaGraphCapture::finish: not capturing");
151 }
152 ✗ auto rec = std::move(impl_->recorder);
153 ✗ CudaGraph g;
154 ✗ g.impl_.reset(new CudaGraph::Impl);
155 ✗ g.impl_->device = impl_->device;
156 ✗ g.impl_->exec = rec->finish();
157 ✗ return g;
158 ✗ }
159
160 } // namespace brotensor
161