Intel nGraph で Halide はどのように使われているか？(その3)

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

Intel nGraph で Halide はどのように使われているか？(その2)に続き。

cpu::op::HalideOpは、src/ngraph/runtime/cpu/op/halide_op.hpp にて宣言されています。

                class HalideOp : public ngraph::op::Op
                {
                public:
                    HalideOp(const NodeVector& args,
                             const std::list<std::shared_ptr<Node>>& ops,
                             const element::Type& out_type,
                             const Shape& out_shape);

                    virtual void validate_and_infer_types() override;

                    virtual std::shared_ptr<Node>
                        copy_with_new_args(const NodeVector& new_args) const override;

                    const std::list<std::shared_ptr<Node>>& get_ops() const { return m_ops; }
                private:
                    std::list<std::shared_ptr<Node>> m_ops;
                    element::Type m_output_type;
                    Shape m_output_shape;
                };

実装は、src/ngraph/runtime/cpu/op/halide_op.cpp

shared_ptr<Node> runtime::cpu::op::HalideOp::copy_with_new_args(const NodeVector& new_args) const
{
    return make_shared<HalideOp>(new_args, m_ops, m_output_type, m_output_shape);
}

runtime::cpu::op::HalideOp::HalideOp(const NodeVector& args,
                                     const std::list<std::shared_ptr<Node>>& ops,
                                     const element::Type& out_type,
                                     const Shape& out_shape)
    : Op("HalideOp", check_single_output_args(args))
    , m_ops(ops)
    , m_output_type(out_type)
    , m_output_shape(out_shape)
{
    constructor_validate_and_infer_types();
}

void runtime::cpu::op::HalideOp::validate_and_infer_types()
{
    set_output_type(0, m_output_type, m_output_shape);
}

HalideOpは、src/ngraph/runtime/cpu/builder/halide_op.cppのBuilder::BUILDER_DECL(ngraph::runtime::cpu::op::HalideOp)にて、Halideを使ってコード生成しています。

            void Builder::BUILDER_DECL(ngraph::runtime::cpu::op::HalideOp)
            {
                const ngraph::runtime::cpu::op::HalideOp* hs =
                    static_cast<const ngraph::runtime::cpu::op::HalideOp*>(node);

                const auto& generators = ngraph::runtime::cpu::halide::get_halide_generators();

get_halide_generators が各Op を Halide に変換していきます。詳細は、明日のブログで解析していきます。

                auto& halide_functions = external_function->get_halide_functions();
                auto& subgraph_params = external_function->get_subgraph_params();
                auto& subgraph_param_sizes = external_function->get_subgraph_param_sizes();
                auto& subgraph_param_ptrs = external_function->get_subgraph_param_ptrs();

external_functionの各メソッドを呼んでいます。この部分については、明日のブログで解析していきます。

hs->get_ops() にて、各Op1 を Halide に変換していきます。

                for (const auto& op : hs->get_ops())
                {
                    if (!generators.count(TI(*op)))
                    {
                        throw ngraph_error("Invalid op in halide subgraph");
                    }
                    vector<Halide::Func> inputs;
                    for (const auto& input : op->get_inputs())
                    {
                        auto tensor_name = input.get_output().get_tensor_ptr()->get_name();
                        if (halide_functions.count(tensor_name))
                        {
                            inputs.emplace_back(halide_functions[tensor_name]);
                        }
                        else
                        {
                            subgraph_params[tensor_name] = Halide::ImageParam(Halide::Float(32), 1);
                            subgraph_param_sizes[tensor_name] =
                                shape_size(input.get_output().get_tensor_ptr()->get_shape());
                            subgraph_param_ptrs.emplace(
                                tensor_name, external_function->get_tensor_data(tensor_name));
                            inputs.emplace_back(subgraph_params[tensor_name]);
                        }
                    }
                    halide_functions[op->get_output_tensor_ptr()->get_name()] =
                        generators.at(TI(*op))(inputs);
                }

次に、Halide部分を　Functor にマッピングします。

                auto out_tensor_name = hs->get_ops().back()->get_output_tensor_ptr()->get_name();
                auto& functors = external_function->get_functors();
                auto& out_tensor = external_function->get_tensor_data(out[0].get_name());
                auto& terminal_func = halide_functions[out_tensor_name];
                auto out_size = out[0].get_size();

                auto functor = [&, out_size](CPURuntimeContext* ctx, CPUExecutionContext* ectx) {
                    for (auto& param : subgraph_params)
                    {
                        Halide::Buffer<float> param_buffer(
                            static_cast<float*>(subgraph_param_ptrs.at(param.first).get()),
                            subgraph_param_sizes.at(param.first));
                        param.second.set(param_buffer);
                    }
                    Halide::Buffer<float> out_buffer(static_cast<float*>(out_tensor), out_size);
                    terminal_func.realize(out_buffer);
                };
                functors.emplace_back(functor);
            }