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

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

/* ==================== CorrelationForward  ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(CorrelationForward);
CorrelationForward::CorrelationForward(
        VarNode* data1, VarNode* data2, const Param& param,
        const OperatorNodeConfig& config)
        : Super{data1->owner_graph(), config, "correlation", {data1, data2}} {
    init_megdnn_opr(*this, param);
    mgb_assert(data1->dtype() == data2->dtype());
    mgb_assert(data1->dtype().category() == DTypeCategory::FLOAT);

    add_input({data1, data2});
    output(0)->dtype(data1->dtype());
}

SymbolVar CorrelationForward::make(
        SymbolVar data1, SymbolVar data2, const Param& param,
        const OperatorNodeConfig& config) {
    return data1.insert_single_output_opr<CorrelationForward>(
            data1.node(), data2.node(), param, config);
}

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

/* ==================== CorrelationBackwardData1 ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(CorrelationBackwardData1);
MEGDNN_OPR_INIT3(CorrelationBackwardData1, "correlation_backward_data1", 1, true);

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

size_t CorrelationBackwardData1::get_workspace_size_bytes(
        const TensorShapeArray& inp_shapes, const TensorShapeArray& out_shapes) const {
    TensorLayout diff{inp_shapes[0], input(0)->dtype(), input(0)->format()},
            data1{inp_shapes[1], input(1)->dtype(), input(1)->format()},
            data2{inp_shapes[2], input(2)->dtype(), input(2)->format()},
            grad1{out_shapes[0], output(0)->dtype(), output(0)->format()};
    return megdnn_opr()->get_workspace_in_bytes(diff, data1, data2, grad1);
}

/* ==================== CorrelationBackwardData2 ==================== */
MGB_DYN_TYPE_OBJ_FINAL_IMPL(CorrelationBackwardData2);
MEGDNN_OPR_INIT3(CorrelationBackwardData2, "correlation_backward_data2", 1, true);

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

size_t CorrelationBackwardData2::get_workspace_size_bytes(
        const TensorShapeArray& inp_shapes, const TensorShapeArray& out_shapes) const {
    TensorLayout diff{inp_shapes[0], input(0)->dtype(), input(0)->format()},
            data1{inp_shapes[1], input(1)->dtype(), input(1)->format()},
            data2{inp_shapes[2], input(2)->dtype(), input(2)->format()},
            grad2{out_shapes[0], output(0)->dtype(), output(0)->format()};
    return megdnn_opr()->get_workspace_in_bytes(diff, data1, data2, grad2);
}

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