Function TEST

src/gopt/test/subgraph_extractor.cpp:59–111

}  // namespace

TEST(TestSubGraphExtractor, MultipleOutputs) {
    HostTensorGenerator<> gen;
    auto graph = ComputingGraph::make();

    auto mkvar = [&](const char* name, const TensorShape& shp) {
        return opr::Host2DeviceCopy::make(*graph, gen(shp)).rename(name);
    };

    auto mkcvar = [&](const char* name, const TensorShape& shp) {
        return opr::SharedDeviceTensor::make(*graph, *gen(shp)).rename(name);
    };

    graph->options().graph_opt_level = 0;
    auto x = mkvar("x", {8, 8, 8, 8}), w1 = mkcvar("w1", {4, 8, 3, 3});
    auto y = mkvar("y", {1, 8, 1, 1});
    auto add = x + y;

    opr::Convolution::Param param;
    param.pad_h = param.pad_w = 1;
    auto c1 = opr::Convolution::make(add, w1, param);
    auto w2 = mkcvar("w2", {8, 4, 3, 3});
    auto c2 = opr::ConvolutionBackwardData::make(w2, add, param, {}, {});
    auto sym_var_arr = MultipleInputOutput::make({c1, c2});
    auto z = sym_var_arr[1];
    z = z + (-128);

    using OprList = SubGraphExtractor::OprList;
    static const OprList opr_list = {
            opr::ConvolutionForward::typeinfo(),
            opr::Elemwise::typeinfo(),
            opr::TypeCvt::typeinfo(),
            MultipleInputOutput::typeinfo(),
    };
    SubGraphExtractor extractor(opr_list);
    auto partitions = extractor.extract({z});
    ASSERT_EQ(partitions.size(), 1u);
    // outputs: sym_var_arr[0], z, add
    ASSERT_EQ(partitions[0].output().size(), 3u);
    ASSERT_TRUE(partitions[0].output().count(add.node()) > 0);
    ASSERT_TRUE(partitions[0].output().count(z.node()) > 0);
    ASSERT_TRUE(partitions[0].output().count(sym_var_arr[0].node()) > 0);
    ASSERT_TRUE(partitions[0].output().count(sym_var_arr[1].node()) == 0);
    // inputs: x, y, w1, c2, (-128)
    ASSERT_EQ(partitions[0].input().size(), 5u);
    ASSERT_TRUE(partitions[0].input().count(x.node()) > 0);
    ASSERT_TRUE(partitions[0].input().count(c2.node()) > 0);
    // opr: (x + y) conv1 multi_io, (z - 128)
    ASSERT_EQ(partitions[0].opr_set().size(), 4u);
    ASSERT_TRUE(partitions[0].opr_set().count(add.node()->owner_opr()) > 0);
    ASSERT_TRUE(partitions[0].opr_set().count(c1.node()->owner_opr()) > 0);
    ASSERT_TRUE(partitions[0].opr_set().count(sym_var_arr[0].node()->owner_opr()) > 0);
    ASSERT_TRUE(partitions[0].opr_set().count(z.node()->owner_opr()) > 0);
}

TEST(TestSubGraphExtractor, MultipleReaders) {
    HostTensorGenerator<> gen;