optimizer/sparse_tensor.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include <limits> | ||
| 2 | #include <stdexcept> | ||
| 3 | #include <string> | ||
| 4 | #include <vector> | ||
| 5 | #include "../storage/kv_engine/base_kv.h" | ||
| 6 | |||
| 7 | #define TAG_TYPE uint8_t | ||
| 8 | |||
| 9 | enum TensorType { PARAMETER = 0, MOMENT_1 = 1, MOMENT_2 = 2 }; | ||
| 10 | |||
| 11 | // ponytail: 8-bit tag, widen TAG_TYPE if table_id * tensors_per_table > 255 | ||
| 12 | 12 | inline constexpr int KeyTagShift() { | |
| 13 | return static_cast<int>(sizeof(uint64_t) * 8) - | ||
| 14 | 12 | static_cast<int>(sizeof(TAG_TYPE) * 8); | |
| 15 | } | ||
| 16 | |||
| 17 | 12 | inline TAG_TYPE ExtractKeyTag(uint64_t key) { | |
| 18 | 12 | return static_cast<TAG_TYPE>(key >> KeyTagShift()); | |
| 19 | } | ||
| 20 | |||
| 21 | inline TAG_TYPE | ||
| 22 | 32 | MakeTensorTag(uint64_t table_id, int role, int tensors_per_table) { | |
| 23 | 32 | if (tensors_per_table <= 0 || role < 0 || role >= tensors_per_table) { | |
| 24 | ✗ | throw std::runtime_error("invalid embedding tensor tag role"); | |
| 25 | } | ||
| 26 | 32 | const uint64_t tag = table_id * static_cast<uint64_t>(tensors_per_table) + | |
| 27 | 32 | static_cast<uint64_t>(role); | |
| 28 | 32 | if (tag > std::numeric_limits<TAG_TYPE>::max()) { | |
| 29 | ✗ | throw std::runtime_error("embedding table tag overflow"); | |
| 30 | } | ||
| 31 | 32 | return static_cast<TAG_TYPE>(tag); | |
| 32 | } | ||
| 33 | |||
| 34 | class SparseTensor { | ||
| 35 | private: | ||
| 36 | std::string name; | ||
| 37 | TensorType type; | ||
| 38 | TAG_TYPE tag; | ||
| 39 | std::vector<uint64_t> shape; | ||
| 40 | BaseKV* kv; | ||
| 41 | uint64_t concatKeyAndTag(uint64_t key, TAG_TYPE tag); | ||
| 42 | |||
| 43 | public: | ||
| 44 | 48 | SparseTensor() = default; | |
| 45 | void init(std::string& name, | ||
| 46 | TensorType type, | ||
| 47 | TAG_TYPE tag, | ||
| 48 | std::vector<uint64_t>& shape, | ||
| 49 | BaseKV* kv); | ||
| 50 | void Get(const uint64_t key, std::string& value, unsigned tid); | ||
| 51 | void Put(const uint64_t key, const std::string_view& value, unsigned tid); | ||
| 52 | void BatchGet(const std::vector<uint64_t>& keys, | ||
| 53 | std::vector<base::ConstArray<float>>* values, | ||
| 54 | unsigned tid); | ||
| 55 | int64_t EmbeddingDim() const; | ||
| 56 | 12 | TAG_TYPE Tag() const { return tag; } | |
| 57 | bool ApplySgdUpdateFlat( | ||
| 58 | const base::ConstArray<uint64_t>& keys, | ||
| 59 | const float* grads, | ||
| 60 | int64_t num_rows, | ||
| 61 | int64_t embedding_dim, | ||
| 62 | float learning_rate, | ||
| 63 | unsigned tid, | ||
| 64 | BaseKV::UpdateFlatStats* stats = nullptr); | ||
| 65 | }; | ||
| 66 |