Function expand_elem_chain

src/gopt/test/basic_arith.cpp:28–42

! * get all operands of a chain of element-wise oprs of the same mode */

 * get all operands of a chain of element-wise oprs of the same mode
 */
std::vector<VarNode*> expand_elem_chain(VarNode* var, opr::Elemwise::Mode mode) {
    if (var->owner_opr()->same_type<opr::Elemwise>()) {
        auto&& op = var->owner_opr()->cast_final<opr::Elemwise>();
        if (op.param().mode == mode) {
            auto ret = expand_elem_chain(op.input(0), mode);
            for (size_t i = 1; i < op.input().size(); ++i) {
                auto cur = expand_elem_chain(op.input()[i], mode);
                ret.reserve(ret.size() + cur.size());
                ret.insert(ret.end(), cur.begin(), cur.end());
            }
            return ret;
        }
    }
    return {var};
}

SymbolVar fma3(SymbolVar a, SymbolVar b, SymbolVar c) {
    return Elemwise::make({a, b, c}, Mode::FUSE_MUL_ADD3);