megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

/**
 * \file src/opr/impl/dnn/softmax.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/opr/dnn/softmax.h"

#include "megbrain/graph/grad_impl.h"
#include "megbrain/opr/internal/out_shape_by_sym_var.h"
#include "megbrain/opr/utility.h"

#include "../internal/megdnn_opr_wrapper.inl"

using namespace mgb;
using namespace opr;

/* ==================== SoftmaxForward  ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(SoftmaxForward);

SoftmaxForward::SoftmaxForward(
        VarNode* inp, const Param& param, const OperatorNodeConfig& config)
        : Super{inp->owner_graph(), config, "softmax", {inp}} {
    init_megdnn_opr(*this, param);

    add_input({inp});
    output(0)->dtype(inp->dtype());
}

SymbolVar SoftmaxForward::make(
        SymbolVar inp, const Param& param, const OperatorNodeConfig& config) {
    auto out = inp.node()
                       ->owner_graph()
                       ->insert_opr(std::make_unique<SoftmaxForward>(
                               inp.node(), param, config))
                       ->output();

    return out[0];
}

void SoftmaxForward::get_output_var_shape(
        const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const {
    out_shape[0] = inp_shape[0];
}

size_t SoftmaxForward::get_workspace_size_bytes(
        const TensorShapeArray& input_shapes,
        const TensorShapeArray& output_shapes) const {
    return megdnn_opr()->get_workspace_in_bytes(
            {input_shapes[0], input(0)->dtype(), input(0)->format()},
            {output_shapes[0], output(0)->dtype(), output(0)->format()});
}

void SoftmaxForward::scn_do_execute() {
    megdnn_opr()->exec(
            input(0)->dev_tensor().as_megdnn(), output(0)->dev_tensor().as_megdnn(),
            intl::get_megdnn_workspace_from_var(output().back()));
}

#if MGB_ENABLE_GRAD
MGB_IMPL_OPR_GRAD(SoftmaxForward) {
    SymbolVar grad = SoftmaxBackward::make(opr.output(0), out_grad[0], opr.param());

    return grad.node();
}
#endif

// /* ==================== SoftmaxBackward ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(SoftmaxBackward);

SoftmaxBackward::SoftmaxBackward(
        VarNode* src, VarNode* diff, const Param& param,
        const OperatorNodeConfig& config)
        : Super({src->owner_graph(), config, "Softmax_backward", {src, diff}}, 0,
                true) {
    init_megdnn_opr(*this, param);
    add_input({src, diff});
}

SymbolVar SoftmaxBackward::make(
        SymbolVar src, SymbolVar diff, const Param& param,
        const OperatorNodeConfig& config) {
    auto out = src.node()
                       ->owner_graph()
                       ->insert_opr(std::make_unique<SoftmaxBackward>(
                               src.node(), diff.node(), param, config))
                       ->output();
    return out[0];
}

void SoftmaxBackward::init_output_static_infer_desc() {
    using namespace cg::static_infer;
    auto&& mgr = owner_graph()->static_infer_manager();
    mgr.register_shape_infer(output(0), ShapeInferDesc::make_identity(input(0)));
    this->init_output_static_infer_desc_workspace(false);
}

void SoftmaxBackward::init_output_dtype() {
    output(0)->dtype(input(0)->dtype());
}

size_t SoftmaxBackward::get_workspace_size_bytes(
        const TensorShapeArray& input_shapes,
        const TensorShapeArray& output_shapes) const {
    return megdnn_opr()->get_workspace_in_bytes(
            {input_shapes[0], input(0)->dtype(), input(0)->format()},
            {input_shapes[1], input(1)->dtype(), input(1)->format()},
            {output_shapes[0], output(0)->dtype(), output(0)->format()});
}

void SoftmaxBackward::scn_do_execute() {
    megdnn_opr()->exec(
            input(0)->dev_tensor().as_megdnn(), input(1)->dev_tensor().as_megdnn(),
            output(0)->dev_tensor().as_megdnn(),
            intl::get_megdnn_workspace_from_var(output().back()));
}

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