framework/pytorch/op_prefetch.cc
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "framework/pytorch/op_prefetch.h" | ||
| 2 | |||
| 3 | #include <torch/extension.h> | ||
| 4 | |||
| 5 | #include "framework/pytorch/op_internal.h" | ||
| 6 | |||
| 7 | namespace recstore { | ||
| 8 | namespace framework { | ||
| 9 | namespace { | ||
| 10 | |||
| 11 | // Async prefetch: returns a unique prefetch id (uint64_t) | ||
| 12 | 2 | int64_t emb_prefetch_torch(const torch::Tensor& keys) { | |
| 13 | 2 | CheckKeys(keys); | |
| 14 | |||
| 15 | 2 | auto op = GetKVClientOp(); | |
| 16 | 2 | torch::Tensor cpu_keys = ToCpuTensor(keys); | |
| 17 | 2 | base::RecTensor rec_keys = ToRecTensor(cpu_keys, base::DataType::UINT64); | |
| 18 | // Dummy values tensor (unused by backend prefetch implementation) | ||
| 19 | 2 | auto dummy_vals = torch::empty({0, 0}, keys.options().dtype(torch::kFloat32)); | |
| 20 | 2 | base::RecTensor rec_vals = ToRecTensor(dummy_vals, base::DataType::FLOAT32); | |
| 21 | 2 | uint64_t pid = op->EmbPrefetch(rec_keys, rec_vals); | |
| 22 | 2 | return static_cast<int64_t>(pid); | |
| 23 | 2 | } | |
| 24 | |||
| 25 | // Wait for prefetch and return result tensor [N, embedding_dim] on CPU | ||
| 26 | torch::Tensor | ||
| 27 | 2 | emb_wait_result_torch(int64_t prefetch_id, int64_t embedding_dim) { | |
| 28 | 2 | TORCH_CHECK(embedding_dim > 0, "Embedding dimension must be positive"); | |
| 29 | 2 | auto kv_op = GetConcreteKVClientOp(); | |
| 30 | // RDMA result consumption waits internally and can borrow the slot payload. | ||
| 31 | // For a single RPC, a prior wait would materialize the borrowed payload. | ||
| 32 | 2 | if (kv_op->CurrentPSBackend() != "rdma") { | |
| 33 | 2 | kv_op->WaitForPrefetch(static_cast<uint64_t>(prefetch_id)); | |
| 34 | } | ||
| 35 | auto owned = std::make_shared<base::RecTensor>( | ||
| 36 | 4 | std::vector<int64_t>{0, embedding_dim}, base::DataType::FLOAT32); | |
| 37 | 2 | kv_op->GetPretchResult(static_cast<uint64_t>(prefetch_id), *owned); | |
| 38 | auto options = | ||
| 39 | 2 | torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU); | |
| 40 | 2 | const int64_t L = owned->dim() == 2 ? owned->shape(0) : 0; | |
| 41 | 2 | if (L == 0) { | |
| 42 | ✗ | return torch::empty({0, embedding_dim}, options); | |
| 43 | } | ||
| 44 | return torch::from_blob( | ||
| 45 | 2 | owned->data(), {L, embedding_dim}, [owned](void*) {}, options); | |
| 46 | 2 | } | |
| 47 | |||
| 48 | } // namespace | ||
| 49 | |||
| 50 | 16 | void RegisterPrefetchOps(torch::Library& m) { | |
| 51 | 16 | m.def("emb_prefetch", emb_prefetch_torch); | |
| 52 | 16 | m.def("emb_wait_result", emb_wait_result_torch); | |
| 53 | 16 | } | |
| 54 | |||
| 55 | } // namespace framework | ||
| 56 | } // namespace recstore | ||
| 57 |