淺談 Halide 新的實用排程指令 - compute_with

今年二月中時, Halide 釋出了一個新版本, 其中第一個列出的新特性是

Scheduling:

• New scheduling directive: compute_with

也就新增了這篇所要討論的 compute_with 的排程指令, 目前沒有相關的文件說明這個排程指令的用法, 只能從 Halide 原始碼中僅有的正確性測試範例來了解其目的與使用的方式.

首先要複習一下這一類既有的排程指令:

• compute_root()
• compute_at()
• store_root()
• store_at()

這類排程指令主要在於控制中間產物在過程中, 其數值於 loop 中計算與存放的位置. 然而不同的函數 Func 間是各自獨立排程與計算的. 當各 Func 中使用的資料來源不同, 這可能不是個問題, 然而當 Func 間有著類似或相同的輸入, 與計算的 loop 結構, 這就造成了 loop 浪費以及 cache locality 的低落.

以下為 Halide 測試中項目中的一個

f(x, y) = x + y;

g(x, y) = x - y;

h(x, y) = f(x - 1, y + 1) + g(x + 2, y - 2);

f.compute_root();

g.compute_root();

Var xo("xo"), xi("xi");

f.split(x, xo, xi, 7);

g.split(x, xo, xi, 7);

 g.compute_with(f, xo);

由於 f 與 g 都使用在相同的變數值域下計算, 若 f, g 都個別處理, 則浪費的可以重複使用的 loop 數值, 這時可以透過將 g 利用 f 的 loop 依附著指定的變數上作計算. 也就等同於使用一個 x, y loop 來將 f, g 同時計算出來. 由於 f, g 的計算很簡單, 這裡主要是節省了 loop iteration 的次數

在影像處理中, 這樣類似的情況很多, 這裡以著名的 Guided Filter 為例

這裡在一開頭的需要的4個以 box average 方式計算的數值函數就可以應用: mean(I), mean(p), corr(I) 與 corr(Ip)
此四個在 Halide 可以定義為 Func 函數, 並透過指定的變數來計算:

Var x, y;
RDom r(-(radius/2), radius);

Func meanI_x, meanP_x, corrI_x, corrIP_x;
Func meanI, meanP, corrI, corrIP;
Expr area = r*r;
meanI_x(x, y) = sum(I(x+r, y));
meanP_x(x, y) = sum(P(x+r, y));
corrI_x(x, y) = sum(I(x+r, y)*I(x+r, y));
corrIP_x(x, y) = sum(I(x+r, y)*P(x+r, y));

meanI(x, y) = sum(meanI_x(x, y+r))/area;
meanP(x, y) = sum(meanP_x(x, y+r))/area;
corrI(x, y) = sum(corrI_x(x, y+r))/area;
corrIP(x, y) = sum(corrIP_x(x, y+r))/area;

meanP_x.compute_with(meanI_x, x);
corrI_x.compute_with(meanI_x, x);
corrIP_x.compute_with(meanI_x, x);

// meanI, meanP, corrI, corrIP can be schedule with compute_with, too

如此可以即可更快速有效地計算出所需要的中間產物