GCC Code Coverage Report

chrono_batch_render/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
164 / 181
90.6%
Branches:
32 / 50
64.0%
Functions:
17 / 21
81.0%
python/src/device.cpp
Line Branch Exec Source
1 // SPDX-License-Identifier: BSD-3-Clause
2 #include <cstddef>
3 #include <cstdint>
4 #include <optional>
5 #include <string>
6 #include <vector>
7
8 #include <nanobind/ndarray.h>
9 #include <nanobind/stl/optional.h>
10 #include <nanobind/stl/string.h>
11 #include <nanobind/stl/vector.h>
12
13 #include "bindings.h"
14 #include "chrono_batch_render/util/log.h"
15
16 namespace nb = nanobind;
17
18 namespace chrono::batch_render::python {
19
20 using device::AdapterInfo;
21 using device::AdapterType;
22 using device::Backend;
23 using device::Device;
24 using device::DType;
25 using device::Sampling;
26 using device::Tensor;
27 using device::Traversal;
28
29 260 nb::object Numpy(std::vector<std::byte> bytes, const std::vector<size_t>& shape, DLDataType dtype) {
30 260 auto* owned = new std::vector<std::byte>(std::move(bytes));
31 260 nb::capsule owner(
32
1/2
✓ Branch 0 taken 260 times.
✗ Branch 1 not taken.
 260 owned, [](void* held) noexcept { delete static_cast<std::vector<std::byte>*>(held); });
33 260 return nb::cast(nb::ndarray<nb::numpy>(owned->data(), shape.size(), shape.data(), owner,
34
2/2
✓ Branch 1 taken 260 times.
✓ Branch 4 taken 260 times.
 260 nullptr, {dtype.code, dtype.bits, dtype.lanes}));
35 260 }
36
37 namespace {
38
39 2 void RegisterEnums(nb::module_& module) {
40 2 nb::enum_<Backend>(module, "Backend")
41 2 .value("VULKAN", Backend::Vulkan)
42 2 .value("CUDA", Backend::Cuda)
43 2 .value("HIP", Backend::Hip)
44 2 .value("METAL", Backend::Metal)
45 2 .value("CPU", Backend::Cpu);
46
47 2 nb::enum_<Traversal>(module, "Traversal")
48 2 .value("AUTO", Traversal::Auto)
49 2 .value("HARDWARE", Traversal::Hardware)
50 2 .value("SOFTWARE", Traversal::Software);
51
52 2 nb::enum_<Sampling>(module, "Sampling")
53 2 .value("AUTO", Sampling::Auto)
54 2 .value("HARDWARE", Sampling::Hardware)
55 2 .value("SOFTWARE", Sampling::Software);
56
57 2 nb::enum_<AdapterType>(module, "AdapterType")
58 2 .value("DISCRETE", AdapterType::Discrete)
59 2 .value("INTEGRATED", AdapterType::Integrated)
60 2 .value("VIRTUAL", AdapterType::Virtual)
61 2 .value("CPU", AdapterType::Cpu)
62 2 .value("OTHER", AdapterType::Other);
63
64 2 nb::enum_<LogLevel>(module, "LogLevel")
65 2 .value("DEBUG", LogLevel::Debug)
66 2 .value("INFO", LogLevel::Info)
67 2 .value("WARNING", LogLevel::Warning)
68 2 .value("ERROR", LogLevel::Error)
69 2 .value("OFF", LogLevel::Off);
70
71 2 nb::enum_<DType>(module, "DType")
72 2 .value("UINT8", DType::UInt8)
73 2 .value("UINT32", DType::UInt32)
74 2 .value("FLOAT32", DType::Float32)
75 2 .value("FLOAT64", DType::Float64)
76 2 .value("INT32", DType::Int32)
77 2 .value("UINT64", DType::UInt64)
78 2 .value("INT64", DType::Int64);
79 2 }
80
81 2 void RegisterAdapters(nb::module_& module) {
82 2 nb::class_<AdapterInfo>(module, "AdapterInfo")
83 2 .def_ro("backend", &AdapterInfo::backend)
84 2 .def_ro("name", &AdapterInfo::name)
85 2 .def_ro("alias", &AdapterInfo::alias)
86 2 .def_ro("ordinal", &AdapterInfo::ordinal)
87 2 .def_ro("type", &AdapterInfo::type)
88 2 .def_ro("architecture", &AdapterInfo::architecture)
89 2 .def_ro("memory_bytes", &AdapterInfo::memoryBytes)
90 2 .def_ro("unsupported", &AdapterInfo::unsupported)
91 2 .def_prop_ro("api_version",
92 ✗ [](const AdapterInfo& adapter) { return adapter.apiVersion.ToString(); })
93 2 .def_prop_ro("driver_version",
94 ✗ [](const AdapterInfo& adapter) { return adapter.driverVersion.ToString(); })
95 2 .def_prop_ro("hardware_traversal",
96
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 ✗ [](const AdapterInfo& adapter) { return adapter.features.hardwareTraversal; })
97 2 .def("__repr__", [](const AdapterInfo& adapter) {
98
0/1
✗ Branch 2 not taken.
 ✗ return "AdapterInfo(" + adapter.alias + ": " + adapter.name + ")";
99 });
100
101 2 module.def("list_adapters", &device::EnumerateAdapters);
102 2 module.def("backend_name",
103 1 [](Backend backend) { return std::string(device::BackendName(backend)); });
104 2 }
105
106 2 void RegisterDevice(nb::module_& module) {
107 2 nb::class_<Device>(module, "Device")
108
1/1
✓ Branch 1 taken 2 times.
 2 .def(nb::new_([](std::optional<Backend> backend, std::string adapter, Traversal traversal,
109 Sampling sampling, bool validation, bool synchronous, bool presentation) {
110 1 nb::gil_scoped_release release;
111 1 return Device::Create({.backend = backend,
112 .adapter = std::move(adapter),
113 .traversal = traversal,
114 .sampling = sampling,
115 .validation = validation,
116 .synchronous = synchronous,
117 1 .presentation = presentation});
118
1/1
✓ Branch 3 taken 1 time.
 2 }),
119
2/2
✓ Branch 1 taken 2 times.
✓ Branch 4 taken 2 times.
 6 nb::kw_only(), nb::arg("backend") = nb::none(), nb::arg("adapter") = "",
120
2/2
✓ Branch 1 taken 2 times.
✓ Branch 4 taken 2 times.
 6 nb::arg("traversal") = Traversal::Auto, nb::arg("sampling") = Sampling::Auto,
121
2/2
✓ Branch 1 taken 2 times.
✓ Branch 4 taken 2 times.
 6 nb::arg("validation") = false, nb::arg("synchronous") = false,
122
2/2
✓ Branch 1 taken 2 times.
✓ Branch 4 taken 2 times.
 6 nb::arg("presentation") = false)
123 2 .def_prop_ro("adapter", &Device::Adapter)
124 2 .def_prop_ro("traversal", &Device::GetTraversal)
125 4 .def_prop_ro("sampling", &Device::GetSampling)
126
1/1
✓ Branch 1 taken 2 times.
 2 .def_prop_ro("shared_device",
127 ✗ [](Device& device) -> std::optional<nb::tuple> {
128
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ if (const std::optional<DLDevice> shared = device.SharedDevice()) {
129 ✗ return nb::make_tuple(int(shared->device_type), shared->device_id);
130 }
131 ✗ return std::nullopt;
132 })
133 2 .def(
134 "create_tensor",
135 1 [](Device& device, DType dtype, std::vector<int64_t> shape, std::string label) {
136 1 nb::gil_scoped_release release;
137
1/1
✓ Branch 2 taken 1 time.
 1 return device.CreateTensor(dtype, std::move(shape), std::move(label));
138 1 },
139
1/1
✓ Branch 1 taken 2 times.
 6 nb::arg("dtype"), nb::arg("shape"), nb::arg("label") = "tensor")
140 2 .def(
141 "copy",
142 3 [](Device& device, const Tensor& src, const Tensor& dst,
143 std::optional<device::Operation> after) {
144 3 nb::gil_scoped_release release;
145
2/3
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
✓ Branch 3 taken 3 times.
 6 return device.Copy(src, dst, after.value_or(device::Operation{}));
146 3 },
147 4 nb::arg("src"), nb::arg("dst"), nb::arg("after") = nb::none(), nb::keep_alive<0, 1>())
148 2 .def(
149 "import_tensor",
150 8 [](Device& device, nb::handle source, std::optional<uintptr_t> stream) {
151 8 const TakenTensor taken = TakeTensor(source, stream);
152 8 const std::shared_ptr<void> held = Hold(taken.managed);
153 8 nb::gil_scoped_release release;
154
2/2
✓ Branch 1 taken 8 times.
✓ Branch 4 taken 8 times.
 16 return device.Import(taken.managed->dl_tensor, taken.producer);
155
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 16 },
156 4 nb::arg("source"), nb::arg("stream") = nb::none(), nb::keep_alive<0, 2>())
157 2 .def(
158 "follow",
159 2 [](Device& device, std::optional<uintptr_t> stream) {
160
0/1
✗ Branch 1 not taken.
 ✗ return device.Follow(ToStream(stream));
161 },
162 4 nb::arg("stream") = nb::none(), nb::call_guard<nb::gil_scoped_release>(),
163 nb::keep_alive<0, 1>())
164 2 .def("__repr__",
165 ✗ [](const Device& device) { return "Device(" + device.Adapter().name + ")"; });
166
167 2 nb::class_<device::Operation>(module, "Operation")
168 4 .def_prop_ro("ready", &device::Operation::Ready, nb::call_guard<nb::gil_scoped_release>())
169 2 .def(
170 "wait",
171 2 [](const device::Operation& operation, std::optional<uintptr_t> stream) {
172
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 11 times.
 11 if (stream) {
173 ✗ operation.Wait(ToStream(stream));
174 } else {
175 11 operation.Wait();
176 }
177 11 },
178 4 nb::arg("stream") = nb::none(), nb::call_guard<nb::gil_scoped_release>());
179 2 }
180
181 2 void RegisterTensor(nb::module_& module) {
182 2 nb::class_<Tensor>(module, "Tensor")
183 4 .def_prop_ro("dtype", &Tensor::Type)
184 2 .def_prop_ro("shape",
185 2 [](const Tensor& tensor) {
186 2 return std::vector<int64_t>(tensor.Shape().begin(), tensor.Shape().end());
187 })
188 4 .def_prop_ro("address", &Tensor::Address)
189 2 .def_prop_ro("device",
190 ✗ [](const Tensor& tensor) { return Ref<Device>(&tensor.GetDevice()); })
191 2 .def("__len__",
192
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ [](const Tensor& tensor) { return tensor.Shape().empty() ? 0 : tensor.Shape()[0]; })
193 2 .def("__getitem__",
194 2 [](const Tensor& tensor, int64_t index) {
195
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 4 const int64_t size = tensor.Shape().empty() ? 0 : tensor.Shape()[0];
196
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (index < -size || index >= size) {
197
0/1
✗ Branch 2 not taken.
 ✗ throw nb::index_error();
198 }
199
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 return tensor.Select(index < 0 ? index + size : index);
200 })
201 2 .def(
202 "__dlpack__",
203 2 [](const Tensor& tensor, std::optional<int64_t> stream,
204 std::optional<nb::tuple> maxVersion, std::optional<nb::tuple> dlDevice,
205 std::optional<bool> copy) {
206
0/1
✗ Branch 3 not taken.
 4 return ExportTensor(tensor, stream, maxVersion, dlDevice, copy);
207 },
208 4 nb::kw_only(), nb::arg("stream") = nb::none(), nb::arg("max_version") = nb::none(),
209 4 nb::arg("dl_device") = nb::none(), nb::arg("copy") = nb::none())
210 2 .def("__dlpack_device__", &TensorDevice)
211 2 .def("numpy", [](const Tensor& tensor) {
212 222 std::vector<std::byte> bytes;
213 222 {
214 222 nb::gil_scoped_release release;
215
1/1
✓ Branch 1 taken 222 times.
 444 bytes = tensor.ToHostBytes();
216 ✗ }
217
2/2
✓ Branch 1 taken 222 times.
✓ Branch 5 taken 222 times.
 444 return Numpy(std::move(bytes), {tensor.Shape().begin(), tensor.Shape().end()},
218
3/3
✓ Branch 1 taken 222 times.
✓ Branch 5 taken 222 times.
✓ Branch 8 taken 222 times.
 444 DataType(tensor.Type()));
219 222 });
220 2 }
221
222 2 void RegisterBuild(nb::module_& module) {
223 2 nb::class_<BuildInfo>(module, "BuildInfo")
224 2 .def_prop_ro("version",
225 1 [](const BuildInfo& build) { return std::string(build.fullVersion); })
226 2 .def_prop_ro("commit", [](const BuildInfo& build) { return std::string(build.commit); })
227 2 .def_ro("modified", &BuildInfo::modified)
228 2 .def_prop_ro("commit_date",
229 ✗ [](const BuildInfo& build) { return std::string(build.commitDate); })
230 2 .def_prop_ro("build_date",
231 ✗ [](const BuildInfo& build) { return std::string(build.buildDate); });
232 2 module.def("build_info", &GetBuildInfo, nb::rv_policy::reference);
233
234 2 nb::class_<device::BuildProgress>(module, "BuildProgress")
235 2 .def_ro("queued", &device::BuildProgress::queued)
236 2 .def_ro("finished", &device::BuildProgress::finished)
237 2 .def_ro("cached", &device::BuildProgress::cached)
238 2 .def_ro("building", &device::BuildProgress::building)
239 4 .def_prop_ro("done", &device::BuildProgress::Done);
240 2 module.def("kernel_builds", &device::CurrentBuilds);
241 2 module.def("get_log_level", &GetLogLevel);
242 2 module.def("set_log_level", &SetLogLevel, nb::arg("level"));
243 2 }
244
245 }
246
247 2 void BindDevice(nb::module_& module) {
248 2 RegisterBuild(module);
249 2 RegisterEnums(module);
250 2 RegisterAdapters(module);
251 2 RegisterDevice(module);
252 2 RegisterTensor(module);
253 2 }
254
255 }
256