GCC Code Coverage Report

src/
Coverage:
low: ≥ 0%
medium: ≥ 75.0%
high: ≥ 90.0%
Lines:
16 / 18
88.9%
Branches:
-
-%
Functions:
7 / 8
87.5%
framework/op.h
Line Branch Exec Source
1 #pragma once
2
3 #include "base/tensor.h"
4 #include "framework/common/op_runtime_support.h"
5 #include "framework/cache/cache_aware_op.h"
6 #include "framework/common/sparse_gradient_commit_runtime.h"
7 #include "ps/base/base_client.h"
8 #include <cstddef>
9 #include <memory>
10 #include <string>
11 #include <mutex>
12 #include <unordered_map>
13 #include <vector>
14
15 using base::RecTensor;
16
17 namespace recstore {
18 enum class InitStrategyType { Normal, Uniform, Xavier, Zero };
19 struct LocalShmFlatGetHandle;
20
21 struct InitStrategy {
22 InitStrategy() = delete;
23 InitStrategyType type;
24
25 // Optional fields depending on type
26 float mean = 0.0f;
27 float std = 1.0f;
28 float lower = -1.0f;
29 float upper = 1.0f;
30
31 12 InitStrategy(InitStrategyType t) : type(t) {}
32
33 2 static InitStrategy Normal(float mean, float std) {
34 2 InitStrategy s(InitStrategyType::Normal);
35 2 s.mean = mean;
36 2 s.std = std;
37 2 return s;
38 }
39
40 2 static InitStrategy Uniform(float lower, float upper) {
41 2 InitStrategy s(InitStrategyType::Uniform);
42 2 s.lower = lower;
43 2 s.upper = upper;
44 2 return s;
45 }
46
47 2 static InitStrategy Xavier() {
48 2 return InitStrategy(InitStrategyType::Xavier);
49 }
50 6 static InitStrategy Zero() { return InitStrategy(InitStrategyType::Zero); }
51 };
52 class CommonOp {
53 public:
54 // keys: uint64_t tensor with shape [N]
55 // values: emb.dtype tensor with shape [N, D]
56
57 36 CommonOp() = default;
58
59 virtual void EmbInit(const RecTensor& keys, const RecTensor& init_values) = 0;
60 virtual void EmbInit(const RecTensor& keys, const InitStrategy& strategy) = 0;
61
62 // Core KV APIs (sync)
63 virtual void
64 EmbRead(const RecTensor& keys, RecTensor& values) = 0; // sync read
65 virtual void
66 EmbWrite(const RecTensor& keys, const RecTensor& values) = 0; // sync write
67
68 virtual bool
69 EmbExists(const RecTensor& keys) = 0; // not urgent, optional existence check
70 virtual void
71 EmbDelete(const RecTensor& keys) = 0; // not urgent, optional deletion
72
73 // Optional Gradient Hook (can be omitted if optimizer is outside)
74 virtual void
75 EmbUpdate(const RecTensor& keys, const RecTensor& grads) = 0; // not urgent
76 virtual void EmbUpdate(const std::string& table_name,
77 const RecTensor& keys,
78 const RecTensor& grads) = 0;
79
80 virtual int InitEmbeddingTable(const std::string& table_name,
81 const EmbeddingTableConfig& config) = 0;
82
83 // Prefetch & write (async)
84 virtual uint64_t
85 EmbPrefetch(const RecTensor& keys,
86 const RecTensor& values) = 0; // async prefetch, returns a unique
87 // ID to track the prefetch status.
88 virtual bool IsPrefetchDone(
89 uint64_t prefetch_id) = 0; // returns true if the prefetch identified by
90 // prefetch_id is complete.
91 virtual void WaitForPrefetch(
92 uint64_t prefetch_id) = 0; // blocks until the prefetch identified by
93 // prefetch_id is complete.
94 virtual void GetPretchResult(uint64_t prefetch_id, RecTensor& values) = 0;
95
96 virtual uint64_t
97 EmbWriteAsync(const RecTensor& keys,
98 const RecTensor& values) = 0; // async write, returns a unique
99 // ID to track the write status.
100 virtual bool
101 IsWriteDone(uint64_t write_id) = 0; // returns true if the asynchronous write
102 // identified by write_id is complete.
103 virtual void
104 WaitForWrite(uint64_t write_id) = 0; // blocks until the asynchronous write
105 // identified by write_id is complete.
106
107 // Persistence
108 virtual void SaveToFile(const std::string& path) = 0; // not urgent
109 virtual void LoadFromFile(const std::string& path) = 0; // not urgent
110 virtual bool
111 SaveCheckpoint(const std::string& path, const std::string& metadata) = 0;
112 virtual bool
113 LoadCheckpoint(const std::string& path, const std::string& metadata) = 0;
114
115 36 virtual ~CommonOp() = default;
116 };
117
118 class KVClientOp : public CommonOp {
119 public:
120 KVClientOp();
121
122 void EmbInit(const base::RecTensor& keys,
123 const base::RecTensor& init_values) override;
124 void EmbInit(const base::RecTensor& keys,
125 const InitStrategy& strategy) override;
126 void EmbRead(const base::RecTensor& keys, base::RecTensor& values) override;
127 void EmbWrite(const base::RecTensor& keys,
128 const base::RecTensor& values) override;
129 void EmbUpdate(const base::RecTensor& keys,
130 const base::RecTensor& grads) override;
131 void EmbUpdate(const std::string& table_name,
132 const base::RecTensor& keys,
133 const base::RecTensor& grads) override;
134 int64_t SparseGradientCommit(const std::string& table,
135 SparseGradientCommitHeader header,
136 const base::RecTensor& keys,
137 const base::RecTensor& grads);
138 int64_t SparseGradientCommitAsync(
139 const std::string& table,
140 SparseGradientCommitHeader header,
141 const base::RecTensor& keys,
142 const base::RecTensor& grads,
143 std::function<void()> ready);
144 void WaitSparseGradientCommit(int64_t id);
145 void CloseSparseGradientCommit(const std::string& table,
146 SparseGradientCommitHeader header);
147 uint64_t EmbUpdateAsync(const std::string& table_name,
148 const base::RecTensor& keys,
149 const base::RecTensor& grads,
150 framework::CacheUpdateActions cache = {});
151 void WaitForEmbUpdate(uint64_t update_id);
152 void EmbUpdateCached(const std::string& table,
153 const base::RecTensor& keys,
154 const base::RecTensor& grads,
155 const framework::CacheUpdateActions& cache,
156 bool local = false);
157 void EmbWriteCached(const base::RecTensor& keys,
158 const base::RecTensor& values,
159 const framework::CacheUpdateActions& cache,
160 bool per_key = false);
161 int InitEmbeddingTableCached(const std::string& table_name,
162 const EmbeddingTableConfig& config,
163 const framework::CacheUpdateActions& cache);
164 bool LoadCheckpointCached(const std::string& path,
165 const std::string& metadata,
166 const framework::CacheUpdateActions& cache);
167 int InitEmbeddingTable(const std::string& table_name,
168 const EmbeddingTableConfig& config) override;
169 bool EmbExists(const base::RecTensor& keys) override;
170 void EmbDelete(const base::RecTensor& keys) override;
171 uint64_t EmbPrefetch(const base::RecTensor& keys,
172 const base::RecTensor& values) override;
173 bool IsPrefetchDone(uint64_t prefetch_id) override;
174 void WaitForPrefetch(uint64_t prefetch_id) override;
175 void GetPretchResult(uint64_t prefetch_id, RecTensor& values) override;
176 uint64_t EmbWriteAsync(const base::RecTensor& keys,
177 const base::RecTensor& values) override;
178 bool IsWriteDone(uint64_t write_id) override;
179 void WaitForWrite(uint64_t write_id) override;
180 void SaveToFile(const std::string& path) override;
181 void LoadFromFile(const std::string& path) override;
182 bool
183 SaveCheckpoint(const std::string& path, const std::string& metadata) override;
184 bool
185 LoadCheckpoint(const std::string& path, const std::string& metadata) override;
186 void SetPSConfig(const std::string& host, int port);
187 void SetPSBackend(const std::string& backend);
188 std::string CurrentPSBackend() const;
189 void RequireLocalBackend() const;
190 ✗ bool SupportsBorrowedLocalLookup() const {
191 ✗ return ps_backend_name_ == "local_shm";
192 }
193 void LocalLookupFlat(const base::RecTensor& keys, base::RecTensor& values);
194 std::shared_ptr<base::RecTensor>
195 AcquireLocalLookupFlat(const base::RecTensor& keys, int64_t embedding_dim);
196 int SubmitLocalLookupFlat(const base::RecTensor& keys,
197 int64_t embedding_dim,
198 LocalShmFlatGetHandle* handle);
199 int WaitLocalLookupFlat(LocalShmFlatGetHandle* handle);
200 void ReleaseLocalLookupFlat(LocalShmFlatGetHandle* handle);
201 bool GetLocalLookupFlatPayloadRegion(const void** base, std::size_t* bytes);
202 void LocalUpdateFlat(const std::string& table_name,
203 const base::RecTensor& keys,
204 const base::RecTensor& grads);
205
206 private:
207 SparseGradientCommitRuntime sparse_commits_;
208 framework::PendingCacheUpdates pending_updates_;
209 int64_t embedding_dim_;
210 std::string ps_backend_name_ = "unknown";
211 static BasePSClient* ps_client_;
212 static std::unique_ptr<BasePSClient> ps_client_holder_;
213
214 #ifdef USE_FAKE_KVCLIENT
215 std::unordered_map<uint64_t, std::vector<float>> store_;
216 std::mutex mtx_;
217 float learning_rate_;
218 std::unordered_map<uint64_t, std::vector<std::vector<float>>>
219 prefetch_results_;
220 uint64_t next_prefetch_id_ = 1;
221 #endif
222 };
223
224 std::shared_ptr<CommonOp> GetKVClientOp();
225
226 } // namespace recstore
227