megenginelite-sys 1.8.2

/**
 * \file src/gopt/test/helper.h
 * 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.
 */

#pragma once

#include "megbrain/test/helper.h"

#include "megbrain/gopt/framework.h"
#include "megbrain/opr/io.h"

namespace mgb {
//! make an opr that reads \p x; only used for test
SymbolVar opr_reader_for_test(SymbolVar x);

template <class Pass>
class TestGoptBasicArithPass : public ::testing::Test {
protected:
    HostTensorGenerator<> gen;
    std::shared_ptr<ComputingGraph> graph = ComputingGraph::make();

    SymbolVar mkvar(
            const char* name, const TensorShape& shp = {1},
            CompNode cn = CompNode::load("xpu0")) {
        return opr::Host2DeviceCopy::make(*graph, gen(shp), cn).rename(name);
    }

    SymbolVar mkcvar(
            const char* name, const TensorShape& shp = {1},
            CompNode cn = CompNode::load("xpu0")) {
        return opr::SharedDeviceTensor::make(*graph, *gen(shp), cn).rename(name);
    }

    template <typename... Args>
    SymbolVarArray run_opt(const SymbolVarArray& inp, Args&&... args) {
        return gopt::GraphOptimizer{}
                .add_pass<Pass>(std::forward<Args>(args)...)
                .apply({{inp}})
                .endpoint_vars();
    }

    template <bool check_ne = true, typename... Args>
    void check(SymbolVar expect, SymbolVar inp, Args&&... args) {
        if (check_ne) {
            ASSERT_NE(expect.node(), inp.node());
        } else {
            ASSERT_EQ(expect, inp);
        }
        SymbolVar get;
        unpack_vector(run_opt({inp}, std::forward<Args>(args)...), get);
        ASSERT_EQ(expect, get);

        // test multiple readers
        unpack_vector(
                gopt::GraphOptimizer{}
                        .add_pass<Pass>(std::forward<Args>(args)...)
                        .apply({{inp + opr_reader_for_test(inp)}})
                        .endpoint_vars(),
                get);

        ASSERT_EQ(expect + opr_reader_for_test(expect), get);
    }
};
}  // namespace mgb

#define TEST_PASS(pass, name)                                  \
    using TestGopt##pass = TestGoptBasicArithPass<gopt::pass>; \
    TEST_F(TestGopt##pass, name)

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