对torchac codec的详解

torchac-backend.cpp出处

/** This code is from https://github.com/fab-jul/torchac/blob/master/torchac/backend/torchac_backend.cpp, the core codec part of torchac 
    A lot of thanks for the author's generous work.
**/

/** Encode symbols `sym` with CDF represented by `cdf_ptr`. NOTE: this is not exposted to python. */
py::bytes encode(
        const cdf_ptr& cdf_ptr, //概率分布函数
        const torch::Tensor& sym){

    OutCacheString out_cache;

    uint32_t low = 0;
    uint32_t high = 0xFFFFFFFFU; //初始编码区间
    uint64_t pending_bits = 0; 
    //将不能立即确定的编码位延后，从而提升效率，时间复杂度最好为O(n)

    const int precision = 16;

    /* cdf是序列式的，即序列中每个位置有Lp个符号的累积概率，
     第i个位置的概率分布反映为cdf[i * Lp]到cdf[(i + 1) * Lp - 1] */
    const cdf_t* cdf = cdf_ptr.data;
    const int N_sym = cdf_ptr.N_sym;
    const int Lp = cdf_ptr.Lp;
    const int max_symbol = Lp - 2;

    auto sym_ = sym.accessor<int16_t, 1>();

    for (int i = 0; i < N_sym; ++i) {
        const int16_t sym_i = sym_[i];

        const uint64_t span = static_cast<uint64_t>(high) - static_cast<uint64_t>(low) + 1; //当前区间跨度

        const int offset = i * Lp;
        // Left boundary is at offset + sym_i
        const uint32_t c_low = cdf[offset + sym_i]; //符号概率区间的下界
        // Right boundary is at offset + sym_i + 1, except for the `max_symbol`
        // For which we hardcode the maxvalue. So if e.g.
        // L == 4, it means that Lp == 5, and the allowed symbols are
        // {0, 1, 2, 3}. The max symbol is thus Lp - 2 == 3. It's probability
        // is then given by c_max - cdf[-2].
        const uint32_t c_high = sym_i == max_symbol ? 0x10000U : cdf[offset + sym_i + 1];//符号概率区间的上界

        high = (low - 1) + ((span * static_cast<uint64_t>(c_high)) >> precision); //按照符号概率细分编码区间
        low =  (low)     + ((span * static_cast<uint64_t>(c_low))  >> precision);

        while (true) { // 所有的移位和位操作都是为了重归一化以保持数值精度
            // pending比特是防止数据一直无法确定输出使得区间越来越小直到精度不足，编码失败
            // pending_bits决定为当前决定的最高bit的否定，如0111、1000，体现出它接近0.5的区间特点
            if (high < 0x80000000U) {
                // 概率区间小于0.5时，最高位为0
                out_cache.append_bit_and_pending(0, pending_bits);
                low <<= 1;
                high <<= 1; //根据最高位确定了为0，把low、high都左移一位，即去掉了最高位的0，继续决定
                high |= 1; // 新增的一位填成1
            } else if (low >= 0x80000000U) {
                // 概率区间大于0.5，最高位为1
                out_cache.append_bit_and_pending(1, pending_bits);
                low <<= 1;
                high <<= 1;
                high |= 1;
            } else if (low >= 0x40000000U && high < 0xC0000000U) {
                // 概率区间在0.25到0.75之间，需要考察后续位，无法立即决定，需要后续位来进一步确定区间
                // 0x40000000U分割00与01，0xC0000000U分割10与11
                pending_bits++;
                low <<= 1;
                low &= 0x7FFFFFFF; // low左移1位后的最高位已经在比较中被考虑过了，所以清除
                high <<= 1;
                high |= 0x80000001; // 确保high的最高位为1，新增的一位填成1
                // 两个位操作将0.25-0.75区间重映射到了0-1
            } else {
                break;
            }
        }
    }

    pending_bits += 1;

    if (pending_bits) {
        if (low < 0x40000000U) {
            out_cache.append_bit_and_pending(0, pending_bits);
        } else {
            out_cache.append_bit_and_pending(1, pending_bits);
        }
    }

    out_cache.flush();

#ifdef VERBOSE
    std::chrono::steady_clock::time_point end= std::chrono::steady_clock::now();
    std::cout << "Time difference (sec) = " << (std::chrono::duration_cast<std::chrono::microseconds>(end - begin).count()) /1000000.0 <<std::endl;
#endif

    return py::bytes(out_cache.out);
}

torch::Tensor decode(
        const cdf_ptr& cdf_ptr,
        const std::string& in) {

#ifdef VERBOSE
    std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
#endif

    const cdf_t* cdf = cdf_ptr.data;
    const int N_sym = cdf_ptr.N_sym;  // To know the # of syms to decode. Is encoded in the stream!
    const int Lp = cdf_ptr.Lp;  // To calculate offset
    const int max_symbol = Lp - 2;

    // 16 bit!
    auto out = torch::empty({N_sym}, torch::kShort);
    auto out_ = out.accessor<int16_t, 1>();

    uint32_t low = 0;
    uint32_t high = 0xFFFFFFFFU;
    uint32_t value = 0;
    const uint32_t c_count = 0x10000U;
    const int precision = 16;

    InCacheString in_cache(in);
    in_cache.initialize(value);

    for (int i = 0; i < N_sym; ++i) {
        const uint64_t span = static_cast<uint64_t>(high) - static_cast<uint64_t>(low) + 1;
        // always < 0x10000 ???
        // 01编码得到概率点
        const uint16_t count = ((static_cast<uint64_t>(value) - static_cast<uint64_t>(low) + 1) * c_count - 1) / span;
        
        const int offset = i * Lp;
        // 概率点二分查找得到第一个符号
        auto sym_i = binsearch(cdf, count, (cdf_t)max_symbol, offset);

        out_[i] = (int16_t)sym_i;

        if (i == N_sym-1) {
            break;
        }

        // 找到第一个符号后，更新解码区间到它对应的子区间（缩小区间
        const uint32_t c_low = cdf[offset + sym_i]; 
        const uint32_t c_high = sym_i == max_symbol ? 0x10000U : cdf[offset + sym_i + 1];

        high = (low - 1) + ((span * static_cast<uint64_t>(c_high)) >> precision);
        low =  (low)     + ((span * static_cast<uint64_t>(c_low))  >> precision);

        // 重归一化以放大区间为下一个符号的解码做准备
        while (true) {
            if (low >= 0x80000000U || high < 0x80000000U) { // 区间较小，重归一化以放大区间，消费比特流
                low <<= 1;
                high <<= 1;
                high |= 1;
                in_cache.get(value); // 从比特流读取下一位
            } else if (low >= 0x40000000U && high < 0xC0000000U) {
                /**
                 * 0100 0000 ... <= value <  1100 0000 ...
                 * <=>
                 * 0100 0000 ... <= value <= 1011 1111 ...
                 * <=>
                 * value starts with 01 or 10.
                 * 01 - 01 == 00  |  10 - 01 == 01
                 * i.e., with shifts
                 * 01A -> 0A  or  10A -> 1A, i.e., discard 2SB as it's all the same while we are in
                 *    near convergence
                 */
                low <<= 1;
                low &= 0x7FFFFFFFU;  // make MSB 0
                high <<= 1;
                high |= 0x80000001U;  // add 1 at the end, retain MSB = 1
                value -= 0x40000000U; // 重归一化
                in_cache.get(value);
            } else { // 区间足够大，可以解码下一个符号
                break;
            }
        }
    }

#ifdef VERBOSE
    std::chrono::steady_clock::time_point end= std::chrono::steady_clock::now();
    std::cout << "Time difference (sec) = " << (std::chrono::duration_cast<std::chrono::microseconds>(end - begin).count()) /1000000.0 <<std::endl;
#endif

    return out;
}