elemwise.cpp 11.5 KB
Newer Older
1 2 3 4
/**
 * \file imperative/src/impl/ops/elemwise.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8 9 10 11
 *
 * 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.
 */

M
Megvii Engine Team 已提交
12
#include "megbrain/imperative/opr_utility.h"
13
#include "megbrain/imperative/ops/autogen.h"
14
#include "megbrain/imperative/utils/stats.h"
15
#include "megbrain/opr/basic_arith.h"
16
#include "megbrain/opr/utility.h"
17

18
#include "../blob_manager_impl.h"
M
Megvii Engine Team 已提交
19 20
#include "../dnn_op_helper.h"
#include "../op_trait.h"
21 22 23 24 25 26 27 28 29 30 31

namespace mgb {
namespace imperative {

namespace {

std::shared_ptr<OpDef> make_from_op_node(cg::OperatorNodeBase* node_) {
    auto* node = &node_->cast_final_safe<opr::Elemwise>();
    return Elemwise::make(node->param().mode);
}

M
Megvii Engine Team 已提交
32
auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
33
    auto&& elemwise_opr = def.cast_final_safe<Elemwise>();
34
    OperatorNodeConfig config{elemwise_opr.make_name()};
35
    return opr::Elemwise::make(inputs, elemwise_opr.mode, config);
36 37
}

38
std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible(
M
Megvii Engine Team 已提交
39
        const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) {
40
    auto&& op_def = def.cast_final_safe<Elemwise>();
41
    auto trait = megdnn::Elemwise::ModeTrait::from_mode(op_def.mode);
M
Megvii Engine Team 已提交
42 43 44
    mgb_assert(
            inputs.size() == trait.arity, "%s expects %u inputs; got %zu actually",
            trait.name, trait.arity, inputs.size());
45 46 47
    TensorShapeArray inp_shapes;
    DType out_dt;
    CompNode out_cn;
M
Megvii Engine Team 已提交
48 49
    for (size_t i = 0; i < inputs.size(); ++i) {
        auto&& t = inputs[i];
50 51 52 53 54 55 56 57 58 59 60 61 62
        if (!i) {
            out_cn = t.comp_node;
            out_dt = t.layout.dtype;
        } else {
            mgb_assert(t.comp_node == out_cn);
            mgb_assert(t.layout.dtype == out_dt);
        }
        if (t.layout.ndim > 0) {
            inp_shapes.push_back(t.layout);
        } else {
            TensorLayout out_layout;
            out_layout.ndim = 0;
            out_layout.dtype = out_dt;
63
            return {{{out_layout, out_cn}}, false};
64 65
        }
    }
66 67 68
    // copy from megdnn::ElemwiseForward::check_dtype
    switch (out_dt.category()) {
        case DTypeCategory::FLOAT:
M
Megvii Engine Team 已提交
69
            mgb_assert(trait.allow_float, "unsupport mode %s for float\n", trait.name);
70 71
            break;
        case DTypeCategory::INT:
M
Megvii Engine Team 已提交
72
            mgb_assert(trait.allow_int, "unsupport mode %s for int\n", trait.name);
73 74
            break;
        case DTypeCategory::BOOL:
M
Megvii Engine Team 已提交
75
            mgb_assert(trait.allow_bool, "unsupport mode %s for bool\n", trait.name);
76 77 78 79 80 81
            break;
        default:
            // Quantized Dtype could also be handled by this op,
            // but scales need to be the same.
            break;
    }
82 83

    auto&& out_shape = opr::Elemwise::get_output_var_shape(op_def.mode, inp_shapes);
84
    return {{{TensorLayout(out_shape, out_dt, inputs[0].layout.format), out_cn}}, true};
85 86
}

87
DispatchMode decide_dispatch_mode(
M
Megvii Engine Team 已提交
88
        const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) {
89 90 91
    bool host_computable = true;
    constexpr int size_threshhold = TensorShape::MAX_NDIM;
    for (auto&& inp : inputs) {
M
Megvii Engine Team 已提交
92 93
        if (inp.value.empty() || inp.value.layout().ndim == 0 ||
            inp.value.layout().total_nr_elems() > size_threshhold) {
94 95 96 97 98 99 100 101
            host_computable = false;
            break;
        }
    }
    return host_computable ? DEFAULT_CPU : KERNEL;
}

void apply_on_device_tensornd(
M
Megvii Engine Team 已提交
102
        const OpDef& def, const SmallVector<DeviceTensorND>& inputs,
103
        SmallVector<DeviceTensorND>* outputs) {
104
    auto&& op_def = def.cast_final_safe<Elemwise>();
105
    auto&& trait = megdnn::Elemwise::ModeTrait::from_mode(op_def.mode);
M
Megvii Engine Team 已提交
106 107 108 109 110
    mgb_assert(
            inputs.size() == trait.arity, "%s expects %u inputs; got %zu actually",
            trait.name, trait.arity, inputs.size());
    auto&& dnn_opr =
            opr::intl::create_megdnn_opr<megdnn::Elemwise>(inputs[0].comp_node());
111 112 113 114
    opr::Elemwise::perform(op_def.mode, (*outputs)[0], inputs, dnn_opr);
}

SmallVector<TensorPtr> apply_on_physical_tensor(
M
Megvii Engine Team 已提交
115
        const OpDef& def, const SmallVector<TensorPtr>& inputs) {
116
    auto&& op_def = def.cast_final_safe<Elemwise>();
117
    SmallVector<DeviceTensorND> inp_tensornds(inputs.size());
118
    TensorShapeArray inp_shapes(inputs.size());
M
Megvii Engine Team 已提交
119
    for (unsigned i = 0; i < inputs.size(); ++i) {
120
        inp_tensornds[i] = inputs[i]->dev_tensor();
121
        inp_shapes[i] = inputs[i]->layout();
122
    }
M
Megvii Engine Team 已提交
123 124 125
    TensorShape shape = opr::Elemwise::get_output_var_shape(op_def.mode, inp_shapes);
    DeviceTensorND out = BlobManager::inst()->alloc_workspace_with_defrag(
            inp_tensornds[0].comp_node(), {shape, inp_tensornds[0].layout().dtype});
126
    SmallVector<DeviceTensorND> oup_tensornds = {out};
127 128
    apply_on_device_tensornd(def, inp_tensornds, &oup_tensornds);
    return {Tensor::make(oup_tensornds[0])};
129 130
}

M
Megvii Engine Team 已提交
131 132
MGB_DEFINE_OPR_CLASS(
        ForceInplaceElemwise,
133
        cg::SingleCNOperatorNodeBaseT<opr::mixin::MegDNNOprHolder>) // {
134
public:
135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
    struct Param {
        using Mode = megdnn::Elemwise::Param::Mode;
        Mode mode;
        size_t inplace_index;
    };
    using Mode = Param::Mode;
    ForceInplaceElemwise(
            const VarNodeArray& inputs, Param param, OperatorNodeConfig config = {})
            : Super(inputs[0]->owner_graph(), config, "device_add_update", inputs),
              m_param{param} {
        for (auto* input : inputs) {
            add_input({input});
        }
        add_output(None)
                ->set_fwd_in2out_writable_force(input(param.inplace_index))
                .add_flag(VarNode::Flag::NO_MEM_RECLAIM);
151
    }
152 153 154 155 156 157 158 159
    static SymbolVar make(const VarNodeArray& inputs, Param param) {
        return SymbolVar{inputs[0]}.insert_single_output_opr<ForceInplaceElemwise>(
                inputs, param);
    }
    static cg::OperatorNodeBase* shallow_copy(
            const serialization::OprShallowCopyContext& ctx,
            const cg::OperatorNodeBase& opr_, const VarNodeArray& inputs,
            const OperatorNodeConfig& config);
M
Megvii Engine Team 已提交
160

161
protected:
162 163 164 165
    NodeProp* do_make_node_prop() const override {
        auto ret = Super::do_make_node_prop();
        ret->add_flag(NodeProp::Flag::FORCE_UPDATE_INPUT_VAR);
        return ret;
166
    }
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187
    void create_megdnn_opr() override {
        auto opr = DnnOprCaller<megdnn::Elemwise>::create_operator(comp_node());
        opr->param().mode = m_param.mode;
        set_megdnn_opr(std::move(opr));
    }
    void scn_do_execute() override {
        auto to_dnnnd = [&](auto* var) { return var->dev_tensor().as_megdnn(); };
        megdnn::TensorNDArray inputs_dnnnd;
        for (auto* input : input()) {
            inputs_dnnnd.push_back(to_dnnnd(input));
        }
        mgb_assert(
                input(m_param.inplace_index)
                        ->contain_flag(VarNode::Flag::NO_SYS_MEM_ALLOC),
                "ForceInplaceElemwise cannot be applied in internal tensor");
        auto* out_dest = output(0);
        auto* opr = static_cast<megdnn::Elemwise*>(megdnn_opr());
        opr->exec(std::move(inputs_dnnnd), to_dnnnd(out_dest));
    }
    void init_output_static_infer_desc() override {
        using namespace cg::static_infer;
M
Megvii Engine Team 已提交
188

189 190 191
        owner_graph()->static_infer_manager().register_shape_infer(
                output(0), ShapeInferDesc::make_identity(input(m_param.inplace_index)));
    }
192 193

private:
194 195 196 197
    Param m_param;
    void record_execute_deps(ExecDependencyArray& deps) override {
        record_megdnn_opr(deps);
    }
198 199 200 201 202
};

MGB_DYN_TYPE_OBJ_FINAL_IMPL(ForceInplaceElemwise);

cg::OperatorNodeBase* ForceInplaceElemwise::shallow_copy(
M
Megvii Engine Team 已提交
203 204 205 206
        const serialization::OprShallowCopyContext& ctx,
        const cg::OperatorNodeBase& opr_, const VarNodeArray& inputs,
        const OperatorNodeConfig& config) {
    auto&& opr = opr_.cast_final_safe<ForceInplaceElemwise>();
207
    auto* graph = ctx.owner_graph(opr, inputs);
M
Megvii Engine Team 已提交
208 209
    return graph->insert_opr(
            std::make_unique<ForceInplaceElemwise>(inputs, opr.m_param, config));
210 211 212 213 214
}

MGB_REG_OPR_SHALLOW_COPY(ForceInplaceElemwise, ForceInplaceElemwise::shallow_copy);

cg::OperatorNodeBase* apply_inplace_add_on_var_node(
M
Megvii Engine Team 已提交
215 216
        const OpDef& def, const VarNodeArray& inputs) {
    auto dest = inputs[0], delta = inputs[1], alpha = inputs[2], beta = inputs[3];
217
    auto mode = ForceInplaceElemwise::Param::Mode::FUSE_MUL_ADD4;
M
Megvii Engine Team 已提交
218 219 220
    return ForceInplaceElemwise::make({alpha, dest, beta, delta}, {mode, 1})
            .node()
            ->owner_opr();
221 222 223
}

SmallVector<TensorPtr> apply_inplace_add_on_physical_tensor(
M
Megvii Engine Team 已提交
224 225
        const OpDef& def, const SmallVector<TensorPtr>& inputs) {
    mgb_assert(
226
            inputs[0]->blob().use_count() == 1 && inputs[0]->blob()->storage().unique(),
227
            "This inplace modification may change the elements of other tensors. "
M
Megvii Engine Team 已提交
228 229 230
            "Please set MEGENGINE_INPLACE_UPDATE to 0 to ensure the program runs "
            "correctly.");
    auto dest = inputs[0], delta = inputs[1], alpha = inputs[2], beta = inputs[3];
231 232 233 234
    auto tensor_to_scalar = [](const TensorPtr& tensor) -> float {
        return *tensor->get_value().ptr<float>();
    };
    DnnOprCaller<megdnn::AddUpdate> caller{dest->comp_node()};
M
Megvii Engine Team 已提交
235
    caller.op->param() = {tensor_to_scalar(alpha), tensor_to_scalar(beta)};
236
    caller.op->exec(dest->dev_tensor().as_megdnn(), delta->dev_tensor().as_megdnn());
M
Megvii Engine Team 已提交
237
    return {std::make_shared<Tensor>(dest->blob(), dest->offset(), dest->layout())};
238 239 240
}

std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_inplace_add_output_attrs_fallible(
M
Megvii Engine Team 已提交
241
        const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) {
242 243
    mgb_assert(inputs.size() == 4, "invalid input number for inplace_add");
    CompNode cn;
M
Megvii Engine Team 已提交
244
    for (auto&& input : inputs) {
245 246 247 248 249 250
        if (!cn.valid()) {
            cn = input.comp_node;
        } else {
            mgb_assert(input.comp_node == cn, "inputs should be in same comp_node");
        }
    }
M
Megvii Engine Team 已提交
251
    auto dest = inputs[0], delta = inputs[1], alpha = inputs[2], beta = inputs[3];
252 253
    bool succeed = dest.layout.ndim != 0;
    if (succeed) {
M
Megvii Engine Team 已提交
254 255 256 257 258 259 260 261 262
        mgb_assert(
                delta.layout.ndim == 0 || dest.layout.eq_shape(delta.layout),
                "dest and delta must have same shape");
        mgb_assert(
                alpha.layout.ndim == 0 || alpha.layout.eq_shape({1}),
                "alpha should be scalar");
        mgb_assert(
                beta.layout.ndim == 0 || beta.layout.eq_shape({1}),
                "beta should be scalar");
263 264 265
    }
    mgb_assert(alpha.layout.dtype == dtype::Float32(), "alpha should be float32");
    mgb_assert(beta.layout.dtype == dtype::Float32(), "beta should be float32");
266 267
    // inplace op result's desc value is changed
    return {{{dest.layout, dest.comp_node}}, succeed};
268 269
}

270
OP_TRAIT_REG(Elemwise, Elemwise, opr::Elemwise)
M
Megvii Engine Team 已提交
271 272 273 274 275 276 277
        .make_from_op_node(make_from_op_node)
        .decide_dispatch_mode(decide_dispatch_mode)
        .apply_on_var_node(apply_on_var_node)
        .infer_output_attrs_fallible(infer_output_attrs_fallible)
        .apply_on_device_tensornd(apply_on_device_tensornd)
        .apply_on_physical_tensor(apply_on_physical_tensor)
        .fallback();
278 279

OP_TRAIT_REG(InplaceAdd, InplaceAdd, opr::AddUpdate)
M
Megvii Engine Team 已提交
280 281 282 283 284
        .apply_on_var_node(apply_inplace_add_on_var_node)
        .apply_on_physical_tensor(apply_inplace_add_on_physical_tensor)
        .infer_output_attrs_fallible(infer_inplace_add_output_attrs_fallible)
        .fallback();
}  // anonymous namespace
285 286 287 288 289

}  // namespace imperative
}  // namespace mgb

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