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

/**
 * \file src/opr/impl/dnn/roi_align.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/roi_align.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;

/* ==================== ROIAlignForward  ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(ROIAlignForward);
ROIAlignForward::ROIAlignForward(
        VarNode* src, VarNode* rois, const Param& param,
        const OperatorNodeConfig& config)
        : Super{src->owner_graph(), config, "roi_align", {src, rois}} {
    init_megdnn_opr(*this, param);
    mgb_assert(src->dtype() == dtype::Float32());
    add_input({src, rois});
    output(0)->dtype(dtype::Float32());
    output(1)->dtype(dtype::Int32());
}

SymbolVar ROIAlignForward::make(
        SymbolVar src, SymbolVar rois, const Param& param,
        const OperatorNodeConfig& config) {
    return src.insert_single_output_opr<ROIAlignForward>(
            src.node(), rois.node(), param, config);
}

#if MGB_ENABLE_GRAD
MGB_IMPL_OPR_GRAD(ROIAlignForward) {
    if (wrt_idx == 0) {
        // wrt src
        SymbolVar grad = ROIAlignBackward::make(
                out_grad[0], opr.input(0), opr.input(1), opr.output(1), opr.param());
        return grad.node();
    } else {
        mgb_assert(wrt_idx == 1);
        return nullptr;
    }
}
#endif

/* ==================== ROIAlignBackward ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(ROIAlignBackward);
MEGDNN_OPR_INIT4(ROIAlignBackward, "roi_align_backward", 1, true);

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

size_t ROIAlignBackward::get_workspace_size_bytes(
        const TensorShapeArray& inp_shapes, const TensorShapeArray& out_shapes) const {
    TensorLayout diff{inp_shapes[0], input(0)->dtype(), input(0)->format()},
            rois{inp_shapes[2], input(2)->dtype(), input(2)->format()},
            index{inp_shapes[3], input(3)->dtype(), input(3)->format()},
            grad{out_shapes[0], output(0)->dtype(), output(0)->format()};
    return megdnn_opr()->get_workspace_in_bytes(diff, rois, index, grad);
}

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