Method init_dispatch_table

src/gopt/impl/basic_arith/trans.cpp:352–408

}

void FinalArithTransformPass::Impl::init_dispatch_table() {
    /*
     * Note: each rule takes var on original graph as input
     */
    auto add_dispatcher = [&](Mode mode) -> DispatchEntry& {
        auto&& vec = m_dispatch_table[mode];
        vec.emplace_back();
        return vec.back();
    };

    auto add_dispatcher_with_name =
            [&](Mode mode, const char* name) -> DispatchEntry::first_type& {
        auto&& ret = add_dispatcher(mode);
        ret.second = name;
        return ret.first;
    };

#define REG(_mode, _name)                                                \
    add_dispatcher_with_name(Mode::_mode, mgb_cstr_log(_name)) = [this]( \
            const VarNodeArray& inp) -> SymbolVar

#define REG_THIS(_mode, _fn)        \
    add_dispatcher(Mode::_mode) = { \
            std::bind(&Impl::_fn, this, std::placeholders::_1), mgb_cstr_log(#_fn)}

    REG_THIS(ADD, merge_negate<MergeNegateAddTrait>);
    REG_THIS(MUL, merge_negate<MergeNegateMulTrait>);

    REG(POW, "powc and exp merge") {
        auto exp_maybe = SymbolVar{inp[1]}.as_immutable_scalar_require_shape();
        if (!exp_maybe.valid()) {
            return {};
        }
        float exp = exp_maybe->get_cast<float>();
        VarNode* base = m_rewriter.get_var(inp[0]);
        Elemwise* base_pow;
        if ((base_pow = as_elem_opr(base, Mode::POW)) && can_replace_var(base)) {
            // powc(pow(x, a), b) => pow(x, a * b); a is not const scalar
            VarNode* exp_new;
            VarNode* exp_old = base_pow->input(1);
            if (almost_equal(exp, -1.f)) {
                // handle reciprocal
                exp_new = get_neg_repl<Mode::ADD>(exp_old, true);
                if (!exp_new) {
                    exp_new = opr::negate(exp_old).node();
                }
            } else {
                exp_new = (SymbolVar{exp_old} * exp).node();
            }
            return opr::pow(base_pow->input(0), exp_new);
        }
        return opr::PowC::make(base, exp);
    };

#undef REG
#undef REG_THIS
}