GCC Code Coverage Report

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