megenginelite-sys 1.8.2

/**
 * \file src/gopt/impl/basic_arith/trans.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "megbrain/gopt/basic_arith.h"
#include "megbrain/opr/basic_arith_wrapper.h"
#include "megbrain/serialization/serializer.h"

//! TODO: here has to be know some megdnn::opr when there is produced midout.h
//! fix it if there is another graceful way.
#include "megdnn/oprs.h"

#include "megbrain/utils/hash_ct.h"
#include "midout.h"

MIDOUT_DECL(megbrain_trans)
#define MIDOUT_B(tag) MIDOUT_BEGIN(megbrain_trans, midout_iv(MGB_HASH_STR(tag))) {
#define MIDOUT_E \
    }            \
    MIDOUT_END();

using namespace mgb;
using namespace gopt;

namespace {
/*!
 * \brief helper for implementing term-rewriting for elemwise oprs
 */
class ElemwiseRewriteImplHelper {
    void on_opr(cg::OperatorNodeBase* opr);

protected:
    using Elemwise = opr::Elemwise;
    using Mode = Elemwise::Mode;

    /*!
     * \brief a node in elemwise chain
     *
     * An elemwise chain is a flattened tree represented as a postfix
     * expression to preserve original tree structure, consisting of
     * elemwise oprs
     */
    struct ElemwiseChainNode;
    using ElemwiseChain = std::vector<ElemwiseChainNode>;

    const Pass& m_pass;
    OptState& m_opt_state;
    SubGraph::Rewriter m_rewriter;
    UniqReaderCheck m_uniq_reader_check;

    ElemwiseRewriteImplHelper(const Pass& pass, OptState& opt_state)
            : m_pass{pass},
              m_opt_state{opt_state},
              m_rewriter{opt_state.graph().make_rewriter()},
              m_uniq_reader_check{opt_state.graph()} {}

    ~ElemwiseRewriteImplHelper() noexcept = default;

    /*!
     * \brief callback when an Elemwise operator is visited
     * \param elem the operator on ORIGINAL graph
     */
    virtual void on_opr_elemwise(Elemwise* elem) = 0;

    //! whether a var node can be replaced
    bool can_replace_var(VarNode* var) const { return m_uniq_reader_check(var); }

    /*!
     * \brief run the rewriter
     *
     * call on_opr_elemwise() for all elemwise oprs and apply m_rewriter
     */
    void run_elemwise_rewriter();

    /*!
     * \brief extract an elemwise chain where all internal nodes share
     *      the given mode and are replaceable
     *
     * Note: must consult can_replace_var() before replacing a leaf var
     */
    ElemwiseChain extract_elemwise_chain(VarNode* endpoint, Mode mode);

    //! reconstruct a var from given elemwise chain
    VarNode* reconstruct_elemwise_chain(const ElemwiseChain& chain);
};
}  // namespace

/* ================ ElemwiseRewriteImplHelper ================ */
struct ElemwiseRewriteImplHelper::ElemwiseChainNode {
    enum class Type { LEAF, INTERNAL };
    Type type;
    union {
        VarNode* leaf;
        Mode mode;
    } data;

    static ElemwiseChainNode make_leaf(VarNode* var) {
        ElemwiseChainNode ret;
        ret.type = Type::LEAF;
        ret.data.leaf = var;
        return ret;
    }

    static ElemwiseChainNode make_internal(Mode mode) {
        ElemwiseChainNode ret;
        ret.type = Type::INTERNAL;
        ret.data.mode = mode;
        return ret;
    }
};

void ElemwiseRewriteImplHelper::run_elemwise_rewriter() {
    using namespace std::placeholders;
    m_opt_state.graph().iter(std::bind(&ElemwiseRewriteImplHelper::on_opr, this, _1));
    m_rewriter.apply_inplace();
}

void ElemwiseRewriteImplHelper::on_opr(OperatorNodeBase* opr) {
    m_uniq_reader_check.update_on_opr_auto_replace(
            opr, m_rewriter.auto_replace_outputs(opr));

    if (auto elem = try_cast_as_op<Elemwise>(opr)) {
        on_opr_elemwise(elem);
    }
}

ElemwiseRewriteImplHelper::ElemwiseChain ElemwiseRewriteImplHelper::
        extract_elemwise_chain(VarNode* endpoint, Mode mode) {
    ElemwiseChain ret;
    auto check_internal = [mode, this](VarNode* var) -> bool {
        return as_elem_opr(var, mode) && can_replace_var(var);
    };
    auto on_leaf = [&ret](VarNode* var) {
        ret.push_back(ElemwiseChainNode::make_leaf(var));
    };
    auto on_internal_finish = [&ret](OperatorNodeBase* opr) {
        ;
        ret.push_back(ElemwiseChainNode::make_internal(
                opr->cast_final_safe<Elemwise>().param().mode));
    };
    visit_opr_tree(endpoint, check_internal, on_leaf, on_internal_finish);
    return ret;
}

VarNode* ElemwiseRewriteImplHelper::reconstruct_elemwise_chain(
        const ElemwiseChain& chain) {
    VarNodeArray stack;
    SymbolVarArray tmp_inp;
    for (auto&& i : chain) {
        if (i.type == ElemwiseChainNode::Type::LEAF) {
            stack.push_back(i.data.leaf);
        } else {
            mgb_assert(i.type == ElemwiseChainNode::Type::INTERNAL);
            auto mode = i.data.mode;
            auto arity = megdnn::Elemwise::ModeTrait::from_mode(mode).arity;
            mgb_assert(arity <= stack.size());
            tmp_inp.resize(arity);
            for (size_t i = 0; i < arity; ++i) {
                tmp_inp[i] = stack[stack.size() - arity + i];
            }
            stack.resize(stack.size() - arity);
            stack.push_back(Elemwise::make(tmp_inp, mode).node());
        }
    }
    mgb_assert(stack.size() == 1);
    return stack[0];
}

/* ================ ArithMulDistributePass ================ */

class ArithMulDistributePass::Impl final : public ElemwiseRewriteImplHelper {
    //! total size reduced by prev try_distribute() call
    size_t m_eliminated_computing;

    void on_opr_elemwise(Elemwise* elem) override;

    /*!
     * \brief try to distribute \p mul over terms in \p add_endpoint
     *
     * The given vars must reside on original graph.
     *
     * \return transformed var for \p add_endpoint times \p mul, or nullptr if
     *      failed
     */
    VarNode* try_distribute(VarNode* add_endpoint, VarNode* mul);

public:
    Impl(const ArithMulDistributePass& pass, OptState& opt_state)
            : ElemwiseRewriteImplHelper(pass, opt_state) {
        run_elemwise_rewriter();
    }
};

void ArithMulDistributePass::Impl::on_opr_elemwise(Elemwise* elem) {
    if (elem->param().mode != Mode::MUL)
        return;
    auto i0 = elem->input(0), i1 = elem->input(1), out = elem->output(0);
    auto &&shp0 = i0->shape(), &&shp1 = i1->shape(), &&oshp = out->shape();
    auto sz0 = shp0.total_nr_elems(), sz1 = shp1.total_nr_elems();
    if (!oshp.ndim || sz0 == sz1 || !(oshp.eq_shape(shp0) || oshp.eq_shape(shp1))) {
        return;
    }
    if (sz0 < sz1) {
        std::swap(i0, i1);
    }

    if (auto end = try_distribute(i0, i1)) {
        m_rewriter.replace_var(
                out, end,
                mgb_ssprintf_log("%zu less elemwise-computing", m_eliminated_computing)
                        .c_str());
    }
}

VarNode* ArithMulDistributePass::Impl::try_distribute(
        VarNode* add_endpoint, VarNode* mul) {
    TensorShapeArray check_compatible_inp(2);
    check_compatible_inp[0] = mul->shape();
    auto shape_compatible = [&](VarNode* var) {
        check_compatible_inp[1] = var->shape();
        TensorShape tshp;
        megdnn::Elemwise::deduce_shape(check_compatible_inp, tshp);
        return tshp.eq_shape(var->shape());
    };

    mul = m_rewriter.get_var(mul);
    auto add_chain = extract_elemwise_chain(add_endpoint, Mode::ADD);

    // mul chain, combine position in mul chain
    std::vector<std::pair<ElemwiseChain, size_t>> terms;

    using Type = ElemwiseChainNode::Type;

    if (add_chain.size() < 3)
        return nullptr;

    m_eliminated_computing = add_endpoint->shape().total_nr_elems();
    for (auto&& term : add_chain) {
        if (term.type != Type::LEAF)
            continue;
        auto mul_chain = extract_elemwise_chain(term.data.leaf, Mode::MUL);
        size_t best_pos = 0, best_size = m_eliminated_computing;
        // find smallest compatible var in mul_chain
        for (size_t i = 0; i < mul_chain.size(); ++i) {
            if (mul_chain[i].type == Type::LEAF) {
                auto var = m_rewriter.get_var(mul_chain[i].data.leaf);
                mul_chain[i].data.leaf = var;
                if (shape_compatible(var)) {
                    auto size = var->shape().total_nr_elems();
                    if (size < best_size) {
                        best_size = size;
                        best_pos = i;
                    }
                }
            }
        }
        if (best_size == m_eliminated_computing) {
            return nullptr;
        }
        m_eliminated_computing -= best_size;
        terms.push_back({{}, best_pos});
        terms.back().first = std::move(mul_chain);
    }

    auto mul_chain_iter = terms.begin();
    for (auto&& term : add_chain) {
        if (term.type != Type::LEAF)
            continue;
        auto&& var = mul_chain_iter->first[mul_chain_iter->second].data.leaf;
        var = (SymbolVar{var} * mul).node();
        term.data.leaf = reconstruct_elemwise_chain(mul_chain_iter->first);
        ++mul_chain_iter;
    }
    mgb_assert(mul_chain_iter == terms.end());

    return reconstruct_elemwise_chain(add_chain);
}

const char* ArithMulDistributePass::name() const {
    return mgb_cstr_log("mul_distribute");
}

void ArithMulDistributePass::apply(OptState& opt) const {
    MIDOUT_B("ArithMulDistributePass::apply")
    Impl{*this, opt};
    MIDOUT_E
}

/* ================ FinalArithTransformPass ================ */

class FinalArithTransformPass::Impl final : public ElemwiseRewriteImplHelper,
                                            public NonCopyableObj {
    using DispatchEntry =
            std::pair<thin_function<SymbolVar(const VarNodeArray&)>, const char*>;

    //! for merge_negate() with ADD/SUB modes
    struct MergeNegateAddTrait;
    //! for merge_negate() with MUL/TRUE_DIV modes
    struct MergeNegateMulTrait;

    ThinHashMap<Mode, std::vector<DispatchEntry>> m_dispatch_table;

    /*!
     * \brief get neg src; also set var to current var
     * \tparam mode either Mode::ADD or Mode::MUL, to define the inv-group
     * \param[in,out] var input original var, and output currently replaced var
     * \return x if var == -x and can_replace_var(var); nullptr otherwise
     */
    template <Mode mode>
    VarNode* get_neg_repl(VarNode*& var, bool require_replaceable) const;

    /*!
     * \brief try to decompose var as a * b if var is replaceable
     *
     * Also requires shapes of a and b are known.
     *
     * \param[in,out] var input original var, and output replaced var
     * \param[out] a decomposed first term; nullptr if can not decompose
     * \param[out] b decomposed second term; nullptr if can not decompose
     */
    void as_replaceable_mul(VarNode*& var, VarNode*& a, VarNode*& b);

    /*!
     * \brief merge negate operators like (-a) + (-b) -> -(a+b)
     * \tparam Trait provides ADD, SUB and neg()
     */
    template <class Trait>
    SymbolVar merge_negate(const VarNodeArray& inp);

    void init_dispatch_table();
    void on_opr_elemwise(Elemwise* elem) override;

public:
    Impl(const FinalArithTransformPass& pass, OptState& opt_state)
            : ElemwiseRewriteImplHelper(pass, opt_state) {
        init_dispatch_table();
        run_elemwise_rewriter();
    }
};

void FinalArithTransformPass::Impl::on_opr_elemwise(Elemwise* elem) {
    auto mode = elem->param().mode;
    auto&& iter = m_dispatch_table.find(mode);
    if (iter != m_dispatch_table.end()) {
        for (auto&& dispatch : iter->second) {
            auto repl = dispatch.first(elem->input()).node();
            if (repl) {
                auto src = elem->output(0);
                m_rewriter.replace_var(src, repl, dispatch.second);
                return;
            }
        }
    }
}

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
}

/* ---------------- merge_negate ---------------- */
template <class Trait>
SymbolVar FinalArithTransformPass::Impl::merge_negate(const VarNodeArray& inp) {
    VarNode *i0 = inp[0], *i1 = inp[1], *neg0 = get_neg_repl<Trait::ADD>(i0, false),
            *neg1 = get_neg_repl<Trait::ADD>(i1, false);
    // always replace neg (do not check unique reader) since this does not
    // introduce new opr

    auto add = [](SymbolVar a, SymbolVar b) {
        return opr::Elemwise::make({a, b}, Trait::ADD);
    };
    auto sub = [](SymbolVar a, SymbolVar b) {
        return opr::Elemwise::make({a, b}, Trait::SUB);
    };

    if (!neg0 && !neg1)
        return {};
    if (neg0 && neg1)
        return Trait::neg(add(neg0, neg1));
    if (neg0)
        return sub(i1, neg0);
    return sub(i0, neg1);
}

struct FinalArithTransformPass::Impl::MergeNegateAddTrait {
    static constexpr Mode ADD = Mode::ADD, SUB = Mode::SUB;

    static SymbolVar neg(SymbolVar x) { return -x; }
};

struct FinalArithTransformPass::Impl::MergeNegateMulTrait {
    static constexpr Mode ADD = Mode::MUL, SUB = Mode::TRUE_DIV;

    static SymbolVar neg(SymbolVar x) { return opr::PowC::make(x, -1); }
};

/* ---------------- helpers ---------------- */
template <opr::Elemwise::Mode mode>
VarNode* FinalArithTransformPass::Impl::get_neg_repl(
        VarNode*& var, bool require_replaceable) const {
    auto new_var = m_rewriter.get_var(var);
    VarNode* ret = nullptr;
    if (!require_replaceable || can_replace_var(var)) {
        ret = check_is_group_inverse_opr<mode>(new_var);
    }
    var = new_var;
    return ret;
}

void FinalArithTransformPass::Impl::as_replaceable_mul(
        VarNode*& var, VarNode*& a, VarNode*& b) {
    a = b = nullptr;
    auto new_var = m_rewriter.get_var(var);
    Elemwise* elem = nullptr;
    if (var->shape().ndim && can_replace_var(var) &&
        (elem = as_elem_opr(new_var->owner_opr(), Mode::MUL))) {
        a = elem->input(0);
        b = elem->input(1);
    }
    var = new_var;
}

const char* FinalArithTransformPass::name() const {
    return mgb_cstr_log("final_arith_transform");
}

void FinalArithTransformPass::apply(OptState& opt) const {
    MIDOUT_B("FinalArithTransformPass::apply")
    Impl{*this, opt};
    MIDOUT_E
}

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}