Function TEST

src/tensorrt/test/opr_replace.cpp:33–83

using namespace tensorrt;

TEST(TestTensorRTReplace, Basic) {
    REQUIRE_GPU(1);
    HostTensorGenerator<> gen;
    std::shared_ptr<HostTensorND> host_x, host_w, host_b, host_w_full, host_o;
    std::shared_ptr<ComputingGraph> graph;

    host_x = gen({5, 8, 28, 28});
    host_w = gen({2, 16, 4, 3, 3});
    host_b = gen({1, 32, 1, 1});
    host_w_full = gen({32 * 13 * 13, 10});
    host_o = gen({1, 10});

    graph = ComputingGraph::make();
    using ConvParam = megdnn::Convolution::Param;
    ConvParam conv_param1;
    conv_param1.sparse = ConvParam::Sparse::GROUP;
    auto x = Host2DeviceCopy::make(*graph, host_x),
         w_conv1 = SharedDeviceTensor::make(*graph, *host_w),
         b_conv1 = SharedDeviceTensor::make(*graph, *host_b),
         f_conv1 = opr::Convolution::make(x, w_conv1, conv_param1),
         y_conv1 = f_conv1 + b_conv1;
    using PoolParam = megdnn::Pooling::Param;
    PoolParam pool_param1;
    pool_param1.mode = PoolParam::Mode::MAX;
    pool_param1.window_h = 2;
    pool_param1.window_w = 2;
    auto y_reshape = y_conv1.reshape({5, 32, 26, 26});
    auto y_pooling1 = opr::Pooling::make(y_reshape, pool_param1);
    auto w_full1 = SharedDeviceTensor::make(*graph, *host_w_full);
    auto x_full1 = y_pooling1.reshape({5, 32 * 13 * 13});
    auto y_full1 = opr::MatrixMul::make(x_full1, w_full1);
    auto o = SharedDeviceTensor::make(*graph, *host_o);
    auto out = y_full1 + o;

    SymbolVar out_trt;
    unpack_vector(
            gopt::GraphOptimizer{}
                    .add_pass<gopt::TensorRTReplacePass>()
                    .apply({{out}})
                    .endpoint_vars(),
            out_trt);
    HostTensorND host_z1;
    HostTensorND host_z2;

    ASSERT_NE(out_trt.node(), out.node());
    auto func = graph->compile(
            {make_callback_copy(out, host_z1), make_callback_copy(out_trt, host_z2)});
    func->execute();

    MGB_ASSERT_TENSOR_NEAR(host_z1, host_z2, 1e-3);
}

TEST(TensorRTReplacePass, MatrixMul) {
    REQUIRE_GPU(1);