chrono_batch_render/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
134 / 180
74.4%
Branches:
89 / 129
69.0%
Functions:
22 / 26
84.6%

device/tensor.cpp
Line Branch Exec Source
1 // SPDX-License-Identifier: BSD-3-Clause
2 #include "chrono_batch_render/device/tensor.h"
3
4 #include <algorithm>
5 #include <cstring>
6 #include <memory>
7 #include <optional>
8
9 #include "chrono_batch_render/device/buffer.h"
10 #include "chrono_batch_render/device/device.h"
11 #include "chrono_batch_render/device/queue.h"
12 #include "chrono_batch_render/util/error.h"
13
14 namespace chrono::batch_render::device {
15 namespace {
16
17 70951 std::vector<int64_t> ContiguousStrides(std::span<const int64_t> shape) {
18 70951 std::vector<int64_t> strides(shape.size());
19 70951 int64_t stride = 1;
20
2/2
✓ Branch 0 taken 141508 times.
✓ Branch 1 taken 70951 times.
212459 for (size_t i = shape.size(); i-- > 0;) {
21 141508 strides[i] = stride;
22 141508 stride *= shape[i];
23 }
24 70951 return strides;
25 }
26
27 71033 int64_t LastElement(std::span<const int64_t> shape, std::span<const int64_t> strides) {
28 71033 int64_t last = 0;
29
2/2
✓ Branch 0 taken 141711 times.
✓ Branch 1 taken 71033 times.
212744 for (size_t i = 0; i < shape.size(); ++i) {
30
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 141711 times.
141711 last += (std::max<int64_t>(shape[i], 1) - 1) * strides[i];
31 }
32 71033 return last;
33 }
34
35 8564 bool Packed(std::span<const int64_t> shape, std::span<const int64_t> strides) {
36 8564 int64_t packed = 1;
37
2/2
✓ Branch 0 taken 29607 times.
✓ Branch 1 taken 8557 times.
38164 for (size_t axis = shape.size(); axis-- > 0;) {
38
4/4
✓ Branch 0 taken 21510 times.
✓ Branch 1 taken 8097 times.
✓ Branch 2 taken 21503 times.
✓ Branch 3 taken 7 times.
29607 if (shape[axis] > 1 && strides[axis] != packed) {
39 return false;
40 }
41 29600 packed *= shape[axis];
42 }
43 return true;
44 }
45
46 template <size_t kBytes>
47 13 std::byte* CopyStrided(std::byte* to, const std::byte* from, int64_t count, int64_t step,
48 uint64_t bytes) {
49 13 const uint64_t each = kBytes == 0 ? bytes : kBytes;
50
6/10
std::byte* chrono::batch_render::device::(anonymous namespace)::CopyStrided<0ul>(std::byte*, std::byte const*, long, long, unsigned long):
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
std::byte* chrono::batch_render::device::(anonymous namespace)::CopyStrided<1ul>(std::byte*, std::byte const*, long, long, unsigned long):
✗ Branch 0 not taken.
✗ Branch 1 not taken.
std::byte* chrono::batch_render::device::(anonymous namespace)::CopyStrided<2ul>(std::byte*, std::byte const*, long, long, unsigned long):
✗ Branch 0 not taken.
✗ Branch 1 not taken.
std::byte* chrono::batch_render::device::(anonymous namespace)::CopyStrided<4ul>(std::byte*, std::byte const*, long, long, unsigned long):
✓ Branch 0 taken 27 times.
✓ Branch 1 taken 11 times.
std::byte* chrono::batch_render::device::(anonymous namespace)::CopyStrided<8ul>(std::byte*, std::byte const*, long, long, unsigned long):
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 1 time.
44 for (int64_t n = 0; n < count; ++n, from += step, to += each) {
51 31 std::memcpy(to, from, kBytes == 0 ? bytes : kBytes);
52 }
53 13 return to;
54 }
55
56 5 void Gather(std::span<const std::byte> source, std::span<const int64_t> shape,
57 std::span<const int64_t> strides, uint32_t elementBytes, int64_t elements,
58 std::byte* out) {
59 5 size_t outer = shape.size();
60 5 int64_t run = 1;
61
5/6
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 3 times.
✓ Branch 5 taken 5 times.
9 while (outer > 0 && (shape[outer - 1] == 1 || strides[outer - 1] == run)) {
62 4 run *= shape[--outer];
63 }
64 5 const uint64_t runBytes = uint64_t(run) * elementBytes;
65 5 const size_t last = outer - 1;
66 5 const int64_t count = shape[last];
67 5 const int64_t step = strides[last] * elementBytes;
68 5 std::vector<int64_t> index(last, 0);
69 5 std::byte* to = out;
70 5 const std::byte* const end = out + uint64_t(elements) * elementBytes;
71
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 5 times.
18 while (to < end) {
72 int64_t at = 0;
73
2/2
✓ Branch 0 taken 17 times.
✓ Branch 1 taken 13 times.
30 for (size_t i = 0; i < last; ++i) {
74 17 at += index[i] * strides[i];
75 }
76
3/5
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 11 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 1 time.
13 const std::byte* from = source.data() + at * elementBytes;
77
3/5
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 11 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 1 time.
13 switch (runBytes) {
78 ✗ case 1:
79 ✗ to = CopyStrided<1>(to, from, count, step, runBytes);
80 ✗ break;
81 ✗ case 2:
82 ✗ to = CopyStrided<2>(to, from, count, step, runBytes);
83 ✗ break;
84 11 case 4:
85 11 to = CopyStrided<4>(to, from, count, step, runBytes);
86 11 break;
87 1 case 8:
88 1 to = CopyStrided<8>(to, from, count, step, runBytes);
89 1 break;
90 1 default:
91 1 to = CopyStrided<0>(to, from, count, step, runBytes);
92 }
93
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 5 times.
18 for (size_t i = last; i-- > 0;) {
94
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 8 times.
13 if (++index[i] < shape[i]) {
95 break;
96 }
97 5 index[i] = 0;
98 }
99 }
100 5 }
101
102 }
103
104 81176 Tensor::Tensor() = default;
105
2/2
✓ Branch 2 taken 75521 times.
✓ Branch 3 taken 82683 times.
158204 Tensor::~Tensor() = default;
106
2/4
✓ Branch 2 taken 837 times.
✓ Branch 5 taken 837 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
837 Tensor::Tensor(const Tensor& other) = default;
107 5163 Tensor::Tensor(Tensor&& other) noexcept = default;
108
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 3860 times.
3860 Tensor& Tensor::operator=(const Tensor& other) = default;
109
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 2017 times.
2057 Tensor& Tensor::operator=(Tensor&& other) noexcept = default;
110
111 71028 Tensor::Tensor(Ref<Buffer> buffer, uint64_t offset, DType dtype, std::vector<int64_t> shape,
112 71028 std::vector<int64_t> strides)
113
2/2
✓ Branch 0 taken 70951 times.
✓ Branch 1 taken 77 times.
71028 : m_buffer(std::move(buffer)), m_offset(offset), m_dtype(dtype), m_shape(std::move(shape)),
114 71028 m_strides(std::move(strides)) {
115
2/2
✓ Branch 0 taken 70951 times.
✓ Branch 1 taken 77 times.
71028 if (m_strides.empty()) {
116
1/1
✓ Branch 1 taken 70951 times.
70951 m_strides = ContiguousStrides(m_shape);
117 }
118
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 71028 times.
71028 if (m_strides.size() != m_shape.size()) {
119
0/1
✗ Branch 2 not taken.
✗ throw Error("Tensor of {} axes with {} strides", m_shape.size(), m_strides.size());
120 }
121
122 71028 const uint64_t end =
123 71028 m_offset + (uint64_t(LastElement(m_shape, m_strides)) + 1) * ElementBytes(m_dtype);
124
3/5
✓ Branch 1 taken 71028 times.
✓ Branch 3 taken 71028 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 71028 times.
71028 if (Elements() > 0 && end > m_buffer->Bytes()) {
125
0/1
✗ Branch 1 not taken.
✗ throw Error("Tensor runs to byte {} of buffer \"{}\", which has {}", end, m_buffer->Label(),
126
0/1
✗ Branch 2 not taken.
✗ m_buffer->Bytes());
127 }
128 71028 }
129
130 138274 Device& Tensor::GetDevice() const {
131
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 138274 times.
138274 if (!m_buffer) {
132
0/2
✗ Branch 2 not taken.
✗ Branch 5 not taken.
✗ throw Error("An empty tensor is on no device");
133 }
134 138274 return m_buffer->GetDevice();
135 }
136
137 71075 uint64_t Tensor::Address() const {
138 71075 return m_buffer->Address() + m_offset;
139 }
140
141 83642 int64_t Tensor::Elements() const {
142 83642 int64_t elements = 1;
143
2/2
✓ Branch 0 taken 185423 times.
✓ Branch 1 taken 83642 times.
269065 for (const int64_t size : m_shape) {
144 185423 elements *= size;
145 }
146
1/2
✓ Branch 0 taken 83642 times.
✗ Branch 1 not taken.
83642 return Empty() ? 0 : elements;
147 }
148
149 43 Tensor Tensor::Select(int64_t index) const {
150
4/6
✓ Branch 0 taken 43 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 43 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 41 times.
43 if (m_shape.empty() || index < 0 || index >= m_shape[0]) {
151 2 throw Error("Index {} of a tensor of {} along its first axis", index,
152
2/3
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
2 m_shape.empty() ? 0 : m_shape[0]);
153 }
154 82 return Tensor(m_buffer, m_offset + uint64_t(index * m_strides[0]) * ElementBytes(m_dtype),
155
1/1
✓ Branch 1 taken 41 times.
82 m_dtype, {m_shape.begin() + 1, m_shape.end()},
156
2/2
✓ Branch 2 taken 41 times.
✓ Branch 6 taken 41 times.
82 {m_strides.begin() + 1, m_strides.end()});
157 }
158
159 244 std::vector<std::byte> Tensor::ToHostBytes() const {
160
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 244 times.
244 if (Elements() == 0) {
161 ✗ return {};
162 }
163
2/2
✓ Branch 0 taken 243 times.
✓ Branch 1 taken 1 time.
244 if (Packed(m_shape, m_strides)) {
164 243 const Download download = ToHostAsync();
165
1/1
✓ Branch 1 taken 243 times.
243 download.Wait();
166
1/1
✓ Branch 1 taken 243 times.
243 const std::span<const std::byte> bytes = download.Bytes();
167
1/1
✓ Branch 1 taken 243 times.
243 return {bytes.begin(), bytes.end()};
168 243 }
169 1 std::vector<std::byte> out(uint64_t(Elements()) * ElementBytes(m_dtype));
170
1/1
✓ Branch 1 taken 1 time.
1 ReadInto(out.data());
171 1 return out;
172 1 }
173
174 4040 void Tensor::ReadInto(void* out) const {
175 4040 const int64_t elements = Elements();
176 4040 const uint32_t size = ElementBytes(m_dtype);
177
2/2
✓ Branch 0 taken 4035 times.
✓ Branch 1 taken 5 times.
4040 if (Packed(m_shape, m_strides)) {
178 4035 const Download download = ToHostAsync();
179
1/1
✓ Branch 1 taken 4035 times.
4035 download.Wait();
180
1/1
✓ Branch 1 taken 4035 times.
4035 std::memcpy(out, download.Bytes().data(), uint64_t(elements) * size);
181 4035 return;
182 4035 }
183 5 const uint64_t spanBytes = (uint64_t(LastElement(m_shape, m_strides)) + 1) * size;
184 5 const Download download =
185 5 m_buffer->GetDevice().GetQueue().ReadAsync(m_buffer->Slice(m_offset, spanBytes));
186
1/1
✓ Branch 1 taken 5 times.
5 download.Wait();
187
2/2
✓ Branch 1 taken 5 times.
✓ Branch 4 taken 5 times.
5 Gather(download.Bytes(), m_shape, m_strides, size, elements, static_cast<std::byte*>(out));
188 5 }
189
190 4281 Download Tensor::ToHostAsync() const {
191
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 4280 times.
4281 if (Empty()) {
192
2/2
✓ Branch 2 taken 1 time.
✓ Branch 5 taken 1 time.
2 throw Error("An empty tensor has nothing to read");
193 }
194
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 4279 times.
4280 if (!Packed(m_shape, m_strides)) {
195
1/1
✓ Branch 2 taken 1 time.
2 throw UnsupportedError("ToHostAsync reads contiguous tensors only");
196 }
197 4279 return m_buffer->GetDevice().GetQueue().ReadAsync(
198 4279 m_buffer->Slice(m_offset, uint64_t(Elements()) * ElementBytes(m_dtype)));
199 }
200
201 1 DLTensor Tensor::DLPack() const {
202 1 const std::optional<DLDevice> shared =
203
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
1 m_buffer ? m_buffer->GetDevice().SharedDevice() : std::nullopt;
204
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
1 if (!shared) {
205
1/1
✓ Branch 2 taken 1 time.
2 throw UnsupportedError("This tensor's memory isn't shared with CUDA or HIP");
206 }
207 ✗ const ExternalAddress external = m_buffer->External(shared->device_type);
208 ✗ const uint8_t bits = uint8_t(ElementBytes(m_dtype) * 8);
209
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ const auto code = uint8_t(m_dtype == DType::Float32 || m_dtype == DType::Float64 ? kDLFloat
210 ✗ : m_dtype == DType::Int32 || m_dtype == DType::Int64 ? kDLInt
211 : kDLUInt);
212 ✗ DLTensor tensor{};
213 ✗ tensor.data = reinterpret_cast<void*>(external.address + m_offset);
214 ✗ tensor.device = external.device;
215
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ tensor.ndim = int32_t(m_shape.size());
216 ✗ tensor.dtype = {code, bits, 1};
217 ✗ tensor.shape = const_cast<int64_t*>(m_shape.data());
218
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ tensor.strides = m_strides.empty() ? nullptr : const_cast<int64_t*>(m_strides.data());
219 ✗ return tensor;
220 }
221
222 ✗ DLManagedTensorVersioned* Tensor::Export() const {
223 ✗ struct Exported {
224 DLManagedTensorVersioned managed{};
225 Tensor tensor;
226 };
227
228 ✗ auto exported = std::make_unique<Exported>();
229
0/1
✗ Branch 1 not taken.
✗ exported->tensor = *this;
230
0/1
✗ Branch 1 not taken.
✗ DLManagedTensorVersioned& managed = exported->managed;
231
0/1
✗ Branch 1 not taken.
✗ managed.dl_tensor = exported->tensor.DLPack();
232 ✗ managed.version = {DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION};
233 ✗ managed.manager_ctx = exported.get();
234 ✗ managed.deleter = [](DLManagedTensorVersioned* self) {
235
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ delete static_cast<Exported*>(self->manager_ctx);
236 ✗ };
237 ✗ return &exported.release()->managed;
238 ✗ }
239
240 }
241