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

/**
 * \file src/opr/impl/internal/out_shape_by_sym_var.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/internal/out_shape_by_sym_var.h"
#include "megbrain/opr/io.h"
#include "megbrain/utils/arith_helper.h"

using namespace mgb;
using namespace opr;
using namespace mixin;

/* ===================== OutshapeBySymvarOpr ===================== */

OutshapeBySymvarOpr::~OutshapeBySymvarOpr() = default;

void OutshapeBySymvarOpr::mixin_outshape_by_symvar_enable(
        OperatorNodeBase& opr, size_t nr_shape_inp, size_t hostval_inp_start) {
    mgb_assert(!m_enable_out_shape_by_symbol_var);
    mgb_assert(
            hostval_inp_start >= nr_shape_inp &&
            hostval_inp_start < opr.input().size());
    m_nr_shape_inp = nr_shape_inp;
    m_hostval_inp_start = hostval_inp_start;
    m_enable_out_shape_by_symbol_var = true;

    m_shape_infer_info.shape_inp_shp.resize(nr_shape_inp);
    m_shape_infer_info.shpval_inp_val.resize(opr.input().size() - hostval_inp_start);
}

void OutshapeBySymvarOpr::mixin_init_output_static_infer_desc(OperatorNodeBase& opr) {
    using namespace cg::static_infer;
    DepVal deps;
    for (size_t i = 0; i < m_nr_shape_inp; ++i)
        deps.push_back({opr.input(i), DepType::SHAPE});
    for (size_t i = m_hostval_inp_start; i < opr.input().size(); ++i)
        deps.push_back({opr.input(i), DepType::VALUE});

    auto infer_shape = [&opr, this](TensorShape& dest, const InpVal&) {
        outshape_by_symvar_do_get_output_shape(
                dest, mixin_outshape_by_symvar_get_shape_infer_info(opr));
        return true;
    };
    opr.owner_graph()->static_infer_manager().register_shape_infer(
            opr.output(0), {SourceType::DEP, deps, infer_shape});
}

const OutshapeBySymvarOpr::ShapeInferInfo& OutshapeBySymvarOpr::
        mixin_outshape_by_symvar_get_shape_infer_info(
                const OperatorNodeBase& opr) const {
    auto&& mgr = opr.owner_graph()->static_infer_manager();

    for (size_t i = 0; i < m_nr_shape_inp; ++i) {
        m_shape_infer_info.shape_inp_shp[i] = mgr.infer_shape(opr.input(i));
    }

    for (size_t i = m_hostval_inp_start; i < opr.input().size(); ++i) {
        m_shape_infer_info.shpval_inp_val[i - m_hostval_inp_start] =
                &mgr.infer_value(opr.input(i));
    }
    return m_shape_infer_info;
}

void OutshapeBySymvarOpr::mixin_outshape_by_symvar_reset_node_dep_type(
        const OperatorNodeBase& opr, NodeProp* prop) const {
    SmallVector<NodeProp::DepType> dt(opr.input().size(), NodeProp::DepType::DEV_VALUE);
    for (size_t i = m_hostval_inp_start; i < opr.input().size(); ++i) {
        dt[i] = NodeProp::DepType::HOST_VALUE;
    }
    prop->reset_dep_type(opr.input(), dt);
}

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