以 halide 實作 block matching 演算法 - block matching with Halide language

由於 multi-frame 與 multi-view 的影像處理需求的提升,  image alignment 的需求日益提升, 而相關的演算也很多, 像是基於 pyramid-based, 9-steps, hexagon or diamond search. 而大小的部份也有著 match / apply 不同大小的 non-local means, 然而這些演算的基本還是源自於特定 search window 的 block matching. 本文重點在於以 Halide language 描述 matching algorithm, 至於 scheduling 的優化不在此文重點. (或是擇日分享)

Step 0 - 輸入格式 (input format - single channel as example)
首先假設要處理的是 image 的某個 channel, multi-channel 的影像則僅是套用到多個 frame 或是以 channel 輸入, 同常有一張是 base image 這裡以 in 代表, 另一張為 reference image 以 ref 代表.

Step 1 - 邊界處理, 用來避免 search block 範圍超過邊界 (boundary handling, to avoid error of part of target block is out of frame)

Func input = BoundaryConditions::repeat_edge(in);
Func refer = BoundaryConditions::repeat_edge(ref);

Step 2 - 用以代表 block 內所有 input 與 reference 的 pixel 的座標 (setup the RDom that represent all pixels of corresponding input & refer block)

// the position in motion vector output table

Var bid_x, bid_y;

// variable for motion vector

Var mv_x, mv_y;

RDom rblk(0, BLOCK_SIZE, 0, BLOCK_SIZE)
// position in input frame
Expr blk_x = (bid_x * BLOCK_SIZE) + rblk.x;

Expr blk_y = (bid_y * BLOCK_SIZE) + rblk.y;
// position in reference frame

Expr ref_x = blk_x + mv_x;
Expr ref_y = blk_y + mv_y;

Step 3 - 建立比對標準, 這裡以 SAD 為例 ( matching criteria calculation (Sum of Absolute Difference(SAD) as example))

Expr in_val = i16(input(blk_x, blk_y);

Expr ref_val = i16(refer(ref_x, ref_y));

Func sad;

sad(mv_x, mv_y, bid_x, bid_y) = sum(abs(in_val - ref_val));

Step 4 - 搜尋與輸出, 這裡是透過 RDom + argmin 搭配來做比對搜尋並取得結果 (search and output, here we use RDom + argmin for matching and finding)

// setup search window
RDom search_win(-range/2, range, -range/2, range);

// argmin will find the one in search_window which get minimal sad.

Tuple me = argmin(sad(search_win.x, search_win.y, bid_x, bid_y));

// output the motion estimation result to MV table

Var ch;

out_mv(bid_x, bid_y, ch) = me[ch];

比較不直覺的地方在於兩個 RDom 套用的層次不同, 第一個 RDom 用來計算 matching criteria, 二個 RDom 是用來作為 motion estimation 尋找 mv, 總之這是以 search window 的方式作窮舉的 block matching 定義, 然而可以延伸與改進用在一開始所說的各種 matching 演算法上

撰寫各 x86 CPU 平台適用的 compiler vector extension code

在套用 gcc / clang 的 -ftree-vectorize or compiler vector extension 時來產生適當的 SIMD optimization 時,  同時會產生一個問題 - 當編譯時套用的 architecture 參數(e.g.: -mavx2) 所使用的指令集與運作時的 CPU 不一致時可能會出現 SIGILL (illegal instruction) 的問題. 過往可能會採用 CPU feature detection + dispatching code 的撰寫, 大致上邏輯如下:

// AVX-512 version

void __foo_avx512(void){  ... }

// AVX version

void __foo_avx2(void){ ... }

// SSE4 version

void __foo_sse4(void){  ... }

// C version

void __foo(void){  ... }

void foo(void)

{

    if(_avx512_available())

        __foo_avx512();

    else if(_avx2_available())

        __foo_avx2();

    else if(_sse4_available())

        __foo_sse4();

    else

        __foo();

}

這是過去長久以來很平常的作法, 而最近研究在接觸 Clear Linux, 研究了一下 Intel 採用的方式, 發現了有趣的事(這些 Intel 都有在 2016 年的簡報 - Improving Linux Performance with GCC latest technologies 中說明), 主要是 gcc 4.8 中已新增了 Function Multi-Versioning (FMV)的功能, 能透過函數屬性更方便地做到原本的目的, 以上述的例子來說就會變成:

// AVX-512 version

__attribute__ ((target ("avx512f")))

void foo(void){  ... }

// AVX2 version

__attribute__ ((target ("avx2")))

void foo(void){ ... }

// SSE4 version

__attribute__ ((target ("sse4.2")))

void foo(void){  ... }

// C version

__attribute__ ((target ("default")))

void foo(void){  ... }

使用 FMV 功能的好處首先所有的 function 都是相同的名稱, 僅是屬性中的 target 不同, 而GCC 會自動產生偵測 CPU 功能並從中選擇對應 foo 的 code, 這點可以省去維護上面 dispatching code 的問題(若有多個函數,  增刪都會是很冗長的工作). 使用 __attribute__ ((target ("TARGET_NAME"))) 的語法方式目前已被 gcc / icc / clang 所採用.

然而並不是所有人都有心力以 SIMD ISA 去實作版本, 這時複製多份相同的 C code 來做不同 SIMD ISA 的 auto-vectorization 似乎不是很實際的方式, 因此 gcc 進一步提供了有趣的屬性功能 - target_clones:

// C version

__attribute__ ((target_clones ("avx512f", "avx2", "sse4.2", "default")))

void foo(void){  ... }

而 GCC FMV 功能亦能夠與 compiler vector extension 合併使用, 因此以 compiler vector 實作後函數可以使用上述 target_clones 的方式, 如此 compiler 會自動針對不同的 SIMD ISA 生出多版本的, 以個人常用來教學的 tiled 8x8 gemm 來說就會像是:

#define TSIZE 8
#if defined (__clang__)
typedef float vfloat __attribute__((ext_vector_type(TSIZE)));
#else
typedef float vfloat __attribute__ ((vector_size (TSIZE*4)));
#endif

__attribute__ ((target_clones ("avx512f", "avx2", "avx", "sse4.2", "default")))
void gemm_vec(float *a, int sa, float *b, int sb, float *c, int sc)
{
    vfloat vb[TSIZE];
        for(int y = 0; y < TSIZE; y++){
            vb[y] = *((vfloat*)(b + sb*y));
        }
        for(int y = 0; y < TSIZE; y++){
            vfloat vc = *((vfloat*)(c + sc*y));
            vfloat va = *((vfloat*)(a + sa*y));
                for(int x = 0; x < TSIZE; x++){
            vc += va[x] * vb[x];
                }
                *((vfloat*)(c + sc*y)) = vc;
        }
}

接著以 gcc 編譯:

$ gcc -O3 -shared mm.c -o libgemm.so

透過 objdump 觀察

$ objdump -t libgemm.so

可以看到中間有如下的輸出:

0000000000001720 l     F .text    0000000000000c96 gemm_vec.default.4
00000000000023c0 l     F .text    00000000000005df gemm_vec.avx512f.0
00000000000029a0 l     F .text    00000000000005df gemm_vec.avx2.1
0000000000002f80 l     F .text    000000000000060c gemm_vec.avx.2
0000000000003590 l     F .text    0000000000000c95 gemm_vec.sse4_2.3

此外值得一提的是在 2019 年初 Clear Linux 計劃中 Intel 釋出了一個 FMV patch generator, 透過 make-fmv-patch 這個工具能針對  C/C++ code 自動產生對應的 FMV 屬性 patch, 套用後重新編譯即可產生一體適用的 SIMD optimized libraries / executables