1 /*
2 * Copyright (C) 2018 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "annotator/quantization.h"
18
19 #include "utils/base/logging.h"
20
21 namespace libtextclassifier3 {
22 namespace {
DequantizeValue(int num_sparse_features,int quantization_bias,float multiplier,int value)23 float DequantizeValue(int num_sparse_features, int quantization_bias,
24 float multiplier, int value) {
25 return 1.0 / num_sparse_features * (value - quantization_bias) * multiplier;
26 }
27
DequantizeAdd8bit(const float * scales,const uint8 * embeddings,int bytes_per_embedding,const int num_sparse_features,const int bucket_id,float * dest,int dest_size)28 void DequantizeAdd8bit(const float* scales, const uint8* embeddings,
29 int bytes_per_embedding, const int num_sparse_features,
30 const int bucket_id, float* dest, int dest_size) {
31 static const int kQuantizationBias8bit = 128;
32 const float multiplier = scales[bucket_id];
33 for (int k = 0; k < dest_size; ++k) {
34 dest[k] +=
35 DequantizeValue(num_sparse_features, kQuantizationBias8bit, multiplier,
36 embeddings[bucket_id * bytes_per_embedding + k]);
37 }
38 }
39
DequantizeAddNBit(const float * scales,const uint8 * embeddings,int bytes_per_embedding,int num_sparse_features,int quantization_bits,int bucket_id,float * dest,int dest_size)40 void DequantizeAddNBit(const float* scales, const uint8* embeddings,
41 int bytes_per_embedding, int num_sparse_features,
42 int quantization_bits, int bucket_id, float* dest,
43 int dest_size) {
44 const int quantization_bias = 1 << (quantization_bits - 1);
45 const float multiplier = scales[bucket_id];
46 for (int i = 0; i < dest_size; ++i) {
47 const int bit_offset = i * quantization_bits;
48 const int read16_offset = bit_offset / 8;
49
50 uint16 data = embeddings[bucket_id * bytes_per_embedding + read16_offset];
51 // If we are not at the end of the embedding row, we can read 2-byte uint16,
52 // but if we are, we need to only read uint8.
53 if (read16_offset < bytes_per_embedding - 1) {
54 data |= embeddings[bucket_id * bytes_per_embedding + read16_offset + 1]
55 << 8;
56 }
57 int value = (data >> (bit_offset % 8)) & ((1 << quantization_bits) - 1);
58 dest[i] += DequantizeValue(num_sparse_features, quantization_bias,
59 multiplier, value);
60 }
61 }
62 } // namespace
63
CheckQuantizationParams(int bytes_per_embedding,int quantization_bits,int output_embedding_size)64 bool CheckQuantizationParams(int bytes_per_embedding, int quantization_bits,
65 int output_embedding_size) {
66 if (bytes_per_embedding * 8 / quantization_bits < output_embedding_size) {
67 return false;
68 }
69
70 return true;
71 }
72
DequantizeAdd(const float * scales,const uint8 * embeddings,int bytes_per_embedding,int num_sparse_features,int quantization_bits,int bucket_id,float * dest,int dest_size)73 bool DequantizeAdd(const float* scales, const uint8* embeddings,
74 int bytes_per_embedding, int num_sparse_features,
75 int quantization_bits, int bucket_id, float* dest,
76 int dest_size) {
77 if (quantization_bits == 8) {
78 DequantizeAdd8bit(scales, embeddings, bytes_per_embedding,
79 num_sparse_features, bucket_id, dest, dest_size);
80 } else if (quantization_bits != 8) {
81 DequantizeAddNBit(scales, embeddings, bytes_per_embedding,
82 num_sparse_features, quantization_bits, bucket_id, dest,
83 dest_size);
84 } else {
85 TC3_LOG(ERROR) << "Unsupported quantization_bits: " << quantization_bits;
86 return false;
87 }
88
89 return true;
90 }
91
92 } // namespace libtextclassifier3
93