megenginelite-sys 1.8.2

/**
 * \file src/gopt/impl/gtrans.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/gtrans.h"
#include "megbrain/opr/basic_arith.h"
#include "megbrain/opr/basic_arith_wrapper.h"
#include "megbrain/opr/blas.h"
#include "megbrain/opr/dnn/convolution.h"
#include "megbrain/opr/tensor_manip.h"

using namespace mgb;
using namespace gopt;
using namespace opr;

namespace {
//! check whether *w* has shape-1 on non-channel axes
bool check_conv_brd_shp(VarNode* w) {
    auto bshp = w->shape();
    if (!bshp.ndim)
        return false;
    for (size_t i = 0; i < bshp.ndim; ++i) {
        if (i + 3 != bshp.ndim && bshp.shape[i] != 1) {
            // only allow non-broadcasting axis in channel
            return false;
        }
    }
    return true;
}

bool normalize_matmul_shape(const TensorShape& src, TensorShape& dst) {
    if (!src.ndim)
        return false;
    dst = src;
    if (src.ndim == 1) {
        ++dst.ndim;
        dst[1] = dst[0];
        dst.shape[0] = 1;
    }
    mgb_assert(dst.ndim == 2);
    return true;
}
}  // anonymous namespace

/* ================ BinaryTrans20 ================ */

class BinaryTrans20::Rule {
protected:
    ~Rule() = default;

public:
    virtual const char* desc() = 0;
    virtual std::pair<Typeinfo*, Typeinfo*> types() = 0;
    virtual VarNode* apply(
            VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
            VarNode* a, VarNode* b, VarNode* c) = 0;
};

GTransResult BinaryTrans20::apply(
        OperatorNodeBase* fop, bool swap_fop_inp, bool swap_gop_inp) const {
    mgb_assert(fop->input().size() == 2);
    auto ab = fop->input(0), c = fop->input(1);
    if (swap_fop_inp) {
        mgb_assert(is_commutable_binary(fop));
        std::swap(ab, c);
    }
    auto gop = ab->owner_opr();
    mgb_assert(gop->input().size() == 2);
    auto a = gop->input(0), b = gop->input(1);
    if (swap_gop_inp) {
        mgb_assert(is_commutable_binary(gop));
        std::swap(a, b);
    }

    auto iter = m_rules.find({fop->dyn_typeinfo(), gop->dyn_typeinfo()});
    GTransResultItem ret;
    if (iter != m_rules.end()) {
        ret.result = iter->second->apply(ret.internal.data(), fop, gop, a, b, c);
        if (ret.result) {
            ret.msg = iter->second->desc();
            return ret;
        }
    }
    return None;
}

//! register a single rule class to given trans object
#define BINARY_TRANS_20_REG_RULE(_cls, t)                                              \
    static _cls _cls##_ins;                                                            \
    do {                                                                               \
        auto ir =                                                                      \
                t.m_rules                                                              \
                        .insert({static_cast<Rule&>(_cls##_ins).types(), &_cls##_ins}) \
                        .second;                                                       \
        mgb_assert(ir);                                                                \
    } while (0)

class BinaryTrans20::AssociativeRuleReg {
    class ElemArith final : public Rule {
        using Mode = Elemwise::Mode;
        const char* desc() override {
            return mgb_cstr_log("elem(elem(x,w1),w2)->elem(x,elem(w1,w2))");
        }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {Elemwise::typeinfo(), Elemwise::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            auto bshp = b->shape(), cshp = c->shape();
            if (!bshp.ndim || !cshp.ndim)
                return nullptr;

            auto&& elem0 = fop->cast_final_safe<Elemwise>();
            auto&& elem1 = gop->cast_final_safe<Elemwise>();
            auto mode = elem0.param().mode;
            if (mode != elem1.param().mode)
                return nullptr;

            if (mode != Mode::ADD && mode != Mode::MUL && mode != Mode::MAX &&
                mode != Mode::MIN) {
                return nullptr;
            }

            auto bcshp = Elemwise::get_output_var_shape(mode, {bshp, cshp});
            if (!bcshp.eq_shape(bshp) && !bcshp.eq_shape(cshp)) {
                // do not allow broadcast
                return nullptr;
            }

            return Elemwise::make(
                           {a, internal[0] = Elemwise::make({b, c}, mode).node()}, mode)
                    .node();
        }
    };

    class ConvMul final : public Rule {
        const char* desc() override { return mgb_cstr_log("conv(x*k,w)->conv(x,w*k)"); }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {Convolution::typeinfo(), Elemwise::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(gop, Elemwise::Mode::MUL) || !check_conv_brd_shp(b))
                return nullptr;

            auto&& orig_conv = fop->cast_final_safe<Convolution>();

            SymbolVar k1 = b;
            if (orig_conv.param().sparse == opr::Convolution::Param::Sparse::GROUP &&
                !k1.shape().is_scalar()) {
                // group convolution with non-scalar multiplicand
                auto one = k1.make_scalar_dt(1);
                k1 = k1.reshape(Concat::make(
                        {GetVarShape::make(c, 0), one, GetVarShape::make(c, 2), one,
                         one},
                        0));
            }
            return Convolution::make(
                           a, internal[0] = (k1 * SymbolVar{c}).node(),
                           orig_conv.param(), orig_conv.execution_policy())
                    .node();
        }
    };
    class MatmulMul final : public Rule {
        const char* desc() override {
            return mgb_cstr_log("matmul(x*k,w)->matmul(x,w*k)");
        }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {MatrixMul::typeinfo(), Elemwise::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(gop, Elemwise::Mode::MUL))
                return nullptr;
            auto&& mm = fop->cast_final_safe<opr::MatrixMul>();
            // axis that must be broadcasting
            TensorShape bshp;
            if (!normalize_matmul_shape(b->shape(), bshp) ||
                bshp[!!mm.param().transposeA] != 1)
                return nullptr;

            SymbolVar tb{b};
            tb = tb.flatten();
            if (mm.param().transposeB) {
                tb = opr::Dimshuffle::make(tb, {-1, 0});
            } else {
                tb = opr::Dimshuffle::make(tb, {0, -1});
            }
            return MatrixMul::make(a, internal[0] = (tb * c).node(), mm.param()).node();
        }
    };
    class MulConv final : public Rule {
        const char* desc() override { return mgb_cstr_log("conv(x,w)*k->conv(x,w*k)"); }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {Elemwise::typeinfo(), Convolution::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(fop, Elemwise::Mode::MUL) || !check_conv_brd_shp(c))
                return nullptr;

            SymbolVar k1{c};

            auto&& orig_conv = gop->cast_final_safe<Convolution>();
            if (orig_conv.param().sparse == opr::Convolution::Param::Sparse::GROUP) {
                auto one = k1.make_scalar_dt(1);
                auto tshp = opr::Concat::make(
                        {GetVarShape::make(b, 0), GetVarShape::make(b, 1), one, one,
                         one},
                        0);
                if (k1.shape().is_scalar()) {
                    k1 = k1.broadcast(tshp);
                } else {
                    k1 = k1.reshape(tshp);
                }

            } else {
                // reshape to [-1, 1, 1, 1]
                k1 = Reshape::make(k1, TensorShape{1, 1, 1, 1}, 0);
            }

            return Convolution::make(
                           a, internal[0] = (k1 * b).node(), orig_conv.param(),
                           orig_conv.execution_policy())
                    .node();
        }
    };
    class MulMatmul final : public Rule {
        const char* desc() override {
            return mgb_cstr_log("matmul(x,w)*k->matmul(x,w*k)");
        }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {Elemwise::typeinfo(), MatrixMul::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(fop, Elemwise::Mode::MUL))
                return nullptr;
            TensorShape cshp;
            if (!normalize_matmul_shape(c->shape(), cshp) || cshp[0] != 1)
                return nullptr;

            auto&& mm = gop->cast_final_safe<opr::MatrixMul>();

            SymbolVar tc{c};
            tc = tc.flatten();
            if (mm.param().transposeB) {
                tc = opr::Dimshuffle::make(tc, {0, -1});
            } else {
                tc = opr::Dimshuffle::make(tc, {-1, 0});
            }
            return MatrixMul::make(a, internal[0] = (tc * b).node(), mm.param()).node();
        }
    };

public:
    AssociativeRuleReg(BinaryTrans20& t) {
        BINARY_TRANS_20_REG_RULE(ElemArith, t);
        BINARY_TRANS_20_REG_RULE(ConvMul, t);
        BINARY_TRANS_20_REG_RULE(MatmulMul, t);
        BINARY_TRANS_20_REG_RULE(MulConv, t);
        BINARY_TRANS_20_REG_RULE(MulMatmul, t);
    }
};

class BinaryTrans20::DistributiveAddRuleReg {
    class ConvAdd final : public Rule {
        const char* desc() override {
            return mgb_cstr_log("conv(x+b,w)->conv(x,w)+b1");
        }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {Convolution::typeinfo(), Elemwise::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(gop, Elemwise::Mode::ADD) || !check_conv_brd_shp(b))
                return nullptr;

            auto&& orig_conv = fop->cast_final_safe<Convolution>();
            auto&& param = orig_conv.param();
            if (param.pad_h || param.pad_w)
                return nullptr;

            internal[0] =
                    Convolution::make(a, c, param, orig_conv.execution_policy()).node();

            if (param.sparse == opr::Convolution::Param::Sparse::GROUP) {
                // group conv
                SymbolVar b1 = b, one = b1.make_scalar(1);
                b1 = b1.reshape(Concat::make(
                        {GetVarShape::make(c, 0), one, GetVarShape::make(c, 2), one,
                         one},
                        0));
                b1 = b1 * c;
                b1 = reduce_sum(
                        b1, Concat::make(
                                    {GetVarShape::make(c, 0), GetVarShape::make(c, 1),
                                     one, one, one},
                                    0));
                b1 = Reshape::make(b1, TensorShape{1, 1, 1, 1}, 1);
                return (b1 + internal[0]).node();
            }

            // dense conv
            SymbolVar b0{b}, w{c},
                    b1 = b0 * w, ochl = GetVarShape::make(w, 0),
                    b1_tshp = Concat::make({ochl, ochl.make_scalar_dt(1)}, 0);
            b1 = Reshape::make(b1, b1_tshp, 1);
            b1 = Reduce::make(b1, {Reduce::Param::Mode::SUM}, b1_tshp);
            return (Dimshuffle::make(b1, {-1, 0, 1, -1}) + internal[0]).node();
        }
    };
    class MatmulAdd final : public Rule {
        const char* desc() override {
            return mgb_cstr_log("matmul(x+b,w)->conv(x,w)+b1");
        }

        std::pair<Typeinfo*, Typeinfo*> types() override {
            return {MatrixMul::typeinfo(), Elemwise::typeinfo()};
        }

        VarNode* apply(
                VarNode** internal, OperatorNodeBase* fop, OperatorNodeBase* gop,
                VarNode* a, VarNode* b, VarNode* c) override {
            if (!as_elem_opr(gop, Elemwise::Mode::ADD))
                return nullptr;
            TensorShape bshp;
            auto&& mm = fop->cast_final_safe<MatrixMul>();
            if (!normalize_matmul_shape(b->shape(), bshp) ||
                bshp[!!mm.param().transposeA] != 1)
                return nullptr;
            auto&& cshp = c->shape();
            if (!cshp.ndim)
                return nullptr;

            bshp[!mm.param().transposeA] = cshp[!!mm.param().transposeB];

            auto bias = MatrixMul::make(SymbolVar{b}.broadcast(bshp), c, mm.param());
            if (bias.shape().ndim)
                mgb_assert(bias.shape()[0] == 1);
            internal[0] = MatrixMul::make(a, c, mm.param()).node();
            return (bias + internal[0]).node();
        }
    };

public:
    DistributiveAddRuleReg(BinaryTrans20& t) {
        BINARY_TRANS_20_REG_RULE(ConvAdd, t);
        BINARY_TRANS_20_REG_RULE(MatmulAdd, t);
    }
};

BinaryTrans20& BinaryTrans20::associtive() {
    static BinaryTrans20 trans;
    static AssociativeRuleReg rule{trans};
    return trans;
}

BinaryTrans20& BinaryTrans20::distributive_add() {
    static BinaryTrans20 trans;
    static DistributiveAddRuleReg rule{trans};
    return trans;
}

/* ================ misc  ================ */

namespace mgb {
namespace gopt {

template <>
VarNode* check_is_group_inverse_opr<Elemwise::Mode::MUL>(SymbolVar x) {
    auto opr = x.node()->owner_opr();
    auto elem = as_elem_opr(opr, Elemwise::Mode::POW);
    if (!elem) {
        if (auto powc = try_cast_as_op<opr::PowC>(opr)) {
            if (almost_equal(powc->param().exp, -1.f)) {
                return powc->input(0);
            }
        }
        return nullptr;
    }
    auto exp = SymbolVar{elem->input(1)}.as_immutable_scalar_require_shape();
    if (exp.valid() && almost_equal(exp->get_cast<float>(), -1.f)) {
        return opr->input(0);
    }
    return nullptr;
}

}  // namespace gopt
}  // namespace mgb

VarNode* gopt::elemwise_reduce_var_list(
        const VarNodeArray& vars, opr::Elemwise::Mode mode, VarNodeArray* mid_results) {
    mgb_assert(!vars.empty());
    VarNode* s = vars[0];
    for (size_t i = 1; i < vars.size(); ++i) {
        s = Elemwise::make({s, vars[i]}, mode).node();
        if (mid_results)
            mid_results->push_back(s);
    }
    return s;
}

VarNode* gopt::get_opr_single_output_var(OperatorNodeBase* opr) {
    VarNode* ret = nullptr;
    for (auto i : opr->output()) {
        if (!i->contain_flag(VarNode::Flag::VOLATILE_CONTENT)) {
            if (ret)
                return nullptr;
            ret = i;
        }
    }
    return ret;
}

void gopt::visit_opr_tree(
        VarNode* endpoint, const thin_function<bool(VarNode*)>& check_internal,
        const thin_function<void(VarNode*)>& on_leaf,
        const thin_function<void(OperatorNodeBase*)>& on_internal_finish,
        bool allow_multi_cn) {
    struct StackFrame {
        using VarNodeArrayPtr = VarNode* const*;
        VarNodeArrayPtr var0, var1;
        OperatorNodeBase* reader_opr;
    };
    std::vector<StackFrame> stack;
    stack.push_back({&endpoint, &endpoint + 1, nullptr});
    while (!stack.empty()) {
        auto&& top = stack.back();
        if (top.var0 == top.var1) {
            if (top.reader_opr && on_internal_finish) {
                on_internal_finish(top.reader_opr);
            }
            stack.pop_back();
            continue;
        }
        VarNode* var = *(top.var0++);
        if (check_internal(var)) {
            auto opr = var->owner_opr();
            mgb_assert(
                    var == opr->output(0),
                    "bad check_internal() provided to visit_opr_tree");
            auto&& inp = opr->input();
            if (!allow_multi_cn) {
                bool multi_cn = false;
                for (auto i : inp) {
                    if (i->comp_node() != var->comp_node()) {
                        multi_cn = true;
                        break;
                    }
                }
                if (multi_cn) {
                    if (on_leaf)
                        on_leaf(var);
                    continue;
                }
            }
            stack.push_back({inp.data(), inp.data() + inp.size(), opr});
        } else {
            if (on_leaf)
                on_leaf(var);
        }
    }
}

VarNodeArray gopt::extract_opr_leaves(
        VarNode* endpoint, const std::function<bool(OperatorNodeBase*)>& pred,
        bool allow_multi_cn) {
    VarNodeArray ret;
    auto check_internal = [&](VarNode* var) -> bool { return pred(var->owner_opr()); };
    auto on_leaf = [&ret](VarNode* var) { ret.push_back(var); };
    visit_opr_tree(endpoint, check_internal, on_leaf);
    return ret;
}

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