Skip to content

ion kit startup guide

Jump to bottom
Takuro IIZUKA edited this page Apr 7, 2021 · 3 revisions

ion-kit のコンセプト

ion-kit とはBuilding Block（後述するようにこれはHalide::Generator）の上にグラフ構造を導入するためのF/W.

ここでは fixstars/ion-kitを利用するに際に必要になる各種情報を, 同リポジトリに含まれているサンプルコードを例に説明していく. ion-kit では構築したグラフの実行形式に2種類存在する.

  • JIT Halide::realizeによって実行
  • AOT Halide::Module::compileにより一度コンパイル済みのグラフを出力し, 別のソースファイルからコンパイル済みのグラフを実行する形式

JIT形式

test/simple_graph_jit.ccを例に ion-kit の使い方の基本を説明する.

Building Blockの定義

まずBuilding Block（以下BB）と呼ばれているクラスを定義する必要がある.

BBとは実装上は, ion/block.hにある通り Halide::Generatorのalias templateとなっていて, BBを定義する際はclass Derived : public ion::BuildingBlock<Derived>というようなクラスを実装する. 例えば, 今回の例では test/test-bb.h で定義れているProducer, ComsumerがBBとなる. BBの内部の実装に関しては, Halide::Generatorに関する公式のチュートリアルやHalide::GeneratorのDetailed Descriptionなどを参照（ここでは割愛）.

Input/Output/GeneratorParam

基本的にBBが持つメンバ変数は以下,

注意するべきこととしては, Input/Outputの定義順序は重要という点. 以下ドキュメントとチュートリアルより引用.

Note that the Inputs and Outputs will appear in the C function call in the order they are declared.

// They'll appear in the signature of our generated // function in the same order as we declare them.

一方, GeneratorParamに関しては順序は無関係.

GeneratorParams are always referenced by name, not position, so their order is irrelevant.

BBの実装例

以上を踏まえた上で実装例としてProducer/Comsumerの定義を以下に示す.

class Producer : public ion::BuildingBlock<Producer> {
public:
    // Inputがないので入力はなく, 出力が一つある
    Output<Halide::Func> output{"output", Int(32), 2};
    // 以下はコンパイル時のパラメータ
    GeneratorParam<std::string> string_param{"string_param", "string value"};
    GeneratorParam<int32_t> v{"v", 42};
    void generate() { /* 略 */ }
    void schedule() { /* 略 */ }
private:
    Halide::Var x, y;
};
// test_producerを登録する
ION_REGISTER_BUILDING_BLOCK(Producer, test_producer);

class Consumer : public ion::BuildingBlock<Consumer> {
public:
    Input<Halide::Func> input{"input", Int(32), 2};
    Input<int32_t> desired_min0{"desired_min0", 0};
    Input<int32_t> desired_extent0{"desired_extent0", 0};
    Input<int32_t> desired_min1{"desired_min1", 0};
    Input<int32_t> desired_extent1{"desired_extent1", 0};
    Input<int32_t> v{"v", 1};

    Output<int> output{"output"};
    void generate() { /* 略 */ }
    void schedule() {}
private:
    Halide::Var x, y;
};
// test_consumerを登録する
ION_REGISTER_BUILDING_BLOCK(Consumer, test_consumer);

ION_REGISTER_BUILDING_BLOCKの部分は, BBのインスタンスを実際に作成する時のために必要. （もう少し詳述すれば, IONにおいてHalide::Pipelineを作成する処理であるion::Builder::build内で, Halide::Internal::GeneratorRegistry::enumerateで登録されている名前を列挙したり, Halide::Internal::GeneratorRegistry::createでインスタンスを作成するために必要ということ）

グラフの定義

Node

Node(ion::Node)はIONで作成される計算グラフのノードであり, BBと一対一に対応しているクラスで, Nodeインスタンスの作成は原則ion::Builder::add経由で行われる.

Nodeは名前, id, 複数のPort(ion::Port)やParam(ion::Param)などを持ち, BBとの関係性は（厳密ではないかもしれないがイメージとしては）以下の通り.

Node BB
Port Input/Output
Param GeneratorParam

Port

Portは名前, 紐づいているNodeのidを持っていて, BBのInput/Outputに対応しているクラスとなっている. Portには特定のNodeに紐づいているか（有効なNode idが設定されているか）そうでないかの種別があり, 前者の状態をbounded, 後者の状態をunoboundedと呼んでいる（公式な名称ではなく, 内部実装のコメントではそうなっている）.

  • ion::Node::operator[]とは?
    後述する実装例などでは, `n["output"]`のようにNode `n`の`"output"`という名前のPortを取得しているような記述が出てくる. 実装上（[ion/node.h#L91-L98](https://github.com/fixstars/ion-kit/blob/master/include/ion/node.h#L91-L98)）は, ```cpp Port operator[](const std::string& key) { return Port(key, impl_->id); } ``` のようになっており, `n["output"]`では`"output"`という名前のNode `n`に紐づいた（bounded）Portのインスタンスを作成している. 従って, Nodeは, そのNodeに対応するBBの`Output`に対応するPortは所有しているわけではない.

PortMap

JITの場合に, Portに対して実際の入力値を与えたり, 出力のバッファを対応付けたりするためのテーブル（のようなもの）.

/**
 * PortMap is used to assign actual value to the port input.
 */
class PortMap {

今回のProducer/Consumerの例では, Consumerの入力Portと出力Portに対して以下のように設定している.

test/simple_graph_jit.cc#L19-L27

PortMap pm;
// Consumerの入力Port(unbounded)に対して入力値を与える
pm.set(min0, 0);
pm.set(extent0, 2);
pm.set(min1, 0);
pm.set(extent1, 2);
pm.set(v, 1);

// Consumerの出力Port(bounded)に出力バッファrを対応付ける
Halide::Buffer<int32_t> r = Halide::Buffer<int32_t>::make_scalar();
pm.set(n["output"], r);

このように設定したPortMapをBuilder::runへ渡すことで, 設定した入力値で動作し出力は設定したバッファに書き込まれる.

少し詳しい話

Port/PortMapの仕組み

実際にはboundedなのかunbounded, Halide::Buffer<T>なのかどうかでPortMat::setの挙動は異なる. PortMapは以下のデータを管理している.

// (Port名, Halide::Expr)のkey, valueペア
// unbounded PortとHalide::Paramが対応する場合
std::unordered_map<std::string, Halide::Expr> param_expr_;
// (Port名, Halide::Func)のkey, valueペア
// unbounded PortとHalide::ImageParamが対応する場合
std::unordered_map<std::string, Halide::Func> param_func_;
// ((node id, Port名), 出力バッファの配列)のkey, valueペア
std::unordered_map<std::tuple<std::string, std::string>,
std::vector<Halide::Buffer<>>> output_buffer_;
// Halide::Paramとその値の対応関係
Halide::ParamMap param_map_;

PortMap::setで追加した時, Portやそれに対応している値によって以下のように扱いが異なる.

  • boundedなPortの場合 output_buffer_に指定したHalide::Buffer<>が追加される
  • unboundedなPortでHalide::Param<T>(スカラー値)の場合 param_map_にHalide::Param<T>とその値との対応関係が追加される param_expr_にPort名と, そのPortに対応するHalide::Exprとの対応関係が追加される
  • unboundedなPortでHalide::ImageParam<T>の場合 param_map_にHalide::ImageParam<T>とその値(Halide::Buffer<T>)との対応関係が追加される param_expr_にPort名と, そのPortに対応するHalide::Funcとの対応関係が追加される

従って, Producer/Consumerの例の場合, PortMapを設定した後は以下の図のようになっている.

このPortMapをBuilder::run(const ion::PortMap&)に渡せば, その内部で以下のようにHalide::Pipeline::realize(RealizationArg, const Target&, const ParamMap&)を呼ぶ.

src/builder.cc#L166-L169

void Builder::run(const ion::PortMap& pm) {
    auto p = build(pm, &outputs_); // pはHalide::Pipeline
    return p.realize(Halide::Realization(outputs_), target_, pm.get_param_map());
}

ここで,

Evaluate this Pipeline into an existing allocated buffer or buffers.

とある通り, この場合はoutputs_という既存のバッファに対してHalide::Pipelineが評価される. この時に, Halide::Pipeline::realizeに渡される, outputs_がPortMapの管理しているoutput_buffer_であり, pm.get_param_map()がparam_map_である. 今回の例の場合は, output_buffer_の中身はHalide::Buffer<int32_t> rであり, param_map_は, min0, extent0, min1, extent1, vが保持しているHalide::Paramとその値の対応関係であったので, 指定した入力値に対してrに出力されるという仕組み.

void Builder::run(const ion::PortMap& pm) {
    auto p = build(pm, &outputs_);
    return p.realize(Halide::Realization(outputs_), target_, pm.get_param_map());
}

Param

Paramは実質的にはkey, valueペアであり, GeneratorParamの名前と値に対応しており, Builder::builder内部でBuildingBlock::set_generator_param_valuesに渡すHalide::Internal::GeneratorParamsMapの作成時になどに使用される（builder.cc#L234-L238）.

グラフの作成例

例として test/simple_graph_jit.ccの場合を以下に示す.

// ConsumerのInputのうちdesired_min0, desired_extent0, desired_min1, desired_extent1, v
// に対応するPort
Port min0{"min0", t}, extent0{"extent0", t}, min1{"min1", t}, extent1{"extent1", t}, v{"v", t};

// ProducerのGeneratorParamのvに対応するParam
Param v41{"v", "41"};

// Builder作成
Builder b;
b.set_target(Halide::get_host_target());

// Node作成
Node n;
// Builder::addへはION_REGISTER_BUILDING_BLOCKで指定した名前と同一のものを渡す
// Paramの追加はNode::set_param
n = b.add("test_producer").set_param(v41);
// Portの追加はNode::operator()
n = b.add("test_consumer")(n["output"], min0, extent0, min1, extent1, v);

Nodeの作成はBuilder::addで行うが, その時に渡す引数の文字列はION_REGISTER_BUILDING_BLOCKで登録した名前と同一のものである必要があり, BBとの対応関係はこの名前で取られる. コメントの繰り返しになるが, NodeへのParamの追加はBuilder::set_paramへion::Paramのインスタンスを渡すことで, Portの追加はBuilder::operator()へion::Portのインスタンスを渡すことで行う.

以上のことを踏まえ, 今回作成したグラフを段階的にイメージ図で再度説明する.

グラフ概略図

実行

これまで作成したグラフの実行は, 入力値と出力バッファを設定したPortMapをBuilder::runへ渡すことでグラフのコンパイルとHalide::Pipelineの評価が行われる.

test/simple_graph_jit.cc#L29

    b.run(pm);

AOT形式

グラフの作成まではJITの場合と同様であるため省略.

AOTではBuilder::runで実行する代わりに, Builder::compileに出力ファイル名を渡すと, 最終的にHalide::Module::compileが呼ばれ, simple_graph.hというヘッダファイルが生成される.

test/simple_graph_compile.cc#L17

b.compile("simple_graph");

グラフの実行の方法は test/simple_graph_run.cc にある通りで, 生成されたsimple_graph.hをインクルードし, simple_graphを実行する. JITではPortMapを利用して入力値や出力先を設定していたのに対して, AOTではsimple_graphの引数としてそれらを渡す.

#include "simple_graph.h"
int main()
{
    Buffer<int32_t> out = Buffer<int32_t>::make_scalar();
    return simple_graph(2, 2, 0, 0, 1, out);
}

External Function

Halideの枠組み外の関数を利用する仕組みで, Func::define_externを使用して, 当該Funcへ外部の定義を設定する.

  • Func::define_extern()の説明

    This lets you define a Func that represents an external pipeline stage. You can, for example, use it to wrap a call to an extern library such as fftw.

実はConsumerでも使用されていて, 下記ではconsumeというC++関数を呼び出している.

test/test-bb.h#L38-L42

// External Function向けの引数リスト
std::vector<ExternFuncArgument> params{in, desired_min0, desired_extent0, desired_min1, desired_extent1, v};
Func consume;
// consumeというFuncへの外部
// Func::define_extern("関数名", 引数, 出力型, 出力次元)
consume.define_extern("consume", params, Int(32), 0);
consume.compute_root();
output() = consume();

外部関数consumeの定義は以下のようになる.

test/test-rt.h#L13-L31

// 最後のoutが出力
// return 0はステータスコード的なもので処理の出力ではない
extern "C"
int consume(halide_buffer_t *in, int desired_min0, int desired_extent0, int desired_min1, int desired_extent1, int32_t v, halide_buffer_t *out) {
    if (in->is_bounds_query()) {
        // Bounds query mode
        // 入力範囲の指定
        in->dim[0].min = desired_min0;
        in->dim[0].extent = desired_extent0;
        in->dim[1].min = desired_min1;
        in->dim[1].extent = desired_extent1;
    } else {
        // 実際の処理
        Halide::Runtime::Buffer<int32_t> ibuf(*in);
        for (int y=0; y<in->dim[1].extent; ++y) {
            for (int x=0; x<in->dim[0].extent; ++x) {
                std::cout << ibuf(x, y) + v << " ";
            }
            std::cout << std::endl;
        }
    }
    return 0;
}

HalideにはBounds Inferenceと呼ばれる出力の範囲と計算式から入力の参照範囲を計算する仕組みがある.

        // Note that consumer was evaluated over a 4x4 box, so Halide
        // automatically inferred that producer was needed over a 5x5
        // box. This is the same 'bounds inference' logic we saw in
        // the previous lesson, where it was used to detect and avoid
        // out-of-bounds reads from an input image.

この機構はHalideの枠組み内でのみ有効であるので, その外側である外部関数内では明示的に境界情報を指定する必要がある. 具体的には, 外部関数が境界を決定するモードで実行された場合（is_bounds_query()が真）は, その処理（この場合consume）が計算すべき範囲（min/extent）を指定する.

Clone this wiki locally