python/device.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file device.cpp | ||
| 3 | /// @brief Python bindings of device columns, shared with PyTorch, Warp and any DLPack library in | ||
| 4 | /// place: `torch.from_dlpack(column)` and `wp.from_dlpack(column)` view Ernie's memory. | ||
| 5 | #include "ernie/python/device.h" | ||
| 6 | |||
| 7 | #include "ernie/device/cuda.h" | ||
| 8 | #include "ernie/device/staging.h" | ||
| 9 | #include "ernie/python/dlpack.h" | ||
| 10 | |||
| 11 | #include <nanobind/stl/string.h> | ||
| 12 | #include <nanobind/stl/string_view.h> | ||
| 13 | #include <nanobind/stl/tuple.h> | ||
| 14 | #include <nanobind/stl/vector.h> | ||
| 15 | |||
| 16 | #include <array> | ||
| 17 | #include <memory> | ||
| 18 | #include <optional> | ||
| 19 | #include <stdexcept> | ||
| 20 | #include <string_view> | ||
| 21 | #include <utility> | ||
| 22 | |||
| 23 | namespace nb = nanobind; | ||
| 24 | |||
| 25 | namespace ernie::python { | ||
| 26 | namespace { | ||
| 27 | |||
| 28 | constexpr std::array<std::pair<std::string_view, DType>, 8> kTypes{{ | ||
| 29 | {"uint8", DType::UInt8}, | ||
| 30 | {"uint16", DType::UInt16}, | ||
| 31 | {"uint32", DType::UInt32}, | ||
| 32 | {"int32", DType::Int32}, | ||
| 33 | {"int64", DType::Int64}, | ||
| 34 | {"float16", DType::Float16}, | ||
| 35 | {"float32", DType::Float32}, | ||
| 36 | {"float64", DType::Float64}, | ||
| 37 | }}; | ||
| 38 | |||
| 39 | ✗ | DType ParseType(std::string_view name) { | |
| 40 | ✗ | for (const auto& [text, type] : kTypes) { | |
| 41 | ✗ | if (text == name) { | |
| 42 | ✗ | return type; | |
| 43 | } | ||
| 44 | } | ||
| 45 | ✗ | throw nb::value_error(("No element type '" + std::string(name) + "'").c_str()); | |
| 46 | ✗ | } | |
| 47 | |||
| 48 | ✗ | std::string_view TypeName(DType type) { | |
| 49 | ✗ | for (const auto& [text, value] : kTypes) { | |
| 50 | ✗ | if (value == type) { | |
| 51 | ✗ | return text; | |
| 52 | } | ||
| 53 | } | ||
| 54 | ✗ | return ""; | |
| 55 | } | ||
| 56 | |||
| 57 | /// A DLPack stream argument as a CUDA stream: None and 1 are the legacy default stream, 2 the | ||
| 58 | /// per-thread one, anything else a `cudaStream_t`. | ||
| 59 | ✗ | CUstream ToStream(const nb::object& stream) { | |
| 60 | ✗ | if (stream.is_none()) { | |
| 61 | return nullptr; | ||
| 62 | } | ||
| 63 | ✗ | const auto value = nb::cast<std::intptr_t>(stream); | |
| 64 | ✗ | return reinterpret_cast<CUstream>(value == 1 ? 0 : value); | |
| 65 | } | ||
| 66 | |||
| 67 | /// A device, with the shared arenas Python's columns live in and its CUDA interop. | ||
| 68 | struct Context { | ||
| 69 | std::unique_ptr<DeviceContext> context; | ||
| 70 | std::unique_ptr<Arenas> shared; | ||
| 71 | std::unique_ptr<CudaInterop> cuda; | ||
| 72 | |||
| 73 | ✗ | explicit Context(const std::string& device) | |
| 74 | ✗ | : context(std::make_unique<DeviceContext>(DeviceContextOptions{.device = device})) { | |
| 75 | ✗ | if (context->Capabilities().externalMemory) { | |
| 76 | ✗ | shared = std::make_unique<Arenas>(*context, MemoryUse::Shared, "python columns"); | |
| 77 | } | ||
| 78 | ✗ | cuda = CudaInterop::Create(*context); | |
| 79 | ✗ | } | |
| 80 | ✗ | ~Context() { | |
| 81 | ✗ | cuda.reset(); | |
| 82 | ✗ | context->WaitIdle(); | |
| 83 | ✗ | } | |
| 84 | |||
| 85 | ✗ | CudaInterop& Cuda() const { | |
| 86 | ✗ | if (!cuda) { | |
| 87 | ✗ | throw nb::buffer_error(("No CUDA device pairs with " + context->Info().name).c_str()); | |
| 88 | } | ||
| 89 | ✗ | return *cuda; | |
| 90 | } | ||
| 91 | ✗ | Arenas& Shared() const { | |
| 92 | ✗ | if (!shared) { | |
| 93 | ✗ | throw nb::buffer_error((context->Info().name + " cannot share memory").c_str()); | |
| 94 | } | ||
| 95 | ✗ | return *shared; | |
| 96 | } | ||
| 97 | }; | ||
| 98 | |||
| 99 | /// A column, and the ticket of CUDA's writes Ernie has yet to see. | ||
| 100 | ✗ | struct Column { | |
| 101 | Context& owner; | ||
| 102 | std::unique_ptr<DeviceColumn> column; | ||
| 103 | /// The tensor a wrapped column shares, which `__dlpack__` hands back as it is. | ||
| 104 | std::optional<Tensor> wrapped; | ||
| 105 | Ticket written; | ||
| 106 | |||
| 107 | ✗ | Column(Context& context, std::unique_ptr<DeviceColumn> device, | |
| 108 | std::optional<Tensor> tensor = std::nullopt) | ||
| 109 | ✗ | : owner(context), column(std::move(device)), wrapped(std::move(tensor)) {} | |
| 110 | |||
| 111 | ✗ | nb::object Dlpack(const nb::object& stream, const nb::object& maxVersion) { | |
| 112 | ✗ | CudaInterop& cuda = owner.Cuda(); | |
| 113 | // -1 asks for no synchronization. | ||
| 114 | ✗ | if (stream.is_none() || nb::cast<std::intptr_t>(stream) != -1) { | |
| 115 | ✗ | for (const Queue* queue : owner.context->Queues()) { | |
| 116 | ✗ | cuda.Wait(ToStream(stream), queue->Last()); | |
| 117 | ✗ | } | |
| 118 | } | ||
| 119 | ✗ | const bool versioned = | |
| 120 | ✗ | !maxVersion.is_none() && nb::cast<std::tuple<int, int>>(maxVersion) >= | |
| 121 | ✗ | std::tuple<int, int>(DLPACK_MAJOR_VERSION, 0); | |
| 122 | ✗ | return Capsule((wrapped ? *wrapped : cuda.View(*column)).ToDlpack(), versioned); | |
| 123 | } | ||
| 124 | |||
| 125 | ✗ | void Assign(const nb::object& source, const nb::object& stream) { | |
| 126 | ✗ | written = owner.Cuda().Assign(*column, TakeTensor(source, stream), ToStream(stream)); | |
| 127 | ✗ | } | |
| 128 | |||
| 129 | ✗ | nb::bytes Read() { | |
| 130 | ✗ | Readback readback(*owner.context, owner.context->Main(), column->Storage(), | |
| 131 | ✗ | column->Range().offset, column->Bytes(), std::span(&written, 1)); | |
| 132 | ✗ | const std::span<const std::byte> bytes = readback.Bytes(); | |
| 133 | ✗ | return nb::bytes(reinterpret_cast<const char*>(bytes.data()), column->Bytes()); | |
| 134 | ✗ | } | |
| 135 | }; | ||
| 136 | |||
| 137 | /// @return The allocations `ernie_cuda_malloc` made, or none without a CUDA driver: PyTorch's | ||
| 138 | /// allocator hook is process-wide, so they are too. | ||
| 139 | ✗ | ShareableAllocations* TorchAllocations() { | |
| 140 | ✗ | static const std::unique_ptr<ShareableAllocations> allocations = | |
| 141 | ✗ | CudaDriver::Get() != nullptr ? std::make_unique<ShareableAllocations>(*CudaDriver::Get()) | |
| 142 | ✗ | : nullptr; | |
| 143 | ✗ | return allocations.get(); | |
| 144 | } | ||
| 145 | |||
| 146 | } // namespace | ||
| 147 | |||
| 148 | 6 | void BindDevice(nb::module_& module) { | |
| 149 | 6 | nb::class_<Context>(module, "DeviceContext", | |
| 150 | "A device: where columns live, and the CUDA device paired with it.") | ||
| 151 | 24 | .def(nb::init<const std::string&>(), nb::arg("device") = "", | |
| 152 | "Opens the best device, or the one whose UUID or name matches `device`.") | ||
| 153 | 6 | .def_prop_ro("name", [](const Context& self) { return self.context->Info().name; }) | |
| 154 | 6 | .def_prop_ro("uuid", [](const Context& self) { return self.context->Info().uuid; }) | |
| 155 | 6 | .def_prop_ro("cuda_ordinal", | |
| 156 | ✗ | [](const Context& self) -> nb::object { | |
| 157 | ✗ | if (!self.cuda) { | |
| 158 | ✗ | return nb::none(); | |
| 159 | } | ||
| 160 | ✗ | return nb::int_(self.cuda->Ordinal()); | |
| 161 | }) | ||
| 162 | 6 | .def_prop_ro("transfers", | |
| 163 | ✗ | [](const Context& self) { | |
| 164 | ✗ | const Transfers transfers = self.context->Transfers(); | |
| 165 | ✗ | nb::dict result; | |
| 166 | ✗ | result["uploaded"] = transfers.uploaded; | |
| 167 | ✗ | result["read_back"] = transfers.readBack; | |
| 168 | ✗ | result["copied"] = transfers.copied; | |
| 169 | ✗ | return result; | |
| 170 | ✗ | }) | |
| 171 | 6 | .def( | |
| 172 | "column", | ||
| 173 | ✗ | [](Context& self, const std::string& name, const std::string& dtype, | |
| 174 | const std::vector<std::int64_t>& shape) { | ||
| 175 | ✗ | if (shape.empty() || shape.front() < 0) { | |
| 176 | ✗ | throw nb::value_error("A column's shape starts with its row count"); | |
| 177 | } | ||
| 178 | ✗ | return new Column(self, | |
| 179 | ✗ | std::make_unique<DeviceColumn>( | |
| 180 | ✗ | self.Shared(), name, ParseType(dtype), | |
| 181 | ✗ | std::vector<std::int64_t>(shape.begin() + 1, shape.end()), | |
| 182 | ✗ | static_cast<std::uint64_t>(shape.front()))); | |
| 183 | }, | ||
| 184 | 6 | nb::arg("name"), nb::arg("dtype"), nb::arg("shape"), nb::keep_alive<0, 1>(), | |
| 185 | "A column of `shape[0]` rows on the device, shared in place through DLPack.") | ||
| 186 | 6 | .def( | |
| 187 | "wrap", | ||
| 188 | ✗ | [](Context& self, const nb::object& tensor, const std::string& name, | |
| 189 | const nb::object& stream) { | ||
| 190 | ✗ | CudaInterop& cuda = self.Cuda(); | |
| 191 | ✗ | Tensor source = TakeTensor(tensor, stream); | |
| 192 | ✗ | std::unique_ptr<DeviceColumn> column = cuda.Wrap(source, name); | |
| 193 | ✗ | if (!column) { | |
| 194 | ✗ | throw nb::buffer_error( | |
| 195 | "The tensor's memory cannot be shared: allocate it with " | ||
| 196 | ✗ | "ernie.cuda_allocator(), or copy it into a column with assign()"); | |
| 197 | } | ||
| 198 | ✗ | auto* wrapped = new Column(self, std::move(column), std::move(source)); | |
| 199 | ✗ | wrapped->written = cuda.Signal(ToStream(stream)); | |
| 200 | ✗ | return wrapped; | |
| 201 | ✗ | }, | |
| 202 | 18 | nb::arg("tensor"), nb::arg("name") = "", nb::arg("stream") = nb::none(), | |
| 203 | nb::keep_alive<0, 1>(), | ||
| 204 | "A column over a CUDA tensor's own memory, in place, when it is exportable."); | ||
| 205 | |||
| 206 | 6 | nb::class_<Column>(module, "DeviceColumn", | |
| 207 | "Rows on the device that CUDA, PyTorch and Warp read and write in place.") | ||
| 208 | 6 | .def_prop_ro("name", [](const Column& self) { return self.column->Name(); }) | |
| 209 | 6 | .def_prop_ro("shape", [](const Column& self) { return self.column->Shape(); }) | |
| 210 | 6 | .def_prop_ro("dtype", [](const Column& self) { return TypeName(self.column->Type()); }) | |
| 211 | 6 | .def( | |
| 212 | "__dlpack__", | ||
| 213 | ✗ | [](Column& self, const nb::object& stream, const nb::object& maxVersion, | |
| 214 | const nb::object&, const nb::object& copy) { | ||
| 215 | ✗ | if (!copy.is_none() && nb::cast<bool>(copy)) { | |
| 216 | ✗ | throw nb::buffer_error("Device columns are shared in place, never copied"); | |
| 217 | } | ||
| 218 | ✗ | return self.Dlpack(stream, maxVersion); | |
| 219 | }, | ||
| 220 | 12 | nb::kw_only(), nb::arg("stream") = nb::none(), nb::arg("max_version") = nb::none(), | |
| 221 | 12 | nb::arg("dl_device") = nb::none(), nb::arg("copy") = nb::none(), | |
| 222 | "A CUDA view of the rows; `stream` waits for Ernie's work on them first.") | ||
| 223 | 6 | .def("__dlpack_device__", | |
| 224 | ✗ | [](const Column& self) { | |
| 225 | ✗ | return std::make_tuple(static_cast<int>(kDLCUDA), self.owner.Cuda().Ordinal()); | |
| 226 | }) | ||
| 227 | 6 | .def( | |
| 228 | "ready", | ||
| 229 | 6 | [](Column& self, const nb::object& stream) { | |
| 230 | ✗ | self.written = self.owner.Cuda().Signal(ToStream(stream)); | |
| 231 | ✗ | }, | |
| 232 | 12 | nb::arg("stream") = nb::none(), | |
| 233 | "Hands the rows back to Ernie once the work enqueued on `stream` so far is done.") | ||
| 234 | 12 | .def("assign", &Column::Assign, nb::arg("tensor"), nb::arg("stream") = nb::none(), | |
| 235 | "Takes a DLPack tensor's values: in place for this column's own view, else by one " | ||
| 236 | "copy on the device.") | ||
| 237 | 6 | .def("read", &Column::Read, "Copies the rows to the host, after CUDA's writes."); | |
| 238 | 6 | } | |
| 239 | |||
| 240 | } // namespace ernie::python | ||
| 241 | |||
| 242 | #if defined(_WIN32) | ||
| 243 | #define ERNIE_ALLOCATOR extern "C" __declspec(dllexport) | ||
| 244 | #else | ||
| 245 | #define ERNIE_ALLOCATOR extern "C" __attribute__((visibility("default"))) | ||
| 246 | #endif | ||
| 247 | |||
| 248 | /// @brief PyTorch's pluggable allocator, allocating each tensor as exportable memory that | ||
| 249 | /// `DeviceContext.wrap` shares in place (`ernie.cuda_allocator()` installs it). | ||
| 250 | ✗ | ERNIE_ALLOCATOR void* ernie_cuda_malloc(std::size_t size, int device, void* stream) { | |
| 251 | ✗ | static_cast<void>(stream); | |
| 252 | ✗ | ernie::ShareableAllocations* allocations = ernie::python::TorchAllocations(); | |
| 253 | ✗ | return allocations != nullptr ? reinterpret_cast<void*>(allocations->Allocate(device, size)) | |
| 254 | ✗ | : nullptr; | |
| 255 | } | ||
| 256 | |||
| 257 | ✗ | ERNIE_ALLOCATOR void ernie_cuda_free(void* pointer, std::size_t size, int device, void* stream) { | |
| 258 | ✗ | static_cast<void>(size); | |
| 259 | ✗ | static_cast<void>(device); | |
| 260 | // Work enqueued on the stream may still use the memory: unmapping waits for none. | ||
| 261 | ✗ | ernie::CudaDriver::Get()->streamSynchronize(static_cast<ernie::CUstream>(stream)); | |
| 262 | ✗ | ernie::python::TorchAllocations()->Free(reinterpret_cast<ernie::CUdeviceptr>(pointer)); | |
| 263 | ✗ | } | |
| 264 |