GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 29.5% 136 / 0 / 461
Functions: 19.1% 9 / 0 / 47
Branches: 18.2% 144 / 0 / 792
framework/pytorch/op_torch.cc
Line Branch Exec Source
1 #include <torch/extension.h>
2
3 #include <cstdlib>
4 #include <cstring>
5 #include <iostream>
6 #include <chrono>
7 #include <mutex>
8 #include <string>
9 #include <unordered_map>
10 #include <unistd.h>
11 #include "base/tensor.h"
12 #include "framework/op.h"
13 #include "ps/local_shm/local_shm_client.h"
14 // Log level: 0=ERROR, 1=WARNING, 2=INFO, 3=DEBUG
15 #include <glog/logging.h>
16
17 #ifdef RECSTORE_ENABLE_GPU_CACHE
18 # include "framework/gpu/gpu_embedding_cache.h"
19 #endif
20
21 #if __has_include(<cuda_runtime_api.h>)
22 # include <ATen/cuda/CUDAContext.h>
23 # include <c10/cuda/CUDAException.h>
24 # include <c10/cuda/CUDAGuard.h>
25 # include <cuda_runtime_api.h>
26 # define RECSTORE_HAS_CUDA_RUNTIME_API 1
27 #else
28 # define RECSTORE_HAS_CUDA_RUNTIME_API 0
29 #endif
30
31 namespace recstore {
32 namespace framework {
33
34 namespace {
35
36 ✗ bool IsLocalFastPathBackend(const std::string& backend) {
37 ✗ return backend == "local_shm" || backend == "hierkv";
38 }
39
40 enum LookupProfileIndex : std::size_t {
41 kLookupTotalMs = 0,
42 kLookupKeysStageMs,
43 kLookupSubmitMs,
44 kLookupWaitMs,
45 kLookupPayloadPinMs,
46 kLookupFallbackCopyMs,
47 kLookupValuesH2DEnqueueMs,
48 kLookupProfileSize,
49 };
50
51 enum UpdateProfileIndex : std::size_t {
52 kUpdateTotalMs = 0,
53 kUpdateKeysStageMs,
54 kUpdateGradsStageMs,
55 kUpdateShmCallMs,
56 kUpdateStageWaitMs,
57 kUpdateProfileSize,
58 };
59
60 thread_local std::vector<double>
61 g_last_local_lookup_flat_profile(kLookupProfileSize, 0.0);
62 thread_local std::vector<double>
63 g_last_local_update_flat_profile(kUpdateProfileSize, 0.0);
64
65 ✗ inline std::chrono::steady_clock::time_point SteadyNow() {
66 ✗ return std::chrono::steady_clock::now();
67 }
68
69 ✗ inline double ElapsedMs(std::chrono::steady_clock::time_point start) {
70 ✗ return std::chrono::duration_cast<std::chrono::duration<double, std::milli>>(
71 ✗ SteadyNow() - start)
72 ✗ .count();
73 }
74
75 ✗ inline void ResetLocalLookupFlatProfile() {
76 ✗ std::fill(g_last_local_lookup_flat_profile.begin(),
77 ✗ g_last_local_lookup_flat_profile.end(),
78 ✗ 0.0);
79 ✗ }
80
81 ✗ inline void ResetLocalUpdateFlatProfile() {
82 ✗ std::fill(g_last_local_update_flat_profile.begin(),
83 ✗ g_last_local_update_flat_profile.end(),
84 ✗ 0.0);
85 ✗ }
86
87 #ifdef RECSTORE_ENABLE_GPU_CACHE
88 constexpr int64_t kGpuCacheBypassMinRows = 1024;
89 constexpr int kGpuCacheLowHitLimit = 1;
90 constexpr double kGpuCacheLowHitRatio = 0.05;
91 thread_local int g_gpu_cache_low_hit_streak = 0;
92 thread_local bool g_gpu_cache_lookup_bypassed = false;
93 thread_local bool g_gpu_cache_lookup_bypass_enabled = true;
94
95 void SafeClearGpuCacheNoThrow();
96
97 void ResetGpuCacheBypassState() {
98 g_gpu_cache_low_hit_streak = 0;
99 g_gpu_cache_lookup_bypassed = false;
100 }
101
102 bool ShouldBypassGpuCacheLookup(int64_t num_keys) {
103 return g_gpu_cache_lookup_bypass_enabled &&
104 num_keys >= kGpuCacheBypassMinRows &&
105 g_gpu_cache_low_hit_streak >= kGpuCacheLowHitLimit;
106 }
107
108 void RecordGpuCacheLookupOutcome(
109 int64_t num_keys, double hit_count, double request_count) {
110 if (num_keys < kGpuCacheBypassMinRows || request_count <= 0.0) {
111 return;
112 }
113 const double hit_ratio = hit_count / request_count;
114 if (hit_ratio < kGpuCacheLowHitRatio) {
115 ++g_gpu_cache_low_hit_streak;
116 } else {
117 g_gpu_cache_low_hit_streak = 0;
118 g_gpu_cache_lookup_bypassed = false;
119 }
120 }
121
122 bool ShouldBypassGpuCacheMaintenance(int64_t num_keys) {
123 return g_gpu_cache_lookup_bypass_enabled &&
124 num_keys >= kGpuCacheBypassMinRows && g_gpu_cache_lookup_bypassed;
125 }
126
127 void MarkGpuCacheLookupBypassed() {
128 if (!g_gpu_cache_lookup_bypassed) {
129 SafeClearGpuCacheNoThrow();
130 g_gpu_cache_low_hit_streak = kGpuCacheLowHitLimit;
131 }
132 g_gpu_cache_lookup_bypassed = true;
133 }
134
135 void EnsureGpuCacheSafeForLookup() {
136 if (g_gpu_cache_lookup_bypassed) {
137 SafeClearGpuCacheNoThrow();
138 ResetGpuCacheBypassState();
139 }
140 }
141
142 void SafeClearGpuCacheNoThrow() {
143 try {
144 gpu::ClearGpuCache();
145 } catch (const std::exception& e) {
146 LOG(WARNING) << "Failed to clear GPU cache: " << e.what();
147 } catch (...) {
148 LOG(WARNING) << "Failed to clear GPU cache: unknown exception";
149 }
150 }
151
152 void SetGpuCacheLookupBypassEnabled(bool enabled) {
153 g_gpu_cache_lookup_bypass_enabled = enabled;
154 if (!enabled) {
155 ResetGpuCacheBypassState();
156 }
157 }
158
159 void MaintainGpuCacheAfterUpdateNoThrow(const torch::Tensor& keys,
160 const torch::Tensor& grads,
161 int64_t embedding_dim) {
162 (void)grads;
163 if (!gpu::IsGpuCacheEnabled()) {
164 return;
165 }
166 if (ShouldBypassGpuCacheMaintenance(keys.numel())) {
167 gpu::ResetLastGpuCacheProfile();
168 return;
169 }
170 if (gpu::CanUseGpuCache(keys, embedding_dim)) {
171 try {
172 gpu::InvalidateGpuCache(keys);
173 return;
174 } catch (const std::exception& e) {
175 LOG(WARNING) << "GPU cache invalidation failed after backend update "
176 "succeeded; clearing cache and continuing: "
177 << e.what();
178 } catch (...) {
179 LOG(WARNING) << "GPU cache invalidation failed after backend update "
180 "succeeded; clearing cache and continuing: "
181 << "unknown exception";
182 }
183 }
184 SafeClearGpuCacheNoThrow();
185 gpu::ResetLastGpuCacheProfile();
186 }
187 #endif
188
189 } // namespace
190
191 static inline base::RecTensor
192 108 ToRecTensor(const torch::Tensor& tensor, base::DataType dtype) {
193 108 std::vector<int64_t> shape;
194
3/4
✓ Branch 1 taken 270 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 162 times.
✓ Branch 4 taken 108 times.
 270 for (int i = 0; i < tensor.dim(); ++i) {
195
2/4
✓ Branch 1 taken 162 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 162 times.
✗ Branch 5 not taken.
 162 shape.push_back(tensor.size(i));
196 }
197
2/4
✓ Branch 1 taken 108 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 108 times.
✗ Branch 5 not taken.
 216 return base::RecTensor(const_cast<void*>(tensor.data_ptr()), shape, dtype);
198 108 }
199
200 ✗ static torch::TensorOptions PinnedCpuOptions(torch::ScalarType dtype) {
201 ✗ return torch::TensorOptions()
202 ✗ .device(torch::kCPU)
203 ✗ .dtype(dtype)
204 ✗ .pinned_memory(true);
205 }
206
207 ✗ static torch::Tensor StageCudaTensorToPinnedCpu(const torch::Tensor& tensor,
208 torch::ScalarType dtype) {
209 ✗ auto cpu_tensor = torch::empty(tensor.sizes(), PinnedCpuOptions(dtype));
210 ✗ cpu_tensor.copy_(tensor.to(dtype), /*non_blocking=*/false);
211 ✗ return cpu_tensor;
212 ✗ }
213
214 static torch::Tensor
215 ✗ StageCudaTensorToPinnedCpuAsyncNoCast(const torch::Tensor& tensor) {
216 auto cpu_tensor =
217 ✗ torch::empty(tensor.sizes(), PinnedCpuOptions(tensor.scalar_type()));
218 ✗ cpu_tensor.copy_(tensor, /*non_blocking=*/true);
219 ✗ return cpu_tensor;
220 ✗ }
221
222 ✗ static void SynchronizeCurrentCudaStreamForTensor(const torch::Tensor& tensor) {
223 #if RECSTORE_HAS_CUDA_RUNTIME_API
224 if (!tensor.is_cuda()) {
225 return;
226 }
227 c10::cuda::CUDAGuard device_guard(tensor.device());
228 C10_CUDA_CHECK(
229 cudaStreamSynchronize(at::cuda::getCurrentCUDAStream().stream()));
230 #else
231 (void)tensor;
232 #endif
233 ✗ }
234
235 ✗ static bool EnsurePinnedLocalShmPayload(const void* ptr, std::size_t bytes) {
236 #if !RECSTORE_HAS_CUDA_RUNTIME_API
237 (void)ptr;
238 (void)bytes;
239 ✗ return false;
240 #else
241 if (ptr == nullptr || bytes == 0) {
242 return false;
243 }
244 const long page_size = ::sysconf(_SC_PAGESIZE);
245 if (page_size <= 0) {
246 return false;
247 }
248 const std::size_t page_bytes = static_cast<std::size_t>(page_size);
249 const uintptr_t raw_begin = reinterpret_cast<uintptr_t>(ptr);
250 const uintptr_t raw_end = raw_begin + bytes;
251 const uintptr_t page_begin =
252 raw_begin & ~(static_cast<uintptr_t>(page_bytes) - 1U);
253 const uintptr_t page_end =
254 (raw_end + page_bytes - 1U) & ~(static_cast<uintptr_t>(page_bytes) - 1U);
255 const std::size_t required_bytes =
256 static_cast<std::size_t>(page_end - page_begin);
257
258 static std::mutex mu;
259 static std::unordered_map<uintptr_t, std::size_t> registered_bytes_by_base;
260 std::lock_guard<std::mutex> guard(mu);
261 const std::size_t existing_bytes = registered_bytes_by_base[page_begin];
262 if (existing_bytes >= required_bytes) {
263 return true;
264 }
265
266 void* register_ptr = reinterpret_cast<void*>(page_begin + existing_bytes);
267 const std::size_t register_bytes = required_bytes - existing_bytes;
268 const cudaError_t err =
269 cudaHostRegister(register_ptr, register_bytes, cudaHostRegisterPortable);
270 if (err != cudaSuccess && err != cudaErrorHostMemoryAlreadyRegistered) {
271 LOG(WARNING) << "cudaHostRegister failed for local_shm payload: "
272 << cudaGetErrorString(err)
273 << " base=" << reinterpret_cast<void*>(page_begin)
274 << " bytes=" << required_bytes;
275 return false;
276 }
277 registered_bytes_by_base[page_begin] = required_bytes;
278 return true;
279 #endif
280 }
281
282 32 torch::Tensor emb_read_torch(const torch::Tensor& keys, int64_t embedding_dim) {
283
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 bool is_cuda = keys.is_cuda();
284
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 auto orig_device = keys.device();
285
286
2/4
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 32 times.
 32 TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
287
2/4
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 32 times.
 32 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
288 "Keys tensor must have dtype int64");
289
2/4
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 32 times.
 32 TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
290
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
 32 TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive");
291
292
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 const int64_t num_keys = keys.size(0);
293
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
 32 if (num_keys == 0) {
294 return torch::empty(
295 ✗ {0, embedding_dim}, torch::TensorOptions().dtype(torch::kFloat32));
296 }
297
298
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 auto op = GetKVClientOp();
299
300 #ifdef RECSTORE_ENABLE_GPU_CACHE
301 gpu::ResetLastGpuCacheProfile();
302 const bool can_use_gpu_cache = gpu::CanUseGpuCache(keys, embedding_dim);
303 const bool bypass_gpu_cache_lookup =
304 can_use_gpu_cache && ShouldBypassGpuCacheLookup(num_keys);
305 if (bypass_gpu_cache_lookup) {
306 MarkGpuCacheLookupBypassed();
307 }
308 if (can_use_gpu_cache && !bypass_gpu_cache_lookup) {
309 EnsureGpuCacheSafeForLookup();
310 try {
311 auto cache_result = gpu::QueryGpuCache(keys, embedding_dim);
312 RecordGpuCacheLookupOutcome(
313 num_keys,
314 static_cast<double>(num_keys - cache_result.missing_count),
315 static_cast<double>(num_keys));
316 if (cache_result.missing_count == 0) {
317 return cache_result.values;
318 }
319
320 auto missing_cpu_values = torch::empty(
321 {cache_result.missing_count, embedding_dim},
322 torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32));
323 base::RecTensor rec_missing_keys = ToRecTensor(
324 cache_result.missing_keys_cpu.contiguous(), base::DataType::UINT64);
325 base::RecTensor rec_missing_values =
326 ToRecTensor(missing_cpu_values, base::DataType::FLOAT32);
327 const auto backend_start = SteadyNow();
328 op->EmbRead(rec_missing_keys, rec_missing_values);
329 gpu::AddGpuCacheBackendLookupMs(ElapsedMs(backend_start));
330
331 auto miss_keys_cuda =
332 cache_result.missing_keys_cpu.to(orig_device, /*non_blocking=*/false);
333 auto miss_values_cuda =
334 missing_cpu_values.to(orig_device, /*non_blocking=*/false);
335 gpu::FillGpuCache(miss_keys_cuda, miss_values_cuda);
336 gpu::ScatterMissValues(&cache_result.values,
337 cache_result.missing_positions_cpu,
338 miss_values_cuda);
339 return cache_result.values;
340 } catch (const std::exception& e) {
341 LOG(WARNING)
342 << "GPU cache emb_read failed; clearing cache and falling back: "
343 << e.what();
344 SafeClearGpuCacheNoThrow();
345 gpu::ResetLastGpuCacheProfile();
346 } catch (...) {
347 LOG(WARNING)
348 << "GPU cache emb_read failed; clearing cache and falling back: "
349 << "unknown exception";
350 SafeClearGpuCacheNoThrow();
351 gpu::ResetLastGpuCacheProfile();
352 }
353 }
354 #endif
355
356
2/6
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 6 taken 32 times.
✗ Branch 7 not taken.
 32 torch::Tensor cpu_keys = is_cuda ? keys.cpu() : keys;
357
358 auto cpu_values = torch::empty(
359
2/4
✓ Branch 2 taken 32 times.
✗ Branch 3 not taken.
✓ Branch 8 taken 32 times.
✗ Branch 9 not taken.
 32 {num_keys, embedding_dim}, torch::TensorOptions().dtype(torch::kFloat32));
360
361
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
362
1/2
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 32 base::RecTensor rec_values = ToRecTensor(cpu_values, base::DataType::FLOAT32);
363
364
1/2
✓ Branch 2 taken 32 times.
✗ Branch 3 not taken.
 32 op->EmbRead(rec_keys, rec_values);
365
366
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
 32 if (is_cuda) {
367 ✗ return cpu_values.to(orig_device);
368 }
369 32 return cpu_values;
370 32 }
371
372 ✗ static std::shared_ptr<KVClientOp> GetConcreteKVClientOp() {
373 ✗ auto op = GetKVClientOp();
374 ✗ auto kv_op = std::dynamic_pointer_cast<KVClientOp>(op);
375 ✗ TORCH_CHECK(kv_op != nullptr, "storage backend is not KVClientOp");
376 ✗ return kv_op;
377 ✗ }
378
379 ✗ static torch::Tensor BackendLocalLookupFlat(
380 const std::shared_ptr<KVClientOp>& kv_op,
381 const torch::Tensor& cpu_keys,
382 const torch::Device& result_device,
383 bool result_on_cuda,
384 int64_t embedding_dim,
385 const std::chrono::steady_clock::time_point& total_start,
386 bool record_profile = true) {
387 ✗ const int64_t num_keys = cpu_keys.size(0);
388 ✗ base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
389 ✗ if (kv_op->CurrentPSBackend() != "local_shm") {
390 auto cpu_values =
391 result_on_cuda
392 ? torch::empty({num_keys, embedding_dim},
393 PinnedCpuOptions(torch::kFloat32))
394 : torch::empty({num_keys, embedding_dim},
395 ✗ torch::TensorOptions()
396 ✗ .device(torch::kCPU)
397 ✗ .dtype(torch::kFloat32));
398 base::RecTensor rec_values =
399 ✗ ToRecTensor(cpu_values, base::DataType::FLOAT32);
400 ✗ kv_op->LocalLookupFlat(rec_keys, rec_values);
401 ✗ if (record_profile) {
402 ✗ g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
403 }
404 ✗ if (result_on_cuda) {
405 ✗ return cpu_values.to(result_device, /*non_blocking=*/true);
406 }
407 ✗ return cpu_values;
408 ✗ }
409
410 ✗ if (!result_on_cuda) {
411 auto cpu_values = torch::empty(
412 {num_keys, embedding_dim},
413 ✗ torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32));
414 base::RecTensor rec_values =
415 ✗ ToRecTensor(cpu_values, base::DataType::FLOAT32);
416 ✗ kv_op->LocalLookupFlat(rec_keys, rec_values);
417 ✗ if (record_profile) {
418 ✗ g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
419 }
420 ✗ return cpu_values;
421 ✗ }
422
423 ✗ LocalShmFlatGetHandle handle;
424 ✗ const auto submit_start = SteadyNow();
425 ✗ TORCH_CHECK(
426 kv_op->SubmitLocalLookupFlat(rec_keys, embedding_dim, &handle) == 0,
427 "Failed to submit local_shm flat lookup.");
428 ✗ if (record_profile) {
429 ✗ g_last_local_lookup_flat_profile[kLookupSubmitMs] = ElapsedMs(submit_start);
430 }
431 ✗ const auto wait_start = SteadyNow();
432 ✗ const int wait_ret = kv_op->WaitLocalLookupFlat(&handle);
433 ✗ if (record_profile) {
434 ✗ g_last_local_lookup_flat_profile[kLookupWaitMs] = ElapsedMs(wait_start);
435 }
436 ✗ if (wait_ret != 0) {
437 ✗ kv_op->ReleaseLocalLookupFlat(&handle);
438 ✗ TORCH_CHECK(false, "Failed to wait for local_shm flat lookup.");
439 }
440 ✗ const float* payload_values = handle.values;
441 ✗ const int64_t payload_rows = handle.num_rows;
442 ✗ const int64_t payload_dim = handle.embedding_dim;
443 ✗ const std::size_t payload_bytes =
444 static_cast<std::size_t>(handle.output_bytes);
445 ✗ const int64_t expected_bytes =
446 ✗ num_keys * embedding_dim * static_cast<int64_t>(sizeof(float));
447 ✗ if (payload_values == nullptr || payload_rows != num_keys ||
448 ✗ payload_dim != embedding_dim ||
449 ✗ static_cast<int64_t>(payload_bytes) != expected_bytes) {
450 ✗ kv_op->ReleaseLocalLookupFlat(&handle);
451 ✗ TORCH_CHECK(false,
452 "local_shm flat lookup returned unexpected payload metadata.");
453 }
454 ✗ const auto pin_start = SteadyNow();
455 const bool payload_is_pinned =
456 ✗ EnsurePinnedLocalShmPayload(payload_values, payload_bytes);
457 ✗ if (record_profile) {
458 ✗ g_last_local_lookup_flat_profile[kLookupPayloadPinMs] =
459 ✗ ElapsedMs(pin_start);
460 }
461 ✗ if (payload_is_pinned) {
462 try {
463 ✗ LocalShmFlatGetHandle handle_for_release = handle;
464 auto cpu_view = torch::from_blob(
465 const_cast<float*>(payload_values),
466 {num_keys, embedding_dim},
467 ✗ [kv_op, handle_for_release](void* /*unused*/) mutable {
468 ✗ kv_op->ReleaseLocalLookupFlat(&handle_for_release);
469 ✗ },
470 ✗ PinnedCpuOptions(torch::kFloat32));
471 ✗ const auto h2d_start = SteadyNow();
472 ✗ auto result = cpu_view.to(result_device, /*non_blocking=*/true);
473 ✗ if (record_profile) {
474 ✗ g_last_local_lookup_flat_profile[kLookupValuesH2DEnqueueMs] =
475 ✗ ElapsedMs(h2d_start);
476 ✗ g_last_local_lookup_flat_profile[kLookupTotalMs] =
477 ✗ ElapsedMs(total_start);
478 }
479 ✗ return result;
480 ✗ } catch (...) {
481 ✗ kv_op->ReleaseLocalLookupFlat(&handle);
482 ✗ throw;
483 ✗ }
484 }
485
486 auto cpu_values = torch::empty(
487 ✗ {num_keys, embedding_dim}, PinnedCpuOptions(torch::kFloat32));
488 ✗ const auto fallback_copy_start = SteadyNow();
489 ✗ std::memcpy(cpu_values.data_ptr<float>(), payload_values, payload_bytes);
490 ✗ if (record_profile) {
491 ✗ g_last_local_lookup_flat_profile[kLookupFallbackCopyMs] =
492 ✗ ElapsedMs(fallback_copy_start);
493 }
494 ✗ kv_op->ReleaseLocalLookupFlat(&handle);
495 ✗ const auto h2d_start = SteadyNow();
496 ✗ auto result = cpu_values.to(result_device, /*non_blocking=*/true);
497 ✗ if (record_profile) {
498 ✗ g_last_local_lookup_flat_profile[kLookupValuesH2DEnqueueMs] =
499 ✗ ElapsedMs(h2d_start);
500 ✗ g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
501 }
502 ✗ return result;
503 ✗ }
504
505 torch::Tensor
506 ✗ local_lookup_flat_torch(const torch::Tensor& keys, int64_t embedding_dim) {
507 ✗ ResetLocalLookupFlatProfile();
508 #ifdef RECSTORE_ENABLE_GPU_CACHE
509 gpu::ResetLastGpuCacheProfile();
510 #endif
511 ✗ const auto total_start = SteadyNow();
512 ✗ const bool is_cuda = keys.is_cuda();
513 ✗ auto orig_device = keys.device();
514
515 ✗ TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
516 ✗ TORCH_CHECK(keys.scalar_type() == torch::kInt64,
517 "Keys tensor must have dtype int64");
518 ✗ TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
519 ✗ TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive");
520
521 ✗ auto kv_op = GetConcreteKVClientOp();
522 ✗ TORCH_CHECK(IsLocalFastPathBackend(kv_op->CurrentPSBackend()),
523 "local_lookup_flat requires local_shm or hierkv backend, but "
524 "current backend is ",
525 kv_op->CurrentPSBackend());
526
527 ✗ const int64_t num_keys = keys.size(0);
528 ✗ if (num_keys == 0) {
529 return torch::empty(
530 ✗ {0, embedding_dim}, torch::TensorOptions().dtype(torch::kFloat32));
531 }
532
533 #ifdef RECSTORE_ENABLE_GPU_CACHE
534 const bool can_use_gpu_cache = gpu::CanUseGpuCache(keys, embedding_dim);
535 const bool bypass_gpu_cache_lookup =
536 can_use_gpu_cache && ShouldBypassGpuCacheLookup(num_keys);
537 if (bypass_gpu_cache_lookup) {
538 MarkGpuCacheLookupBypassed();
539 }
540 if (can_use_gpu_cache && !bypass_gpu_cache_lookup) {
541 EnsureGpuCacheSafeForLookup();
542 try {
543 auto cache_result = gpu::QueryGpuCache(keys, embedding_dim);
544 RecordGpuCacheLookupOutcome(
545 num_keys,
546 static_cast<double>(num_keys - cache_result.missing_count),
547 static_cast<double>(num_keys));
548 if (cache_result.missing_count == 0) {
549 g_last_local_lookup_flat_profile[kLookupTotalMs] =
550 ElapsedMs(total_start);
551 return cache_result.values;
552 }
553
554 const auto backend_start = SteadyNow();
555 auto miss_values = BackendLocalLookupFlat(
556 kv_op,
557 cache_result.missing_keys_cpu.contiguous(),
558 orig_device,
559 /*result_on_cuda=*/false,
560 embedding_dim,
561 total_start);
562 const double backend_ms = ElapsedMs(backend_start);
563 gpu::AddGpuCacheBackendLookupMs(backend_ms);
564 auto miss_keys_cuda =
565 cache_result.missing_keys_cpu.to(orig_device, /*non_blocking=*/false);
566 auto miss_values_cuda =
567 miss_values.to(orig_device, /*non_blocking=*/false);
568 gpu::FillGpuCache(miss_keys_cuda, miss_values_cuda);
569 gpu::ScatterMissValues(&cache_result.values,
570 cache_result.missing_positions_cpu,
571 miss_values_cuda);
572 g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
573 return cache_result.values;
574 } catch (const std::exception& e) {
575 LOG(WARNING)
576 << "GPU cache lookup failed; clearing cache and falling back: "
577 << e.what();
578 SafeClearGpuCacheNoThrow();
579 gpu::ResetLastGpuCacheProfile();
580 } catch (...) {
581 LOG(WARNING)
582 << "GPU cache lookup failed; clearing cache and falling back: "
583 << "unknown exception";
584 SafeClearGpuCacheNoThrow();
585 gpu::ResetLastGpuCacheProfile();
586 }
587 }
588 #endif
589
590 ✗ torch::Tensor cpu_keys = keys;
591 ✗ if (is_cuda) {
592 ✗ const auto stage_start = SteadyNow();
593 ✗ cpu_keys = StageCudaTensorToPinnedCpu(keys, torch::kInt64);
594 ✗ g_last_local_lookup_flat_profile[kLookupKeysStageMs] =
595 ✗ ElapsedMs(stage_start);
596 }
597
598 return BackendLocalLookupFlat(
599 ✗ kv_op, cpu_keys, orig_device, is_cuda, embedding_dim, total_start);
600 ✗ }
601
602
603 // GPU-cache-accelerated flat lookup that works with ANY backend (BRPC, GRPC,
604 // RDMA, local_shm). Cache hits are served from the GPU cache; misses are
605 // fetched via EmbRead and filled back into the cache. This is the forward
606 // path used by the BagPipe controller when the local_shm fast path is
607 // unavailable, so the GPU cache is actually queried instead of bypassed.
608 torch::Tensor
609 ✗ gpu_cache_lookup_flat_torch(const torch::Tensor& keys,
610 int64_t embedding_dim) {
611 ✗ ResetLocalLookupFlatProfile();
612 #ifdef RECSTORE_ENABLE_GPU_CACHE
613 gpu::ResetLastGpuCacheProfile();
614 #endif
615 ✗ const auto total_start = SteadyNow();
616 ✗ const bool is_cuda = keys.is_cuda();
617 ✗ auto orig_device = keys.device();
618
619 ✗ TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
620 ✗ TORCH_CHECK(keys.scalar_type() == torch::kInt64,
621 "Keys tensor must have dtype int64");
622 ✗ TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
623 ✗ TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive");
624
625 ✗ const int64_t num_keys = keys.size(0);
626 ✗ if (num_keys == 0) {
627 return torch::empty(
628 ✗ {0, embedding_dim}, torch::TensorOptions().dtype(torch::kFloat32));
629 }
630
631 #ifdef RECSTORE_ENABLE_GPU_CACHE
632 const bool can_use_gpu_cache = gpu::CanUseGpuCache(keys, embedding_dim);
633 const bool bypass_gpu_cache_lookup =
634 can_use_gpu_cache && ShouldBypassGpuCacheLookup(num_keys);
635 if (bypass_gpu_cache_lookup) {
636 MarkGpuCacheLookupBypassed();
637 }
638 if (can_use_gpu_cache && !bypass_gpu_cache_lookup) {
639 EnsureGpuCacheSafeForLookup();
640 try {
641 auto cache_result = gpu::QueryGpuCache(keys, embedding_dim);
642 RecordGpuCacheLookupOutcome(
643 num_keys,
644 static_cast<double>(num_keys - cache_result.missing_count),
645 static_cast<double>(num_keys));
646 if (cache_result.missing_count == 0) {
647 g_last_local_lookup_flat_profile[kLookupTotalMs] =
648 ElapsedMs(total_start);
649 return cache_result.values;
650 }
651
652 // Fetch misses via EmbRead (works with BRPC / GRPC / RDMA).
653 const auto backend_start = SteadyNow();
654 auto miss_cpu_keys = cache_result.missing_keys_cpu.contiguous();
655 const int64_t miss_count = miss_cpu_keys.size(0);
656 auto miss_cpu_values = torch::empty(
657 {miss_count, embedding_dim},
658 torch::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32));
659 auto op = GetKVClientOp();
660 base::RecTensor rec_miss_keys =
661 ToRecTensor(miss_cpu_keys, base::DataType::UINT64);
662 base::RecTensor rec_miss_values =
663 ToRecTensor(miss_cpu_values, base::DataType::FLOAT32);
664 op->EmbRead(rec_miss_keys, rec_miss_values);
665 gpu::AddGpuCacheBackendLookupMs(ElapsedMs(backend_start));
666
667 auto miss_keys_cuda =
668 miss_cpu_keys.to(orig_device, /*non_blocking=*/false);
669 auto miss_values_cuda =
670 miss_cpu_values.to(orig_device, /*non_blocking=*/false);
671 gpu::FillGpuCache(miss_keys_cuda, miss_values_cuda);
672 gpu::ScatterMissValues(&cache_result.values,
673 cache_result.missing_positions_cpu,
674 miss_values_cuda);
675 g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
676 return cache_result.values;
677 } catch (const std::exception& e) {
678 LOG(WARNING)
679 << "gpu_cache_lookup_flat: cache lookup failed; falling back: "
680 << e.what();
681 SafeClearGpuCacheNoThrow();
682 gpu::ResetLastGpuCacheProfile();
683 } catch (...) {
684 LOG(WARNING)
685 << "gpu_cache_lookup_flat: cache lookup failed; falling back";
686 SafeClearGpuCacheNoThrow();
687 gpu::ResetLastGpuCacheProfile();
688 }
689 }
690 #endif
691
692 // Fallback: direct EmbRead (no GPU cache).
693 ✗ torch::Tensor cpu_keys = keys;
694 ✗ if (is_cuda) {
695 ✗ const auto stage_start = SteadyNow();
696 ✗ cpu_keys = StageCudaTensorToPinnedCpu(keys, torch::kInt64);
697 ✗ g_last_local_lookup_flat_profile[kLookupKeysStageMs] =
698 ✗ ElapsedMs(stage_start);
699 }
700 ✗ auto op = GetKVClientOp();
701 auto cpu_values = torch::empty(
702 ✗ {cpu_keys.size(0), embedding_dim},
703 ✗ is_cuda ? PinnedCpuOptions(torch::kFloat32)
704 ✗ : torch::TensorOptions()
705 ✗ .device(torch::kCPU)
706 ✗ .dtype(torch::kFloat32));
707 ✗ base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
708 ✗ base::RecTensor rec_values = ToRecTensor(cpu_values, base::DataType::FLOAT32);
709 ✗ op->EmbRead(rec_keys, rec_values);
710 ✗ g_last_local_lookup_flat_profile[kLookupTotalMs] = ElapsedMs(total_start);
711 ✗ if (is_cuda) {
712 ✗ return cpu_values.to(orig_device, /*non_blocking=*/true);
713 }
714 ✗ return cpu_values;
715 ✗ }
716
717 // Async prefetch: returns a unique prefetch id (uint64_t)
718 2 int64_t emb_prefetch_torch(const torch::Tensor& keys) {
719
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
720
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
721 "Keys tensor must have dtype int64");
722
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
723
724
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 auto op = GetKVClientOp();
725
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 torch::Tensor cpu_keys = keys;
726
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 if (keys.is_cuda()) {
727 ✗ cpu_keys = keys.cpu();
728 }
729
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
730 // Dummy values tensor (unused by backend prefetch implementation)
731
2/4
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
 2 auto dummy_vals = torch::empty({0, 0}, keys.options().dtype(torch::kFloat32));
732
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 base::RecTensor rec_vals = ToRecTensor(dummy_vals, base::DataType::FLOAT32);
733
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 uint64_t pid = op->EmbPrefetch(rec_keys, rec_vals);
734 2 return static_cast<int64_t>(pid);
735 2 }
736
737 // Wait for prefetch and return result tensor [N, embedding_dim] on CPU
738 torch::Tensor
739 2 emb_wait_result_torch(int64_t prefetch_id, int64_t embedding_dim) {
740
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive");
741
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 auto op = GetKVClientOp();
742
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 op->WaitForPrefetch(static_cast<uint64_t>(prefetch_id));
743 2 std::vector<float> flat_values;
744 2 int64_t L = 0;
745
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 op->GetPretchResultFlat(
746 static_cast<uint64_t>(prefetch_id), &flat_values, &L, embedding_dim);
747 auto options =
748
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
749
1/2
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 2 auto out = torch::empty({L, embedding_dim}, options);
750
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
 2 if (L > 0 && !flat_values.empty()) {
751
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 std::memcpy(out.data_ptr<float>(),
752 2 flat_values.data(),
753 2 static_cast<size_t>(L) * static_cast<size_t>(embedding_dim) *
754 sizeof(float));
755 }
756 4 return out;
757 2 }
758
759 ✗ void emb_update_torch(const torch::Tensor& keys, const torch::Tensor& grads) {
760 ✗ throw std::runtime_error(
761 "emb_update_torch is deprecated. Use the Python-based sparse "
762 ✗ "optimizer.");
763 }
764
765 2 void emb_update_table_torch(const std::string& table_name,
766 const torch::Tensor& keys,
767 const torch::Tensor& grads) {
768
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
 2 TORCH_CHECK(!table_name.empty(), "table_name must be non-empty");
769
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
770
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
771 "Keys tensor must have dtype int64");
772
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
773
774
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(grads.dim() == 2, "Grads tensor must be 2-dimensional");
775
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(grads.scalar_type() == torch::kFloat32,
776 "Grads tensor must have dtype float32");
777
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 TORCH_CHECK(grads.is_contiguous(), "Grads tensor must be contiguous");
778
3/6
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
 2 TORCH_CHECK(keys.size(0) == grads.size(0),
779 "Keys and grads tensors must have the same number of entries");
780
781
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 if (keys.size(0) == 0) {
782 ✗ return;
783 }
784
785
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 auto op = GetKVClientOp();
786
787
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 torch::Tensor cpu_keys = keys;
788
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 torch::Tensor cpu_grads = grads;
789
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 if (keys.is_cuda()) {
790 ✗ cpu_keys = keys.cpu();
791 }
792
2/4
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
 2 if (grads.is_cuda()) {
793 ✗ cpu_grads = grads.cpu();
794 }
795
796
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
797
1/2
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
 2 base::RecTensor rec_grads = ToRecTensor(cpu_grads, base::DataType::FLOAT32);
798
799
1/2
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 op->EmbUpdate(table_name, rec_keys, rec_grads);
800 #ifdef RECSTORE_ENABLE_GPU_CACHE
801 MaintainGpuCacheAfterUpdateNoThrow(keys, grads, grads.size(1));
802 #endif
803 2 }
804
805 ✗ void local_update_flat_torch(const std::string& table_name,
806 const torch::Tensor& keys,
807 const torch::Tensor& grads) {
808 ✗ ResetLocalUpdateFlatProfile();
809 #ifdef RECSTORE_ENABLE_GPU_CACHE
810 gpu::ResetLastGpuCacheProfile();
811 #endif
812 ✗ const auto total_start = SteadyNow();
813 ✗ TORCH_CHECK(!table_name.empty(), "table_name must be non-empty");
814 ✗ TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
815 ✗ TORCH_CHECK(keys.scalar_type() == torch::kInt64,
816 "Keys tensor must have dtype int64");
817 ✗ TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
818
819 ✗ TORCH_CHECK(grads.dim() == 2, "Grads tensor must be 2-dimensional");
820 ✗ TORCH_CHECK(grads.scalar_type() == torch::kFloat32,
821 "Grads tensor must have dtype float32");
822 ✗ TORCH_CHECK(grads.is_contiguous(), "Grads tensor must be contiguous");
823 ✗ TORCH_CHECK(keys.size(0) == grads.size(0),
824 "Keys and grads tensors must have the same number of entries");
825
826 ✗ auto kv_op = GetConcreteKVClientOp();
827 ✗ TORCH_CHECK(IsLocalFastPathBackend(kv_op->CurrentPSBackend()),
828 "local_update_flat requires local_shm or hierkv backend, but "
829 "current backend is ",
830 kv_op->CurrentPSBackend());
831
832 ✗ if (keys.size(0) == 0) {
833 ✗ g_last_local_update_flat_profile[kUpdateTotalMs] = ElapsedMs(total_start);
834 ✗ return;
835 }
836
837 ✗ torch::Tensor cpu_keys = keys;
838 const bool can_async_stage_cuda =
839 ✗ (keys.is_cuda() || grads.is_cuda()) &&
840 ✗ (!keys.is_cuda() || !grads.is_cuda() || keys.device() == grads.device());
841 ✗ bool staged_cuda_async = false;
842 ✗ if (keys.is_cuda()) {
843 ✗ const auto keys_stage_start = SteadyNow();
844 ✗ if (can_async_stage_cuda) {
845 ✗ cpu_keys = StageCudaTensorToPinnedCpuAsyncNoCast(keys);
846 ✗ staged_cuda_async = true;
847 } else {
848 ✗ cpu_keys = StageCudaTensorToPinnedCpu(keys, torch::kInt64);
849 }
850 ✗ g_last_local_update_flat_profile[kUpdateKeysStageMs] =
851 ✗ ElapsedMs(keys_stage_start);
852 }
853 ✗ torch::Tensor cpu_grads = grads;
854 ✗ if (grads.is_cuda()) {
855 ✗ const auto grads_stage_start = SteadyNow();
856 ✗ if (can_async_stage_cuda) {
857 ✗ cpu_grads = StageCudaTensorToPinnedCpuAsyncNoCast(grads);
858 ✗ staged_cuda_async = true;
859 } else {
860 ✗ cpu_grads = StageCudaTensorToPinnedCpu(grads, torch::kFloat32);
861 }
862 ✗ g_last_local_update_flat_profile[kUpdateGradsStageMs] =
863 ✗ ElapsedMs(grads_stage_start);
864 }
865 ✗ if (staged_cuda_async) {
866 ✗ const auto stage_wait_start = SteadyNow();
867 ✗ SynchronizeCurrentCudaStreamForTensor(keys.is_cuda() ? keys : grads);
868 ✗ g_last_local_update_flat_profile[kUpdateStageWaitMs] =
869 ✗ ElapsedMs(stage_wait_start);
870 }
871
872 ✗ base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
873 ✗ base::RecTensor rec_grads = ToRecTensor(cpu_grads, base::DataType::FLOAT32);
874
875 ✗ const auto shm_call_start = SteadyNow();
876 try {
877 ✗ kv_op->LocalUpdateFlat(table_name, rec_keys, rec_grads);
878 ✗ } catch (...) {
879 #ifdef RECSTORE_ENABLE_GPU_CACHE
880 if (gpu::IsGpuCacheEnabled()) {
881 SafeClearGpuCacheNoThrow();
882 gpu::ResetLastGpuCacheProfile();
883 }
884 #endif
885 ✗ throw;
886 ✗ }
887 ✗ g_last_local_update_flat_profile[kUpdateShmCallMs] =
888 ✗ ElapsedMs(shm_call_start);
889
890 #ifdef RECSTORE_ENABLE_GPU_CACHE
891 MaintainGpuCacheAfterUpdateNoThrow(keys, grads, grads.size(1));
892 #endif
893
894 ✗ g_last_local_update_flat_profile[kUpdateTotalMs] = ElapsedMs(total_start);
895 ✗ }
896
897 ✗ std::vector<double> get_last_local_lookup_flat_profile_torch() {
898 ✗ return g_last_local_lookup_flat_profile;
899 }
900
901 ✗ std::vector<double> get_last_local_update_flat_profile_torch() {
902 ✗ return g_last_local_update_flat_profile;
903 }
904
905 ✗ bool warmup_local_lookup_flat_cuda_region_torch() {
906 ✗ auto kv_op = GetConcreteKVClientOp();
907 ✗ const void* payload_base = nullptr;
908 ✗ std::size_t payload_bytes = 0;
909 ✗ if (!kv_op->GetLocalLookupFlatPayloadRegion(&payload_base, &payload_bytes)) {
910 ✗ return false;
911 }
912 ✗ return EnsurePinnedLocalShmPayload(payload_base, payload_bytes);
913 ✗ }
914
915 10 bool init_embedding_table_torch(const std::string& table_name,
916 int64_t num_embeddings,
917 int64_t embedding_dim) {
918
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
 10 TORCH_CHECK(!table_name.empty(), "table_name must be non-empty");
919
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
 10 TORCH_CHECK(num_embeddings > 0, "num_embeddings must be positive");
920
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
 10 TORCH_CHECK(embedding_dim > 0, "embedding_dim must be positive");
921
922 10 EmbeddingTableConfig cfg{};
923 10 cfg.num_embeddings = static_cast<uint64_t>(num_embeddings);
924 10 cfg.embedding_dim = static_cast<uint64_t>(embedding_dim);
925
926
1/2
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
 10 auto op = GetKVClientOp();
927
1/2
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
 10 const bool ok = op->InitEmbeddingTable(table_name, cfg);
928 #ifdef RECSTORE_ENABLE_GPU_CACHE
929 if (ok && gpu::IsGpuCacheEnabled()) {
930 SafeClearGpuCacheNoThrow();
931 gpu::ResetLastGpuCacheProfile();
932 }
933 #endif
934 10 return ok;
935 10 }
936
937 ✗ bool save_checkpoint_torch(const std::string& path,
938 const std::string& metadata) {
939 ✗ TORCH_CHECK(!path.empty(), "checkpoint path must be non-empty");
940 ✗ TORCH_CHECK(!metadata.empty(), "checkpoint metadata must be non-empty");
941 ✗ return GetKVClientOp()->SaveCheckpoint(path, metadata);
942 }
943
944 ✗ bool load_checkpoint_torch(const std::string& path,
945 const std::string& metadata) {
946 ✗ TORCH_CHECK(!path.empty(), "checkpoint path must be non-empty");
947 ✗ TORCH_CHECK(!metadata.empty(), "checkpoint metadata must be non-empty");
948 ✗ const bool ok = GetKVClientOp()->LoadCheckpoint(path, metadata);
949 #ifdef RECSTORE_ENABLE_GPU_CACHE
950 if (ok && gpu::IsGpuCacheEnabled()) {
951 SafeClearGpuCacheNoThrow();
952 gpu::ResetLastGpuCacheProfile();
953 }
954 #endif
955 ✗ return ok;
956 }
957
958 18 void emb_write_torch(const torch::Tensor& keys, const torch::Tensor& values) {
959
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
960
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
961 "Keys tensor must have dtype int64");
962
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
963
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(values.dim() == 2, "Values tensor must be 2-dimensional");
964
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(values.scalar_type() == torch::kFloat32,
965 "Values tensor must have dtype float32");
966
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 TORCH_CHECK(values.is_contiguous(), "Values tensor must be contiguous");
967
3/6
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✓ Branch 4 taken 18 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 18 times.
 18 TORCH_CHECK(keys.size(0) == values.size(0),
968 "Keys and Values tensors must have the same number of entries");
969
970
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 if (keys.size(0) == 0) {
971 ✗ return;
972 }
973
974
1/2
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
 18 auto op = GetKVClientOp();
975
976
1/2
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
 18 torch::Tensor cpu_keys = keys;
977
1/2
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
 18 torch::Tensor cpu_values = values;
978
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 if (keys.is_cuda()) {
979 ✗ cpu_keys = keys.cpu();
980 }
981
2/4
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 18 times.
 18 if (values.is_cuda()) {
982 ✗ cpu_values = values.cpu();
983 }
984
985
1/2
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
 18 base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
986
1/2
✓ Branch 1 taken 18 times.
✗ Branch 2 not taken.
 18 base::RecTensor rec_values = ToRecTensor(cpu_values, base::DataType::FLOAT32);
987
988
1/2
✓ Branch 2 taken 18 times.
✗ Branch 3 not taken.
 18 op->EmbWrite(rec_keys, rec_values);
989 #ifdef RECSTORE_ENABLE_GPU_CACHE
990 if (gpu::IsGpuCacheEnabled()) {
991 SafeClearGpuCacheNoThrow();
992 gpu::ResetLastGpuCacheProfile();
993 }
994 #endif
995 18 }
996
997
998 ✗ void emb_write_values_torch(const torch::Tensor& keys,
999 const torch::Tensor& values) {
1000 // Direct value-set to the PS for a subset of keys, with *per-key* GPU
1001 // cache invalidation (not a full clear). Used by the BagPipe eviction
1002 // writeback path to push locally-updated cache values back to the PS
1003 // without disturbing other cached entries. Mirrors emb_write_torch but
1004 // replaces SafeClearGpuCacheNoThrow() with InvalidateGpuCache(keys).
1005 ✗ TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
1006 ✗ TORCH_CHECK(keys.scalar_type() == torch::kInt64,
1007 "Keys tensor must have dtype int64");
1008 ✗ TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
1009 ✗ TORCH_CHECK(values.dim() == 2, "Values tensor must be 2-dimensional");
1010 ✗ TORCH_CHECK(values.scalar_type() == torch::kFloat32,
1011 "Values tensor must be float32");
1012 ✗ TORCH_CHECK(values.is_contiguous(), "Values tensor must be contiguous");
1013 ✗ TORCH_CHECK(keys.size(0) == values.size(0),
1014 "Keys and Values tensors must have the same number of entries");
1015
1016 ✗ if (keys.size(0) == 0) {
1017 ✗ return;
1018 }
1019
1020 ✗ auto op = GetKVClientOp();
1021
1022 ✗ torch::Tensor cpu_keys = keys;
1023 ✗ torch::Tensor cpu_values = values;
1024 ✗ if (keys.is_cuda()) {
1025 ✗ cpu_keys = keys.cpu();
1026 }
1027 ✗ if (values.is_cuda()) {
1028 ✗ cpu_values = values.cpu();
1029 }
1030
1031 ✗ base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64);
1032 ✗ base::RecTensor rec_values = ToRecTensor(cpu_values, base::DataType::FLOAT32);
1033
1034 ✗ op->EmbWrite(rec_keys, rec_values);
1035 #ifdef RECSTORE_ENABLE_GPU_CACHE
1036 if (gpu::IsGpuCacheEnabled()) {
1037 int64_t embedding_dim = values.size(1);
1038 if (keys.is_cuda() && gpu::CanUseGpuCache(keys, embedding_dim)) {
1039 try {
1040 gpu::InvalidateGpuCache(keys);
1041 } catch (...) {
1042 SafeClearGpuCacheNoThrow();
1043 }
1044 } else {
1045 SafeClearGpuCacheNoThrow();
1046 }
1047 gpu::ResetLastGpuCacheProfile();
1048 }
1049 #endif
1050 ✗ }
1051
1052 ✗ void set_ps_config_torch(const std::string& host, int64_t port) {
1053 ✗ auto kv_op = GetConcreteKVClientOp();
1054 ✗ kv_op->SetPSConfig(host, static_cast<int>(port));
1055 ✗ }
1056
1057 ✗ void set_ps_backend_torch(const std::string& backend) {
1058 ✗ auto kv_op = GetConcreteKVClientOp();
1059 ✗ kv_op->SetPSBackend(backend);
1060 ✗ }
1061
1062 ✗ std::string current_ps_backend_torch() {
1063 ✗ auto kv_op = GetConcreteKVClientOp();
1064 ✗ return kv_op->CurrentPSBackend();
1065 ✗ }
1066
1067 ✗ bool enable_gpu_cache_torch(int64_t capacity, int64_t embedding_dim) {
1068 #ifdef RECSTORE_ENABLE_GPU_CACHE
1069 const bool enabled = gpu::EnableGpuCache(capacity, embedding_dim);
1070 if (enabled) {
1071 ResetGpuCacheBypassState();
1072 }
1073 return enabled;
1074 #else
1075 (void)capacity;
1076 (void)embedding_dim;
1077 ✗ return false;
1078 #endif
1079 }
1080
1081 ✗ void disable_gpu_cache_torch() {
1082 #ifdef RECSTORE_ENABLE_GPU_CACHE
1083 gpu::DisableGpuCache();
1084 ResetGpuCacheBypassState();
1085 #endif
1086 ✗ }
1087
1088 22 void clear_gpu_cache_torch() {
1089 #ifdef RECSTORE_ENABLE_GPU_CACHE
1090 gpu::ClearGpuCache();
1091 ResetGpuCacheBypassState();
1092 #endif
1093 22 }
1094
1095 ✗ void prefill_gpu_cache_torch(const torch::Tensor& keys,
1096 const torch::Tensor& values) {
1097 #ifdef RECSTORE_ENABLE_GPU_CACHE
1098 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
1099 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
1100 "keys must have dtype int64");
1101 TORCH_CHECK(values.dim() == 2, "values must be 2-dimensional");
1102 TORCH_CHECK(values.scalar_type() == torch::kFloat32,
1103 "values must have dtype float32");
1104 TORCH_CHECK(keys.size(0) == values.size(0),
1105 "keys and values must have the same number of rows");
1106 if (keys.numel() == 0) {
1107 return;
1108 }
1109 TORCH_CHECK(keys.is_cuda() || values.is_cuda(),
1110 "prefill_gpu_cache requires keys or values on CUDA");
1111 const auto cache_device = values.is_cuda() ? values.device() : keys.device();
1112 auto keys_cuda = keys.is_cuda() ? keys : keys.to(cache_device);
1113 auto values_cuda = values.is_cuda() ? values : values.to(cache_device);
1114 if (!keys_cuda.is_contiguous()) {
1115 keys_cuda = keys_cuda.contiguous();
1116 }
1117 if (!values_cuda.is_contiguous()) {
1118 values_cuda = values_cuda.contiguous();
1119 }
1120 gpu::FillGpuCache(keys_cuda, values_cuda);
1121 #else
1122 (void)keys;
1123 (void)values;
1124 #endif
1125 ✗ }
1126
1127 ✗ void invalidate_gpu_cache_torch(const torch::Tensor& keys) {
1128 #ifdef RECSTORE_ENABLE_GPU_CACHE
1129 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
1130 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
1131 "keys must have dtype int64");
1132 if (keys.numel() == 0) {
1133 return;
1134 }
1135 TORCH_CHECK(keys.is_cuda(), "invalidate_gpu_cache requires keys on CUDA");
1136 auto keys_cuda = keys;
1137 if (!keys_cuda.is_contiguous()) {
1138 keys_cuda = keys_cuda.contiguous();
1139 }
1140 gpu::InvalidateGpuCache(keys_cuda);
1141 #else
1142 (void)keys;
1143 #endif
1144 ✗ }
1145
1146 ✗ bool apply_sgd_update_gpu_cache_torch(const torch::Tensor& keys,
1147 const torch::Tensor& grads,
1148 double learning_rate) {
1149 #ifdef RECSTORE_ENABLE_GPU_CACHE
1150 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
1151 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
1152 "keys must have dtype int64");
1153 TORCH_CHECK(grads.dim() == 2, "grads must be 2-dimensional");
1154 TORCH_CHECK(grads.scalar_type() == torch::kFloat32,
1155 "grads must have dtype float32");
1156 TORCH_CHECK(keys.size(0) == grads.size(0),
1157 "keys and grads must have the same number of rows");
1158 if (keys.numel() == 0) {
1159 return true;
1160 }
1161 TORCH_CHECK(keys.is_cuda() || grads.is_cuda(),
1162 "apply_sgd_update_gpu_cache requires keys or grads on CUDA");
1163 const auto cache_device = grads.is_cuda() ? grads.device() : keys.device();
1164 auto keys_cuda = keys.is_cuda() ? keys : keys.to(cache_device);
1165 auto grads_cuda = grads.is_cuda() ? grads : grads.to(cache_device);
1166 if (!keys_cuda.is_contiguous()) {
1167 keys_cuda = keys_cuda.contiguous();
1168 }
1169 if (!grads_cuda.is_contiguous()) {
1170 grads_cuda = grads_cuda.contiguous();
1171 }
1172 return gpu::ApplySgdUpdateGpuCache(keys_cuda, grads_cuda, learning_rate);
1173 #else
1174 (void)keys;
1175 (void)grads;
1176 (void)learning_rate;
1177 ✗ return false;
1178 #endif
1179 }
1180
1181 ✗ void set_gpu_cache_lookup_bypass_enabled_torch(bool enabled) {
1182 #ifdef RECSTORE_ENABLE_GPU_CACHE
1183 SetGpuCacheLookupBypassEnabled(enabled);
1184 #else
1185 (void)enabled;
1186 #endif
1187 ✗ }
1188
1189 ✗ bool is_gpu_cache_lookup_bypass_enabled_torch() {
1190 #ifdef RECSTORE_ENABLE_GPU_CACHE
1191 return g_gpu_cache_lookup_bypass_enabled;
1192 #else
1193 ✗ return false;
1194 #endif
1195 }
1196
1197 ✗ bool is_gpu_cache_lookup_bypassed_torch() {
1198 #ifdef RECSTORE_ENABLE_GPU_CACHE
1199 return g_gpu_cache_lookup_bypassed;
1200 #else
1201 ✗ return false;
1202 #endif
1203 }
1204
1205 ✗ void reset_gpu_cache_bypass_state_torch() {
1206 #ifdef RECSTORE_ENABLE_GPU_CACHE
1207 ResetGpuCacheBypassState();
1208 #endif
1209 ✗ }
1210
1211 ✗ std::vector<double> get_last_gpu_cache_profile_torch() {
1212 #ifdef RECSTORE_ENABLE_GPU_CACHE
1213 const auto profile = gpu::GetLastGpuCacheProfile();
1214 return {
1215 profile.query_ms,
1216 profile.backend_lookup_ms,
1217 profile.fill_ms,
1218 profile.update_ms,
1219 profile.hit_count,
1220 profile.invalidate_ms,
1221 profile.request_count,
1222 profile.miss_count,
1223 };
1224 #else
1225 ✗ return {};
1226 #endif
1227 }
1228
1229 // ---- BagPipe-style GPU cache ops (query / update / invalidate / sgd) ----
1230
1231 std::tuple<torch::Tensor, torch::Tensor>
1232 ✗ query_gpu_cache_torch(const torch::Tensor& keys, int64_t embedding_dim) {
1233 #ifdef RECSTORE_ENABLE_GPU_CACHE
1234 if (!gpu::IsGpuCacheEnabled() || keys.numel() == 0) {
1235 auto opts = torch::TensorOptions().dtype(torch::kFloat32);
1236 auto dev = keys.is_cuda() ? keys.device() : torch::kCPU;
1237 return {torch::empty({0, embedding_dim}, opts.device(dev)),
1238 torch::empty({0}, torch::TensorOptions().dtype(torch::kInt64))};
1239 }
1240 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
1241 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
1242 "keys must have dtype int64");
1243 auto keys_contig = keys.is_contiguous() ? keys : keys.contiguous();
1244 auto result = gpu::QueryGpuCache(keys_contig, embedding_dim);
1245 return {result.values, result.missing_keys_cpu};
1246 #else
1247 (void)keys;
1248 (void)embedding_dim;
1249 ✗ auto opts = torch::TensorOptions().dtype(torch::kFloat32);
1250 ✗ return {torch::empty({0, 1}, opts), torch::empty({0}, opts.dtype(torch::kInt64))};
1251 #endif
1252 }
1253
1254 ✗ void update_gpu_cache_torch(const torch::Tensor& keys,
1255 const torch::Tensor& values) {
1256 #ifdef RECSTORE_ENABLE_GPU_CACHE
1257 if (keys.numel() == 0) return;
1258 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
1259 TORCH_CHECK(keys.scalar_type() == torch::kInt64, "keys must be int64");
1260 TORCH_CHECK(values.dim() == 2, "values must be 2-dimensional");
1261 TORCH_CHECK(values.scalar_type() == torch::kFloat32,
1262 "values must be float32");
1263 TORCH_CHECK(keys.size(0) == values.size(0), "row count mismatch");
1264 TORCH_CHECK(keys.is_cuda() || values.is_cuda(),
1265 "update_gpu_cache requires keys or values on CUDA");
1266 const auto dev = values.is_cuda() ? values.device() : keys.device();
1267 auto keys_cuda = keys.is_cuda() ? keys : keys.to(dev);
1268 auto values_cuda = values.is_cuda() ? values : values.to(dev);
1269 if (!keys_cuda.is_contiguous()) keys_cuda = keys_cuda.contiguous();
1270 if (!values_cuda.is_contiguous()) values_cuda = values_cuda.contiguous();
1271 gpu::UpdateGpuCache(keys_cuda, values_cuda);
1272 #else
1273 (void)keys;
1274 (void)values;
1275 #endif
1276 ✗ }
1277
1278 14 TORCH_LIBRARY(recstore_ops, m) {
1279
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_read", emb_read_torch);
1280
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("local_lookup_flat", local_lookup_flat_torch);
1281
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("gpu_cache_lookup_flat", gpu_cache_lookup_flat_torch);
1282
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_update", emb_update_torch);
1283
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_update_table", emb_update_table_torch);
1284
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("local_update_flat", local_update_flat_torch);
1285
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("init_embedding_table", init_embedding_table_torch);
1286
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("save_checkpoint", save_checkpoint_torch);
1287
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("load_checkpoint", load_checkpoint_torch);
1288
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_write", emb_write_torch);
1289
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_write_values", emb_write_values_torch);
1290
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_prefetch", emb_prefetch_torch);
1291
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("emb_wait_result", emb_wait_result_torch);
1292
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("set_ps_config", set_ps_config_torch);
1293
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("set_ps_backend", set_ps_backend_torch);
1294
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("current_ps_backend", current_ps_backend_torch);
1295
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("get_last_local_lookup_flat_profile",
1296 get_last_local_lookup_flat_profile_torch);
1297
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("get_last_local_update_flat_profile",
1298 get_last_local_update_flat_profile_torch);
1299
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("warmup_local_lookup_flat_cuda_region",
1300 warmup_local_lookup_flat_cuda_region_torch);
1301
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("enable_gpu_cache", enable_gpu_cache_torch);
1302
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("disable_gpu_cache", disable_gpu_cache_torch);
1303
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("clear_gpu_cache", clear_gpu_cache_torch);
1304
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("prefill_gpu_cache", prefill_gpu_cache_torch);
1305
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("invalidate_gpu_cache", invalidate_gpu_cache_torch);
1306
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("apply_sgd_update_gpu_cache", apply_sgd_update_gpu_cache_torch);
1307
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("set_gpu_cache_lookup_bypass_enabled",
1308 set_gpu_cache_lookup_bypass_enabled_torch);
1309
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("is_gpu_cache_lookup_bypass_enabled",
1310 is_gpu_cache_lookup_bypass_enabled_torch);
1311
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("is_gpu_cache_lookup_bypassed", is_gpu_cache_lookup_bypassed_torch);
1312
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("reset_gpu_cache_bypass_state", reset_gpu_cache_bypass_state_torch);
1313
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("get_last_gpu_cache_profile", get_last_gpu_cache_profile_torch);
1314
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("query_gpu_cache", query_gpu_cache_torch);
1315
1/2
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 14 m.def("update_gpu_cache", update_gpu_cache_torch);
1316 14 }
1317
1318 } // namespace framework
1319 } // namespace recstore
1320