Optimizing Huffman Decoding

Question

I've been working on implementing a fast Huffman decoder for video decoding. With some help from you, I have been able to produce a decent implementation. However I am still not satisfied with the results. I have done some research both on Stack Overflow and some research papers, however I haven't found anything to further improve my current implementation.

I'm using a lookup-table for decoding. I've removed most of the branches and byte moves (MOVZX). And also aligned my data for optimal cache usage.

I think I've got the decoding itself as fast as it can get, however building the lookup-table (build_tables2) takes as much time as the decoding itself.

My profiler gives me:

• decoding: 52%
• build lookup-table (build_tables2): 45%

In average my decoder is building ~100 tables per frame (of max 256, one for each node in the biggest tree), which is much but I don't see how I could reduce it. 16 bit tables seems out of the question.

Does anyone have any ideas as to how I can speed up building the lookup-table? Maybe I could somehow encode some kind of additional data in the frame/"input" which would help for building the tables?

typedef uint8_t block_type;
static const int bits_per_block = 8;
static const int values_per_block = 256;

struct node_t
{
    node_t* children; // 0 right, 1 left
    int32_t value;    // if leaf then its the value of the current node, if not leaf then its the index of the next lookup-table
    int32_t is_leaf;
};

// NOTE:
// entry is 16 bytes and will be aligned during allocation, thus "values" will always be on 16 byte alignment
struct entry 
{
    uint8_t values[bits_per_block];
    std::array<entry, values_per_block>* next_table;
    int32_t num_values;
};

typedef std::array<entry, values_per_block>     table_t;
typedef std::array<table_t, values_per_block-1> tables_t;

void build_tables2(node_t* nodes, tables_t& tables, int& table_count);

void unpack_tree(void* data, node_t* nodes);

void decode_huff(void* input, uint8_t* dest)
{
    // Initial setup
    std::shared_ptr<node_t> nodes(reinterpret_cast<node_t*>(scalable_aligned_malloc(sizeof(node_t)*(513), 64)), scalable_aligned_free); // cacheline alignment  
    memset_sse(nodes.get(), 0, sizeof(node_t)*512);

    auto data = reinterpret_cast<uint32_t*>(input); 
    size_t table_size   = *(data++); // Size is first 32 bits.
    size_t result_size  = *(data++); // Data size is second 32 bits.

    // Unpack huffman-tree.
    unpack_tree(data, nodes.get());

    if(nodes->children != nullptr)
    {           
        // BUILD LOOKUP-TABLE   

        std::shared_ptr<tables_t> tables(reinterpret_cast<tables_t*>(scalable_aligned_malloc(sizeof(tables_t), 64)), scalable_aligned_free); // cacheline alignment 

        int table_count = 0;

        // Build lookup-table from huffman-tree.
        build_tables2(nodes.get(), *tables, table_count);

        // INIT DECODING

        auto huffman_data = data + table_size/32; 
        auto ptr          = dest;
        auto end          = ptr + result_size;

        entry entry2;
        entry2.next_table = tables->data();

        // NOTE:
        // Use 2x MOVQ instead of memcpy, could use MOVDQA and MOVDQU but we only care about 64 bits, thus we use MOVQ which might be faster?
        // _mm_storel_epi64(reinterpret_cast<__m128i*>(ptr), _mm_loadl_epi64(reinterpret_cast<__m128i*>(entry))); load is aligned, store is unaligned
        // <=>
        // memcpy(ptr, entry, entry->num_values);
        //
        // The decoding loop will always be overwriting some values ahead (which haven't been decoded), this means
        // that the target memory buffer will need to have extra padding in the end to avoid access violations.

        // DECODING

        auto entry = &entry2;
        while(ptr < end)
        {
            auto elem = *(huffman_data++); // Use single mov and shift and masking, instead of 4x movzx/byte-moves

            entry = entry->next_table->data() + (elem & 0xff); // entry = &entry->next_table->at(elem & 0xff);
            _mm_storel_epi64(reinterpret_cast<__m128i*>(ptr), _mm_loadl_epi64(reinterpret_cast<__m128i*>(entry)));          
            ptr += entry->num_values;

            entry = entry->next_table->data() + ((elem >> 8) & 0xff);
            _mm_storel_epi64(reinterpret_cast<__m128i*>(ptr), _mm_loadl_epi64(reinterpret_cast<__m128i*>(entry)));          
            ptr += entry->num_values;

            entry = entry->next_table->data() + ((elem >> 16) & 0xff);
            _mm_storel_epi64(reinterpret_cast<__m128i*>(ptr), _mm_loadl_epi64(reinterpret_cast<__m128i*>(entry)));          
            ptr += entry->num_values;

            entry = entry->next_table->data() + (elem >> 24);
            _mm_storel_epi64(reinterpret_cast<__m128i*>(ptr), _mm_loadl_epi64(reinterpret_cast<__m128i*>(entry)));          
            ptr += entry->num_values;
        }
    }
    else
        memset_sse(dest, static_cast<char>(nodes->value), result_size);
}

void build_tables2(node_t* nodes, tables_t& tables, int& table_count)
{
    // Using a stack instead of function recursion helps the profiler.
    std::stack<node_t*> stack;
    stack.push(nodes);

    while(!stack.empty())
    {
        auto root = stack.top();
        stack.pop();

        auto& table = tables[root->value];

        for(int n = 0; n < values_per_block; ++n)
        {
            auto current = root;

            auto& entry = table[n];

            entry.num_values = 0;

            // Looped and branched version
            //for(int i = 0; i < 8; ++i)
            //{
            //  current = current->children + ((n >> i) & 1);       
            //  if(current->is_leaf)
            //      entry.values[entry.num_values++] = current->value;      
            //}

            // NOTE:
            //*reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            // <=>
            // entry.values[entry.num_values] = current->value;
            //
            // Avoids byte moves, the last write can overwrite the next element in the entry structure, "next_table",
            // this is not a problem since we set "next_table" below and don't care about its current value

            current = current->children + (n & 1); // +0 is left and +1 is right child      
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value; // dummy write if not leaf
            entry.num_values += current->is_leaf; // only increment if it was a leaf

            current = current->children + ((n >> 1) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            current = current->children + ((n >> 2) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            current = current->children + ((n >> 3) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;


            current = current->children + ((n >> 4) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            current = current->children + ((n >> 5) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            current = current->children + ((n >> 6) & 1);       
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            current = current->children + (n >> 7);     
            *reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;
            entry.num_values += current->is_leaf;

            if(!current->is_leaf)
            {
                // a non-leaf, then current->value is the next lookup-table
                if(current->value == 0)
                {
                    current->value = ++table_count;
                    stack.push(current);
                }

                table[n].next_table = &tables[current->value];
            }
            else
                table[n].next_table = &tables[0]; // All bits were used, back to root
        }   
    }
}

Answer 1

Do you have some control over the input encoder? If so it should be providing you with a Huffman tree delta, rather than a whole new tree for every frame. In the same way that the HT algorithm applies similarity rules to encoded data, from frame to frame the HT will have also have a very high degree of similarity.

Also some efficiency and platform concerns: Please rewrite expressions of the form

((n >> #) & 1)

where # is a digit as,

(((n & (1 << #)) != 0) ? 1 : 0)

Compile time versus run time calculation is always going to be faster.

*reinterpret_cast<int32_t*>(entry.values + entry.num_values) = current->value;

Is dependent on the little endian architecture. Prefer correct over fast.

auto* next_value = entry.values; // get the starting value element address
...
*next_value = static_cast<uint_8>(current->value); // set its value
next_value += current->is_leaf;
...
// at the end of the loop unrolling do this calculate num_value 
// from the offset, thus avoiding multiple dereferences of the struct member
entry.num_values = next_value - entry.values;