GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
25 / 75
33.3%
Branches:
21 / 70
30.0%
Functions:
2 / 22
9.1%
framework/pytorch/cache/cache_bindings.cc
Line Branch Exec Source
1 #include "framework/pytorch/cache/cache_bindings.h"
2 #include "framework/pytorch/op_internal.h"
3
4 namespace recstore {
5 namespace framework {
6 namespace {
7
8 torch::Tensor
9 ✗ gpu_cache_lookup_flat_torch(const torch::Tensor& keys, int64_t dim) {
10 ✗ return std::get<0>(CachedReadTorch(keys, dim));
11 }
12 std::tuple<torch::Tensor, torch::Tensor>
13 ✗ gpu_cache_lookup_flat_no_evict_torch(const torch::Tensor& keys, int64_t dim) {
14 ✗ return CachedReadTorch(keys, dim, false, true);
15 }
16 ✗ void emb_write_values_torch(const torch::Tensor& keys,
17 const torch::Tensor& values) {
18
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ CheckRows(keys, values);
19
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
 ✗ if (!keys.numel())
20 ✗ return;
21
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 ✗ auto cpu_keys = keys.cpu(), cpu_values = values.cpu();
22
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
 ✗ GetConcreteKVClientOp()->EmbWriteCached(
23
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ ToRecTensor(cpu_keys, base::DataType::UINT64),
24
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ ToRecTensor(cpu_values, base::DataType::FLOAT32),
25
0/4
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 ✗ CacheActions(keys, values.size(1)),
26 true);
27 ✗ }
28
29 ✗ bool enable_gpu_cache_torch(int64_t capacity, int64_t embedding_dim) {
30 #ifdef RECSTORE_ENABLE_GPU_CACHE
31 return torch_cache::EnableGpuCache(capacity, embedding_dim);
32 #else
33 (void)capacity;
34 (void)embedding_dim;
35 ✗ return false;
36 #endif
37 }
38
39 ✗ void disable_gpu_cache_torch() {
40 #ifdef RECSTORE_ENABLE_GPU_CACHE
41 torch_cache::DisableGpuCache();
42 #endif
43 ✗ }
44
45 22 void clear_gpu_cache_torch() {
46 #ifdef RECSTORE_ENABLE_GPU_CACHE
47 torch_cache::ResetGpuCache();
48 #endif
49 22 }
50
51 ✗ void prefill_gpu_cache_torch(const torch::Tensor& keys,
52 const torch::Tensor& values) {
53 #ifdef RECSTORE_ENABLE_GPU_CACHE
54 CheckRows(keys, values, false);
55 if (!keys.numel()) {
56 return;
57 }
58 auto [k, v] = CudaCacheRows(keys, values);
59 torch_cache::FillGpuCache(k, v);
60 #else
61 (void)keys;
62 (void)values;
63 #endif
64 ✗ }
65
66 ✗ torch::Tensor prefill_gpu_cache_no_evict_torch(const torch::Tensor& keys,
67 const torch::Tensor& values) {
68 #ifdef RECSTORE_ENABLE_GPU_CACHE
69 CheckRows(keys, values, false);
70 if (!keys.numel()) {
71 return torch::empty({0}, torch::TensorOptions().dtype(torch::kBool));
72 }
73 auto [k, v] = CudaCacheRows(keys, values);
74 return torch_cache::FillGpuCacheNoEvict(k, v);
75 #else
76 (void)keys;
77 (void)values;
78 ✗ TORCH_CHECK(false, "GPU cache support is unavailable");
79 #endif
80 }
81
82 ✗ torch::Tensor contains_gpu_cache_torch(const torch::Tensor& keys) {
83 #ifdef RECSTORE_ENABLE_GPU_CACHE
84 return torch_cache::ContainsGpuCache(keys);
85 #else
86 (void)keys;
87 ✗ TORCH_CHECK(false, "GPU cache support is unavailable");
88 #endif
89 }
90
91 ✗ int64_t get_gpu_cache_generation_torch() {
92 #ifdef RECSTORE_ENABLE_GPU_CACHE
93 return static_cast<int64_t>(torch_cache::GetGpuCacheGeneration());
94 #else
95 ✗ return 0;
96 #endif
97 }
98
99 ✗ void invalidate_gpu_cache_torch(const torch::Tensor& keys) {
100 #ifdef RECSTORE_ENABLE_GPU_CACHE
101 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
102 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
103 "keys must have dtype int64");
104 if (keys.numel() == 0) {
105 return;
106 }
107 TORCH_CHECK(keys.is_cuda(), "invalidate_gpu_cache requires keys on CUDA");
108 auto keys_cuda = keys;
109 if (!keys_cuda.is_contiguous()) {
110 keys_cuda = keys_cuda.contiguous();
111 }
112 torch_cache::InvalidateGpuCache(keys_cuda);
113 #else
114 (void)keys;
115 #endif
116 ✗ }
117
118 ✗ torch::Tensor invalidate_gpu_cache_with_mask_torch(const torch::Tensor& keys) {
119 #ifdef RECSTORE_ENABLE_GPU_CACHE
120 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
121 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
122 "keys must have dtype int64");
123 if (keys.numel() == 0) {
124 return torch::empty({0}, torch::TensorOptions().dtype(torch::kBool));
125 }
126 TORCH_CHECK(keys.is_cuda(),
127 "invalidate_gpu_cache_with_mask requires CUDA keys");
128 auto keys_cuda = keys;
129 if (!keys_cuda.is_contiguous()) {
130 keys_cuda = keys_cuda.contiguous();
131 }
132 return torch_cache::InvalidateGpuCacheWithMask(keys_cuda);
133 #else
134 (void)keys;
135 ✗ TORCH_CHECK(false, "GPU cache support is unavailable");
136 #endif
137 }
138
139 std::tuple<torch::Tensor, torch::Tensor>
140 ✗ gpu_cache_lookup_flat_assuming_hits_torch(const torch::Tensor& keys,
141 int64_t embedding_dim) {
142 #ifdef RECSTORE_ENABLE_GPU_CACHE
143 TORCH_CHECK(keys.dim() == 1, "Keys tensor must be 1-dimensional");
144 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
145 "Keys must have dtype int64");
146 TORCH_CHECK(keys.is_contiguous(), "Keys tensor must be contiguous");
147 TORCH_CHECK(keys.is_cuda(),
148 "gpu_cache_lookup_flat_assuming_hits requires CUDA keys");
149 TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive");
150 return torch_cache::LookupGpuCacheAssumingHits(keys, embedding_dim);
151 #else
152 (void)keys;
153 (void)embedding_dim;
154 ✗ TORCH_CHECK(false, "GPU cache support is unavailable");
155 #endif
156 }
157
158 ✗ bool apply_sgd_update_gpu_cache_torch(const torch::Tensor& keys,
159 const torch::Tensor& grads,
160 double learning_rate) {
161 #ifdef RECSTORE_ENABLE_GPU_CACHE
162 CheckRows(keys, grads, false);
163 if (!keys.numel())
164 return true;
165 auto [k, v] = CudaCacheRows(keys, grads);
166 return torch_cache::ApplySgdUpdateGpuCache(k, v, learning_rate);
167 #else
168 (void)keys;
169 (void)grads;
170 (void)learning_rate;
171 ✗ return false;
172 #endif
173 }
174
175 // Best-effort in-place SGD on the GPU cache: value -= lr * grad for keys
176 // present in the cache; missing keys silently skipped. No Query, no
177 // missing report, no device synchronization. This is the cache-side SGD
178 // primitive for controllers (BagPipe) that own their own cache policy.
179 ✗ void apply_sgd_update_gpu_cache_best_effort_torch(
180 const torch::Tensor& keys,
181 const torch::Tensor& grads,
182 double learning_rate) {
183 #ifdef RECSTORE_ENABLE_GPU_CACHE
184 CheckRows(keys, grads, false);
185 if (!keys.numel()) {
186 return;
187 }
188 auto [k, v] = CudaCacheRows(keys, grads);
189 torch_cache::ApplySgdUpdateBestEffortGpuCache(k, v, learning_rate);
190 #else
191 (void)keys;
192 (void)grads;
193 (void)learning_rate;
194 #endif
195 ✗ }
196
197 ✗ void set_gpu_cache_lookup_bypass_enabled_torch(bool enabled) {
198 #ifdef RECSTORE_ENABLE_GPU_CACHE
199 g_cache_aware_op.SetBypassEnabled(enabled);
200 #else
201 (void)enabled;
202 #endif
203 ✗ }
204
205 ✗ bool is_gpu_cache_lookup_bypass_enabled_torch() {
206 #ifdef RECSTORE_ENABLE_GPU_CACHE
207 return g_cache_aware_op.BypassEnabled();
208 #else
209 ✗ return false;
210 #endif
211 }
212
213 ✗ bool is_gpu_cache_lookup_bypassed_torch() {
214 #ifdef RECSTORE_ENABLE_GPU_CACHE
215 return g_cache_aware_op.Bypassed();
216 #else
217 ✗ return false;
218 #endif
219 }
220
221 ✗ void reset_gpu_cache_bypass_state_torch() {
222 #ifdef RECSTORE_ENABLE_GPU_CACHE
223 g_cache_aware_op.Reset();
224 #endif
225 ✗ }
226
227 // ---- BagPipe-style GPU cache ops (query / update / invalidate / sgd) ----
228 std::tuple<torch::Tensor, torch::Tensor>
229 ✗ query_gpu_cache_torch(const torch::Tensor& keys, int64_t embedding_dim) {
230 #ifdef RECSTORE_ENABLE_GPU_CACHE
231 if (!torch_cache::IsGpuCacheEnabled() || keys.numel() == 0) {
232 auto opts = torch::TensorOptions().dtype(torch::kFloat32);
233 auto dev = keys.is_cuda() ? keys.device() : torch::kCPU;
234 return {torch::empty({0, embedding_dim}, opts.device(dev)),
235 torch::empty({0}, torch::TensorOptions().dtype(torch::kInt64))};
236 }
237 TORCH_CHECK(keys.dim() == 1, "keys must be 1-dimensional");
238 TORCH_CHECK(keys.scalar_type() == torch::kInt64,
239 "keys must have dtype int64");
240 auto keys_contig = keys.is_contiguous() ? keys : keys.contiguous();
241 auto result = torch_cache::QueryGpuCache(keys_contig, embedding_dim);
242 return {result.values, result.missing_keys_cpu};
243 #else
244 (void)keys;
245 (void)embedding_dim;
246
0/2
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 ✗ auto opts = torch::TensorOptions().dtype(torch::kFloat32);
247
0/2
✗ Branch 3 not taken.
✗ Branch 4 not taken.
 ✗ return {torch::empty({0, 1}, opts),
248
0/2
✗ Branch 5 not taken.
✗ Branch 6 not taken.
 ✗ torch::empty({0}, opts.dtype(torch::kInt64))};
249 #endif
250 }
251
252 ✗ void update_gpu_cache_torch(const torch::Tensor& keys,
253 const torch::Tensor& values) {
254 #ifdef RECSTORE_ENABLE_GPU_CACHE
255 if (!keys.numel())
256 return;
257 CheckRows(keys, values, false);
258 auto [k, v] = CudaCacheRows(keys, values);
259 torch_cache::UpdateGpuCache(k, v);
260 #else
261 (void)keys;
262 (void)values;
263 #endif
264 ✗ }
265
266 } // namespace
267
268 16 void RegisterGpuCacheOps(torch::Library& m) {
269
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("gpu_cache_lookup_flat", gpu_cache_lookup_flat_torch);
270
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("gpu_cache_lookup_flat_no_evict", gpu_cache_lookup_flat_no_evict_torch);
271
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("emb_write_values", emb_write_values_torch);
272
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("enable_gpu_cache", enable_gpu_cache_torch);
273
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("disable_gpu_cache", disable_gpu_cache_torch);
274
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("clear_gpu_cache", clear_gpu_cache_torch);
275
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("prefill_gpu_cache", prefill_gpu_cache_torch);
276
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("prefill_gpu_cache_no_evict", prefill_gpu_cache_no_evict_torch);
277
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("contains_gpu_cache", contains_gpu_cache_torch);
278
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("get_gpu_cache_generation", get_gpu_cache_generation_torch);
279
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("invalidate_gpu_cache", invalidate_gpu_cache_torch);
280
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("invalidate_gpu_cache_with_mask", invalidate_gpu_cache_with_mask_torch);
281
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("gpu_cache_lookup_flat_assuming_hits",
282 gpu_cache_lookup_flat_assuming_hits_torch);
283
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("apply_sgd_update_gpu_cache", apply_sgd_update_gpu_cache_torch);
284
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("apply_sgd_update_gpu_cache_best_effort",
285 apply_sgd_update_gpu_cache_best_effort_torch);
286
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("set_gpu_cache_lookup_bypass_enabled",
287 set_gpu_cache_lookup_bypass_enabled_torch);
288
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("is_gpu_cache_lookup_bypass_enabled",
289 is_gpu_cache_lookup_bypass_enabled_torch);
290
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("is_gpu_cache_lookup_bypassed", is_gpu_cache_lookup_bypassed_torch);
291
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("reset_gpu_cache_bypass_state", reset_gpu_cache_bypass_state_torch);
292
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("query_gpu_cache", query_gpu_cache_torch);
293
1/2
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
 16 m.def("update_gpu_cache", update_gpu_cache_torch);
294 16 }
295
296 } // namespace framework
297 } // namespace recstore
298