シェーダ言語では、specialization と呼ばれる方法で、
シェーディング処理の高速化を図ることができます.
(ここでは、partial evaluation == specialization とします)
specialization は、同じ計算はキャッシュしておき、
次に使うときはそのキャッシュを使うようにしたり、
変化しない(関数)パラメータは、定数とみなして、
変化する部分だけ再計算するような処理を指します.
グラフィックス(シェーディング)では、
一つのプログラムがピクセルやサンプル点ごとに何回も呼ばれますが、
すべてのパラメータが変化することはなく、いくつかは固定(たとえばマテリアルパラメータなど)
であることが多いです.
このため、グラフィックス(シェーディング)は、
specialization 技術が非常に有効な分野の一つになります.
シェーダの specialization については既存研究がいくつかありますが [1, 2]、
これをやろうとすると、AST(抽象構文器) をいじくったりしないといけません.
そこで、自前でそんなめんどくさい機構をプログラミングするのは面倒なので、
期待のコンパイラインフラである LLVM を使って
partial evaluation(specialization) を行うサンプルを作ってみました.
(web をあさったのですが、[3] くらいしか情報源がなくてちょっと手こずりましたが)
今回のファイル群はこちらにあります.
http://lucille.atso-net.jp/svn/angelina/llvm/partial_eval/
例題
// func.c
int func(int a, int b)
{
return a * a * a * b;
}
このような関数が与えられたとき、実際には a が 1 しか取らないことが分かったとしましょう.
その場合、以下のような簡略化された関数に変換することができます.
int func_specialized(int a, int b)
{
return b; // a = 1
}
この簡略化処理を, LLVM (の C++ API) で行ってみます.
LLVM で partial evaluation
基本的なやり方は [3] の通りです.
入力はビットコード(.bc)ファイルとします.
[準備] ビットコードを読み込んで Module を作る.
1. doIt() を Module を引数にして呼ぶ
2. Module 内の関数を探索し、’func’ Function を探す.
3. 関数の’枠’のコピーを作る(CloneFunctionHeader)
このとき関数の中身のコピーは必要なくて、引数構造のコピーと ValueMap と呼ばれるデータを作る.
ValueMap は関数を Specialize する CloneAndPruneFunctionInto で必要になる.
4. 引数 ‘a’ の ValueMap を定数で置き換え
(引数 a を定数にして関数をコールする、という意味になるようです)
5. CloneAndPruneFunctionInto を、specialize する引数を書き換えた ValueMap で読んで,
specialize された関数を作る.
[終了処理] ビットコードファイルに specialize された関数を書き出し.
関数の specialize を行うコードは、以下のようになります.
Function *CloneFunctionHeader(const Function *F,
DenseMap &ValueMap)
{
std::vector ArgTypes;
// The user might be deleting arguments to the function by specifying them in
// the ValueMap. If so, we need to not add the arguments to the arg ty vector
//
for (Function::const_arg_iterator I = F->arg_begin(), E = F->arg_end();
I != E; ++I)
if (ValueMap.count(I) == 0) // Haven't mapped the argument to anything yet?
ArgTypes.push_back(I->getType());
// Create a new function type...
FunctionType *FTy = FunctionType::get(F->getFunctionType()->getReturnType(),
ArgTypes, F->getFunctionType()->isVarArg());
// Create the new function...
Function *NewF = Function::Create(FTy, F->getLinkage(), F->getName());
// Loop over the arguments, copying the names of the mapped arguments over...
Function::arg_iterator DestI = NewF->arg_begin();
for (Function::const_arg_iterator I = F->arg_begin(), E = F->arg_end();
I != E; ++I)
if (ValueMap.count(I) == 0) { // Is this argument preserved?
DestI->setName(I->getName()); // Copy the name over...
ValueMap[I] = DestI++; // Add mapping to ValueMap
}
return NewF;
}
void
doIt(Module *m, std::ostream *Out)
{
Module::iterator f, e;
// Iterate over functions in the module.
for (f = m->begin(), e = m->end(); f != e; ++f) {
if (!f->isDeclaration()) {
if (f->getNameStr() == "func") break;
}
}
if (f == m->end()) {
cout << "Can't find function 'func' in the module\n";
exit(2);
}
//
// Got function 'func'
//
DenseMap ValueMap;
Function *specializedFunction = CloneFunctionHeader(f, ValueMap);
specializedFunction->setName(f->getNameStr() + "_specialized");
// Iterate over function argments.
for (Function::arg_iterator i = f->arg_begin(),
ie = f->arg_end(); i != ie; ++i) {
Argument *arg = i;
cout << "Arg: " <getNameStr() <getNameStr() == "a") {
// Replace arg 'a' with Const expresison 1.
val = ConstantInt::get(Type::Int32Ty, 1);
ValueMap[i] = val;
}
}
//
// Generate specialized function
//
std::vector Returns;
ClonedCodeInfo SpecializedFunctionInfo;
CloneAndPruneFunctionInto( specializedFunction, // NewFunc
f, // OldFunc
ValueMap, // ValueMap
Returns, // Returns
".", // NameSuffix
&SpecializedFunctionInfo, // CodeInfo
0); // TD
f->print(cout);
cout <print(cout);
// Add specialized function to current module.
m->getFunctionList().push_back(specializedFunction);
}
実験
では実際に確認してみましょう. 上記のコード片を含んだ例題プログラムは partial_eval というコマンドラインプログラムとします.
まず、入力の LLVM ビットコード func.bc を作ります.
$ llvm-gcc -emit-llvm -c func.c
$ llvm-as func.ll -f
$ llvm-dis func.bc
define i32 @func(i32 %a, i32 %b) {
entry:
%a.addr = alloca i32 ; [#uses=4]
%b.addr = alloca i32 ; [#uses=2]
store i32 %a, i32* %a.addr
store i32 %b, i32* %b.addr
%tmp = load i32* %a.addr ; [#uses=1]
%tmp1 = load i32* %a.addr ; [#uses=1]
%mul = mul i32 %tmp, %tmp1 ; [#uses=1]
%tmp2 = load i32* %a.addr ; [#uses=1]
%mul3 = mul i32 %mul, %tmp2 ; [#uses=1]
%tmp4 = load i32* %b.addr ; [#uses=1]
%mul5 = mul i32 %mul3, %tmp4 ; [#uses=1]
ret i32 %mul5
}
specialize を行うサンプルプログラムを走らせます.
(specialize された関数は、同じモジュール(ビットコードファイル)に書かれます)
$ ./partial_eval func.bc -o func.par.bc -f
$ llvm-dis func.par.bc -f
$ cat func.par.ll
define i32 @func(i32 %a, i32 %b) {
entry:
%a.addr = alloca i32 ; [#uses=4]
%b.addr = alloca i32 ; [#uses=2]
store i32 %a, i32* %a.addr
store i32 %b, i32* %b.addr
%tmp = load i32* %a.addr ; [#uses=1]
%tmp1 = load i32* %a.addr ; [#uses=1]
%mul = mul i32 %tmp, %tmp1 ; [#uses=1]
%tmp2 = load i32* %a.addr ; [#uses=1]
%mul3 = mul i32 %mul, %tmp2 ; [#uses=1]
%tmp4 = load i32* %b.addr ; [#uses=1]
%mul5 = mul i32 %mul3, %tmp4 ; [#uses=1]
ret i32 %mul5
}
define i32 @func_specialized(i32 %a, i32 %b) {
entry.:
%a.addr. = alloca i32 ; [#uses=4]
%b.addr. = alloca i32 ; [#uses=2]
store i32 1, i32* %a.addr.
store i32 %b, i32* %b.addr.
%tmp. = load i32* %a.addr. ; [#uses=1]
%tmp1. = load i32* %a.addr. ; [#uses=1]
%mul. = mul i32 %tmp., %tmp1. ; [#uses=1]
%tmp2. = load i32* %a.addr. ; [#uses=1]
%mul3. = mul i32 %mul., %tmp2. ; [#uses=1]
%tmp4. = load i32* %b.addr. ; [#uses=1]
%mul5. = mul i32 %mul3., %tmp4. ; [#uses=1]
ret i32 %mul5.
}
func_specialized が、a=1 として specialize された関数です.
このままではなにも変わっていません.
オプティマイザを走らせてみます.
$ opt -std-compile-opts func.par.bc -o func.opt.bc -f
$ llvm-dis func.opt.bc -f
$ cat func.opt.ll
define i32 @func(i32 %a, i32 %b) nounwind {
entry:
%mul = mul i32 %a, %a ; [#uses=1]
%mul3 = mul i32 %mul, %a ; [#uses=1]
%mul5 = mul i32 %mul3, %b ; [#uses=1]
ret i32 %mul5
}
define i32 @func_specialized(i32 %a, i32 %b) nounwind {
entry.:
ret i32 %b
}
元の関数 func は 3 muls ですが、func_specialize は
期待通り b を返すだけになりました.
おお、ブラボー！！！
まとめ
今回は必要な部分だけ C++ で作って、あとはコマンドラインプログラムの組み合わせでやりましたが、
C++ 側に LLVM の最適化のパスなども取り込んだり、LLVM の JIT を使うようにしてもっと
on-the-fly で partial evaluation(specialization) を行うこともできるでしょう.
LLVM はコンパイラインフラなので、必要な機能はほとんど C++ API から呼び出すことができます.
現実のアプリでは、必ずしてもパラメータが一つの定数になったり、
一つの引数だけが定数になることはないでしょうから、
定数ごと(a=1, 3, 4, ..)や組み合わせごとに specialize された関数を生成することになると思います.
Download:
You can download files used in this post here.
http://lucille.atso-net.jp/svn/angelina/llvm/partial_eval/
References:
[1] the lightspeed automatic interactive lighting preview system
http://people.csail.mit.edu/jrk/lightspeed/
[2] Specializing Shaders
http://research.microsoft.com/research/pubs/view.aspx?pubid=1041
[3] [LLVMdev] Partial evaluation and LLVM
http://lists.cs.uiuc.edu/pipermail/llvmdev/2007-June/009337.html
Syoyo Fujita's Blog 