Function TEST

src/gopt/test/reformat_manager.cpp:12–59

using namespace gopt;

TEST(TestReformatManager, Feature) {
    constexpr size_t N = 16, C = 128, H = 7, W = 7;
    HostTensorGenerator<> gen;
    using ReformatKey = ReformatManager::ReformatKey;
    auto src_format = TensorFormats::NHWC, dst_format = TensorFormats::NCHWc64;
    ReformatKey key{src_format, dst_format};
    auto reformat = ReformatManager::instance().get(key);

    auto graph = ComputingGraph::make();
    graph->options().graph_opt_level = 0;

    auto r = [](VarNode* inp) {
        auto x = SymbolVar(inp);
        auto xshp = opr::GetVarShape::make(x);
        auto cv = [&x](int v) { return x.make_scalar(v); };
        auto sub = [&xshp, &cv](int idx) {
            return opr::IndexAt::make(xshp, {{0, cv(idx)}});
        };
        auto tshp0 =
                opr::Concat::make({sub(0), sub(1), sub(2), sub(3) / 64, cv(64)}, 0);
        auto y0 = opr::Reshape::make(x, tshp0);
        auto y1 = opr::Dimshuffle::make(y0, {0, 3, 1, 2, 4});
        return y1;
    };

    auto mkvar = [&](const char* name, const TensorShape& shp) {
        return opr::Host2DeviceCopy::make(*graph, gen(shp)).rename(name);
    };
    auto x = mkvar("x", {N, H, W, C});
    auto y1 = SymbolVar(reformat({x.node()}));
    auto y2 = r(x.node());
    size_t nr_shapeof = 0;
    size_t nr_reshape = 0;
    cg::DepOprIter{[&nr_shapeof, &nr_reshape](cg::OperatorNodeBase* o) {
        if (o->same_type<opr::GetVarShape>())
            nr_shapeof++;
        if (o->same_type<opr::Reshape>())
            nr_reshape++;
    }}.add(y1.node()->owner_opr());
    ASSERT_EQ(nr_shapeof, 1);
    ASSERT_EQ(nr_reshape, 1);
    HostTensorND t1, t2;
    auto func1 = graph->compile({make_callback_copy(y1, t1)});
    func1->execute();
    auto func2 = graph->compile({make_callback_copy(y2, t2)});
    func2->execute();
    MGB_ASSERT_TENSOR_EQ(t1, t2);
}

TEST(TestReformatManager, Weight) {
    constexpr size_t G = 8, K = 128, C = 128, R = 3, S = 3;