H264/H265 NALU 起始码搜索性能优化(1)

码农心语

发布于 2024-04-09 15:56:33

870

发布于 2024-04-09 15:56:33

文章被收录于专栏：码农心语

在ffmpeg中，在进行h264 rbsp流demux的时候，需要进行starting code的搜索，其采用的方法比较简单，就是不断比较字节流中连续的三个字节，是不是 0x00, 0x00, 0x01，ffmpeg采用如下代码用来找到各个NALU的分界点：

static int find_next_start_code(const uint8_t *buf, const uint8_t *next_avc)
{

    int i = 0;

    if (buf + 3 >= next_avc)

        return next_avc - buf;

    while (buf + i + 3 < next_avc) {

        if (buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1)
            break;

        i++;

    }
    return i + 3;
}

这种方法比较传统，一般来说如果对于性能没有特别的要求，是能够很好满足要求了。

按照ffmpeg的算法，搜索一个内存块中的所有H264 NALU的起始码，并将起始码的位置写入nalu_borders数组中，写出来的代码如下，整个逻辑简单明了：

vector<int> split_simd_ffmpeg(const uint8_t *data, size_t cb)
{
    const uint8_t *p = data;

    const uint8_t *e = data + cb;

    vector<int> list_index;

    while (p < e -3){
        if (p[0] == 0 && p[1] == 0 && p[2] == 1){
            list_index.emplace_back(p - data);
        }
        p++;
    }

    return list_index;
}

本文提出一种更高效的starting code的搜索算法：

void split_nalu(char *video_data, int size,  std::vector<int> &nalu_borders)
{

    while(index + 8 <= (size_t)size){

        uint64_t code = *(uint64_t*)(video_data+index);

        if ((code & 0xFFFFFFull) == 0x010000ull){
            nalu_borders.emplace_back(index);
            index += 3;
            continue;
        }
                    
        if ((code & 0xFFFFFF00ull) == 0x01000000ull){
            nalu_borders.emplace_back(index+1);
            index += 4;
            continue;
        }

                    
        if ((code & 0xFFFFFF0000ull) == 0x0100000000ull){
            nalu_borders.emplace_back(index+2);
            index += 5;
            continue;
        }
                    
        if (((code & 0xFFFFFF000000ull) == 0x010000000000ull){
            nalu_borders.emplace_back(index+3);
            index += 6;
            continue;
        }

        if ((code & 0xFFFFFF00000000ull) ==  0x01000000000000ull){
            nalu_borders.emplace_back(index+4);
            index += 7;
            continue;
        }

        if ((code & 0xFFFFFF0000000000ull) ==  0x0100000000000000ull){
            nalu_borders.emplace_back(index+5);
            index += 8;
            continue;
        }

        if ((code & 0xFFFF000000000000ull) ==  0x0ull){
            index += 6;
        }else if ((code & 0xFF00000000000000ull) ==  0x0ull){
            index += 7;
        }
        else{
            index += 8;
        }
    }

    while(index + 4 < (size_t)size){

        uint32_t code = *(uint64_t*)(video_data+index);

        if ((code & 0x00FFFFFFu) == 0x00010000u){
            nalu_borders.emplace_back(index);
            index += 3;
            continue;
        }

        index++;
    }
}

看上去代码复杂了很多，但是实际测试结果，对比ffmpeg的搜索算法，在x64的CPU的机器上面，性能至少提升了30%。

本算法利用了64位的寄存器预先读取8个字节，然后在循环提内通过掩码与操作进行比对，一个循环中对8个字节进行264起始码的搜索，相比ffmpeg需要不断从内存中load，然后逐字节比对来说，使得搜索性能得到了显著提升。毕竟寄存器的性能远比内存的读写性能要强多了！而且在一个循环里面可以通过编译器和CPU的优化进行指令的多路并发执行。

然后最近几天拼命想利用SIMD优化指令进行搜索的算法，好不容易熟悉了simd的几个指令，调试了老半天，终于写出来如下代码：

    const uint8_t mask_1[32] = { 0xFF, 0xFF, 0xFF, 0x00,
                                 0xFF, 0xFF, 0xFF, 0x00,  
                                 0xFF, 0xFF, 0xFF, 0x00,  
                                 0xFF, 0xFF, 0xFF, 0x00  };

    const uint8_t dest_1[32] = { 0, 0, 1, 0,   0, 0, 1, 0,
                                 0, 0, 1, 0,   0, 0, 1, 0 };


    const uint8_t mask_2[32] = { 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,};

    const uint8_t dest_2[32] = { 0, 0, 0, 1,   0, 0, 0, 1,
                                 0, 0, 0, 1,   0, 0, 0, 1};

vector<int> split_nalu_simd(const uint8_t *video_data, size_t size)
{
    size_t index = 0;
    vector<int> list_index;
    
    op_data od;

    size_t next_block = 31;

    while(index + 16 <= (size_t)size){
        
        __m128i src  = _mm_load_si128((__m128i*)(video_data+index));
        __m128i mask = _mm_load_si128((__m128i*)od.mask_1);
        __m128i dest = _mm_load_si128((__m128i*)od.dest_1);
        __m128i tmp  = _mm_and_si128(src, mask);
        __m128i tmp2 = _mm_cmpeq_epi32(tmp, dest);

        if (_mm_extract_epi32(tmp2, 0) == 0xFFFFFFFFu){
            list_index.emplace_back(index);
        }
        
        if (_mm_extract_epi32(tmp2, 1) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 4);
        }
        
        if (_mm_extract_epi32(tmp2, 2) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 8);
        }
        
        if (_mm_extract_epi32(tmp2, 3) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 12);
        }



        mask = _mm_load_si128((__m128i*)od.mask_2);
        dest = _mm_load_si128((__m128i*)od.dest_2);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);

        if (_mm_extract_epi32(tmp2, 0) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 1);
        }
        
        if (_mm_extract_epi32(tmp2, 1) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 5);
        }
        
        if (_mm_extract_epi32(tmp2, 2) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 9);
        }
        
        if (_mm_extract_epi32(tmp2, 3) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 13);
        }


        src = _mm_srli_si128(src, 2);

        mask = _mm_load_si128((__m128i*)od.mask_1);
        dest = _mm_load_si128((__m128i*)od.dest_1);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);

        if (_mm_extract_epi32(tmp2, 0) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 2);
        }
        
        if (_mm_extract_epi32(tmp2, 1) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 6);
        }
        
        if (_mm_extract_epi32(tmp2, 2) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 10);
        }


        mask = _mm_load_si128((__m128i*)od.mask_2);
        dest = _mm_load_si128((__m128i*)od.dest_2);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);


        if (_mm_extract_epi32(tmp2, 0) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 3);
        }
        
        if (_mm_extract_epi32(tmp2, 1) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 7);
        }
        
        if (_mm_extract_epi32(tmp2, 2) == 0xFFFFFFFFu){
            list_index.emplace_back(index + 11);
        }


        index += 16;
        uint32_t tail = _mm_extract_epi32(src, 3);



        if (tail == 0x0u){
            if (index + 1 <= size && video_data[index] == 0x01){
                list_index.emplace_back(index-2);
                continue;
            }
        }

        if ((tail & 0x0000FF00u) == 0x0u){
            if (index + 2 <= size && video_data[index] == 0x00 && video_data[index+1] == 0x01){
                list_index.emplace_back(index-1);
            }
        }
    }

    while(index + 4 < (size_t)size){

        uint32_t code = *(uint64_t*)(video_data+index);

        if ((code & 0x00FFFFFFu) == 0x00010000u){
            list_index.emplace_back(index);
            index += 3;
            continue;
        }

        index++;
    }    
    
    std::sort(list_index.begin(), list_index.end());

    return list_index;
}

但是测试下来的效果不太理想，大失所望，基本上和上面的split_nalu不相上下，应该是中间用了太多条件判断的分支，抵消了simd的优势，加上插入起始码位置的时候因为并不是按顺序插入的，需要在结尾处进行一次sort操作一定程度上会引起性能的降低，得不偿失。

感觉能够想到的方法还是split_nalu比较靠谱，也不是太复杂。

今天下午，突然来了灵感，修改了上面的simd的代码，减少了if条件判断分支，尽然性能提升1倍，开心！！！！

    const uint8_t mask_1[32] = { 0xFF, 0xFF, 0xFF, 0x00,
                                 0xFF, 0xFF, 0xFF, 0x00,  
                                 0xFF, 0xFF, 0xFF, 0x00,  
                                 0xFF, 0xFF, 0xFF, 0x00  };

    const uint8_t dest_1[32] = { 0, 0, 1, 0,   0, 0, 1, 0,
                                 0, 0, 1, 0,   0, 0, 1, 0 };


    const uint8_t mask_2[32] = { 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,
                                 0x00, 0xFF, 0xFF, 0xFF,};

    const uint8_t dest_2[32] = { 0, 0, 0, 1,   0, 0, 0, 1,
                                 0, 0, 0, 1,   0, 0, 0, 1};

vector<int> split_nalu_simd2(const uint8_t *video_data, size_t size)
{
    size_t index = 0;
    vector<int> list_index;
    
    op_data od;

    size_t next_block = 31;

    while(index + 16 <= (size_t)size){
        
        __m128i src  = _mm_load_si128((__m128i*)(video_data+index));
        __m128i mask = _mm_load_si128((__m128i*)od.mask_1);
        __m128i dest = _mm_load_si128((__m128i*)od.dest_1);
        __m128i tmp  = _mm_and_si128(src, mask);
        __m128i tmp2 = _mm_cmpeq_epi32(tmp, dest);
        
        int m = _mm_movemask_epi8(tmp2);
        if (m != 0){
            if (m & 0x0F){
                list_index.emplace_back(index);
            }
            
            if (m & 0xF0){
                list_index.emplace_back(index + 4);
            }
            
            if (m & 0xF00){
                list_index.emplace_back(index + 8);
            }
            
            if (m & 0xF000){
                list_index.emplace_back(index + 12);
            }
        }

        mask = _mm_load_si128((__m128i*)od.mask_2);
        dest = _mm_load_si128((__m128i*)od.dest_2);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);

        m = _mm_movemask_epi8(tmp2);
        if (m != 0){
            if (m & 0x0F){
                list_index.emplace_back(index + 1);
            }
            
            if (m & 0xF0){
                list_index.emplace_back(index + 5);
            }
            
            if (m & 0xF00){
                list_index.emplace_back(index + 9);
            }
            
            if (m & 0xF000){
                list_index.emplace_back(index + 13);
            }
        }

        src = _mm_srli_si128(src, 2);

        mask = _mm_load_si128((__m128i*)od.mask_1);
        dest = _mm_load_si128((__m128i*)od.dest_1);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);

        m = _mm_movemask_epi8(tmp2);
        if (m != 0){
            if (m & 0x0F){
                list_index.emplace_back(index + 2);
            }
            
            if (m & 0xF0){
                list_index.emplace_back(index + 6);
            }
            
            if (m & 0xF00){
                list_index.emplace_back(index + 10);
            }
        }

        mask = _mm_load_si128((__m128i*)od.mask_2);
        dest = _mm_load_si128((__m128i*)od.dest_2);
        tmp  = _mm_and_si128(src, mask);
        tmp2 = _mm_cmpeq_epi32(tmp, dest);

        m = _mm_movemask_epi8(tmp2);
        if (m != 0){
            if (m & 0x0F){
                list_index.emplace_back(index + 3);
            }
            
            if (m & 0xF0){
                list_index.emplace_back(index + 7);
            }
            
            if (m & 0xF00){
                list_index.emplace_back(index + 11);
            }
        }


        index += 16;
        uint32_t tail = _mm_extract_epi32(src, 3);



        if (tail == 0x0u){
            if (index + 1 <= size && video_data[index] == 0x01){
                list_index.emplace_back(index-2);
                continue;
            }
        }

        if ((tail & 0x0000FF00u) == 0x0u){
            if (index + 2 <= size && video_data[index] == 0x00 && video_data[index+1] == 0x01){
                list_index.emplace_back(index-1);
            }
        }
    }

    while(index + 4 < (size_t)size){

        uint32_t code = *(uint64_t*)(video_data+index);

        if ((code & 0x00FFFFFFu) == 0x00010000u){
            list_index.emplace_back(index);
            index += 3;
            continue;
        }

        index++;
    }    
    
    std::sort(list_index.begin(), list_index.end());

    return list_index;
}

关键是是通过_mm_movemask_epi8将前面比较的结果合并到一个int型的字段m中，然后直接判断m是否都为0，如果为0就可以直接掉过里面的分支逻辑了，大大减少了条件分支，从而提升了性能。

ps：

需要提一下，如果没有开启-O3编译器优化，split_nalu_simd2比split_nalu版本性能差了一倍，开启以后则反回过了，虽然split_nalu性能也提升了，但是split_nalu_simd2提升非常明显。

假设待编译的代码为test.cpp，那么用如下编译指令：

g++ test.cpp -mavx2 -o test -g -O3

-mavx2是必须要加的，否则编译器不能支持simd相关指令。

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2024-03-07，如有侵权请联系 cloudcommunity@tencent.com 删除

list