framework/pytorch/op_internal.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // Tensor/device conversion helpers shared by the PyTorch op adapters. | ||
| 4 | |||
| 5 | #include <torch/extension.h> | ||
| 6 | |||
| 7 | #include <algorithm> | ||
| 8 | #include <chrono> | ||
| 9 | #include <cstddef> | ||
| 10 | #include <cstdint> | ||
| 11 | #include <memory> | ||
| 12 | #include <string> | ||
| 13 | #include <vector> | ||
| 14 | |||
| 15 | #include "base/tensor.h" | ||
| 16 | #include "framework/op.h" | ||
| 17 | #include "framework/common/op_profile.h" | ||
| 18 | #include "framework/cache/cache_profile.h" | ||
| 19 | #include "framework/pytorch/op_cuda_runtime.h" | ||
| 20 | #include "ps/local_shm/local_shm_client.h" | ||
| 21 | |||
| 22 | #include <glog/logging.h> | ||
| 23 | |||
| 24 | #ifdef RECSTORE_ENABLE_GPU_CACHE | ||
| 25 | # include "framework/pytorch/cache/gpu_cache_adapter.h" | ||
| 26 | #endif | ||
| 27 | |||
| 28 | namespace recstore { | ||
| 29 | namespace framework { | ||
| 30 | |||
| 31 | // --- Shared predicates and tensor helpers ---------------------------------- | ||
| 32 | |||
| 33 | inline base::RecTensor | ||
| 34 | 108 | ToRecTensor(const torch::Tensor& tensor, base::DataType dtype) { | |
| 35 | 108 | std::vector<int64_t> shape; | |
| 36 | 270 | for (int i = 0; i < tensor.dim(); ++i) { | |
| 37 | 162 | shape.push_back(tensor.size(i)); | |
| 38 | } | ||
| 39 | 216 | return base::RecTensor(const_cast<void*>(tensor.data_ptr()), shape, dtype); | |
| 40 | 108 | } | |
| 41 | |||
| 42 | ✗ | inline torch::TensorOptions PinnedCpuOptions(torch::ScalarType dtype) { | |
| 43 | ✗ | return torch::TensorOptions() | |
| 44 | ✗ | .device(torch::kCPU) | |
| 45 | ✗ | .dtype(dtype) | |
| 46 | ✗ | .pinned_memory(true); | |
| 47 | } | ||
| 48 | |||
| 49 | ✗ | inline torch::Tensor StageCudaTensorToPinnedCpu(const torch::Tensor& tensor, | |
| 50 | torch::ScalarType dtype) { | ||
| 51 | ✗ | auto cpu_tensor = torch::empty(tensor.sizes(), PinnedCpuOptions(dtype)); | |
| 52 | ✗ | cpu_tensor.copy_(tensor.to(dtype), /*non_blocking=*/false); | |
| 53 | ✗ | return cpu_tensor; | |
| 54 | ✗ | } | |
| 55 | |||
| 56 | inline torch::Tensor | ||
| 57 | ✗ | StageCudaTensorToPinnedCpuAsyncNoCast(const torch::Tensor& tensor) { | |
| 58 | auto cpu_tensor = | ||
| 59 | ✗ | torch::empty(tensor.sizes(), PinnedCpuOptions(tensor.scalar_type())); | |
| 60 | ✗ | cpu_tensor.copy_(tensor, /*non_blocking=*/true); | |
| 61 | ✗ | return cpu_tensor; | |
| 62 | ✗ | } | |
| 63 | |||
| 64 | 42 | inline torch::Tensor ToCpuTensor(const torch::Tensor& tensor) { | |
| 65 | 42 | return tensor.is_cuda() ? tensor.cpu() : tensor; | |
| 66 | } | ||
| 67 | |||
| 68 | 64 | inline std::shared_ptr<KVClientOp> GetConcreteKVClientOp() { | |
| 69 | 64 | auto op = GetKVClientOp(); | |
| 70 | 64 | auto kv_op = std::dynamic_pointer_cast<KVClientOp>(op); | |
| 71 | 64 | TORCH_CHECK(kv_op != nullptr, "storage backend is not KVClientOp"); | |
| 72 | 128 | return kv_op; | |
| 73 | 64 | } | |
| 74 | |||
| 75 | // Tensor adapters for the framework-independent op/cache runtime. | ||
| 76 | void CheckKeys(const torch::Tensor& keys, bool contiguous = true); | ||
| 77 | void CheckRows(const torch::Tensor& keys, | ||
| 78 | const torch::Tensor& values, | ||
| 79 | bool contiguous = true); | ||
| 80 | std::pair<torch::Tensor, torch::Tensor> | ||
| 81 | CudaCacheRows(const torch::Tensor& keys, const torch::Tensor& values); | ||
| 82 | CacheUpdateActions | ||
| 83 | CacheActions(const torch::Tensor& keys = {}, int64_t dim = 0); | ||
| 84 | std::tuple<torch::Tensor, torch::Tensor> CachedReadTorch( | ||
| 85 | const torch::Tensor& keys, | ||
| 86 | int64_t dim, | ||
| 87 | bool local = false, | ||
| 88 | bool no_evict = false); | ||
| 89 | |||
| 90 | } // namespace framework | ||
| 91 | } // namespace recstore | ||
| 92 |