1 /** 2 * Copyright 2021 Huawei Technologies Co., Ltd 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 #ifndef MINDSPORE_LITE_TOOLS_CONVERTER_QUANTIZER_FSEBITSTREAM_H 18 #define MINDSPORE_LITE_TOOLS_CONVERTER_QUANTIZER_FSEBITSTREAM_H 19 #include <cstdint> 20 21 namespace mindspore::lite::quant { 22 class BitStream { 23 public: 24 BitStream() = default; 25 26 ~BitStream() = default; 27 28 public: 29 int Create(int bit_capacity); 30 void Free(); 31 void Empty(); 32 int64_t Pop(uint8_t bit_count); 33 void Push(int64_t state, uint8_t bit_count); 34 void Flush(); 35 GetCurrChunkIndex()36 int32_t GetCurrChunkIndex() { return this->curr_chunk_index_; } GetCurrChunk()37 uint64_t GetCurrChunk() { return this->curr_chunk_; } GetCurrBitCount()38 int8_t GetCurrBitCount() { return this->curr_bit_count_; } GetChunks()39 uint64_t *GetChunks() { return this->chunks_; } GetChunkCount()40 int GetChunkCount() { return this->chunk_count_; } 41 SetCurrChunkIndex(int32_t curr_chunk_index)42 void SetCurrChunkIndex(int32_t curr_chunk_index) { this->curr_chunk_index_ = curr_chunk_index; } SetCurrChunk(uint64_t curr_chunk)43 void SetCurrChunk(uint64_t curr_chunk) { this->curr_chunk_ = curr_chunk; } SetCurrBitCount(int8_t curr_bit_count)44 void SetCurrBitCount(int8_t curr_bit_count) { this->curr_bit_count_ = curr_bit_count; } SetChunks(uint64_t * chunks)45 void SetChunks(uint64_t *chunks) { this->chunks_ = chunks; } SetChunkCount(int chunk_count)46 void SetChunkCount(int chunk_count) { this->chunk_count_ = chunk_count; } 47 48 private: 49 int32_t curr_chunk_index_{-1}; // the index of the next chunk that we will write to 50 uint64_t curr_chunk_{0}; 51 int8_t curr_bit_count_{0}; // the number of bits that are currently written in the register. 52 uint64_t *chunks_{nullptr}; // the actual memory 53 int chunk_count_{0}; // the number of chunks 54 }; 55 } // namespace mindspore::lite::quant 56 #endif // MINDSPORE_LITE_TOOLS_CONVERTER_QUANTIZER_ 57