| | | 1 | | // Copyright 2025 Digital Holography Foundation |
| | | 2 | | // |
| | | 3 | | // Licensed under the Apache License, Version 2.0 (the "License"); |
| | | 4 | | // you may not use this file except in compliance with the License. |
| | | 5 | | // You may obtain a copy of the License at |
| | | 6 | | // |
| | | 7 | | // http://www.apache.org/licenses/LICENSE-2.0 |
| | | 8 | | // |
| | | 9 | | // Unless required by applicable law or agreed to in writing, software |
| | | 10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
| | | 11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| | | 12 | | // See the License for the specific language governing permissions and |
| | | 13 | | // limitations under the License. |
| | | 14 | | |
| | | 15 | | #include "holotask/syncs/pca.hh" |
| | | 16 | | |
| | | 17 | | #include <cuda.h> |
| | | 18 | | #include <nvJitLink.h> |
| | | 19 | | #include <nvtx3/nvtx3.hpp> |
| | | 20 | | #include <windows.h> |
| | | 21 | | |
| | | 22 | | #include <algorithm> |
| | | 23 | | #include <array> |
| | | 24 | | #include <cstddef> |
| | | 25 | | #include <cstdint> |
| | | 26 | | #include <cstdlib> |
| | | 27 | | #include <filesystem> |
| | | 28 | | #include <format> |
| | | 29 | | #include <limits> |
| | | 30 | | #include <list> |
| | | 31 | | #include <memory> |
| | | 32 | | #include <stdexcept> |
| | | 33 | | #include <string> |
| | | 34 | | #include <utility> |
| | | 35 | | #include <vector> |
| | | 36 | | |
| | | 37 | | #include "curaii/cublas.hh" |
| | | 38 | | #include "curaii/cuda.hh" |
| | | 39 | | #include "curaii/cusolver.hh" |
| | | 40 | | #include "curaii/nvrtc.hh" |
| | | 41 | | |
| | | 42 | | #include "logger.hh" |
| | | 43 | | |
| | | 44 | | namespace holotask::syncs { |
| | | 45 | | |
| | | 46 | | namespace { |
| | | 47 | | |
| | | 48 | | // ------------------------------------------------------------------------------------------------- |
| | | 49 | | // PCA data model |
| | | 50 | | // ------------------------------------------------------------------------------------------------- |
| | | 51 | | |
| | | 52 | | // PCA operates on tensors shaped [..., features, height, width]. All leading dimensions are |
| | | 53 | | // flattened into independent batches. |
| | | 54 | | struct PcaLayout { |
| | | 55 | | size_t feature_axis; |
| | | 56 | | int features; |
| | | 57 | | int height; |
| | | 58 | | int width; |
| | | 59 | | int samples; |
| | | 60 | | size_t batches; |
| | | 61 | | |
| | | 62 | | // Strides in elements between consecutive flattened batches. |
| | | 63 | | long long input_batch_stride; |
| | | 64 | | long long matrix_batch_stride; |
| | | 65 | | |
| | 1 | 66 | | explicit PcaLayout(const holoflow::core::TDesc &desc) { |
| | 1 | 67 | | const auto rank = desc.rank(); |
| | | 68 | | |
| | 1 | 69 | | feature_axis = rank - 3; |
| | 1 | 70 | | features = static_cast<int>(desc.shape.at(feature_axis)); |
| | 1 | 71 | | height = static_cast<int>(desc.shape.at(rank - 2)); |
| | 1 | 72 | | width = static_cast<int>(desc.shape.at(rank - 1)); |
| | 1 | 73 | | samples = height * width; |
| | | 74 | | |
| | 1 | 75 | | batches = 1; |
| | 1 | 76 | | for (size_t axis = 0; axis < feature_axis; ++axis) { |
| | 1 | 77 | | batches *= desc.shape.at(axis); |
| | 1 | 78 | | } |
| | | 79 | | |
| | 1 | 80 | | input_batch_stride = static_cast<long long>(features) * samples; |
| | 1 | 81 | | matrix_batch_stride = static_cast<long long>(features) * features; |
| | 1 | 82 | | } |
| | | 83 | | |
| | 1 | 84 | | [[nodiscard]] long long output_batch_stride(int components) const { |
| | 1 | 85 | | return static_cast<long long>(samples) * components; |
| | 1 | 86 | | } |
| | | 87 | | }; |
| | | 88 | | |
| | | 89 | | struct PcaWorkspace { |
| | | 90 | | // The matrix buffer contains covariance matrices before HEEV and eigenvectors afterwards. |
| | | 91 | | curaii::unique_device_ptr<float> matrices; |
| | | 92 | | curaii::unique_device_ptr<float> eigenvalues; |
| | | 93 | | curaii::unique_device_ptr<int> solver_info; |
| | | 94 | | |
| | | 95 | | explicit PcaWorkspace(const PcaLayout &layout) |
| | 1 | 96 | | : matrices(curaii::make_unique_device_ptr<float>( |
| | | 97 | | layout.batches * static_cast<size_t>(layout.features) * layout.features)), |
| | 1 | 98 | | eigenvalues(curaii::make_unique_device_ptr<float>(layout.batches * layout.features)), |
| | 1 | 99 | | solver_info(curaii::make_unique_device_ptr<int>(layout.batches)) {} |
| | | 100 | | }; |
| | | 101 | | |
| | | 102 | | // ------------------------------------------------------------------------------------------------- |
| | | 103 | | // CUDA Graph replay |
| | | 104 | | // ------------------------------------------------------------------------------------------------- |
| | | 105 | | |
| | | 106 | | using PcaGraphKey = std::array<const void *, 2>; |
| | | 107 | | |
| | | 108 | | template <typename Enqueue> cudaGraph_t capture_cuda_graph(cudaStream_t stream, Enqueue &&enqueue) { |
| | | 109 | | cudaGraph_t graph = nullptr; |
| | | 110 | | bool capturing = false; |
| | | 111 | | |
| | | 112 | | try { |
| | | 113 | | CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); |
| | | 114 | | capturing = true; |
| | | 115 | | |
| | | 116 | | std::forward<Enqueue>(enqueue)(); |
| | | 117 | | |
| | | 118 | | CUDA_CHECK(cudaStreamEndCapture(stream, &graph)); |
| | | 119 | | capturing = false; |
| | | 120 | | return graph; |
| | | 121 | | } catch (...) { |
| | | 122 | | if (capturing) { |
| | | 123 | | cudaGraph_t discarded = nullptr; |
| | | 124 | | (void)cudaStreamEndCapture(stream, &discarded); |
| | | 125 | | if (discarded != nullptr) { |
| | | 126 | | CUDA_CHECK_NT(cudaGraphDestroy(discarded)); |
| | | 127 | | } |
| | | 128 | | } else if (graph != nullptr) { |
| | | 129 | | CUDA_CHECK_NT(cudaGraphDestroy(graph)); |
| | | 130 | | } |
| | | 131 | | throw; |
| | | 132 | | } |
| | | 133 | | } |
| | | 134 | | |
| | 1 | 135 | | [[nodiscard]] bool stream_is_capturing(cudaStream_t stream) { |
| | 1 | 136 | | cudaStreamCaptureStatus status = cudaStreamCaptureStatusNone; |
| | 1 | 137 | | CUDA_CHECK(cudaStreamIsCapturing(stream, &status)); |
| | 1 | 138 | | return status != cudaStreamCaptureStatusNone; |
| | 1 | 139 | | } |
| | | 140 | | |
| | | 141 | | class PcaCudaGraph { |
| | | 142 | | public: |
| | 1 | 143 | | explicit PcaCudaGraph(PcaGraphKey key) : key_(key) {} |
| | | 144 | | |
| | 1 | 145 | | ~PcaCudaGraph() noexcept { |
| | 1 | 146 | | if (executable_ != nullptr) { |
| | 1 | 147 | | CUDA_CHECK_NT(cudaGraphExecDestroy(executable_)); |
| | | 148 | | } |
| | 1 | 149 | | } |
| | | 150 | | |
| | | 151 | | PcaCudaGraph(const PcaCudaGraph &) = delete; |
| | | 152 | | PcaCudaGraph &operator=(const PcaCudaGraph &) = delete; |
| | | 153 | | |
| | 1 | 154 | | [[nodiscard]] bool matches(const PcaGraphKey &key) const noexcept { |
| | 1 | 155 | | return executable_ != nullptr && key_ == key; |
| | 1 | 156 | | } |
| | | 157 | | |
| | 1 | 158 | | void launch(cudaStream_t stream) const { CUDA_CHECK(cudaGraphLaunch(executable_, stream)); } |
| | | 159 | | |
| | | 160 | | template <typename Enqueue> void capture(cudaStream_t stream, Enqueue &&enqueue) { |
| | | 161 | | cudaGraph_t graph = capture_cuda_graph(stream, std::forward<Enqueue>(enqueue)); |
| | | 162 | | |
| | | 163 | | try { |
| | | 164 | | CUDA_CHECK(cudaGraphInstantiateWithFlags(&executable_, graph, 0)); |
| | | 165 | | CUDA_CHECK_NT(cudaGraphDestroy(graph)); |
| | | 166 | | } catch (...) { |
| | | 167 | | CUDA_CHECK_NT(cudaGraphDestroy(graph)); |
| | | 168 | | throw; |
| | | 169 | | } |
| | | 170 | | } |
| | | 171 | | |
| | | 172 | | private: |
| | | 173 | | PcaGraphKey key_{}; |
| | 1 | 174 | | cudaGraphExec_t executable_{nullptr}; |
| | | 175 | | }; |
| | | 176 | | |
| | | 177 | | class PcaGraphCache { |
| | | 178 | | public: |
| | 1 | 179 | | void invalidate() { |
| | 1 | 180 | | graphs_.clear(); |
| | 1 | 181 | | capture_enabled_ = true; |
| | 1 | 182 | | } |
| | | 183 | | |
| | 0 | 184 | | void enable_capture() noexcept { capture_enabled_ = true; } |
| | | 185 | | |
| | 1 | 186 | | [[nodiscard]] bool try_launch(const PcaGraphKey &key, cudaStream_t stream) { |
| | 1 | 187 | | const auto graph = std::find_if(graphs_.begin(), graphs_.end(), |
| | | 188 | | [&](const PcaCudaGraph &entry) { return entry.matches(key); }); |
| | | 189 | | |
| | 1 | 190 | | if (graph == graphs_.end()) { |
| | 1 | 191 | | return false; |
| | | 192 | | } |
| | | 193 | | |
| | 1 | 194 | | graph->launch(stream); |
| | 1 | 195 | | graphs_.splice(graphs_.begin(), graphs_, graph); |
| | 1 | 196 | | return true; |
| | 1 | 197 | | } |
| | | 198 | | |
| | | 199 | | template <typename Enqueue> |
| | | 200 | | void try_capture(const PcaGraphKey &key, cudaStream_t stream, Enqueue &&enqueue) { |
| | | 201 | | if (!capture_enabled_) { |
| | | 202 | | return; |
| | | 203 | | } |
| | | 204 | | |
| | | 205 | | bool entry_inserted = false; |
| | | 206 | | try { |
| | | 207 | | graphs_.emplace_front(key); |
| | | 208 | | entry_inserted = true; |
| | | 209 | | |
| | | 210 | | graphs_.front().capture(stream, std::forward<Enqueue>(enqueue)); |
| | | 211 | | |
| | | 212 | | if (graphs_.size() > capacity) { |
| | | 213 | | graphs_.pop_back(); |
| | | 214 | | } |
| | | 215 | | } catch (const std::exception &error) { |
| | | 216 | | if (entry_inserted) { |
| | | 217 | | graphs_.pop_front(); |
| | | 218 | | } |
| | | 219 | | capture_enabled_ = false; |
| | | 220 | | logger()->warn("[Pca] CUDA Graph capture disabled: {}", error.what()); |
| | | 221 | | } |
| | | 222 | | } |
| | | 223 | | |
| | | 224 | | private: |
| | | 225 | | // cuBLAS and cuSOLVER graph nodes embed their buffer addresses. Cache the finite set of rotating |
| | | 226 | | // pipeline buffers instead of attempting to update opaque library nodes. |
| | | 227 | | static constexpr size_t capacity = 128; |
| | | 228 | | |
| | | 229 | | bool capture_enabled_{true}; |
| | | 230 | | std::list<PcaCudaGraph> graphs_; |
| | | 231 | | }; |
| | | 232 | | |
| | | 233 | | // ------------------------------------------------------------------------------------------------- |
| | | 234 | | // Eigensolver abstraction |
| | | 235 | | // ------------------------------------------------------------------------------------------------- |
| | | 236 | | |
| | | 237 | | class Eigensolver { |
| | | 238 | | public: |
| | 1 | 239 | | virtual ~Eigensolver() = default; |
| | | 240 | | |
| | | 241 | | Eigensolver(const Eigensolver &) = delete; |
| | | 242 | | Eigensolver &operator=(const Eigensolver &) = delete; |
| | | 243 | | |
| | | 244 | | [[nodiscard]] virtual bool is_compatible_stream(cudaStream_t stream) const = 0; |
| | | 245 | | |
| | | 246 | | virtual void solve(float *matrices, float *eigenvalues, int *info, cudaStream_t stream) = 0; |
| | | 247 | | |
| | | 248 | | protected: |
| | 1 | 249 | | Eigensolver() = default; |
| | | 250 | | }; |
| | | 251 | | |
| | | 252 | | std::unique_ptr<Eigensolver> make_eigensolver(const PcaLayout &layout, |
| | | 253 | | const PcaWorkspace &workspace, cudaStream_t stream); |
| | | 254 | | |
| | | 255 | | // ------------------------------------------------------------------------------------------------- |
| | | 256 | | // PCA task |
| | | 257 | | // ------------------------------------------------------------------------------------------------- |
| | | 258 | | |
| | | 259 | | class PcaTask final : public holoflow::core::ISyncTask { |
| | | 260 | | public: |
| | | 261 | | PcaTask(const PcaSettings &settings, const holoflow::core::TDesc &input_desc, |
| | | 262 | | const holoflow::core::SyncCreateCtx &ctx) |
| | 1 | 263 | | : settings_(settings), input_desc_(input_desc), layout_(input_desc), stream_(ctx.stream), |
| | 1 | 264 | | workspace_(layout_) { |
| | 1 | 265 | | CUBLAS_CHECK(cublasSetStream(cublas_.get(), stream_)); |
| | 1 | 266 | | eigensolver_ = make_eigensolver(layout_, workspace_, stream_); |
| | 1 | 267 | | } |
| | | 268 | | |
| | 1 | 269 | | [[nodiscard]] bool can_reuse(const holoflow::core::TDesc &input_desc, cudaStream_t stream) const { |
| | 1 | 270 | | return input_desc.shape == input_desc_.shape && input_desc.strides == input_desc_.strides && |
| | | 271 | | input_desc.dtype == input_desc_.dtype && input_desc.mem_loc == input_desc_.mem_loc && |
| | | 272 | | eigensolver_->is_compatible_stream(stream); |
| | 1 | 273 | | } |
| | | 274 | | |
| | 1 | 275 | | void reconfigure(const PcaSettings &settings, cudaStream_t stream) { |
| | 1 | 276 | | if (settings_ != settings) { |
| | 1 | 277 | | settings_ = settings; |
| | 1 | 278 | | graph_cache_.invalidate(); |
| | | 279 | | } |
| | | 280 | | |
| | 1 | 281 | | if (stream_ != stream) { |
| | 0 | 282 | | stream_ = stream; |
| | 0 | 283 | | CUBLAS_CHECK(cublasSetStream(cublas_.get(), stream_)); |
| | | 284 | | |
| | | 285 | | // Graph executables are not tied to the stream on which they were captured. Reusing them is |
| | | 286 | | // valid here because can_reuse() already guarantees that the CUDA context is unchanged. |
| | 0 | 287 | | graph_cache_.enable_capture(); |
| | | 288 | | } |
| | 1 | 289 | | } |
| | | 290 | | |
| | 1 | 291 | | holoflow::core::OpResult execute(holoflow::core::SyncCtx &ctx) override { |
| | 1 | 292 | | nvtx3::scoped_range range("PCA Sync Task"); |
| | | 293 | | |
| | 1 | 294 | | if (stream_ == nullptr || stream_is_capturing(stream_)) { |
| | 1 | 295 | | return enqueue(ctx); |
| | | 296 | | } |
| | | 297 | | |
| | 1 | 298 | | const PcaGraphKey key{ctx.inputs[0].data(), ctx.outputs[0].data()}; |
| | 1 | 299 | | if (graph_cache_.try_launch(key, stream_)) { |
| | 1 | 300 | | return holoflow::core::OpResult::Ok; |
| | | 301 | | } |
| | | 302 | | |
| | 1 | 303 | | const auto result = enqueue(ctx); |
| | 1 | 304 | | if (result == holoflow::core::OpResult::Ok) { |
| | 1 | 305 | | graph_cache_.try_capture(key, stream_, [&]() { (void)enqueue(ctx); }); |
| | | 306 | | } |
| | 1 | 307 | | return result; |
| | 1 | 308 | | } |
| | | 309 | | |
| | | 310 | | private: |
| | 1 | 311 | | holoflow::core::OpResult enqueue(holoflow::core::SyncCtx &ctx) { |
| | 1 | 312 | | auto *input = reinterpret_cast<float *>(ctx.inputs[0].data()); |
| | 1 | 313 | | auto *output = reinterpret_cast<float *>(ctx.outputs[0].data()); |
| | | 314 | | |
| | 1 | 315 | | enqueue_covariance(input); |
| | 1 | 316 | | enqueue_eigendecomposition(); |
| | 1 | 317 | | enqueue_projection(input, output); |
| | | 318 | | |
| | 1 | 319 | | return holoflow::core::OpResult::Ok; |
| | 1 | 320 | | } |
| | | 321 | | |
| | 1 | 322 | | void enqueue_covariance(const float *input) { |
| | 1 | 323 | | nvtx3::scoped_range range("PCA covariance"); |
| | | 324 | | |
| | 1 | 325 | | constexpr float alpha = 1.0f; |
| | 1 | 326 | | constexpr float beta = 0.0f; |
| | | 327 | | |
| | | 328 | | // Independent GEMMs preserve cuBLAS Split-K selection for the large sample dimension. |
| | 1 | 329 | | for (size_t batch = 0; batch < layout_.batches; ++batch) { |
| | 1 | 330 | | const auto input_offset = static_cast<long long>(batch) * layout_.input_batch_stride; |
| | 1 | 331 | | const auto matrix_offset = static_cast<long long>(batch) * layout_.matrix_batch_stride; |
| | | 332 | | |
| | 1 | 333 | | CUBLAS_CHECK(cublasGemmEx(cublas_.get(), CUBLAS_OP_T, CUBLAS_OP_N, layout_.features, |
| | | 334 | | layout_.features, layout_.samples, &alpha, input + input_offset, |
| | | 335 | | CUDA_R_32F, layout_.samples, input + input_offset, CUDA_R_32F, |
| | | 336 | | layout_.samples, &beta, workspace_.matrices.get() + matrix_offset, |
| | | 337 | | CUDA_R_32F, layout_.features, CUBLAS_COMPUTE_32F_FAST_16F, |
| | | 338 | | CUBLAS_GEMM_DEFAULT)); |
| | 1 | 339 | | } |
| | 1 | 340 | | } |
| | | 341 | | |
| | 1 | 342 | | void enqueue_eigendecomposition() { |
| | 1 | 343 | | nvtx3::scoped_range range("PCA eigendecomposition"); |
| | | 344 | | |
| | 1 | 345 | | eigensolver_->solve(workspace_.matrices.get(), workspace_.eigenvalues.get(), |
| | | 346 | | workspace_.solver_info.get(), stream_); |
| | 1 | 347 | | } |
| | | 348 | | |
| | 1 | 349 | | void enqueue_projection(const float *input, float *output) { |
| | 1 | 350 | | nvtx3::scoped_range range("PCA projection"); |
| | | 351 | | |
| | 1 | 352 | | constexpr float alpha = 1.0f; |
| | 1 | 353 | | constexpr float beta = 0.0f; |
| | | 354 | | |
| | 1 | 355 | | const int components = settings_.components(); |
| | | 356 | | |
| | 1 | 357 | | const auto eigenvector_offset = static_cast<long long>(settings_.begin) * layout_.features; |
| | 1 | 358 | | const auto output_stride = layout_.output_batch_stride(components); |
| | | 359 | | |
| | 1 | 360 | | for (size_t batch = 0; batch < layout_.batches; ++batch) { |
| | 1 | 361 | | const auto input_offset = static_cast<long long>(batch) * layout_.input_batch_stride; |
| | 1 | 362 | | const auto matrix_offset = static_cast<long long>(batch) * layout_.matrix_batch_stride; |
| | 1 | 363 | | const auto output_offset = static_cast<long long>(batch) * output_stride; |
| | | 364 | | |
| | 1 | 365 | | CUBLAS_CHECK( |
| | | 366 | | cublasGemmEx(cublas_.get(), CUBLAS_OP_N, CUBLAS_OP_N, layout_.samples, components, |
| | | 367 | | layout_.features, &alpha, input + input_offset, CUDA_R_32F, layout_.samples, |
| | | 368 | | workspace_.matrices.get() + matrix_offset + eigenvector_offset, CUDA_R_32F, |
| | | 369 | | layout_.features, &beta, output + output_offset, CUDA_R_32F, layout_.samples, |
| | | 370 | | CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT)); |
| | 1 | 371 | | } |
| | 1 | 372 | | } |
| | | 373 | | |
| | | 374 | | PcaSettings settings_; |
| | | 375 | | holoflow::core::TDesc input_desc_; |
| | | 376 | | PcaLayout layout_; |
| | | 377 | | cudaStream_t stream_; |
| | | 378 | | |
| | | 379 | | curaii::CublasHandle cublas_; |
| | | 380 | | PcaWorkspace workspace_; |
| | | 381 | | std::unique_ptr<Eigensolver> eigensolver_; |
| | | 382 | | PcaGraphCache graph_cache_; |
| | | 383 | | }; |
| | | 384 | | |
| | | 385 | | // ------------------------------------------------------------------------------------------------- |
| | | 386 | | // CUDA Driver API support |
| | | 387 | | // ------------------------------------------------------------------------------------------------- |
| | | 388 | | |
| | | 389 | | std::string driver_error_message(CUresult result, const char *expression, const char *file, |
| | 0 | 390 | | int line) { |
| | 0 | 391 | | const char *name = nullptr; |
| | 0 | 392 | | const char *message = nullptr; |
| | 0 | 393 | | (void)cuGetErrorName(result, &name); |
| | 0 | 394 | | (void)cuGetErrorString(result, &message); |
| | | 395 | | |
| | 0 | 396 | | return std::format("CUDA Driver error: {} ({})\n expression : {}\n location : {}:{}", |
| | | 397 | | message != nullptr ? message : "unknown", |
| | | 398 | | name != nullptr ? name : std::to_string(static_cast<int>(result)), expression, |
| | | 399 | | file, line); |
| | 0 | 400 | | } |
| | | 401 | | |
| | 1 | 402 | | void driver_check(CUresult result, const char *expression, const char *file, int line) { |
| | 1 | 403 | | if (result == CUDA_SUCCESS) { |
| | 1 | 404 | | return; |
| | | 405 | | } |
| | | 406 | | |
| | 0 | 407 | | const auto message = driver_error_message(result, expression, file, line); |
| | 0 | 408 | | logger()->error("{}", message); |
| | 0 | 409 | | throw std::runtime_error(message); |
| | 1 | 410 | | } |
| | | 411 | | |
| | | 412 | | #define PCA_DRIVER_CHECK(expr) driver_check((expr), #expr, __FILE__, __LINE__) |
| | | 413 | | |
| | | 414 | | class ScopedCudaContext { |
| | | 415 | | public: |
| | 1 | 416 | | explicit ScopedCudaContext(CUcontext context) { |
| | 1 | 417 | | CUcontext current = nullptr; |
| | 1 | 418 | | PCA_DRIVER_CHECK(cuCtxGetCurrent(¤t)); |
| | | 419 | | |
| | 1 | 420 | | if (current != context) { |
| | 0 | 421 | | PCA_DRIVER_CHECK(cuCtxPushCurrent(context)); |
| | 0 | 422 | | pushed_ = true; |
| | | 423 | | } |
| | 1 | 424 | | } |
| | | 425 | | |
| | 1 | 426 | | ~ScopedCudaContext() noexcept { |
| | 1 | 427 | | if (!pushed_) { |
| | 1 | 428 | | return; |
| | | 429 | | } |
| | | 430 | | |
| | 0 | 431 | | CUcontext popped = nullptr; |
| | 0 | 432 | | const auto result = cuCtxPopCurrent(&popped); |
| | 0 | 433 | | if (result != CUDA_SUCCESS) { |
| | 0 | 434 | | logger()->critical("{}", driver_error_message(result, "cuCtxPopCurrent", __FILE__, __LINE__)); |
| | 0 | 435 | | std::abort(); |
| | | 436 | | } |
| | 1 | 437 | | } |
| | | 438 | | |
| | | 439 | | ScopedCudaContext(const ScopedCudaContext &) = delete; |
| | | 440 | | ScopedCudaContext &operator=(const ScopedCudaContext &) = delete; |
| | | 441 | | |
| | | 442 | | private: |
| | 1 | 443 | | bool pushed_{false}; |
| | | 444 | | }; |
| | | 445 | | |
| | 1 | 446 | | CUcontext cuda_context_for_stream(cudaStream_t stream) { |
| | 1 | 447 | | CUcontext context = nullptr; |
| | | 448 | | |
| | 1 | 449 | | if (stream != nullptr) { |
| | 1 | 450 | | PCA_DRIVER_CHECK(cuStreamGetCtx(reinterpret_cast<CUstream>(stream), &context)); |
| | 1 | 451 | | } else { |
| | 1 | 452 | | PCA_DRIVER_CHECK(cuCtxGetCurrent(&context)); |
| | | 453 | | } |
| | | 454 | | |
| | 1 | 455 | | if (context == nullptr) { |
| | 0 | 456 | | throw std::runtime_error("PCA eigensolver requires an active CUDA context"); |
| | | 457 | | } |
| | | 458 | | |
| | 1 | 459 | | return context; |
| | 1 | 460 | | } |
| | | 461 | | |
| | 1 | 462 | | template <typename T> T read_module_constant(CUmodule module, const char *name) { |
| | 1 | 463 | | CUdeviceptr address = 0; |
| | 1 | 464 | | size_t size = 0; |
| | | 465 | | |
| | 1 | 466 | | PCA_DRIVER_CHECK(cuModuleGetGlobal(&address, &size, module, name)); |
| | 1 | 467 | | if (size != sizeof(T)) { |
| | 0 | 468 | | throw std::runtime_error( |
| | | 469 | | std::format("Unexpected size for cuSolverDx module constant {}", name)); |
| | | 470 | | } |
| | | 471 | | |
| | 1 | 472 | | T value{}; |
| | 1 | 473 | | PCA_DRIVER_CHECK(cuMemcpyDtoH(&value, address, sizeof(T))); |
| | 1 | 474 | | return value; |
| | 1 | 475 | | } |
| | | 476 | | |
| | 1 | 477 | | CUfunction module_function(CUmodule module, const char *name) { |
| | 1 | 478 | | CUfunction function = nullptr; |
| | 1 | 479 | | PCA_DRIVER_CHECK(cuModuleGetFunction(&function, module, name)); |
| | 1 | 480 | | return function; |
| | 1 | 481 | | } |
| | | 482 | | |
| | | 483 | | // ------------------------------------------------------------------------------------------------- |
| | | 484 | | // Conventional cuSOLVER eigensolver |
| | | 485 | | // ------------------------------------------------------------------------------------------------- |
| | | 486 | | |
| | | 487 | | class CusolverEigensolver final : public Eigensolver { |
| | | 488 | | public: |
| | | 489 | | CusolverEigensolver(int features, size_t batches, float *matrices, float *eigenvalues, |
| | | 490 | | cudaStream_t stream) |
| | 1 | 491 | | : features_(features), batches_(batches), context_(cuda_context_for_stream(stream)), |
| | 1 | 492 | | handle_(std::make_unique<curaii::CusolverDnHandle>()), |
| | 1 | 493 | | params_(std::make_unique<curaii::CusolverDnParams>()) { |
| | 1 | 494 | | if (batches_ > static_cast<size_t>((std::numeric_limits<int64_t>::max)())) { |
| | 0 | 495 | | throw std::invalid_argument("[Pca] Batch count exceeds the cuSOLVER API limit"); |
| | | 496 | | } |
| | | 497 | | |
| | 1 | 498 | | CUSOLVER_CHECK(cusolverDnSetStream(handle_->get(), stream)); |
| | 1 | 499 | | CUSOLVER_CHECK(cusolverDnXsyevBatched_bufferSize( |
| | | 500 | | handle_->get(), params_->get(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, features_, |
| | | 501 | | CUDA_R_32F, matrices, features_, CUDA_R_32F, eigenvalues, CUDA_R_32F, |
| | | 502 | | &device_workspace_size_, &host_workspace_size_, static_cast<int64_t>(batches_))); |
| | | 503 | | |
| | 1 | 504 | | if (device_workspace_size_ != 0) { |
| | 1 | 505 | | device_workspace_ = curaii::make_unique_device_ptr<std::byte>(device_workspace_size_); |
| | | 506 | | } |
| | 1 | 507 | | host_workspace_.resize(host_workspace_size_); |
| | 1 | 508 | | } |
| | | 509 | | |
| | 1 | 510 | | [[nodiscard]] bool is_compatible_stream(cudaStream_t stream) const override { |
| | 1 | 511 | | return cuda_context_for_stream(stream) == context_; |
| | 1 | 512 | | } |
| | | 513 | | |
| | 1 | 514 | | void solve(float *matrices, float *eigenvalues, int *info, cudaStream_t stream) override { |
| | 1 | 515 | | if (!is_compatible_stream(stream)) { |
| | 0 | 516 | | throw std::invalid_argument( |
| | | 517 | | "[Pca] cuSOLVER eigensolver and PCA stream use different contexts"); |
| | | 518 | | } |
| | | 519 | | |
| | 1 | 520 | | CUSOLVER_CHECK(cusolverDnSetStream(handle_->get(), stream)); |
| | 1 | 521 | | CUSOLVER_CHECK(cusolverDnXsyevBatched( |
| | | 522 | | handle_->get(), params_->get(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, features_, |
| | | 523 | | CUDA_R_32F, matrices, features_, CUDA_R_32F, eigenvalues, CUDA_R_32F, |
| | | 524 | | device_workspace_.get(), device_workspace_size_, host_workspace_.data(), |
| | | 525 | | host_workspace_size_, info, static_cast<int64_t>(batches_))); |
| | 1 | 526 | | } |
| | | 527 | | |
| | | 528 | | private: |
| | | 529 | | int features_; |
| | | 530 | | size_t batches_; |
| | | 531 | | CUcontext context_{nullptr}; |
| | | 532 | | |
| | | 533 | | std::unique_ptr<curaii::CusolverDnHandle> handle_; |
| | | 534 | | std::unique_ptr<curaii::CusolverDnParams> params_; |
| | | 535 | | |
| | | 536 | | curaii::unique_device_ptr<std::byte> device_workspace_; |
| | 1 | 537 | | size_t device_workspace_size_{0}; |
| | | 538 | | std::vector<std::byte> host_workspace_; |
| | 1 | 539 | | size_t host_workspace_size_{0}; |
| | | 540 | | }; |
| | | 541 | | |
| | | 542 | | // ------------------------------------------------------------------------------------------------- |
| | | 543 | | // cuSolverDx runtime compilation |
| | | 544 | | // ------------------------------------------------------------------------------------------------- |
| | | 545 | | |
| | | 546 | | const char *cusolverdx_source(); |
| | | 547 | | |
| | | 548 | | namespace cusolverdx_runtime { |
| | | 549 | | |
| | | 550 | | struct Assets { |
| | | 551 | | std::filesystem::path cuda_include; |
| | | 552 | | std::filesystem::path cusolverdx_include; |
| | | 553 | | std::filesystem::path cutlass_include; |
| | | 554 | | std::filesystem::path cusolverdx_fatbin; |
| | | 555 | | }; |
| | | 556 | | |
| | 0 | 557 | | std::string jitlink_error_log(nvJitLinkHandle linker) { |
| | 0 | 558 | | size_t size = 0; |
| | 0 | 559 | | if (linker == nullptr || nvJitLinkGetErrorLogSize(linker, &size) != NVJITLINK_SUCCESS || |
| | | 560 | | size == 0) { |
| | 0 | 561 | | return {}; |
| | | 562 | | } |
| | | 563 | | |
| | 0 | 564 | | std::vector<char> log(size); |
| | 0 | 565 | | if (nvJitLinkGetErrorLog(linker, log.data()) != NVJITLINK_SUCCESS) { |
| | 0 | 566 | | return {}; |
| | | 567 | | } |
| | | 568 | | |
| | 0 | 569 | | return std::string(log.data()); |
| | 0 | 570 | | } |
| | | 571 | | |
| | | 572 | | void jitlink_check(nvJitLinkHandle linker, nvJitLinkResult result, const char *expression, |
| | 1 | 573 | | const char *file, int line) { |
| | 1 | 574 | | if (result == NVJITLINK_SUCCESS) { |
| | 1 | 575 | | return; |
| | | 576 | | } |
| | | 577 | | |
| | 0 | 578 | | const auto log = jitlink_error_log(linker); |
| | 0 | 579 | | const auto message = std::format("nvJitLink error: {} at {}:{}\n expression : {}\n{}", |
| | | 580 | | static_cast<int>(result), file, line, expression, log); |
| | 0 | 581 | | logger()->error("{}", message); |
| | 0 | 582 | | throw std::runtime_error(message); |
| | 1 | 583 | | } |
| | | 584 | | |
| | | 585 | | #define PCA_NVJITLINK_CHECK(linker, expr) jitlink_check((linker), (expr), #expr, __FILE__, __LINE__) |
| | | 586 | | |
| | 1 | 587 | | std::filesystem::path executable_directory() { |
| | 1 | 588 | | std::vector<wchar_t> buffer(MAX_PATH); |
| | | 589 | | |
| | 1 | 590 | | while (true) { |
| | 1 | 591 | | const auto length = |
| | | 592 | | GetModuleFileNameW(nullptr, buffer.data(), static_cast<DWORD>(buffer.size())); |
| | | 593 | | |
| | 1 | 594 | | if (length == 0) { |
| | 0 | 595 | | throw std::runtime_error( |
| | | 596 | | std::format("[Pca] Failed to locate the executable (Windows error {})", GetLastError())); |
| | | 597 | | } |
| | | 598 | | |
| | 1 | 599 | | if (length < buffer.size() - 1) { |
| | 1 | 600 | | return std::filesystem::path(std::wstring(buffer.data(), length)).parent_path(); |
| | | 601 | | } |
| | | 602 | | |
| | 0 | 603 | | buffer.resize(buffer.size() * 2); |
| | 0 | 604 | | } |
| | 1 | 605 | | } |
| | | 606 | | |
| | 1 | 607 | | bool assets_exist(const Assets &assets) { |
| | 1 | 608 | | return std::filesystem::is_directory(assets.cuda_include) && |
| | | 609 | | std::filesystem::is_regular_file(assets.cuda_include / "cuda_runtime.h") && |
| | | 610 | | std::filesystem::is_directory(assets.cuda_include / "cccl") && |
| | | 611 | | std::filesystem::is_regular_file(assets.cusolverdx_include / "cusolverdx.hpp") && |
| | | 612 | | std::filesystem::is_regular_file(assets.cusolverdx_include / "cusolverdx_io.hpp") && |
| | | 613 | | std::filesystem::is_directory(assets.cutlass_include) && |
| | | 614 | | std::filesystem::is_regular_file(assets.cusolverdx_fatbin); |
| | 1 | 615 | | } |
| | | 616 | | |
| | 1 | 617 | | Assets locate_assets() { |
| | 1 | 618 | | const auto installed_root = |
| | | 619 | | (executable_directory() / HOLOFLOW_NVRTC_ASSET_RELATIVE_DIR).lexically_normal(); |
| | | 620 | | |
| | | 621 | | const Assets installed{ |
| | 1 | 622 | | installed_root / "cuda/include", |
| | 1 | 623 | | installed_root / "mathdx/include", |
| | 1 | 624 | | installed_root / "mathdx/cutlass/include", |
| | 1 | 625 | | installed_root / "mathdx/lib/libcusolverdx.fatbin", |
| | | 626 | | }; |
| | | 627 | | |
| | 1 | 628 | | if (assets_exist(installed)) { |
| | 0 | 629 | | return installed; |
| | | 630 | | } |
| | | 631 | | |
| | | 632 | | const Assets build{ |
| | 1 | 633 | | HOLOFLOW_CUDA_INCLUDE_DIR, |
| | 1 | 634 | | HOLOFLOW_CUSOLVERDX_INCLUDE_DIR, |
| | 1 | 635 | | HOLOFLOW_CUSOLVERDX_CUTLASS_INCLUDE_DIR, |
| | 1 | 636 | | HOLOFLOW_CUSOLVERDX_FATBIN, |
| | | 637 | | }; |
| | | 638 | | |
| | 1 | 639 | | if (assets_exist(build)) { |
| | 1 | 640 | | return build; |
| | | 641 | | } |
| | | 642 | | |
| | 0 | 643 | | throw std::runtime_error(std::format( |
| | | 644 | | "[Pca] cuSolverDx runtime-compilation assets are missing. Expected installed assets under " |
| | | 645 | | "'{}' (CUDA headers, cuSolverDx/CommonDx headers, CUTLASS headers, and " |
| | | 646 | | "libcusolverdx.fatbin); build-tree fallback is also unavailable", |
| | | 647 | | installed_root.string())); |
| | 1 | 648 | | } |
| | | 649 | | |
| | | 650 | | std::vector<char> compile_lto_ir(int features, int solver_sm, int architecture, |
| | 1 | 651 | | const Assets &assets) { |
| | 1 | 652 | | curaii::NvrtcProgram program(cusolverdx_source(), "pca_heev_kernel.cu"); |
| | | 653 | | |
| | 1 | 654 | | std::vector<std::string> options = { |
| | | 655 | | "--std=c++17", |
| | | 656 | | "--device-as-default-execution-space", |
| | | 657 | | "-dlto", |
| | | 658 | | "--relocatable-device-code=true", |
| | | 659 | | std::format("--gpu-architecture=sm_{}", architecture), |
| | | 660 | | std::format("-DPCA_FEATURES={}", features), |
| | | 661 | | std::format("-DPCA_SOLVER_SM={}", solver_sm), |
| | | 662 | | std::format("--include-path={}", assets.cuda_include.string()), |
| | | 663 | | std::format("--include-path={}", (assets.cuda_include / "cccl").string()), |
| | | 664 | | std::format("--include-path={}", assets.cusolverdx_include.string()), |
| | | 665 | | std::format("--include-path={}", assets.cutlass_include.string()), |
| | | 666 | | }; |
| | | 667 | | |
| | 1 | 668 | | std::vector<const char *> option_ptrs; |
| | 1 | 669 | | option_ptrs.reserve(options.size()); |
| | 1 | 670 | | for (const auto &option : options) { |
| | 1 | 671 | | option_ptrs.push_back(option.c_str()); |
| | 1 | 672 | | } |
| | | 673 | | |
| | 1 | 674 | | const auto compile_result = |
| | | 675 | | nvrtcCompileProgram(program.get(), static_cast<int>(option_ptrs.size()), option_ptrs.data()); |
| | | 676 | | |
| | 1 | 677 | | if (compile_result != NVRTC_SUCCESS) { |
| | 0 | 678 | | size_t log_size = 0; |
| | 0 | 679 | | NVRTC_CHECK(nvrtcGetProgramLogSize(program.get(), &log_size)); |
| | | 680 | | |
| | 0 | 681 | | std::vector<char> log(log_size); |
| | 0 | 682 | | NVRTC_CHECK(nvrtcGetProgramLog(program.get(), log.data())); |
| | 0 | 683 | | logger()->error("[Pca] NVRTC compilation log:\n{}", log.data()); |
| | 0 | 684 | | NVRTC_CHECK(compile_result); |
| | 0 | 685 | | } |
| | | 686 | | |
| | 1 | 687 | | size_t lto_size = 0; |
| | 1 | 688 | | NVRTC_CHECK(nvrtcGetLTOIRSize(program.get(), <o_size)); |
| | | 689 | | |
| | 1 | 690 | | std::vector<char> lto_ir(lto_size); |
| | 1 | 691 | | NVRTC_CHECK(nvrtcGetLTOIR(program.get(), lto_ir.data())); |
| | 1 | 692 | | return lto_ir; |
| | 1 | 693 | | } |
| | | 694 | | |
| | | 695 | | std::vector<char> link_cubin(const std::vector<char> <o_ir, int architecture, |
| | 1 | 696 | | const Assets &assets) { |
| | 1 | 697 | | nvJitLinkHandle linker = nullptr; |
| | 1 | 698 | | const auto arch_option = std::format("-arch=sm_{}", architecture); |
| | 1 | 699 | | const char *options[] = {"-lto", arch_option.c_str()}; |
| | | 700 | | |
| | 1 | 701 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkCreate(&linker, 2, options)); |
| | | 702 | | |
| | | 703 | | try { |
| | 1 | 704 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkAddFile(linker, NVJITLINK_INPUT_FATBIN, |
| | | 705 | | assets.cusolverdx_fatbin.string().c_str())); |
| | 1 | 706 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkAddData(linker, NVJITLINK_INPUT_LTOIR, |
| | | 707 | | const_cast<char *>(lto_ir.data()), lto_ir.size(), |
| | | 708 | | "pca_heev_lto_ir")); |
| | 1 | 709 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkComplete(linker)); |
| | | 710 | | |
| | 1 | 711 | | size_t cubin_size = 0; |
| | 1 | 712 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkGetLinkedCubinSize(linker, &cubin_size)); |
| | | 713 | | |
| | 1 | 714 | | std::vector<char> cubin(cubin_size); |
| | 1 | 715 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkGetLinkedCubin(linker, cubin.data())); |
| | 1 | 716 | | PCA_NVJITLINK_CHECK(linker, nvJitLinkDestroy(&linker)); |
| | 1 | 717 | | return cubin; |
| | 0 | 718 | | } catch (...) { |
| | 0 | 719 | | if (linker != nullptr) { |
| | 0 | 720 | | (void)nvJitLinkDestroy(&linker); |
| | | 721 | | } |
| | 0 | 722 | | throw; |
| | 0 | 723 | | } |
| | 1 | 724 | | } |
| | | 725 | | |
| | 1 | 726 | | std::vector<char> compile_cubin(int features, int solver_sm, int architecture) { |
| | 1 | 727 | | const auto assets = locate_assets(); |
| | 1 | 728 | | const auto lto_ir = compile_lto_ir(features, solver_sm, architecture, assets); |
| | 1 | 729 | | return link_cubin(lto_ir, architecture, assets); |
| | 1 | 730 | | } |
| | | 731 | | |
| | | 732 | | } // namespace cusolverdx_runtime |
| | | 733 | | |
| | | 734 | | // ------------------------------------------------------------------------------------------------- |
| | | 735 | | // cuSolverDx eigensolver |
| | | 736 | | // ------------------------------------------------------------------------------------------------- |
| | | 737 | | |
| | | 738 | | struct DxKernel { |
| | | 739 | | CUfunction function{nullptr}; |
| | | 740 | | unsigned int block_x{0}; |
| | | 741 | | unsigned int shared_memory{0}; |
| | | 742 | | }; |
| | | 743 | | |
| | | 744 | | class CusolverDxEigensolver final : public Eigensolver { |
| | | 745 | | public: |
| | 1 | 746 | | CusolverDxEigensolver(int features, size_t batches, cudaStream_t stream) |
| | 1 | 747 | | : features_(features), batches_(batches) { |
| | 1 | 748 | | CUDA_CHECK(cudaGetDevice(&device_)); |
| | | 749 | | |
| | 1 | 750 | | int major = 0; |
| | 1 | 751 | | int minor = 0; |
| | 1 | 752 | | CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_)); |
| | 1 | 753 | | CUDA_CHECK(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_)); |
| | | 754 | | |
| | 1 | 755 | | architecture_ = major * 10 + minor; |
| | 1 | 756 | | const int solver_sm = architecture_ * 10; |
| | | 757 | | |
| | 1 | 758 | | PCA_DRIVER_CHECK(cuInit(0)); |
| | 1 | 759 | | context_ = cuda_context_for_stream(stream); |
| | | 760 | | |
| | 1 | 761 | | ScopedCudaContext context_guard(context_); |
| | | 762 | | |
| | 1 | 763 | | const auto cubin = cusolverdx_runtime::compile_cubin(features_, solver_sm, architecture_); |
| | 1 | 764 | | PCA_DRIVER_CHECK(cuModuleLoadDataEx(&module_, cubin.data(), 0, nullptr, nullptr)); |
| | | 765 | | |
| | | 766 | | try { |
| | 1 | 767 | | load_kernel_configuration(); |
| | 1 | 768 | | configure_shared_memory(); |
| | 0 | 769 | | } catch (...) { |
| | 0 | 770 | | (void)cuModuleUnload(module_); |
| | 0 | 771 | | module_ = nullptr; |
| | 0 | 772 | | throw; |
| | 0 | 773 | | } |
| | 1 | 774 | | } |
| | | 775 | | |
| | 1 | 776 | | ~CusolverDxEigensolver() noexcept override { |
| | 1 | 777 | | if (module_ == nullptr) { |
| | 0 | 778 | | return; |
| | | 779 | | } |
| | | 780 | | |
| | | 781 | | try { |
| | 1 | 782 | | ScopedCudaContext context_guard(context_); |
| | 1 | 783 | | const auto result = cuModuleUnload(module_); |
| | | 784 | | |
| | 1 | 785 | | if (result != CUDA_SUCCESS) { |
| | 0 | 786 | | logger()->critical("{}", |
| | | 787 | | driver_error_message(result, "cuModuleUnload", __FILE__, __LINE__)); |
| | 0 | 788 | | std::abort(); |
| | | 789 | | } |
| | 1 | 790 | | } catch (...) { |
| | 0 | 791 | | std::abort(); |
| | 0 | 792 | | } |
| | 1 | 793 | | } |
| | | 794 | | |
| | 1 | 795 | | [[nodiscard]] bool is_compatible_stream(cudaStream_t stream) const override { |
| | 1 | 796 | | return cuda_context_for_stream(stream) == context_; |
| | 1 | 797 | | } |
| | | 798 | | |
| | 1 | 799 | | void solve(float *matrices, float *eigenvalues, int *info, cudaStream_t stream) override { |
| | 1 | 800 | | if (!is_compatible_stream(stream)) { |
| | 0 | 801 | | throw std::invalid_argument( |
| | | 802 | | "[Pca] cuSolverDx eigensolver and PCA stream use different contexts"); |
| | | 803 | | } |
| | | 804 | | |
| | 1 | 805 | | const size_t full_batches = (batches_ / batches_per_block_) * batches_per_block_; |
| | 1 | 806 | | if (full_batches != 0) { |
| | 1 | 807 | | launch_batched(matrices, eigenvalues, info, full_batches, stream); |
| | | 808 | | } |
| | | 809 | | |
| | 1 | 810 | | const size_t tail_batches = batches_ - full_batches; |
| | 1 | 811 | | if (tail_batches != 0) { |
| | 1 | 812 | | launch_tail(matrices, eigenvalues, info, full_batches, tail_batches, stream); |
| | | 813 | | } |
| | 1 | 814 | | } |
| | | 815 | | |
| | | 816 | | private: |
| | 1 | 817 | | void load_kernel_configuration() { |
| | 1 | 818 | | batched_kernel_ = { |
| | | 819 | | .function = module_function(module_, "pca_heev_batched"), |
| | | 820 | | .block_x = read_module_constant<unsigned int>(module_, "pca_batched_block_x"), |
| | | 821 | | .shared_memory = read_module_constant<unsigned int>(module_, "pca_batched_shared_memory"), |
| | | 822 | | }; |
| | | 823 | | |
| | 1 | 824 | | tail_kernel_ = { |
| | | 825 | | .function = module_function(module_, "pca_heev_tail"), |
| | | 826 | | .block_x = read_module_constant<unsigned int>(module_, "pca_tail_block_x"), |
| | | 827 | | .shared_memory = read_module_constant<unsigned int>(module_, "pca_tail_shared_memory"), |
| | | 828 | | }; |
| | | 829 | | |
| | 1 | 830 | | batches_per_block_ = read_module_constant<unsigned int>(module_, "pca_batches_per_block"); |
| | 1 | 831 | | } |
| | | 832 | | |
| | 1 | 833 | | void configure_shared_memory() { |
| | 1 | 834 | | int max_shared_memory = 0; |
| | 1 | 835 | | CUDA_CHECK(cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlockOptin, |
| | | 836 | | device_)); |
| | | 837 | | |
| | 1 | 838 | | const bool uses_batched = batches_ >= batches_per_block_; |
| | 1 | 839 | | const bool uses_tail = batches_ < batches_per_block_ || (batches_ % batches_per_block_) != 0; |
| | | 840 | | |
| | 1 | 841 | | if (uses_batched && |
| | | 842 | | batched_kernel_.shared_memory > static_cast<unsigned int>(max_shared_memory)) { |
| | 0 | 843 | | throw_shared_memory_error(batched_kernel_.shared_memory, max_shared_memory); |
| | | 844 | | } |
| | | 845 | | |
| | 1 | 846 | | if (uses_tail && tail_kernel_.shared_memory > static_cast<unsigned int>(max_shared_memory)) { |
| | 0 | 847 | | throw_shared_memory_error(tail_kernel_.shared_memory, max_shared_memory); |
| | | 848 | | } |
| | | 849 | | |
| | 1 | 850 | | if (uses_batched) { |
| | 1 | 851 | | PCA_DRIVER_CHECK(cuFuncSetAttribute(batched_kernel_.function, |
| | | 852 | | CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, |
| | | 853 | | static_cast<int>(batched_kernel_.shared_memory))); |
| | | 854 | | } |
| | | 855 | | |
| | 1 | 856 | | if (uses_tail) { |
| | 1 | 857 | | PCA_DRIVER_CHECK(cuFuncSetAttribute(tail_kernel_.function, |
| | | 858 | | CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, |
| | | 859 | | static_cast<int>(tail_kernel_.shared_memory))); |
| | | 860 | | } |
| | 1 | 861 | | } |
| | | 862 | | |
| | 0 | 863 | | [[noreturn]] void throw_shared_memory_error(unsigned int required, int max_shared_memory) const { |
| | 0 | 864 | | throw std::runtime_error( |
| | | 865 | | std::format("[Pca] cuSolverDx HEEV for {} features requires {} bytes of dynamic " |
| | | 866 | | "shared memory, but device sm_{} provides {} bytes", |
| | | 867 | | features_, required, architecture_, max_shared_memory)); |
| | 0 | 868 | | } |
| | | 869 | | |
| | | 870 | | void launch_batched(float *matrices, float *eigenvalues, int *info, size_t full_batches, |
| | 1 | 871 | | cudaStream_t stream) const { |
| | 1 | 872 | | CUdeviceptr matrix_arg = reinterpret_cast<CUdeviceptr>(matrices); |
| | 1 | 873 | | CUdeviceptr eigenvalues_arg = reinterpret_cast<CUdeviceptr>(eigenvalues); |
| | 1 | 874 | | CUdeviceptr info_arg = reinterpret_cast<CUdeviceptr>(info); |
| | 1 | 875 | | auto batch_arg = static_cast<unsigned int>(full_batches); |
| | 1 | 876 | | void *arguments[] = {&matrix_arg, &eigenvalues_arg, &info_arg, &batch_arg}; |
| | | 877 | | |
| | 1 | 878 | | PCA_DRIVER_CHECK(cuLaunchKernel(batched_kernel_.function, |
| | | 879 | | static_cast<unsigned int>(full_batches / batches_per_block_), 1, |
| | | 880 | | 1, batched_kernel_.block_x, 1, 1, batched_kernel_.shared_memory, |
| | | 881 | | reinterpret_cast<CUstream>(stream), arguments, nullptr)); |
| | 1 | 882 | | } |
| | | 883 | | |
| | | 884 | | void launch_tail(float *matrices, float *eigenvalues, int *info, size_t full_batches, |
| | 1 | 885 | | size_t tail_batches, cudaStream_t stream) const { |
| | 1 | 886 | | const size_t matrix_offset = full_batches * static_cast<size_t>(features_) * features_; |
| | 1 | 887 | | const size_t value_offset = full_batches * static_cast<size_t>(features_); |
| | | 888 | | |
| | 1 | 889 | | CUdeviceptr matrix_arg = reinterpret_cast<CUdeviceptr>(matrices + matrix_offset); |
| | 1 | 890 | | CUdeviceptr eigenvalues_arg = reinterpret_cast<CUdeviceptr>(eigenvalues + value_offset); |
| | 1 | 891 | | CUdeviceptr info_arg = reinterpret_cast<CUdeviceptr>(info + full_batches); |
| | 1 | 892 | | auto batch_arg = static_cast<unsigned int>(tail_batches); |
| | 1 | 893 | | void *arguments[] = {&matrix_arg, &eigenvalues_arg, &info_arg, &batch_arg}; |
| | | 894 | | |
| | 1 | 895 | | PCA_DRIVER_CHECK(cuLaunchKernel(tail_kernel_.function, static_cast<unsigned int>(tail_batches), |
| | | 896 | | 1, 1, tail_kernel_.block_x, 1, 1, tail_kernel_.shared_memory, |
| | | 897 | | reinterpret_cast<CUstream>(stream), arguments, nullptr)); |
| | 1 | 898 | | } |
| | | 899 | | |
| | | 900 | | int features_; |
| | | 901 | | size_t batches_; |
| | 1 | 902 | | int device_{0}; |
| | 1 | 903 | | int architecture_{0}; |
| | 1 | 904 | | CUcontext context_{nullptr}; |
| | 1 | 905 | | CUmodule module_{nullptr}; |
| | | 906 | | |
| | | 907 | | DxKernel batched_kernel_; |
| | | 908 | | DxKernel tail_kernel_; |
| | 1 | 909 | | unsigned int batches_per_block_{1}; |
| | | 910 | | }; |
| | | 911 | | |
| | | 912 | | // ------------------------------------------------------------------------------------------------- |
| | | 913 | | // Eigensolver selection |
| | | 914 | | // ------------------------------------------------------------------------------------------------- |
| | | 915 | | |
| | | 916 | | constexpr int cusolverdx_max_features_exclusive = 256; |
| | | 917 | | |
| | | 918 | | std::unique_ptr<Eigensolver> make_eigensolver(const PcaLayout &layout, |
| | 1 | 919 | | const PcaWorkspace &workspace, cudaStream_t stream) { |
| | 1 | 920 | | if (layout.features < cusolverdx_max_features_exclusive) { |
| | | 921 | | try { |
| | 1 | 922 | | return std::make_unique<CusolverDxEigensolver>(layout.features, layout.batches, stream); |
| | 0 | 923 | | } catch (const std::exception &error) { |
| | 0 | 924 | | logger()->warn("[Pca] cuSolverDx initialization failed for depth {}: {}\n" |
| | | 925 | | "[Pca] Falling back to the conventional cuSOLVER eigensolver", |
| | | 926 | | layout.features, error.what()); |
| | 0 | 927 | | } |
| | | 928 | | } |
| | | 929 | | |
| | | 930 | | try { |
| | 1 | 931 | | return std::make_unique<CusolverEigensolver>(layout.features, layout.batches, |
| | | 932 | | workspace.matrices.get(), |
| | | 933 | | workspace.eigenvalues.get(), stream); |
| | 0 | 934 | | } catch (const std::exception &error) { |
| | 0 | 935 | | logger()->error("[Pca] Failed to initialize any GPU eigensolver for depth {}: {}", |
| | | 936 | | layout.features, error.what()); |
| | | 937 | | |
| | 0 | 938 | | throw std::runtime_error(std::format("PCA could not initialize a GPU eigensolver for depth {}. " |
| | | 939 | | "See the terminal log for details.", |
| | | 940 | | layout.features)); |
| | 0 | 941 | | } |
| | 1 | 942 | | } |
| | | 943 | | |
| | | 944 | | // ------------------------------------------------------------------------------------------------- |
| | | 945 | | // Embedded cuSolverDx device program |
| | | 946 | | // ------------------------------------------------------------------------------------------------- |
| | | 947 | | |
| | 1 | 948 | | const char *cusolverdx_source() { |
| | | 949 | | static constexpr char source[] = R"cusolverdx( |
| | | 950 | | #include <cusolverdx.hpp> |
| | | 951 | | #include <cusolverdx_io.hpp> |
| | | 952 | | |
| | | 953 | | using namespace cusolverdx; |
| | | 954 | | |
| | | 955 | | using Base = decltype(Size<PCA_FEATURES>() + Precision<float>() + Type<type::real>() + |
| | | 956 | | Function<heev>() + FillMode<fill_mode::lower>() + |
| | | 957 | | Arrangement<arrangement::col_major>() + |
| | | 958 | | Job<job::overwrite_vectors>() + SM<PCA_SOLVER_SM>() + Block()); |
| | | 959 | | using BatchedSolver = |
| | | 960 | | decltype(Base() + BatchesPerBlock<Base::suggested_batches_per_block>()); |
| | | 961 | | using TailSolver = decltype(Base() + BatchesPerBlock<1>()); |
| | | 962 | | |
| | | 963 | | extern "C" __constant__ unsigned int pca_batched_block_x = BatchedSolver::block_dim.x; |
| | | 964 | | extern "C" __constant__ unsigned int pca_batched_shared_memory = |
| | | 965 | | BatchedSolver::shared_memory_size; |
| | | 966 | | extern "C" __constant__ unsigned int pca_batches_per_block = |
| | | 967 | | BatchedSolver::batches_per_block; |
| | | 968 | | extern "C" __constant__ unsigned int pca_tail_block_x = TailSolver::block_dim.x; |
| | | 969 | | extern "C" __constant__ unsigned int pca_tail_shared_memory = TailSolver::shared_memory_size; |
| | | 970 | | |
| | | 971 | | template <class Solver> |
| | | 972 | | __device__ void solve(float *covariance, float *eigenvalues, int *info, |
| | | 973 | | const unsigned int batches) { |
| | | 974 | | constexpr unsigned int m = Solver::m_size; |
| | | 975 | | constexpr unsigned int batches_per_block = Solver::batches_per_block; |
| | | 976 | | constexpr unsigned int lda_shared = Solver::lda; |
| | | 977 | | constexpr unsigned int matrix_elements = m * m; |
| | | 978 | | |
| | | 979 | | const unsigned int batch = blockIdx.x * batches_per_block; |
| | | 980 | | if (batch >= batches) { |
| | | 981 | | return; |
| | | 982 | | } |
| | | 983 | | |
| | | 984 | | extern __shared__ __align__(16) cusolverdx::byte shared_memory[]; |
| | | 985 | | auto [matrix_shared, eigenvalues_shared, workspace_shared] = |
| | | 986 | | cusolverdx::shared_memory::slice<float, float, float>( |
| | | 987 | | shared_memory, alignof(float), lda_shared * m * batches_per_block, alignof(float), |
| | | 988 | | m * batches_per_block, alignof(float)); |
| | | 989 | | |
| | | 990 | | float *matrix_global = covariance + matrix_elements * batch; |
| | | 991 | | float *eigenvalues_global = eigenvalues + m * batch; |
| | | 992 | | |
| | | 993 | | cusolverdx::copy_2d<Solver, m, m, arrangement::col_major, batches_per_block>( |
| | | 994 | | matrix_global, m, matrix_shared, lda_shared); |
| | | 995 | | __syncthreads(); |
| | | 996 | | |
| | | 997 | | Solver().execute(matrix_shared, lda_shared, eigenvalues_shared, workspace_shared, info + batch); |
| | | 998 | | |
| | | 999 | | cusolverdx::copy_2d<Solver, m, 1, arrangement::col_major, batches_per_block>( |
| | | 1000 | | eigenvalues_shared, m, eigenvalues_global, m); |
| | | 1001 | | __syncthreads(); |
| | | 1002 | | cusolverdx::copy_2d<Solver, m, m, arrangement::col_major, batches_per_block>( |
| | | 1003 | | matrix_shared, lda_shared, matrix_global, m); |
| | | 1004 | | } |
| | | 1005 | | |
| | | 1006 | | extern "C" __global__ void pca_heev_batched(float *covariance, float *eigenvalues, int *info, |
| | | 1007 | | const unsigned int batches) { |
| | | 1008 | | solve<BatchedSolver>(covariance, eigenvalues, info, batches); |
| | | 1009 | | } |
| | | 1010 | | |
| | | 1011 | | extern "C" __global__ void pca_heev_tail(float *covariance, float *eigenvalues, int *info, |
| | | 1012 | | const unsigned int batches) { |
| | | 1013 | | solve<TailSolver>(covariance, eigenvalues, info, batches); |
| | | 1014 | | } |
| | | 1015 | | )cusolverdx"; |
| | | 1016 | | |
| | 1 | 1017 | | return source; |
| | 1 | 1018 | | } |
| | | 1019 | | |
| | | 1020 | | #undef PCA_NVJITLINK_CHECK |
| | | 1021 | | #undef PCA_DRIVER_CHECK |
| | | 1022 | | |
| | | 1023 | | } // namespace |
| | | 1024 | | |
| | | 1025 | | // ------------------------------------------------------------------------------------------------- |
| | | 1026 | | // JSON serialization |
| | | 1027 | | // ------------------------------------------------------------------------------------------------- |
| | | 1028 | | |
| | 1 | 1029 | | void to_json(nlohmann::json &j, const PcaSettings &settings) { |
| | 1 | 1030 | | j = nlohmann::json{ |
| | | 1031 | | {"begin", settings.begin}, |
| | | 1032 | | {"end", settings.end}, |
| | | 1033 | | }; |
| | 1 | 1034 | | } |
| | | 1035 | | |
| | 1 | 1036 | | void from_json(const nlohmann::json &j, PcaSettings &settings) { |
| | 1 | 1037 | | j.at("begin").get_to(settings.begin); |
| | 1 | 1038 | | j.at("end").get_to(settings.end); |
| | 1 | 1039 | | } |
| | | 1040 | | |
| | | 1041 | | // ------------------------------------------------------------------------------------------------- |
| | | 1042 | | // Factory methods |
| | | 1043 | | // ------------------------------------------------------------------------------------------------- |
| | | 1044 | | |
| | | 1045 | | holoflow::core::InferResult PcaFactory::infer(std::span<const holoflow::core::TDesc> input_descs, |
| | 1 | 1046 | | const nlohmann::json &jsettings) const { |
| | | 1047 | | const auto check = [&](bool condition, const std::string &message) { |
| | | 1048 | | if (!condition) { |
| | | 1049 | | logger()->error("[PcaFactory::infer] error: {}", message); |
| | | 1050 | | throw std::invalid_argument("PcaFactory inference error: " + message); |
| | | 1051 | | } |
| | | 1052 | | }; |
| | | 1053 | | |
| | 1 | 1054 | | const auto settings = jsettings.get<PcaSettings>(); |
| | | 1055 | | |
| | 1 | 1056 | | check(input_descs.size() == 1, "expected exactly one input"); |
| | | 1057 | | |
| | 1 | 1058 | | const auto &input_desc = input_descs.front(); |
| | 1 | 1059 | | check(input_desc.rank() >= 3, "expected input rank >= 3"); |
| | 1 | 1060 | | check(input_desc.dtype == holoflow::core::DType::F32, "PCA currently supports F32 input only"); |
| | 1 | 1061 | | check(input_desc.mem_loc == holoflow::core::MemLoc::Device, "expected input in device memory"); |
| | 1 | 1062 | | check(settings.begin < settings.end, "expected begin < end"); |
| | 1 | 1063 | | check(settings.begin >= 0, "expected begin >= 0"); |
| | | 1064 | | |
| | 1 | 1065 | | const PcaLayout layout(input_desc); |
| | 1 | 1066 | | check(settings.end <= layout.features, "expected end <= n_features"); |
| | | 1067 | | |
| | 1 | 1068 | | auto output_shape = input_desc.shape; |
| | 1 | 1069 | | output_shape.at(layout.feature_axis) = static_cast<size_t>(settings.components()); |
| | | 1070 | | |
| | 1 | 1071 | | holoflow::core::TDesc output_desc(output_shape, input_desc.dtype, input_desc.mem_loc); |
| | | 1072 | | |
| | 1 | 1073 | | return holoflow::core::InferResult{ |
| | | 1074 | | .input_descs = {input_desc}, |
| | | 1075 | | .output_descs = {output_desc}, |
| | | 1076 | | .in_place = {}, |
| | | 1077 | | .owned_inputs = {false}, |
| | | 1078 | | .owned_outputs = {false}, |
| | | 1079 | | .kind = holoflow::core::TaskKind::Sync, |
| | | 1080 | | }; |
| | 1 | 1081 | | } |
| | | 1082 | | |
| | | 1083 | | std::unique_ptr<holoflow::core::ISyncTask> |
| | | 1084 | | PcaFactory::create(std::span<const holoflow::core::TDesc> input_descs, |
| | | 1085 | | const nlohmann::json &jsettings, |
| | 1 | 1086 | | const holoflow::core::SyncCreateCtx &ctx) const { |
| | 1 | 1087 | | this->infer(input_descs, jsettings); |
| | | 1088 | | |
| | 1 | 1089 | | return std::make_unique<PcaTask>(jsettings.get<PcaSettings>(), input_descs.front(), ctx); |
| | 1 | 1090 | | } |
| | | 1091 | | |
| | | 1092 | | std::unique_ptr<holoflow::core::ISyncTask> |
| | | 1093 | | PcaFactory::update(std::unique_ptr<holoflow::core::ISyncTask> old_task, |
| | | 1094 | | std::span<const holoflow::core::TDesc> input_descs, |
| | | 1095 | | const nlohmann::json &jsettings, |
| | 1 | 1096 | | const holoflow::core::SyncCreateCtx &ctx) const { |
| | 1 | 1097 | | this->infer(input_descs, jsettings); |
| | | 1098 | | |
| | 1 | 1099 | | const auto settings = jsettings.get<PcaSettings>(); |
| | 1 | 1100 | | const auto &input_desc = input_descs.front(); |
| | | 1101 | | |
| | 1 | 1102 | | auto *pca = dynamic_cast<PcaTask *>(old_task.get()); |
| | 1 | 1103 | | if (pca != nullptr && pca->can_reuse(input_desc, ctx.stream)) { |
| | 1 | 1104 | | pca->reconfigure(settings, ctx.stream); |
| | 1 | 1105 | | return old_task; |
| | | 1106 | | } |
| | | 1107 | | |
| | 0 | 1108 | | return create(input_descs, jsettings, ctx); |
| | 1 | 1109 | | } |
| | | 1110 | | |
| | | 1111 | | } // namespace holotask::syncs |