framework/pytorch/cache/tensor_adapter.cc
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "framework/pytorch/op_internal.h" | ||
| 2 | |||
| 3 | namespace recstore::framework { | ||
| 4 | |||
| 5 | 54 | void CheckKeys(const torch::Tensor& keys, bool contiguous) { | |
| 6 | 54 | TORCH_CHECK(keys.dim() == 1 && keys.scalar_type() == torch::kInt64 && | |
| 7 | (!contiguous || keys.is_contiguous()), | ||
| 8 | "Keys must be contiguous 1-D int64"); | ||
| 9 | 54 | } | |
| 10 | |||
| 11 | 20 | void CheckRows(const torch::Tensor& keys, | |
| 12 | const torch::Tensor& values, | ||
| 13 | bool contiguous) { | ||
| 14 | 20 | CheckKeys(keys, contiguous); | |
| 15 | 20 | TORCH_CHECK(values.dim() == 2 && values.scalar_type() == torch::kFloat32 && | |
| 16 | (!contiguous || values.is_contiguous()) && | ||
| 17 | values.size(0) == keys.numel(), | ||
| 18 | "Values must be contiguous 2-D float32 with one row per key"); | ||
| 19 | 20 | } | |
| 20 | |||
| 21 | std::pair<torch::Tensor, torch::Tensor> | ||
| 22 | ✗ | CudaCacheRows(const torch::Tensor& keys, const torch::Tensor& values) { | |
| 23 | ✗ | TORCH_CHECK(keys.is_cuda() || values.is_cuda(), | |
| 24 | "Cache rows require a CUDA tensor"); | ||
| 25 | ✗ | auto device = values.is_cuda() ? values.device() : keys.device(); | |
| 26 | ✗ | return {(keys.is_cuda() ? keys : keys.to(device)).contiguous(), | |
| 27 | ✗ | (values.is_cuda() ? values : values.to(device)).contiguous()}; | |
| 28 | } | ||
| 29 | |||
| 30 | 30 | CacheUpdateActions CacheActions(const torch::Tensor& keys, int64_t dim) { | |
| 31 | #ifdef RECSTORE_ENABLE_GPU_CACHE | ||
| 32 | return {[] { return torch_cache::IsGpuCacheEnabled(); }, | ||
| 33 | [keys, dim] { | ||
| 34 | return keys.defined() && torch_cache::CanUseGpuCache(keys, dim); | ||
| 35 | }, | ||
| 36 | [keys] { torch_cache::InvalidateGpuCache(keys); }, | ||
| 37 | [] { | ||
| 38 | torch_cache::ClearGpuCache(); | ||
| 39 | gpu::ResetLastGpuCacheProfile(); | ||
| 40 | }}; | ||
| 41 | #else | ||
| 42 | 30 | return {}; | |
| 43 | #endif | ||
| 44 | } | ||
| 45 | |||
| 46 | namespace { | ||
| 47 | // Allocation and transport-buffer wrapping only. The common op owns the lease | ||
| 48 | // and validates its payload. The CUDA event retains it through the H2D copy. | ||
| 49 | ✗ | torch::Tensor ReadLocalValues(const std::shared_ptr<KVClientOp>& op, | |
| 50 | const torch::Tensor& cpu_keys, | ||
| 51 | int64_t dim, | ||
| 52 | const torch::Device& device) { | ||
| 53 | ✗ | auto keys = ToRecTensor(cpu_keys, base::DataType::UINT64); | |
| 54 | ✗ | if (device.is_cuda() && op->SupportsBorrowedLocalLookup()) { | |
| 55 | ✗ | auto lease = op->AcquireLocalLookupFlat(keys, dim); | |
| 56 | ✗ | auto start = SteadyNow(); | |
| 57 | ✗ | const auto bytes = lease->num_elements() * sizeof(float); | |
| 58 | ✗ | const bool pinned = EnsurePinnedLocalShmPayload(lease->data(), bytes); | |
| 59 | ✗ | g_last_local_lookup_flat_profile[kLookupPayloadPinMs] = ElapsedMs(start); | |
| 60 | auto view = torch::from_blob( | ||
| 61 | lease->data(), | ||
| 62 | ✗ | {cpu_keys.numel(), dim}, | |
| 63 | ✗ | [op, lease](void*) {}, | |
| 64 | ✗ | torch::TensorOptions().dtype(torch::kFloat32).pinned_memory(pinned)); | |
| 65 | ✗ | torch::Tensor staging = view; | |
| 66 | ✗ | if (!pinned) { | |
| 67 | ✗ | start = SteadyNow(); | |
| 68 | ✗ | staging = torch::empty(view.sizes(), PinnedCpuOptions(torch::kFloat32)); | |
| 69 | ✗ | staging.copy_(view); | |
| 70 | ✗ | g_last_local_lookup_flat_profile[kLookupFallbackCopyMs] = | |
| 71 | ✗ | ElapsedMs(start); | |
| 72 | } | ||
| 73 | ✗ | start = SteadyNow(); | |
| 74 | ✗ | auto result = staging.to(device, /*non_blocking=*/true); | |
| 75 | // from_blob storage is not owned by the PyTorch pinned allocator. Retain | ||
| 76 | // the transport lease explicitly until this stream finishes reading it. | ||
| 77 | ✗ | RetainUntilCudaComplete(result, [op, lease] {}); | |
| 78 | ✗ | g_last_local_lookup_flat_profile[kLookupValuesH2DEnqueueMs] = | |
| 79 | ✗ | ElapsedMs(start); | |
| 80 | ✗ | return result; | |
| 81 | ✗ | } | |
| 82 | auto values = torch::empty( | ||
| 83 | ✗ | {cpu_keys.numel(), dim}, | |
| 84 | ✗ | device.is_cuda() ? PinnedCpuOptions(torch::kFloat32) | |
| 85 | ✗ | : torch::TensorOptions().dtype(torch::kFloat32)); | |
| 86 | ✗ | auto out = ToRecTensor(values, base::DataType::FLOAT32); | |
| 87 | ✗ | op->LocalLookupFlat(keys, out); | |
| 88 | ✗ | return values.to(device, /*non_blocking=*/true); | |
| 89 | ✗ | } | |
| 90 | } // namespace | ||
| 91 | |||
| 92 | 32 | std::tuple<torch::Tensor, torch::Tensor> CachedReadTorch( | |
| 93 | const torch::Tensor& keys, int64_t dim, bool local, bool no_evict) { | ||
| 94 | 32 | CheckKeys(keys); | |
| 95 | 32 | TORCH_CHECK(dim > 0, "Embedding dimension must be positive"); | |
| 96 | 32 | ResetLocalLookupFlatProfile(); | |
| 97 | 32 | const auto start = SteadyNow(); | |
| 98 | 32 | std::shared_ptr<KVClientOp> op; | |
| 99 | 32 | if (local) { | |
| 100 | ✗ | op = GetConcreteKVClientOp(); | |
| 101 | ✗ | op->RequireLocalBackend(); | |
| 102 | } | ||
| 103 | auto resident = | ||
| 104 | ✗ | no_evict ? torch::ones({keys.numel()}, keys.options().dtype(torch::kBool)) | |
| 105 | 32 | : torch::Tensor(); | |
| 106 | 32 | if (!keys.numel()) | |
| 107 | return { | ||
| 108 | ✗ | torch::empty({0, dim}, | |
| 109 | ✗ | no_evict ? keys.options().dtype(torch::kFloat32) | |
| 110 | ✗ | : torch::TensorOptions().dtype(torch::kFloat32)), | |
| 111 | ✗ | resident}; | |
| 112 | 32 | auto read = [&](const torch::Tensor& cpu_keys, const torch::Device& device) { | |
| 113 | 32 | if (!op) | |
| 114 | 32 | op = GetConcreteKVClientOp(); | |
| 115 | 32 | if (local) | |
| 116 | ✗ | return ReadLocalValues(op, cpu_keys, dim, device); | |
| 117 | auto values = torch::empty( | ||
| 118 | 32 | {cpu_keys.numel(), dim}, | |
| 119 | 32 | device.is_cuda() ? PinnedCpuOptions(torch::kFloat32) | |
| 120 | 96 | : torch::TensorOptions().dtype(torch::kFloat32)); | |
| 121 | 32 | auto k = ToRecTensor(cpu_keys, base::DataType::UINT64); | |
| 122 | 32 | auto v = ToRecTensor(values, base::DataType::FLOAT32); | |
| 123 | 32 | op->EmbRead(k, v); | |
| 124 | 32 | return values.to(device, /*non_blocking=*/true); | |
| 125 | 32 | }; | |
| 126 | 32 | torch::Tensor output; | |
| 127 | 32 | CacheLookupActions io; | |
| 128 | 32 | io.read_all = [&] { | |
| 129 | 32 | auto stage_start = SteadyNow(); | |
| 130 | auto cpu_keys = | ||
| 131 | 32 | keys.is_cuda() ? StageCudaTensorToPinnedCpu(keys, torch::kInt64) : keys; | |
| 132 | 32 | if (keys.is_cuda()) | |
| 133 | ✗ | g_last_local_lookup_flat_profile[kLookupKeysStageMs] = | |
| 134 | ✗ | ElapsedMs(stage_start); | |
| 135 | 32 | output = read(cpu_keys, keys.device()); | |
| 136 | 64 | }; | |
| 137 | #ifdef RECSTORE_ENABLE_GPU_CACHE | ||
| 138 | gpu::ResetLastGpuCacheProfile(); | ||
| 139 | torch_cache::GpuCacheLookupResult cached; | ||
| 140 | torch::Tensor miss_keys, miss_values, inserted; | ||
| 141 | io.usable = torch_cache::CanUseGpuCache(keys, dim); | ||
| 142 | io.query = [&] { | ||
| 143 | cached = torch_cache::QueryGpuCache(keys, dim); | ||
| 144 | output = cached.values; | ||
| 145 | return cached.missing_count; | ||
| 146 | }; | ||
| 147 | io.read_missing = [&] { | ||
| 148 | auto begin = SteadyNow(); | ||
| 149 | auto values = read(cached.missing_keys_cpu.contiguous(), torch::kCPU); | ||
| 150 | gpu::AddGpuCacheBackendLookupMs(ElapsedMs(begin)); | ||
| 151 | miss_keys = cached.missing_keys_cpu.to(keys.device()); | ||
| 152 | miss_values = values.to(keys.device()); | ||
| 153 | }; | ||
| 154 | io.fill = [&] { | ||
| 155 | if (no_evict) | ||
| 156 | inserted = torch_cache::FillGpuCacheNoEvict(miss_keys, miss_values); | ||
| 157 | else | ||
| 158 | torch_cache::FillGpuCache(miss_keys, miss_values); | ||
| 159 | }; | ||
| 160 | io.scatter = [&] { | ||
| 161 | torch_cache::ScatterMissValues( | ||
| 162 | &output, cached.missing_positions_cpu, miss_values); | ||
| 163 | if (no_evict) | ||
| 164 | resident.index_put_( | ||
| 165 | {cached.missing_positions_cpu.to(keys.device())}, inserted); | ||
| 166 | }; | ||
| 167 | io.clear = [] { | ||
| 168 | torch_cache::ClearGpuCache(); | ||
| 169 | gpu::ResetLastGpuCacheProfile(); | ||
| 170 | }; | ||
| 171 | #endif | ||
| 172 | 32 | g_cache_aware_op.Lookup(keys.numel(), io, no_evict); | |
| 173 | 32 | g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(start); | |
| 174 | 32 | return {output, resident}; | |
| 175 | 32 | } | |
| 176 | |||
| 177 | } // namespace recstore::framework | ||
| 178 |