contrib/faiss/faiss/impl/ProductQuantizer-inl.h
Line | Count | Source |
1 | | /* |
2 | | * Copyright (c) Meta Platforms, Inc. and affiliates. |
3 | | * |
4 | | * This source code is licensed under the MIT license found in the |
5 | | * LICENSE file in the root directory of this source tree. |
6 | | */ |
7 | | |
8 | | namespace faiss { |
9 | | |
10 | | inline PQEncoderGeneric::PQEncoderGeneric( |
11 | | uint8_t* code, |
12 | | int nbits, |
13 | | uint8_t offset) |
14 | 4 | : code(code), offset(offset), nbits(nbits), reg(0) { |
15 | 4 | assert(nbits <= 64); |
16 | 4 | if (offset > 0) { |
17 | 0 | reg = (*code & ((1 << offset) - 1)); |
18 | 0 | } |
19 | 4 | } |
20 | | |
21 | 8 | inline void PQEncoderGeneric::encode(uint64_t x) { |
22 | 8 | reg |= (uint8_t)(x << offset); |
23 | 8 | x >>= (8 - offset); |
24 | 8 | if (offset + nbits >= 8) { |
25 | 0 | *code++ = reg; |
26 | |
|
27 | 0 | for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) { |
28 | 0 | *code++ = (uint8_t)x; |
29 | 0 | x >>= 8; |
30 | 0 | } |
31 | |
|
32 | 0 | offset += nbits; |
33 | 0 | offset &= 7; |
34 | 0 | reg = (uint8_t)x; |
35 | 8 | } else { |
36 | 8 | offset += nbits; |
37 | 8 | } |
38 | 8 | } |
39 | | |
40 | 4 | inline PQEncoderGeneric::~PQEncoderGeneric() { |
41 | 4 | if (offset > 0) { |
42 | 4 | *code = reg; |
43 | 4 | } |
44 | 4 | } |
45 | | |
46 | 0 | inline PQEncoder8::PQEncoder8(uint8_t* code, int nbits) : code(code) { |
47 | 0 | assert(8 == nbits); |
48 | 0 | } |
49 | | |
50 | 0 | inline void PQEncoder8::encode(uint64_t x) { |
51 | 0 | *code++ = (uint8_t)x; |
52 | 0 | } |
53 | | |
54 | | inline PQEncoder16::PQEncoder16(uint8_t* code, int nbits) |
55 | 0 | : code((uint16_t*)code) { |
56 | 0 | assert(16 == nbits); |
57 | 0 | } |
58 | | |
59 | 0 | inline void PQEncoder16::encode(uint64_t x) { |
60 | 0 | *code++ = (uint16_t)x; |
61 | 0 | } |
62 | | |
63 | | inline PQDecoderGeneric::PQDecoderGeneric(const uint8_t* code, int nbits) |
64 | 6 | : code(code), |
65 | 6 | offset(0), |
66 | 6 | nbits(nbits), |
67 | 6 | mask((1ull << nbits) - 1), |
68 | 6 | reg(0) { |
69 | 6 | assert(nbits <= 64); |
70 | 6 | } |
71 | | |
72 | 12 | inline uint64_t PQDecoderGeneric::decode() { |
73 | 12 | if (offset == 0) { |
74 | 6 | reg = *code; |
75 | 6 | } |
76 | 12 | uint64_t c = (reg >> offset); |
77 | | |
78 | 12 | if (offset + nbits >= 8) { |
79 | 0 | uint64_t e = 8 - offset; |
80 | 0 | ++code; |
81 | 0 | for (int i = 0; i < (nbits - (8 - offset)) / 8; ++i) { |
82 | 0 | c |= ((uint64_t)(*code++) << e); |
83 | 0 | e += 8; |
84 | 0 | } |
85 | |
|
86 | 0 | offset += nbits; |
87 | 0 | offset &= 7; |
88 | 0 | if (offset > 0) { |
89 | 0 | reg = *code; |
90 | 0 | c |= ((uint64_t)reg << e); |
91 | 0 | } |
92 | 12 | } else { |
93 | 12 | offset += nbits; |
94 | 12 | } |
95 | | |
96 | 12 | return c & mask; |
97 | 12 | } |
98 | | |
99 | 0 | inline PQDecoder8::PQDecoder8(const uint8_t* code, int nbits_in) : code(code) { |
100 | 0 | assert(8 == nbits_in); |
101 | 0 | } |
102 | | |
103 | 0 | inline uint64_t PQDecoder8::decode() { |
104 | 0 | return (uint64_t)(*code++); |
105 | 0 | } |
106 | | |
107 | | inline PQDecoder16::PQDecoder16(const uint8_t* code, int nbits_in) |
108 | 0 | : code((uint16_t*)code) { |
109 | 0 | assert(16 == nbits_in); |
110 | 0 | } |
111 | | |
112 | 0 | inline uint64_t PQDecoder16::decode() { |
113 | 0 | return (uint64_t)(*code++); |
114 | 0 | } |
115 | | |
116 | | } // namespace faiss |