GCC Code Coverage Report

Directory: src/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 50.8% 67 / 0 / 132
Functions: 47.1% 8 / 0 / 17
Branches: 34.9% 37 / 0 / 106
ps/rdma/rdma_protocol.h
Line Branch Exec Source
1 #pragma once
2
3 #include <algorithm>
4 #include <array>
5 #include <atomic>
6 #include <cstddef>
7 #include <cstdint>
8 #include <cstring>
9 #include <limits>
10 #include <string>
11 #include <string_view>
12
13 #include "base/array.h"
14 #include "base/flatc.h"
15 #include "base/log.h"
16 #include "ps/base/parameters.h"
17 #include "ps/rdma/rdma_status.h"
18
19 namespace petps {
20
21 inline constexpr std::uint32_t kRcProtocolMagic = 0x52435053;
22 inline constexpr std::uint16_t kRcProtocolVersion = 1;
23 inline constexpr std::size_t kTableNameBytes = 64;
24 inline constexpr std::uint32_t kRcSlotReady = 1;
25 inline constexpr std::uint32_t kRcSlotDone = 2;
26 inline constexpr std::uint32_t kRcFlagGetDirectSg = 1U << 0;
27 inline constexpr std::uint32_t kRcFlagGetAllowFallbackCopy = 1U << 1;
28
29 enum class RcOp : std::uint16_t {
30 kGet = 1,
31 kPut = 2,
32 kUpdate = 3,
33 kInitTable = 4,
34 kUpdateFlat = 5,
35 };
36
37 enum class RcHashMethod : std::uint8_t {
38 kCityHash = 1,
39 kSimpleMod = 2,
40 };
41
42 struct alignas(64) RequestDescriptor {
43 std::uint32_t magic = kRcProtocolMagic;
44 std::uint16_t version = kRcProtocolVersion;
45 std::uint16_t op = static_cast<std::uint16_t>(RcOp::kGet);
46 std::uint64_t seq = 0; // Monotonic lane-local request sequence.
47 std::uint32_t shard_id = 0; // Logical shard targeted by this RPC.
48 std::uint32_t client_id = 0; // Logical client owner of this lane.
49 std::uint32_t qp_index = 0; // Lane index within the client.
50 std::uint32_t key_count = 0; // Number of keys in the payload.
51 std::uint32_t value_size = 0; // Row size in bytes for GET responses.
52 std::uint32_t embedding_dim = 0; // Row size expressed as float count.
53 std::uint32_t payload_offset = 0; // Offset from slot base to payload.
54 std::uint32_t payload_bytes = 0; // Bytes occupied by the payload.
55 std::uint32_t response_bytes = 0; // Bytes expected in the response payload.
56 std::uint32_t reserved0 = 0;
57 std::uint64_t client_response_addr =
58 0; // Optional client response address for verbs RC.
59 std::uint32_t client_response_rkey =
60 0; // Optional client response remote key.
61 std::uint32_t client_status_rkey = 0; // Optional client status remote key.
62 std::uint64_t client_status_addr =
63 0; // Optional client status address for verbs RC.
64 std::uint32_t flags = 0; // Op-specific protocol flags.
65 std::uint32_t reserved1 = 0;
66 std::array<char, kTableNameBytes>
67 table_name{}; // Optional logical table name.
68 };
69
70 struct alignas(64) CommitWord {
71 std::atomic<std::uint64_t> seq{0}; // Mirrors RequestDescriptor::seq.
72 std::atomic<std::uint32_t> state{
73 0}; // READY/DONE state published by client/server.
74 std::uint32_t checksum_or_reserved =
75 0; // Reserved for future integrity checks.
76 };
77
78 struct alignas(64) StatusWord {
79 std::atomic<std::uint64_t> seq{0}; // Mirrors the request seq completed here.
80 std::atomic<std::uint32_t> state{
81 0}; // DONE when the server has finished writing.
82 std::int32_t status = 0; // RpcStatus value returned by the server.
83 std::uint32_t response_bytes = 0; // Payload bytes valid in the response slot.
84 std::uint32_t reserved = 0;
85 };
86
87 static_assert(sizeof(RequestDescriptor) == 192, "RequestDescriptor size");
88 static_assert(alignof(RequestDescriptor) == 64, "RequestDescriptor align");
89 static_assert(alignof(CommitWord) == 64, "CommitWord align");
90 static_assert(alignof(StatusWord) == 64, "StatusWord align");
91
92 ✗ inline std::size_t Align64(std::size_t value) {
93 ✗ return (value + 63U) & ~std::size_t{63U};
94 }
95
96 2 inline std::size_t GetKeysPerRpcByResponseBudget(
97 std::size_t value_size, std::size_t mtu_bytes, std::size_t response_mtu) {
98
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 2 if (value_size == 0 || mtu_bytes == 0 || response_mtu == 0) {
99 ✗ return 0;
100 }
101 2 return (mtu_bytes * response_mtu) / value_size;
102 }
103
104 ✗ inline std::size_t GetRequestBytes(std::size_t key_count) {
105 ✗ return key_count * sizeof(std::uint64_t);
106 }
107
108 inline std::size_t
109 ✗ GetResponseBytes(std::size_t key_count, std::size_t value_size) {
110 ✗ return key_count * value_size;
111 }
112
113 inline std::size_t
114 ✗ FixedSlotResponseBytes(std::size_t key_count, std::size_t value_size) {
115 ✗ return GetResponseBytes(key_count, value_size) + sizeof(std::int32_t);
116 }
117
118 ✗ inline std::size_t InitTablePayloadBytes() { return sizeof(std::uint64_t) * 2; }
119
120 ✗ inline std::size_t PutPayloadBudget(std::size_t request_slot_bytes) {
121 ✗ if (request_slot_bytes <=
122 ✗ Align64(sizeof(RequestDescriptor)) + Align64(sizeof(CommitWord))) {
123 ✗ return 0;
124 }
125 ✗ return request_slot_bytes - Align64(sizeof(RequestDescriptor)) -
126 ✗ Align64(sizeof(CommitWord));
127 }
128
129 inline std::size_t ParameterReaderBytes(const ParameterCompressReader& reader) {
130 return static_cast<std::size_t>(reader.byte_size());
131 }
132
133 4 inline std::size_t PutPayloadBytes(
134 const std::vector<std::uint64_t>& keys,
135 const std::vector<std::vector<float>>& values,
136 std::string* payload,
137 std::string* error = nullptr) {
138
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (payload == nullptr) {
139 ✗ if (error != nullptr) {
140 ✗ *error = "payload buffer is null";
141 }
142 ✗ return 0;
143 }
144
1/2
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (keys.size() != values.size()) {
145 ✗ if (error != nullptr) {
146 ✗ *error = "keys and values size mismatch";
147 }
148 ✗ return 0;
149 }
150
151
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 ParameterCompressor compressor;
152
2/2
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 4 times.
 12 for (std::size_t i = 0; i < keys.size(); ++i) {
153 8 ParameterPack pack;
154 8 pack.key = keys[i];
155 8 pack.dim = static_cast<int>(values[i].size());
156 8 pack.emb_data = values[i].data();
157
1/2
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
 8 compressor.AddItem(pack, nullptr);
158 }
159
160 4 payload->clear();
161
1/2
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
 4 compressor.ToBlock(payload);
162 4 return payload->size();
163 4 }
164
165 2 inline std::size_t UpdatePayloadBytes(
166 const std::vector<std::uint64_t>& keys,
167 const std::vector<std::vector<float>>& values,
168 std::string* payload,
169 std::string* error = nullptr) {
170 2 return PutPayloadBytes(keys, values, payload, error);
171 }
172
173 18 inline std::size_t FlatUpdatePayloadBytes(
174 std::size_t key_count, std::size_t embedding_dim) {
175 18 constexpr std::size_t kKeyBytes = sizeof(std::uint64_t);
176 18 constexpr std::size_t kFloatBytes = sizeof(float);
177 18 const std::size_t max_size = std::numeric_limits<std::size_t>::max();
178
3/4
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 16 times.
 18 if (embedding_dim == 0 || embedding_dim > (max_size - kKeyBytes) / kFloatBytes) {
179 2 return 0;
180 }
181 16 const std::size_t row_bytes = kKeyBytes + embedding_dim * kFloatBytes;
182
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 2 times.
 16 return key_count > max_size / row_bytes ? 0 : key_count * row_bytes;
183 }
184
185 4 inline std::size_t UpdatePayloadBytesFlat(
186 base::ConstArray<std::uint64_t> keys,
187 const float* values,
188 std::size_t embedding_dim,
189 std::string* payload,
190 std::string* error = nullptr) {
191
4/8
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 3 taken 4 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
 4 if (payload == nullptr || (keys.Size() > 0 && values == nullptr)) {
192 ✗ if (error != nullptr) {
193 ✗ *error = "payload buffer or values is null";
194 }
195 ✗ return 0;
196 }
197
198 const std::size_t payload_bytes =
199 4 FlatUpdatePayloadBytes(keys.Size(), embedding_dim);
200
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (payload_bytes == 0) {
201 ✗ if (error != nullptr) {
202 ✗ *error = "invalid flat update shape";
203 }
204 ✗ return 0;
205 }
206 4 const std::size_t key_bytes = keys.Size() * sizeof(std::uint64_t);
207 4 const std::size_t value_bytes = keys.Size() * embedding_dim * sizeof(float);
208 4 payload->resize(payload_bytes);
209
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (key_bytes > 0) {
210 4 std::memcpy(payload->data(), keys.Data(), key_bytes);
211 4 std::memcpy(payload->data() + key_bytes, values, value_bytes);
212 }
213 4 return payload->size();
214 }
215
216 4 inline std::size_t PackFlatUpdatePayloadGather(
217 const std::uint64_t* keys,
218 const float* values,
219 std::size_t num_rows,
220 std::size_t embedding_dim,
221 const std::size_t* row_indices,
222 std::size_t row_count,
223 void* payload,
224 std::size_t payload_capacity,
225 std::string* error = nullptr) {
226
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if ((row_count > 0 &&
227
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
 4 (keys == nullptr || values == nullptr || row_indices == nullptr)) ||
228 payload == nullptr) {
229 ✗ if (error != nullptr) {
230 ✗ *error = "flat update gather input is null";
231 }
232 ✗ return 0;
233 }
234 const std::size_t payload_bytes =
235 4 FlatUpdatePayloadBytes(row_count, embedding_dim);
236
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (payload_bytes == 0 || payload_bytes > payload_capacity) {
237 ✗ if (error != nullptr) {
238 ✗ *error = "flat update gather payload does not fit";
239 }
240 ✗ return 0;
241 }
242
243 4 auto* payload_keys = static_cast<std::uint64_t*>(payload);
244 4 auto* payload_values = reinterpret_cast<float*>(
245 4 static_cast<char*>(payload) + row_count * sizeof(std::uint64_t));
246 4 const std::size_t row_bytes = embedding_dim * sizeof(float);
247
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 for (std::size_t row = 0; row < row_count; ++row) {
248 6 const std::size_t source_row = row_indices[row];
249
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (source_row >= num_rows) {
250
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (error != nullptr) {
251 2 *error = "flat update gather row index is out of range";
252 }
253 2 return 0;
254 }
255 4 payload_keys[row] = keys[source_row];
256 4 std::memcpy(
257 4 payload_values + row * embedding_dim,
258 4 values + source_row * embedding_dim,
259 row_bytes);
260 }
261 2 return payload_bytes;
262 }
263
264 ✗ inline bool CopyTableName(std::string_view table_name,
265 std::array<char, kTableNameBytes>* storage) {
266 ✗ if (storage == nullptr || table_name.size() >= kTableNameBytes) {
267 ✗ return false;
268 }
269 ✗ storage->fill('\0');
270 ✗ std::memcpy(storage->data(), table_name.data(), table_name.size());
271 ✗ return true;
272 }
273
274 inline std::string_view
275 ✗ DescriptorTableName(const RequestDescriptor& descriptor) {
276 ✗ return std::string_view(
277 descriptor.table_name.data(),
278 ✗ std::find(
279 ✗ descriptor.table_name.begin(), descriptor.table_name.end(), '\0') -
280 ✗ descriptor.table_name.begin());
281 }
282
283 ✗ inline bool ValidateRequestDescriptor(
284 const RequestDescriptor& descriptor,
285 std::size_t request_slot_bytes,
286 std::size_t response_slot_bytes,
287 std::string* error = nullptr) {
288 ✗ if (descriptor.magic != kRcProtocolMagic) {
289 ✗ if (error != nullptr) {
290 ✗ *error = "bad request magic";
291 }
292 ✗ return false;
293 }
294 ✗ if (descriptor.version != kRcProtocolVersion) {
295 ✗ if (error != nullptr) {
296 ✗ *error = "bad request version";
297 }
298 ✗ return false;
299 }
300 ✗ if (descriptor.payload_offset < sizeof(RequestDescriptor) ||
301 ✗ static_cast<std::size_t>(descriptor.payload_offset) +
302 ✗ descriptor.payload_bytes >
303 request_slot_bytes) {
304 ✗ if (error != nullptr) {
305 ✗ *error = "request payload exceeds slot capacity";
306 }
307 ✗ return false;
308 }
309 ✗ if (descriptor.response_bytes > response_slot_bytes) {
310 ✗ if (error != nullptr) {
311 ✗ *error = "response exceeds slot capacity";
312 }
313 ✗ return false;
314 }
315 ✗ return true;
316 }
317
318 2 inline void ResetStatusWord(StatusWord* status, std::uint64_t seq) {
319 2 status->status = static_cast<std::int32_t>(RpcStatus::kPending);
320 2 status->response_bytes = 0;
321 2 status->seq.store(seq, std::memory_order_release);
322 2 status->state.store(0, std::memory_order_release);
323 2 }
324
325 6 inline bool StatusWordDone(const StatusWord& status, std::uint64_t seq) {
326
2/2
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 2 times.
 16 return status.state.load(std::memory_order_acquire) == kRcSlotDone &&
327
2/2
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
 14 status.seq.load(std::memory_order_acquire) == seq;
328 }
329
330 } // namespace petps
331