/root/doris/contrib/faiss/faiss/impl/LookupTableScaler.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 | | #pragma once |
9 | | |
10 | | #include <cstdint> |
11 | | #include <cstdlib> |
12 | | |
13 | | #include <faiss/utils/simdlib.h> |
14 | | |
15 | | /******************************************* |
16 | | * The Scaler objects are used to specialize the handling of the |
17 | | * norm components in Additive quantizer fast-scan. |
18 | | ********************************************/ |
19 | | |
20 | | namespace faiss { |
21 | | |
22 | | /// no-op handler |
23 | | struct DummyScaler { |
24 | | static constexpr int nscale = 0; |
25 | | |
26 | 0 | inline simd32uint8 lookup(const simd32uint8&, const simd32uint8&) const { |
27 | 0 | FAISS_THROW_MSG("DummyScaler::lookup should not be called."); |
28 | 0 | return simd32uint8(0); |
29 | 0 | } |
30 | | |
31 | 0 | inline simd16uint16 scale_lo(const simd32uint8&) const { |
32 | 0 | FAISS_THROW_MSG("DummyScaler::scale_lo should not be called."); |
33 | 0 | return simd16uint16(0); |
34 | 0 | } |
35 | | |
36 | 0 | inline simd16uint16 scale_hi(const simd32uint8&) const { |
37 | 0 | FAISS_THROW_MSG("DummyScaler::scale_hi should not be called."); |
38 | 0 | return simd16uint16(0); |
39 | 0 | } |
40 | | |
41 | | #ifdef __AVX512F__ |
42 | | inline simd64uint8 lookup(const simd64uint8&, const simd64uint8&) const { |
43 | | FAISS_THROW_MSG("DummyScaler::lookup should not be called."); |
44 | | return simd64uint8(0); |
45 | | } |
46 | | |
47 | | inline simd32uint16 scale_lo(const simd64uint8&) const { |
48 | | FAISS_THROW_MSG("DummyScaler::scale_lo should not be called."); |
49 | | return simd32uint16(0); |
50 | | } |
51 | | |
52 | | inline simd32uint16 scale_hi(const simd64uint8&) const { |
53 | | FAISS_THROW_MSG("DummyScaler::scale_hi should not be called."); |
54 | | return simd32uint16(0); |
55 | | } |
56 | | #endif |
57 | | |
58 | | template <class dist_t> |
59 | | inline dist_t scale_one(const dist_t&) const { |
60 | | FAISS_THROW_MSG("DummyScaler::scale_one should not be called."); |
61 | | return 0; |
62 | | } |
63 | | }; |
64 | | |
65 | | /// consumes 2x4 bits to encode a norm as a scalar additive quantizer |
66 | | /// the norm is scaled because its range if larger than other components |
67 | | struct NormTableScaler { |
68 | | static constexpr int nscale = 2; |
69 | | int scale_int; |
70 | | simd16uint16 scale_simd; |
71 | | |
72 | 0 | explicit NormTableScaler(int scale) : scale_int(scale), scale_simd(scale) {} |
73 | | |
74 | | inline simd32uint8 lookup(const simd32uint8& lut, const simd32uint8& c) |
75 | 0 | const { |
76 | 0 | return lut.lookup_2_lanes(c); |
77 | 0 | } |
78 | | |
79 | 0 | inline simd16uint16 scale_lo(const simd32uint8& res) const { |
80 | 0 | return simd16uint16(res) * scale_simd; |
81 | 0 | } |
82 | | |
83 | 0 | inline simd16uint16 scale_hi(const simd32uint8& res) const { |
84 | 0 | return (simd16uint16(res) >> 8) * scale_simd; |
85 | 0 | } |
86 | | |
87 | | #ifdef __AVX512F__ |
88 | | inline simd64uint8 lookup(const simd64uint8& lut, const simd64uint8& c) |
89 | | const { |
90 | | return lut.lookup_4_lanes(c); |
91 | | } |
92 | | |
93 | | inline simd32uint16 scale_lo(const simd64uint8& res) const { |
94 | | auto scale_simd_wide = simd32uint16(scale_simd, scale_simd); |
95 | | return simd32uint16(res) * scale_simd_wide; |
96 | | } |
97 | | |
98 | | inline simd32uint16 scale_hi(const simd64uint8& res) const { |
99 | | auto scale_simd_wide = simd32uint16(scale_simd, scale_simd); |
100 | | return (simd32uint16(res) >> 8) * scale_simd_wide; |
101 | | } |
102 | | #endif |
103 | | |
104 | | // for non-SIMD implem 2, 3, 4 |
105 | | template <class dist_t> |
106 | 0 | inline dist_t scale_one(const dist_t& x) const { |
107 | 0 | return x * scale_int; |
108 | 0 | } Unexecuted instantiation: _ZNK5faiss15NormTableScaler9scale_oneIfEET_RKS2_ Unexecuted instantiation: _ZNK5faiss15NormTableScaler9scale_oneIhEET_RKS2_ |
109 | | }; |
110 | | |
111 | | } // namespace faiss |