stage/views/staged.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/stage/views/staged.h" | ||
| 3 | |||
| 4 | #include <algorithm> | ||
| 5 | |||
| 6 | namespace chrono::batch_render { | ||
| 7 | |||
| 8 | namespace { | ||
| 9 | |||
| 10 | 14158 | bool InHostMemory(DLDevice device) { | |
| 11 | 14158 | return device.device_type == kDLCPU || device.device_type == kDLCUDAHost || | |
| 12 | 14158 | device.device_type == kDLROCMHost || device.device_type == kDLCUDAManaged; | |
| 13 | } | ||
| 14 | |||
| 15 | 14158 | std::optional<uint64_t> ContiguousBytes(const DLTensor& source) { | |
| 16 | 14158 | uint64_t elements = 1; | |
| 17 | 42473 | for (int axis = source.ndim; axis-- > 0;) { | |
| 18 | 28315 | const int64_t side = source.shape[axis]; | |
| 19 | 28315 | if (side < 0) { | |
| 20 | ✗ | return std::nullopt; | |
| 21 | } | ||
| 22 | 28315 | if (source.strides != nullptr && side > 1 && source.strides[axis] != int64_t(elements)) { | |
| 23 | ✗ | return std::nullopt; | |
| 24 | } | ||
| 25 | 28315 | elements *= uint64_t(side); | |
| 26 | } | ||
| 27 | 14158 | const uint64_t bits = uint64_t(source.dtype.bits) * std::max<uint16_t>(source.dtype.lanes, 1); | |
| 28 | 14158 | if (bits % 8 != 0) { | |
| 29 | ✗ | return std::nullopt; | |
| 30 | } | ||
| 31 | 14158 | return elements * (bits / 8); | |
| 32 | } | ||
| 33 | |||
| 34 | } | ||
| 35 | |||
| 36 | 14158 | std::optional<StagedTensor> StagedTensor::Stage(device::Device& device, const DLTensor& source, | |
| 37 | Stream producer) { | ||
| 38 | 14158 | const std::optional<uint64_t> bytes = ContiguousBytes(source); | |
| 39 | 14158 | if (!bytes) { | |
| 40 | ✗ | return std::nullopt; | |
| 41 | } | ||
| 42 | 14158 | StagedTensor staged; | |
| 43 | 14158 | staged.m_shape.assign(source.shape, source.shape + source.ndim); | |
| 44 | 14158 | staged.m_dtype = source.dtype; | |
| 45 | 14158 | const auto* data = static_cast<const std::byte*>(source.data) + source.byte_offset; | |
| 46 | 14158 | if (InHostMemory(source.device)) { | |
| 47 | 14158 | if (producer.IsCompute()) { | |
| 48 | ✗ | return std::nullopt; | |
| 49 | } | ||
| 50 | 14158 | staged.m_host.assign(data, data + *bytes); | |
| 51 | 14158 | return staged; | |
| 52 | } | ||
| 53 | ✗ | const std::optional<DLDevice> shared = device.SharedDevice(); | |
| 54 | ✗ | if (!producer.IsCompute() || !shared || shared->device_type != source.device.device_type) { | |
| 55 | ✗ | return std::nullopt; | |
| 56 | } | ||
| 57 | ✗ | staged.m_buffer = device.StageOnStream(uint64_t(data), *bytes, producer, "deferred write"); | |
| 58 | ✗ | const device::ExternalAddress external = staged.m_buffer->External(source.device.device_type); | |
| 59 | ✗ | staged.m_address = external.address; | |
| 60 | ✗ | staged.m_device = external.device; | |
| 61 | ✗ | staged.m_producer = producer; | |
| 62 | ✗ | return staged; | |
| 63 | 14158 | } | |
| 64 | |||
| 65 | 14158 | DLTensor StagedTensor::View() const { | |
| 66 | 14158 | DLTensor view{}; | |
| 67 | 28316 | view.data = | |
| 68 | 14158 | m_buffer ? reinterpret_cast<void*>(m_address) : const_cast<std::byte*>(m_host.data()); | |
| 69 | 14158 | view.device = m_device; | |
| 70 | 14158 | view.ndim = int32_t(m_shape.size()); | |
| 71 | 14158 | view.dtype = m_dtype; | |
| 72 | 14158 | view.shape = const_cast<int64_t*>(m_shape.data()); | |
| 73 | 14158 | return view; | |
| 74 | } | ||
| 75 | |||
| 76 | } | ||
| 77 |