ps/base/parameters.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | #include <algorithm> | ||
| 3 | #include <cstring> | ||
| 4 | #include <cstdint> | ||
| 5 | #include <vector> | ||
| 6 | |||
| 7 | #include "base/flatc.h" | ||
| 8 | |||
| 9 | #pragma pack(push, 1) | ||
| 10 | struct ParameterCompressItem { | ||
| 11 | uint64_t key; | ||
| 12 | int dim; | ||
| 13 | |||
| 14 | ✗ | static int GetSize(int dim) { | |
| 15 | ✗ | return sizeof(ParameterCompressItem) + dim * sizeof(float); | |
| 16 | } | ||
| 17 | 10 | const float* data() const { return embedding; } | |
| 18 | |||
| 19 | int byte_size() const { | ||
| 20 | return sizeof(ParameterCompressItem) + dim * sizeof(float); | ||
| 21 | } | ||
| 22 | float embedding[0]; // this must be the tail | ||
| 23 | }; | ||
| 24 | #pragma pack(pop) | ||
| 25 | |||
| 26 | template <> | ||
| 27 | struct Pack<ParameterCompressItem> { | ||
| 28 | static constexpr const bool implemented = true; | ||
| 29 | uint64_t key = 0; | ||
| 30 | int dim = 0; | ||
| 31 | const float* emb_data = nullptr; | ||
| 32 | Pack<ParameterCompressItem>() = default; | ||
| 33 | ✗ | Pack<ParameterCompressItem>(uint64_t key, int dim, const float* emb_data) | |
| 34 | ✗ | : key(key), dim(dim), emb_data(emb_data) {} | |
| 35 | void CompressAppend(std::string* output) const; | ||
| 36 | }; | ||
| 37 | |||
| 38 | using ParameterPack = Pack<ParameterCompressItem>; | ||
| 39 | using ParameterCompressor = FlatItemCompressor<ParameterCompressItem>; | ||
| 40 | using ParameterCompressReader = FlatItemCompressReader<ParameterCompressItem>; | ||
| 41 | |||
| 42 | 234 | inline void CopyCompressItemsToFlat( | |
| 43 | const ParameterCompressReader* reader, | ||
| 44 | float* dst, | ||
| 45 | int64_t embedding_dim, | ||
| 46 | size_t dst_row_offset = 0) { | ||
| 47 | 234 | if (reader == nullptr || dst == nullptr || embedding_dim <= 0) { | |
| 48 | ✗ | return; | |
| 49 | } | ||
| 50 | 1488 | for (int index = 0; index < reader->item_size(); ++index) { | |
| 51 | 1254 | auto item = reader->item(index); | |
| 52 | 1254 | float* row = dst + (dst_row_offset + static_cast<size_t>(index)) * | |
| 53 | 1254 | static_cast<size_t>(embedding_dim); | |
| 54 | 1254 | if (item->dim == 0) { | |
| 55 | 60 | std::memset(row, 0, static_cast<size_t>(embedding_dim) * sizeof(float)); | |
| 56 | 60 | continue; | |
| 57 | } | ||
| 58 | const int64_t copy_d = | ||
| 59 | 1194 | std::min<int64_t>(embedding_dim, static_cast<int64_t>(item->dim)); | |
| 60 | 1194 | std::memcpy( | |
| 61 | 1194 | row, item->embedding, static_cast<size_t>(copy_d) * sizeof(float)); | |
| 62 | 1194 | if (copy_d < embedding_dim) { | |
| 63 | ✗ | std::memset(row + copy_d, | |
| 64 | 0, | ||
| 65 | ✗ | static_cast<size_t>(embedding_dim - copy_d) * sizeof(float)); | |
| 66 | } | ||
| 67 | } | ||
| 68 | } | ||
| 69 | |||
| 70 | 188 | inline void CompressEmbeddingRows( | |
| 71 | ParameterCompressor* compressor, | ||
| 72 | const uint64_t* keys, | ||
| 73 | const float* values, | ||
| 74 | size_t n, | ||
| 75 | int64_t embedding_dim, | ||
| 76 | std::vector<std::string>* blocks = nullptr) { | ||
| 77 | 1874 | for (size_t i = 0; i < n; ++i) { | |
| 78 | 1686 | ParameterPack pack; | |
| 79 | 1686 | pack.key = keys[i]; | |
| 80 | 1686 | pack.dim = static_cast<int>(embedding_dim); | |
| 81 | 1686 | pack.emb_data = values + i * embedding_dim; | |
| 82 | 1686 | compressor->AddItem(pack, blocks); | |
| 83 | } | ||
| 84 | 188 | } | |
| 85 |