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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

/**
 * \file src/opr/test/internal.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/io.h"
#include "megbrain/test/helper.h"

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

using namespace mgb;
using opr::intl::WorkspaceLimitGetter;

namespace {

//! forward unchanged value and set m_infer_called flag
MGB_DEFINE_OPR_CLASS(WorkspaceLimitGetterOpr, cg::SingleCNOperatorNodeBase) // {
public:
    using InferShapeCallback = thin_function<void()>;

    WorkspaceLimitGetterOpr(VarNode* inp, const InferShapeCallback& cb)
            : Super{inp->owner_graph(), {}, "workspace_limit_getter_opr", {inp}},
              m_infer_shape_callback{cb} {
        add_input({inp});
        add_output(None);
    }

    static SymbolVar make(SymbolVar inp, const InferShapeCallback& cb) {
        return inp.insert_single_output_opr<WorkspaceLimitGetterOpr>(inp.node(), cb);
    }

private:
    InferShapeCallback m_infer_shape_callback;

    void scn_do_execute() override {
        output(0)->dev_tensor().copy_from_fixlayout(input(0)->dev_tensor());
    }

    void init_output_static_infer_desc() override {
        using namespace cg::static_infer;
        auto infer_shp = [this](TensorShape& dest, const InpVal& inp) {
            dest = inp.val.at(0).shape();
            m_infer_shape_callback();
            return true;
        };
        auto&& mgr = owner_graph()->static_infer_manager();
        auto ivar = input(0), ovar = output(0);
        auto wk_var = WorkspaceLimitGetter::register_to_graph(owner_graph());
        ASSERT_NE(nullptr, wk_var);
        mgr.register_shape_infer(
                ovar, {SourceType::DEP,
                       {{ivar, DepType::SHAPE}, {wk_var, DepType::VALUE}},
                       infer_shp});
    }
};
MGB_DYN_TYPE_OBJ_FINAL_IMPL(WorkspaceLimitGetterOpr);

void run_test(bool dynamic) {
    HostTensorGenerator<> gen;
    auto graph = ComputingGraph::make();

    if (dynamic) {
        graph->options().force_dynamic_alloc = true;
    }

    auto x = opr::SharedDeviceTensor::make(*graph, *gen({23}));

    int infer_shape_nr_call = 0;
    auto infer_shape_callback = [&]() {
        ++infer_shape_nr_call;
        if (infer_shape_nr_call < 3) {
            ASSERT_TRUE(WorkspaceLimitGetter::is_prealloc_run(graph.get()));
        } else {
            ASSERT_FALSE(WorkspaceLimitGetter::is_prealloc_run(graph.get()));
            auto wk = WorkspaceLimitGetter::get_workspace_limit(
                    graph.get(), x.node()->comp_node(), 123);
            ASSERT_GT(wk, 0u);
            ASSERT_LE(wk, 123u);
            return;
        }
    };

    auto y = WorkspaceLimitGetterOpr::make(x, infer_shape_callback);
    ASSERT_EQ(1, infer_shape_nr_call);

    graph->compile({{x, {}}})->execute();
    ASSERT_EQ(1, infer_shape_nr_call);

    auto func1 = graph->compile({{y, {}}});
    ASSERT_EQ(1, infer_shape_nr_call);
    func1->execute();
    ASSERT_EQ(3, infer_shape_nr_call);

    func1->execute();
    ASSERT_EQ(3, infer_shape_nr_call);
}

}  // namespace

TEST(TestOprInternal, WorkspaceLimitGetter) {
    run_test(false);
}

TEST(TestOprInternal, WorkspaceLimitGetterDynamic) {
    run_test(true);
}

TEST(TestOprInternal, WorkspaceLimitGetterWithoutOpt) {
    HostTensorGenerator<> gen;
    auto graph = ComputingGraph::make();
    graph->options().seq_opt.enable_mem_reuse_alloc = false;

    auto x = opr::SharedDeviceTensor::make(*graph, *gen({23}));

    int infer_shape_nr_call = 0;
    auto infer_shape_callback = [&]() {
        ++infer_shape_nr_call;
        ASSERT_FALSE(WorkspaceLimitGetter::is_prealloc_run(graph.get()));
        auto wk = WorkspaceLimitGetter::get_workspace_limit(
                graph.get(), x.node()->comp_node(), 123);
        ASSERT_GT(wk, 0u);
        ASSERT_LE(wk, 123u);
    };

    auto y = WorkspaceLimitGetterOpr::make(x, infer_shape_callback);
    ASSERT_EQ(1, infer_shape_nr_call);

    graph->compile({{x, {}}})->execute();
    ASSERT_EQ(1, infer_shape_nr_call);

    auto func1 = graph->compile({{y, {}}});
    ASSERT_EQ(1, infer_shape_nr_call);
    func1->execute();
    ASSERT_EQ(2, infer_shape_nr_call);

    func1->execute();
    ASSERT_EQ(2, infer_shape_nr_call);
}

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