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 | 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 | 212744 | for (size_t i = 0; i < shape.size(); ++i) { | |
| 30 | 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 | 38164 | for (size_t axis = shape.size(); axis-- > 0;) { | |
| 38 | 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 | 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 | 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 | 18 | while (to < end) { | |
| 72 | int64_t at = 0; | ||
| 73 | 30 | for (size_t i = 0; i < last; ++i) { | |
| 74 | 17 | at += index[i] * strides[i]; | |
| 75 | } | ||
| 76 | 13 | const std::byte* from = source.data() + at * elementBytes; | |
| 77 | 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 | 18 | for (size_t i = last; i-- > 0;) { | |
| 94 | 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 | 158204 | Tensor::~Tensor() = default; | |
| 106 | 837 | Tensor::Tensor(const Tensor& other) = default; | |
| 107 | 5163 | Tensor::Tensor(Tensor&& other) noexcept = default; | |
| 108 | 3860 | Tensor& Tensor::operator=(const Tensor& other) = default; | |
| 109 | 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 | 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 | 71028 | if (m_strides.empty()) { | |
| 116 | 70951 | m_strides = ContiguousStrides(m_shape); | |
| 117 | } | ||
| 118 | 71028 | if (m_strides.size() != m_shape.size()) { | |
| 119 | ✗ | 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 | 71028 | if (Elements() > 0 && end > m_buffer->Bytes()) { | |
| 125 | ✗ | throw Error("Tensor runs to byte {} of buffer \"{}\", which has {}", end, m_buffer->Label(), | |
| 126 | ✗ | m_buffer->Bytes()); | |
| 127 | } | ||
| 128 | 71028 | } | |
| 129 | |||
| 130 | 138274 | Device& Tensor::GetDevice() const { | |
| 131 | 138274 | if (!m_buffer) { | |
| 132 | ✗ | 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 | 269065 | for (const int64_t size : m_shape) { | |
| 144 | 185423 | elements *= size; | |
| 145 | } | ||
| 146 | 83642 | return Empty() ? 0 : elements; | |
| 147 | } | ||
| 148 | |||
| 149 | 43 | Tensor Tensor::Select(int64_t index) const { | |
| 150 | 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 | 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 | 82 | m_dtype, {m_shape.begin() + 1, m_shape.end()}, | |
| 156 | 82 | {m_strides.begin() + 1, m_strides.end()}); | |
| 157 | } | ||
| 158 | |||
| 159 | 244 | std::vector<std::byte> Tensor::ToHostBytes() const { | |
| 160 | 244 | if (Elements() == 0) { | |
| 161 | ✗ | return {}; | |
| 162 | } | ||
| 163 | 244 | if (Packed(m_shape, m_strides)) { | |
| 164 | 243 | const Download download = ToHostAsync(); | |
| 165 | 243 | download.Wait(); | |
| 166 | 243 | const std::span<const std::byte> bytes = download.Bytes(); | |
| 167 | 243 | return {bytes.begin(), bytes.end()}; | |
| 168 | 243 | } | |
| 169 | 1 | std::vector<std::byte> out(uint64_t(Elements()) * ElementBytes(m_dtype)); | |
| 170 | 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 | 4040 | if (Packed(m_shape, m_strides)) { | |
| 178 | 4035 | const Download download = ToHostAsync(); | |
| 179 | 4035 | download.Wait(); | |
| 180 | 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 | 5 | download.Wait(); | |
| 187 | 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 | 4281 | if (Empty()) { | |
| 192 | 2 | throw Error("An empty tensor has nothing to read"); | |
| 193 | } | ||
| 194 | 4280 | if (!Packed(m_shape, m_strides)) { | |
| 195 | 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 | m_buffer ? m_buffer->GetDevice().SharedDevice() : std::nullopt; | |
| 204 | 1 | if (!shared) { | |
| 205 | 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 | ✗ | 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 | ✗ | tensor.ndim = int32_t(m_shape.size()); | |
| 216 | ✗ | tensor.dtype = {code, bits, 1}; | |
| 217 | ✗ | tensor.shape = const_cast<int64_t*>(m_shape.data()); | |
| 218 | ✗ | 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 | ✗ | exported->tensor = *this; | |
| 230 | ✗ | DLManagedTensorVersioned& managed = exported->managed; | |
| 231 | ✗ | 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 | ✗ | delete static_cast<Exported*>(self->manager_ctx); | |
| 236 | ✗ | }; | |
| 237 | ✗ | return &exported.release()->managed; | |
| 238 | ✗ | } | |
| 239 | |||
| 240 | } | ||
| 241 |