chore: import upstream snapshot with attribution

test/cpp/fluid/tensorrt/CMakeLists.txt

@@ -0,0 +1,58 @@
+if(NOT WIN32)
+  nv_test(
+    test_tensorrt_engine_op
+    SRCS tensorrt_engine_op_test.cc
+    DEPS tensorrt_engine_op analysis python)
+else()
+  get_property(paddle_lib GLOBAL PROPERTY PADDLE_LIB_NAME)
+  nv_test(
+    test_tensorrt_engine_op
+    SRCS tensorrt_engine_op_test.cc
+    DEPS tensorrt_engine_op analysis)
+endif()
+
+if(WITH_ONNXRUNTIME AND WIN32)
+  # Copy onnxruntime for some c++ test in Windows, since the test will
+  # be build only in CI, so suppose the generator in Windows is Ninja.
+  copy_onnx(test_tensorrt_engine_op)
+endif()
+
+if(WIN32)
+  nv_test(
+    test_tensorrt
+    SRCS test_tensorrt.cc
+    DEPS phi common dynload_tensorrt)
+else()
+  nv_test(
+    test_tensorrt
+    SRCS test_tensorrt.cc
+    DEPS phi common)
+endif()
+if(WIN32)
+  nv_test(
+    test_tensorrt_engine
+    SRCS test_engine.cc test_dynamic_engine.cc
+    DEPS phi common dynload_tensorrt tensorrt_engine tensorrt_plugin)
+elseif(WITH_CINN)
+  nv_test(
+    test_tensorrt_engine
+    SRCS test_engine.cc test_dynamic_engine.cc
+    DEPS phi common tensorrt_engine tensorrt_plugin python)
+else()
+  nv_test(
+    test_tensorrt_engine
+    SRCS test_engine.cc test_dynamic_engine.cc
+    DEPS phi common tensorrt_engine tensorrt_plugin python)
+endif()
+nv_test(
+  test_arg_mapping_context
+  SRCS test_arg_mapping_context.cc
+  DEPS phi tensorrt_plugin_arg_mapping_context)
+
+if(WITH_ONNXRUNTIME AND WIN32)
+  # Copy onnxruntime for some c++ test in Windows, since the test will
+  # be build only in CI, so suppose the generator in Windows is Ninja.
+  copy_onnx(test_tensorrt_engine)
+endif()
+
+add_subdirectory(plugin)

test/cpp/fluid/tensorrt/plugin/CMakeLists.txt

@@ -0,0 +1,19 @@
+if(WIN32)
+  nv_test(
+    test_split_plugin
+    SRCS test_split_plugin.cc
+    DEPS paddle_framework ${GLOB_OPERATOR_DEPS} tensorrt_plugin
+         dynload_tensorrt)
+else()
+  nv_test(
+    test_split_plugin
+    SRCS test_split_plugin.cc
+    DEPS paddle_framework ${GLOB_OPERATOR_DEPS} tensorrt_plugin)
+endif()
+
+if(NOT WIN32)
+  nv_test(
+    test_fused_token_prune_plugin
+    SRCS test_fused_token_prune_plugin.cc
+    DEPS paddle_framework ${GLOB_OPERATOR_DEPS} tensorrt_plugin)
+endif()

test/cpp/fluid/tensorrt/plugin/test_fused_token_prune_plugin.cc

@@ -0,0 +1,41 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/inference/tensorrt/plugin/fused_token_prune_op_plugin.h"
+
+namespace paddle::inference::tensorrt::plugin {
+
+TEST(fused_token_prune_op_plugin, test_plugin) {
+  FusedTokenPrunePluginDynamic plugin(true,
+                                      /*keep_first_token*/ false,
+                                      /*keep_order*/ true,
+                                      /*flag_varseqlen*/ false);
+  plugin.initialize();
+  plugin.getPluginType();
+  plugin.getNbOutputs();
+  size_t buf_size = plugin.getSerializationSize();
+  std::vector<char> buf(buf_size);
+  plugin.serialize(buf.data());
+}
+
+TEST(fused_token_prune_op_plugin, test_plugin_creator) {
+  FusedTokenPrunePluginDynamicCreator creator;
+  creator.getFieldNames();
+  creator.createPlugin("test", nullptr);
+  creator.setPluginNamespace("test");
+}
+
+}  // namespace paddle::inference::tensorrt::plugin

test/cpp/fluid/tensorrt/plugin/test_split_plugin.cc

@@ -0,0 +1,60 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h"
+
+namespace paddle::inference::tensorrt::plugin {
+
+TEST(split_op_plugin, test_plugin) {
+  int axis = 1;
+  std::vector<int> output_lengths{1, 1};
+  bool with_fp16 = false;
+  std::vector<nvinfer1::DataType> input_types{nvinfer1::DataType::kFLOAT};
+  std::vector<nvinfer1::Dims> input_dims;
+
+  SplitPlugin sp_plugin(axis, output_lengths, with_fp16);
+  nvinfer1::Dims in_dims;
+  in_dims.nbDims = 4;
+  input_dims.push_back(in_dims);
+  sp_plugin.configurePlugin(input_dims.data(),
+                            1,
+                            nullptr,
+                            2,
+                            input_types.data(),
+                            nullptr,
+                            nullptr,
+                            nullptr,
+                            nvinfer1::PluginFormat::kLINEAR,
+                            4);
+  sp_plugin.initialize();
+  sp_plugin.getPluginType();
+  sp_plugin.canBroadcastInputAcrossBatch(0);
+  sp_plugin.getNbOutputs();
+  auto clone_plugin = sp_plugin.clone();
+  clone_plugin->setPluginNamespace("test");
+  clone_plugin->destroy();
+  sp_plugin.getOutputDataType(0, input_types.data(), 1);
+  sp_plugin.terminate();
+}
+
+TEST(split_op_plugin, test_plugin_creator) {
+  SplitPluginCreator creator;
+  creator.getFieldNames();
+  creator.createPlugin("test", nullptr);
+  creator.setPluginNamespace("test");
+}
+
+}  // namespace paddle::inference::tensorrt::plugin

test/cpp/fluid/tensorrt/tensorrt_engine_op_test.cc

@@ -0,0 +1,317 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/tensorrt/tensorrt_engine_op.h"
+
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/framework/block_desc.h"
+#include "paddle/fluid/framework/lod_tensor.h"
+#include "paddle/fluid/framework/op_desc.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/inference/analysis/helper.h"
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/convert/ut_helper.h"
+#include "paddle/phi/common/data_type.h"
+
+USE_NO_KERNEL_OP(tensorrt_engine);
+namespace paddle {
+namespace operators {
+
+namespace {
+void CreateCUDATensor(framework::Scope* scope,
+                      const std::string& name,
+                      const std::vector<int64_t>& shape) {
+  auto* var = scope->Var(name);
+  auto* tensor = var->GetMutable<phi::DenseTensor>();
+  auto dims = common::make_ddim(shape);
+  tensor->Resize(dims);
+  phi::GPUPlace place;
+  phi::GPUContext ctx(place);
+  ctx.SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                       .GetAllocator(place, ctx.stream())
+                       .get());
+  ctx.PartialInitWithAllocator();
+  inference::tensorrt::RandomizeTensor(tensor, place, ctx);
+}
+
+void AddTensorToBlockDesc(framework::proto::BlockDesc* block,
+                          const std::string& name,
+                          const std::vector<int64_t>& shape) {
+  using framework::proto::VarType;
+  auto* var = block->add_vars();
+  framework::VarDesc desc(name);
+  desc.SetType(VarType::DENSE_TENSOR);
+  desc.SetDataType(VarType::FP32);
+  desc.SetShape(shape);
+  *var = *desc.Proto();
+}
+
+}  // namespace
+
+using inference::analysis::SetAttr;
+
+void DynamicShapeTest(bool allow_build_at_runtime) {
+  framework::ProgramDesc program;
+  auto* block_ = program.Proto()->add_blocks();
+  block_->set_idx(0);
+  block_->set_parent_idx(-1);
+
+  LOG(INFO) << "create block desc";
+  framework::BlockDesc block_desc(&program, block_);
+  LOG(INFO) << "create elementwise_add op";
+  auto* elementwise_add0 = block_desc.AppendOp();
+  elementwise_add0->SetType("elementwise_add");
+  elementwise_add0->SetInput("X",
+                             std::vector<std::string>({"x"}));  // 2 x 4 x 4 x 4
+  elementwise_add0->SetInput("Y",
+                             std::vector<std::string>({"y"}));  // 1 x 4 x 1 x 1
+  elementwise_add0->SetOutput(
+      "Out", std::vector<std::string>({"z"}));  // 2 x 4 x 4 x 4
+  elementwise_add0->SetAttr("axis", static_cast<int32_t>(0));
+
+  LOG(INFO) << "create elementwise_add op";
+  auto* elementwise_add1 = block_desc.AppendOp();
+  elementwise_add1->SetType("elementwise_add");
+  elementwise_add1->SetInput("X",
+                             std::vector<std::string>({"z"}));  // 2 x 4 x 4 x 4
+  elementwise_add1->SetInput(
+      "Y", std::vector<std::string>({"y0"}));  // 1 x 4 x 4 x 4
+  elementwise_add1->SetOutput(
+      "Out", std::vector<std::string>({"z0"}));  // 2 x 4 x 4 x 4
+  elementwise_add1->SetAttr("axis", static_cast<int32_t>(0));
+
+  inference::tensorrt::OpTeller::Global().SetOpConverterType(
+      elementwise_add0, inference::tensorrt::OpConverterType::Default);
+  inference::tensorrt::OpTeller::Global().SetOpConverterType(
+      elementwise_add1, inference::tensorrt::OpConverterType::Default);
+
+  // Set inputs' variable shape in BlockDesc
+  // the batch size is 2, so the dims of 'x' is {2, 4}
+  AddTensorToBlockDesc(block_, "x", std::vector<int64_t>({2, 4, 4, 4}));
+  AddTensorToBlockDesc(block_, "y", std::vector<int64_t>({1, 4, 1, 1}));
+  AddTensorToBlockDesc(block_, "y0", std::vector<int64_t>({1, 4, 4, 4}));
+  AddTensorToBlockDesc(block_, "z", std::vector<int64_t>({2, 4, 4, 4}));
+  AddTensorToBlockDesc(block_, "z0", std::vector<int64_t>({2, 4, 4, 4}));
+
+  // It is wired, need to copy manually.
+  *block_->add_ops() = *elementwise_add0->Proto();
+  *block_->add_ops() = *elementwise_add1->Proto();
+
+  ASSERT_EQ(block_->ops_size(), 2);
+
+  LOG(INFO) << "create tensorrt op desc";
+  framework::OpDesc engine_op_desc(nullptr);
+  engine_op_desc.SetType("tensorrt_engine");
+  engine_op_desc.SetInput("Xs", std::vector<std::string>({"x"}));
+  engine_op_desc.SetOutput("Ys", std::vector<std::string>({"z0"}));
+  engine_op_desc.SetAttr("origin_outputs_dtype", std::vector<int>{5});
+
+  engine_op_desc.SetBlockAttr("sub_block", &block_desc);
+  engine_op_desc.SetAttr("max_batch_size", static_cast<int>(2));
+  engine_op_desc.SetAttr("workspace_size", static_cast<int64_t>(1 << 20));
+  engine_op_desc.SetAttr("parameters", std::vector<std::string>({}));
+  engine_op_desc.SetAttr("engine_key", std::string("a_engine"));
+  engine_op_desc.SetAttr("calibration_engine_key",
+                         std::string("a_calib_engine"));
+  engine_op_desc.SetAttr("predictor_id", 1);
+  engine_op_desc.SetAttr("calibration_data", std::string(""));
+  engine_op_desc.SetAttr("enable_int8", static_cast<bool>(false));
+  engine_op_desc.SetAttr("enable_fp16", static_cast<bool>(false));
+  engine_op_desc.SetAttr("use_calib_mode", static_cast<bool>(false));
+  engine_op_desc.SetAttr("output_name_mapping",
+                         std::vector<std::string>({"z0"}));
+  engine_op_desc.SetAttr("origin_output_rank", std::vector<int>({2}));
+  engine_op_desc.SetAttr("subgraph", std::string(block_->SerializeAsString()));
+  engine_op_desc.SetAttr("engine_serialized_data", std::string(""));
+  int device_id = 0;
+  engine_op_desc.SetAttr("gpu_device_id", device_id);
+  engine_op_desc.SetAttr("shape_range_info_path", std::string(""));
+  engine_op_desc.SetAttr("model_opt_cache_dir", std::string(""));
+  engine_op_desc.SetAttr("allow_build_at_runtime", allow_build_at_runtime);
+  engine_op_desc.SetAttr("use_static_engine", false);
+  engine_op_desc.SetAttr("with_dynamic_shape", false);
+  engine_op_desc.SetAttr("context_memory_sharing", true);
+  engine_op_desc.SetAttr("disable_trt_plugin_fp16", false);
+  engine_op_desc.SetAttr("enable_low_precision_io", false);
+  engine_op_desc.SetAttr("use_inspector", false);
+  engine_op_desc.SetAttr("engine_info_path", std::string(""));
+  engine_op_desc.SetAttr("use_dla", false);
+  engine_op_desc.SetAttr("dla_core", 0);
+
+  LOG(INFO) << "create engine op";
+  auto engine_op = framework::OpRegistry::CreateOp(engine_op_desc);
+  LOG(INFO) << "engine_op " << engine_op.get();
+
+  framework::Scope scope;
+  phi::GPUPlace place;
+  phi::GPUContext ctx(place);
+  ctx.SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                       .GetAllocator(place, ctx.stream())
+                       .get());
+  ctx.PartialInitWithAllocator();
+  // Prepare variables.
+  if (allow_build_at_runtime)
+    CreateCUDATensor(&scope, "x", std::vector<int64_t>({32, 4, 4, 4}));
+  else
+    CreateCUDATensor(&scope, "x", std::vector<int64_t>({2, 4, 4, 4}));
+  CreateCUDATensor(&scope, "y", std::vector<int64_t>({1, 4, 1, 1}));
+
+  CreateCUDATensor(&scope, "y0", std::vector<int64_t>({1, 4, 4, 4}));
+  CreateCUDATensor(&scope, "z0", std::vector<int64_t>({2, 4, 4, 4}));
+
+  // Execute them.
+  LOG(INFO) << "engine_op run";
+  engine_op->Run(scope, place);
+}
+
+TEST(TensorRTEngineOp, manual) { DynamicShapeTest(false); }
+void Execute(int batch_size, int input_dim, int output_dim, int nlayers = 1) {
+  framework::ProgramDesc program;
+  framework::Scope scope;
+  phi::GPUPlace place;
+  phi::GPUContext ctx(place);
+  ctx.SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                       .GetAllocator(place, ctx.stream())
+                       .get());
+  ctx.PartialInitWithAllocator();
+
+  auto* block_ = program.Proto()->add_blocks();
+  block_->set_idx(0);
+  block_->set_parent_idx(-1);
+
+  using shape_t = std::vector<int64_t>;
+
+  LOG(INFO) << "create block desc";
+  framework::BlockDesc block_desc(&program, block_);
+
+  auto AddFCLayer = [&](const std::string& x_name,
+                        const std::string& y_name,
+                        const std::string& z_name,
+                        bool x_created,
+                        const shape_t& x_shape,
+                        const shape_t& y_shape,
+                        const shape_t& z_shape) {
+    LOG(INFO) << "create matrix_multiply op";
+    auto* matrix_multiply = block_desc.AppendOp();
+    matrix_multiply->SetType("matrix_multiply");
+    matrix_multiply->SetInput("X", std::vector<std::string>({x_name}));
+    matrix_multiply->SetInput("Y", std::vector<std::string>({y_name}));
+    matrix_multiply->SetOutput("Out", std::vector<std::string>({z_name}));
+
+    // Set inputs' variable shape in BlockDesc
+    if (!x_created) {
+      AddTensorToBlockDesc(
+          block_, x_name, std::vector<int64_t>({batch_size, input_dim, 1, 1}));
+    }
+    AddTensorToBlockDesc(
+        block_, y_name, std::vector<int64_t>({input_dim, output_dim}));
+    AddTensorToBlockDesc(
+        block_, z_name, std::vector<int64_t>({batch_size, output_dim}));
+
+    // Prepare variables.
+    if (!x_created) {
+      CreateCUDATensor(&scope, x_name, std::vector<int64_t>(x_shape));
+    }
+    CreateCUDATensor(&scope, y_name, std::vector<int64_t>(y_shape));
+    CreateCUDATensor(&scope, z_name, std::vector<int64_t>(z_shape));
+
+    // It is wired, need to copy manually.
+    *block_->add_ops() = *matrix_multiply->Proto();
+  };
+
+  // Test with 4 layer FC
+  AddFCLayer("x0",
+             "y0",
+             "z0",
+             false,
+             {batch_size, input_dim},
+             {input_dim, output_dim},
+             {batch_size, output_dim});
+  AddFCLayer("z0",
+             "y1",
+             "z1",
+             true,
+             {},
+             {output_dim, output_dim},
+             {batch_size, output_dim});
+  AddFCLayer("z1",
+             "y2",
+             "z2",
+             true,
+             {},
+             {output_dim, output_dim},
+             {batch_size, output_dim});
+  AddFCLayer("z2",
+             "y3",
+             "z3",
+             true,
+             {},
+             {output_dim, output_dim},
+             {batch_size, output_dim});
+
+  LOG(INFO) << "create tensorrt op desc";
+  framework::OpDesc engine_op_desc(nullptr);
+  engine_op_desc.SetType("tensorrt_engine");
+  engine_op_desc.SetInput("Xs", std::vector<std::string>({"x0"}));
+  engine_op_desc.SetOutput("Ys", std::vector<std::string>({"z3"}));
+
+  engine_op_desc.SetBlockAttr("sub_block", &block_desc);
+  engine_op_desc.SetAttr("max_batch_size", static_cast<int>(batch_size));
+  engine_op_desc.SetAttr("workspace_size", static_cast<int64_t>(1 << 20));
+  engine_op_desc.SetAttr("parameters",
+                         std::vector<std::string>({"y0", "y1", "y2", "y3"}));
+  engine_op_desc.SetAttr("engine_key", std::string("b_engine"));
+  engine_op_desc.SetAttr("calibration_engine_key",
+                         std::string("b_calib_engine"));
+  engine_op_desc.SetAttr("predictor_id", 1);
+  engine_op_desc.SetAttr("calibration_data", std::string(""));
+  engine_op_desc.SetAttr("enable_int8", static_cast<bool>(false));
+  engine_op_desc.SetAttr("enable_fp16", static_cast<bool>(false));
+  engine_op_desc.SetAttr("use_calib_mode", static_cast<bool>(false));
+  engine_op_desc.SetAttr("output_name_mapping",
+                         std::vector<std::string>({"z3"}));
+  engine_op_desc.SetAttr("origin_output_rank", std::vector<int>({2}));
+  engine_op_desc.SetAttr("subgraph", std::string(block_->SerializeAsString()));
+  engine_op_desc.SetAttr("engine_serialized_data", std::string(""));
+  int device_id = 0;
+  engine_op_desc.SetAttr("gpu_device_id", device_id);
+  engine_op_desc.SetAttr("shape_range_info_path", std::string(""));
+  engine_op_desc.SetAttr("model_opt_cache_dir", std::string(""));
+  engine_op_desc.SetAttr("allow_build_at_runtime", false);
+  engine_op_desc.SetAttr("use_static_engine", false);
+  engine_op_desc.SetAttr("with_dynamic_shape", false);
+  engine_op_desc.SetAttr("context_memory_sharing", true);
+  engine_op_desc.SetAttr("disable_trt_plugin_fp16", false);
+  engine_op_desc.SetAttr("enable_low_precision_io", false);
+  engine_op_desc.SetAttr("use_inspector", false);
+  engine_op_desc.SetAttr("engine_info_path", std::string(""));
+  engine_op_desc.SetAttr("use_dla", false);
+  engine_op_desc.SetAttr("dla_core", 0);
+
+  auto engine_op = framework::OpRegistry::CreateOp(engine_op_desc);
+
+  // Execute them.
+  engine_op->Run(scope, place);
+}
+
+// Test with a larger FC layer.
+// TEST(TensorRTEngineOp, matrix_multiply) { Execute(40, 28, 28); }
+
+}  // namespace operators
+}  // namespace paddle
+
+USE_TRT_CONVERTER(elementwise_add_weight)

test/cpp/fluid/tensorrt/test_arg_mapping_context.cc

@@ -0,0 +1,124 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+
+#include "paddle/fluid/framework/op_desc.h"
+#include "paddle/fluid/inference/tensorrt/plugin_arg_mapping_context.h"
+
+namespace paddle::inference::tensorrt {
+
+TEST(ArgMappingContextTest, BasicFunction) {
+  paddle::framework::proto::OpDesc op;
+  op.set_type("imaged_op");
+  auto *input_var = op.add_inputs();
+  input_var->set_parameter("X");
+  *input_var->add_arguments() = "input";
+
+  auto *output_var = op.add_outputs();
+  output_var->set_parameter("Out");
+  *output_var->add_arguments() = "output";
+
+  auto *attr = op.add_attrs();
+  attr->set_name("int_attr");
+  attr->set_type(paddle::framework::proto::AttrType::INT);
+  attr->set_i(1);
+
+  attr = op.add_attrs();
+  attr->set_name("float_attr");
+  attr->set_type(paddle::framework::proto::AttrType::FLOAT);
+  attr->set_f(1.0);
+
+  attr = op.add_attrs();
+  attr->set_name("string_attr");
+  attr->set_type(paddle::framework::proto::AttrType::STRING);
+  attr->set_s("1");
+
+  attr = op.add_attrs();
+  attr->set_name("bool_attr");
+  attr->set_type(paddle::framework::proto::AttrType::BOOLEAN);
+  attr->set_b(true);
+
+  attr = op.add_attrs();
+  attr->set_name("ints_attr");
+  attr->set_type(paddle::framework::proto::AttrType::INTS);
+  attr->add_ints(1);
+  attr->add_ints(2);
+
+  attr = op.add_attrs();
+  attr->set_name("floats_attr");
+  attr->set_type(paddle::framework::proto::AttrType::FLOATS);
+  attr->add_floats(1.0);
+  attr->add_floats(2.0);
+
+  attr = op.add_attrs();
+  attr->set_name("strings_attr");
+  attr->set_type(paddle::framework::proto::AttrType::STRINGS);
+  attr->add_strings("1");
+  attr->add_strings("2");
+
+  attr = op.add_attrs();
+  attr->set_name("bools_attr");
+  attr->set_type(paddle::framework::proto::AttrType::BOOLEANS);
+  attr->add_bools(true);
+  attr->add_bools(true);
+
+  framework::OpDesc op_desc(op, nullptr);
+  PluginArgumentMappingContext context(&op_desc);
+
+  EXPECT_EQ(context.HasInput("X"), true);
+  EXPECT_EQ(context.HasOutput("Out"), true);
+  EXPECT_EQ(context.HasAttr("int_attr"), true);
+
+  int int_attr = any_cast<int>(context.Attr("int_attr"));
+  EXPECT_EQ(int_attr, 1);
+
+  float float_attr = any_cast<float>(context.Attr("float_attr"));
+  EXPECT_EQ(float_attr, 1);
+
+  std::string string_attr = any_cast<std::string>(context.Attr("string_attr"));
+  EXPECT_EQ(string_attr, "1");
+
+  bool bool_attr = any_cast<bool>(context.Attr("bool_attr"));
+  EXPECT_EQ(bool_attr, true);
+
+  std::vector<int> ints_attr =
+      any_cast<std::vector<int>>(context.Attr("ints_attr"));
+  EXPECT_EQ(ints_attr[0], 1);
+  EXPECT_EQ(ints_attr[1], 2);
+
+  std::vector<float> floats_attr =
+      any_cast<std::vector<float>>(context.Attr("floats_attr"));
+  EXPECT_EQ(floats_attr[0], 1.0);
+  EXPECT_EQ(floats_attr[1], 2.0);
+
+  std::vector<std::string> strings_attr =
+      any_cast<std::vector<std::string>>(context.Attr("strings_attr"));
+  EXPECT_EQ(strings_attr[0], "1");
+  EXPECT_EQ(strings_attr[1], "2");
+
+  std::vector<bool> bools_attr =
+      any_cast<std::vector<bool>>(context.Attr("bools_attr"));
+  EXPECT_EQ(bools_attr[0], true);
+  EXPECT_EQ(bools_attr[1], true);
+
+  EXPECT_EQ(context.InputSize("X"), true);
+  EXPECT_EQ(context.OutputSize("Out"), true);
+  EXPECT_EQ(context.IsDenseTensorInput("X"), true);
+  EXPECT_EQ(context.IsDenseTensorInputs("X"), true);
+
+  EXPECT_EQ(context.IsDenseTensorOutput("Out"), true);
+}
+
+}  // namespace paddle::inference::tensorrt

test/cpp/fluid/tensorrt/test_engine.cc

@@ -0,0 +1,335 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle::inference::tensorrt {
+
+class TensorRTEngineTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    ctx_ = new phi::GPUContext(phi::GPUPlace(0));
+    ctx_->SetAllocator(paddle::memory::allocation::AllocatorFacade::Instance()
+                           .GetAllocator(phi::GPUPlace(0), ctx_->stream())
+                           .get());
+    ctx_->SetHostAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::CPUPlace())
+            .get());
+    ctx_->SetZeroAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetZeroAllocator(phi::GPUPlace(0))
+            .get());
+    ctx_->SetHostZeroAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetZeroAllocator(phi::CPUPlace())
+            .get());
+    ctx_->SetPinnedAllocator(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAllocator(phi::GPUPinnedPlace())
+            .get());
+    ctx_->PartialInitWithAllocator();
+
+    TensorRTEngine::ConstructionParams params;
+    params.max_batch_size = 10;
+    params.max_workspace_size = 1 << 10;
+    params.with_dynamic_shape = true;
+    engine_ = std::make_unique<TensorRTEngine>(params);
+    engine_->InitNetwork();
+  }
+
+  void PrepareInputOutput(const std::vector<float> &input,
+                          std::vector<int> output_shape) {
+    paddle::framework::TensorFromVector(input, *ctx_, &input_);
+    output_.Resize(common::make_ddim(output_shape));
+  }
+
+  void GetOutput(std::vector<float> *output) {
+    paddle::framework::TensorToVector(output_, *ctx_, output);
+  }
+
+ protected:
+  phi::DenseTensor input_;
+  phi::DenseTensor output_;
+  std::unique_ptr<TensorRTEngine> engine_ = nullptr;
+  phi::GPUContext *ctx_ = nullptr;
+};
+
+TEST_F(TensorRTEngineTest, add_layer) {
+  const int size = 1;
+
+  std::vector<float> raw_weight = {2.};  // Weight in CPU memory.
+  std::vector<float> raw_bias = {3.};
+
+  std::vector<void *> buffers(2);  // TRT binded inputs
+
+  LOG(INFO) << "create weights";
+  TensorRTEngine::Weight weight(
+      nvinfer1::DataType::kFLOAT, raw_weight.data(), size);
+  TensorRTEngine::Weight bias(
+      nvinfer1::DataType::kFLOAT, raw_bias.data(), size);
+  auto *x = engine_->DeclareInput(
+      "x", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, 1, 1});
+  auto *weight_layer = TRT_ENGINE_ADD_LAYER(
+      engine_, Constant, nvinfer1::Dims3{1, 1, 1}, weight.get());
+  auto *bias_layer = TRT_ENGINE_ADD_LAYER(
+      engine_, Constant, nvinfer1::Dims3{1, 1, 1}, bias.get());
+  auto *matmul_layer =
+      TRT_ENGINE_ADD_LAYER(engine_,
+                           MatrixMultiply,
+                           *x,
+                           nvinfer1::MatrixOperation::kNONE,
+                           *weight_layer->getOutput(0),
+                           nvinfer1::MatrixOperation::kTRANSPOSE);
+  PADDLE_ENFORCE_NOT_NULL(
+      matmul_layer,
+      common::errors::InvalidArgument(
+          "The TRT MatrixMultiply layer cannot be null. There is something "
+          "wrong with the TRT network building and layer creation."));
+  auto *add_layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                         ElementWise,
+                                         *matmul_layer->getOutput(0),
+                                         *bias_layer->getOutput(0),
+                                         nvinfer1::ElementWiseOperation::kSUM);
+  PADDLE_ENFORCE_NOT_NULL(
+      add_layer,
+      common::errors::InvalidArgument(
+          "The TRT elementwise layer cannot be null. There is something wrong "
+          "with the TRT network building and layer creation."));
+
+  engine_->DeclareOutput(add_layer, 0, "y");
+  LOG(INFO) << "freeze network";
+  engine_->FreezeNetwork();
+#if IS_TRT_VERSION_GE(8600)
+  ASSERT_EQ(engine_->engine()->getNbIOTensors(), 2);
+#else
+  ASSERT_EQ(engine_->engine()->getNbBindings(), 2);
+#endif
+
+  // fill in real data
+  std::vector<float> x_v = {1234};
+  std::vector<float> y_cpu;
+  PrepareInputOutput(x_v, {1});
+
+  auto *x_v_gpu_data = input_.mutable_data<float>(ctx_->GetPlace());
+  auto *y_gpu_data = output_.mutable_data<float>(ctx_->GetPlace());
+
+  buffers[0] = reinterpret_cast<void *>(x_v_gpu_data);
+  buffers[1] = reinterpret_cast<void *>(y_gpu_data);
+
+  LOG(INFO) << "to execute";
+  engine_->Execute(1, &buffers, ctx_->stream());
+
+  LOG(INFO) << "to get output";
+  GetOutput(&y_cpu);
+
+  LOG(INFO) << "to checkout output";
+  ASSERT_EQ(y_cpu[0], x_v[0] * 2 + 3);
+}
+
+TEST_F(TensorRTEngineTest, add_layer_multi_dim) {
+  // Weight in CPU memory.
+  // It seems tensorrt FC use col-major: [[1.0, 3.3], [1.1, 4.4]]
+  // instead of row-major, which is [[1.0, 1.1], [3.3, 4.4]]
+  std::vector<float> raw_weight = {1.0, 1.1, 3.3, 4.4};
+  std::vector<float> raw_bias = {1.3, 2.4};
+  std::vector<void *> buffers(2);  // TRT binded inputs
+
+  TensorRTEngine::Weight weight(
+      nvinfer1::DataType::kFLOAT, raw_weight.data(), 4);
+  TensorRTEngine::Weight bias(nvinfer1::DataType::kFLOAT, raw_bias.data(), 2);
+  auto *x = engine_->DeclareInput(
+      "x", nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, 1, 2});
+  auto *weight_layer = TRT_ENGINE_ADD_LAYER(
+      engine_, Constant, nvinfer1::Dims3{1, 2, 2}, weight.get());
+  auto *bias_layer = TRT_ENGINE_ADD_LAYER(
+      engine_, Constant, nvinfer1::Dims3{1, 1, 2}, bias.get());
+  auto *matmul_layer =
+      TRT_ENGINE_ADD_LAYER(engine_,
+                           MatrixMultiply,
+                           *x,
+                           nvinfer1::MatrixOperation::kNONE,
+                           *weight_layer->getOutput(0),
+                           nvinfer1::MatrixOperation::kTRANSPOSE);
+  PADDLE_ENFORCE_NOT_NULL(
+      matmul_layer,
+      common::errors::InvalidArgument(
+          "The TRT MatrixMultiply layer cannot be null. There is something "
+          "wrong with the TRT network building and layer creation."));
+  auto *add_layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                         ElementWise,
+                                         *matmul_layer->getOutput(0),
+                                         *bias_layer->getOutput(0),
+                                         nvinfer1::ElementWiseOperation::kSUM);
+  PADDLE_ENFORCE_NOT_NULL(
+      add_layer,
+      common::errors::InvalidArgument(
+          "The TRT elementwise layer cannot be null. There is something wrong "
+          "with the TRT network building and layer creation."));
+
+  engine_->DeclareOutput(add_layer, 0, "y");
+  engine_->FreezeNetwork();
+#if IS_TRT_VERSION_GE(8600)
+  ASSERT_EQ(engine_->engine()->getNbIOTensors(), 2);
+#else
+  ASSERT_EQ(engine_->engine()->getNbBindings(), 2);
+#endif
+
+  // fill in real data
+  std::vector<float> x_v = {1.0, 2.0};
+  std::vector<float> y_cpu;
+  PrepareInputOutput(x_v, {2});
+
+  auto *x_v_gpu_data = input_.mutable_data<float>(ctx_->GetPlace());
+  auto *y_gpu_data = output_.mutable_data<float>(ctx_->GetPlace());
+
+  buffers[0] = reinterpret_cast<void *>(x_v_gpu_data);
+  buffers[1] = reinterpret_cast<void *>(y_gpu_data);
+
+  engine_->Execute(1, &buffers, ctx_->stream());
+
+  LOG(INFO) << "to get output";
+  GetOutput(&y_cpu);
+
+  auto dims = engine_->GetITensor("y")->getDimensions();
+  ASSERT_EQ(dims.nbDims, 3);
+  ASSERT_EQ(dims.d[0], 1);
+  ASSERT_EQ(dims.d[1], 1);
+  ASSERT_EQ(dims.d[2], 2);
+
+  ASSERT_EQ(y_cpu[0], 4.5);
+  ASSERT_EQ(y_cpu[1], 14.5);
+}
+
+TEST_F(TensorRTEngineTest, test_conv2d) {
+  // Weight in CPU memory.
+  std::vector<float> raw_weight = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
+  std::vector<float> raw_bias = {0};
+  std::vector<void *> buffers(2);  // TRT binded inputs
+
+  TensorRTEngine::Weight weight(
+      nvinfer1::DataType::kFLOAT, raw_weight.data(), 9);
+  TensorRTEngine::Weight bias(nvinfer1::DataType::kFLOAT, raw_bias.data(), 1);
+  auto *x = engine_->DeclareInput(
+      "x", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{2, 1, 3, 3});
+  auto *conv_layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                          ConvolutionNd,
+                                          *x,
+                                          1,
+                                          nvinfer1::DimsHW{3, 3},
+                                          weight.get(),
+                                          bias.get());
+  PADDLE_ENFORCE_NOT_NULL(conv_layer,
+                          common::errors::InvalidArgument(
+                              "TRT convolution layer building failed."));
+  conv_layer->setStrideNd(nvinfer1::Dims2{1, 1});
+  conv_layer->setPaddingNd(nvinfer1::Dims2{1, 1});
+
+  engine_->DeclareOutput(conv_layer, 0, "y");
+  engine_->FreezeNetwork();
+#if IS_TRT_VERSION_GE(8600)
+  ASSERT_EQ(engine_->engine()->getNbIOTensors(), 2);
+#else
+  ASSERT_EQ(engine_->engine()->getNbBindings(), 2);
+#endif
+
+  // fill in real data
+  std::vector<float> x_v = {1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0,
+                            1.0};
+  std::vector<float> y_cpu;
+  PrepareInputOutput(x_v, {18});
+
+  auto *x_v_gpu_data = input_.mutable_data<float>(ctx_->GetPlace());
+  auto *y_gpu_data = output_.mutable_data<float>(ctx_->GetPlace());
+
+  buffers[0] = reinterpret_cast<void *>(x_v_gpu_data);
+  buffers[1] = reinterpret_cast<void *>(y_gpu_data);
+
+  engine_->Execute(2, &buffers, ctx_->stream());
+
+  LOG(INFO) << "to get output";
+  GetOutput(&y_cpu);
+
+  ASSERT_EQ(y_cpu[0], 4.0);
+  ASSERT_EQ(y_cpu[1], 6.0);
+}
+
+TEST_F(TensorRTEngineTest, test_pool2d) {
+  // Weight in CPU memory.
+  auto *x = engine_->DeclareInput(
+      "x", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{2, 1, 2, 2});
+
+  std::vector<void *> buffers(2);  // TRT binded inputs
+  nvinfer1::PoolingType pool_t = nvinfer1::PoolingType::kAVERAGE;
+  auto *pool_layer = TRT_ENGINE_ADD_LAYER(
+      engine_, PoolingNd, *x, pool_t, nvinfer1::DimsHW{2, 2});
+
+  PADDLE_ENFORCE_NOT_NULL(
+      pool_layer,
+      common::errors::InvalidArgument("TRT pooling layer building failed."));
+  pool_layer->setStrideNd(nvinfer1::Dims2{1, 1});
+  pool_layer->setPaddingNd(nvinfer1::Dims2{0, 0});
+
+  engine_->DeclareOutput(pool_layer, 0, "y");
+  engine_->FreezeNetwork();
+#if IS_TRT_VERSION_GE(8600)
+  ASSERT_EQ(engine_->engine()->getNbIOTensors(), 2);
+#else
+  ASSERT_EQ(engine_->engine()->getNbBindings(), 2);
+#endif
+
+  // fill in real data
+  std::vector<float> x_v = {1.0, 2.0, 5.0, 0.0, 2.0, 3.0, 5.0, 10.0};
+  std::vector<float> y_cpu;
+  PrepareInputOutput(x_v, {2});
+
+  auto *x_v_gpu_data = input_.mutable_data<float>(ctx_->GetPlace());
+  auto *y_gpu_data = output_.mutable_data<float>(ctx_->GetPlace());
+
+  buffers[0] = reinterpret_cast<void *>(x_v_gpu_data);
+  buffers[1] = reinterpret_cast<void *>(y_gpu_data);
+
+  engine_->SetAllNodesLowerToTrt(true);
+  engine_->Execute(2, &buffers, ctx_->stream());
+
+  LOG(INFO) << "to get output";
+  GetOutput(&y_cpu);
+
+  ASSERT_EQ(y_cpu[0], 2.0);
+  ASSERT_EQ(y_cpu[1], 5.0);
+}
+
+}  // namespace paddle::inference::tensorrt

test/cpp/fluid/tensorrt/test_tensorrt.cc

@@ -0,0 +1,213 @@
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <cuda_runtime_api.h>
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+
+#include "NvInfer.h"
+#include "paddle/fluid/inference/tensorrt/helper.h"
+#include "paddle/phi/backends/dynload/tensorrt.h"
+
+namespace dy = phi::dynload;
+
+class Logger : public nvinfer1::ILogger {
+ public:
+  void log(nvinfer1::ILogger::Severity severity,
+           const char* msg) TRT_NOEXCEPT override {
+    switch (severity) {
+      case Severity::kINFO:
+        LOG(INFO) << msg;
+        break;
+      case Severity::kWARNING:
+        LOG(WARNING) << msg;
+        break;
+      case Severity::kINTERNAL_ERROR:
+      case Severity::kERROR:
+        LOG(ERROR) << msg;
+        break;
+      default:
+        break;
+    }
+  }
+};
+
+class ScopedWeights {
+ public:
+  explicit ScopedWeights(float value)
+      : value_(value), w{nvinfer1::DataType::kFLOAT, &value_, 1} {}
+  const nvinfer1::Weights& get() { return w; }
+
+ private:
+  float value_;
+  nvinfer1::Weights w;
+};
+
+// The following two API are implemented in TensorRT's header file, cannot load
+// from the dynamic library. So create our own implementation and directly
+// trigger the method from the dynamic library.
+nvinfer1::IBuilder* createInferBuilder(nvinfer1::ILogger* logger) {
+  return static_cast<nvinfer1::IBuilder*>(
+      dy::createInferBuilder_INTERNAL(logger, NV_TENSORRT_VERSION));
+}
+nvinfer1::IRuntime* createInferRuntime(nvinfer1::ILogger* logger) {
+  return static_cast<nvinfer1::IRuntime*>(
+      dy::createInferRuntime_INTERNAL(logger, NV_TENSORRT_VERSION));
+}
+
+const char* kInputTensor = "input";
+const char* kOutputTensor = "output";
+
+// Creates a network to compute y = 2x + 3
+nvinfer1::IHostMemory* CreateNetwork() {
+  Logger logger;
+  // Create the engine.
+  nvinfer1::IBuilder* builder = createInferBuilder(&logger);
+  auto config = builder->createBuilderConfig();
+  ScopedWeights weights(2.);
+  ScopedWeights bias(3.);
+
+  nvinfer1::INetworkDefinition* network = builder->createNetworkV2(
+      1U << static_cast<uint32_t>(
+          nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
+  // Add the input
+  auto input = network->addInput(
+      kInputTensor, nvinfer1::DataType::kFLOAT, nvinfer1::Dims3{1, 1, 1});
+  EXPECT_NE(input, nullptr);
+  // Add the constant layer for weight
+  auto weight_tensor =
+      network->addConstant(nvinfer1::Dims3{1, 1, 1}, weights.get())
+          ->getOutput(0);
+  // Add the constant layer for bias
+  auto bias_tensor =
+      network->addConstant(nvinfer1::Dims3{1, 1, 1}, bias.get())->getOutput(0);
+  // Add the hidden layer.
+  auto matmul_layer =
+      network->addMatrixMultiply(*input,
+                                 nvinfer1::MatrixOperation::kNONE,
+                                 *weight_tensor,
+                                 nvinfer1::MatrixOperation::kTRANSPOSE);
+  auto add_layer =
+      network->addElementWise(*matmul_layer->getOutput(0),
+                              *bias_tensor,
+                              nvinfer1::ElementWiseOperation::kSUM);
+  EXPECT_NE(add_layer, nullptr);
+  // Mark the output.
+  auto output = add_layer->getOutput(0);
+  output->setName(kOutputTensor);
+  network->markOutput(*output);
+#if IS_TRT_VERSION_GE(8300)
+  config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 10);
+#else
+  config->setMaxWorkspaceSize(1 << 10);
+#endif
+#if IS_TRT_VERSION_GE(8600)
+  nvinfer1::IHostMemory* model =
+      builder->buildSerializedNetwork(*network, *config);
+  EXPECT_NE(model, nullptr);
+#else
+  auto* engine = builder->buildEngineWithConfig(*network, *config);
+  EXPECT_NE(engine, nullptr);
+  // Serialize the engine to create a model, then close.
+  nvinfer1::IHostMemory* model = engine->serialize();
+  delete engine;
+#endif
+  delete network;
+  delete builder;
+  return model;
+}
+
+void Execute(nvinfer1::IExecutionContext* context,
+             const float* input,
+             float* output) {
+  const nvinfer1::ICudaEngine& engine = context->getEngine();
+  // Two binds, input and output
+  cudaStream_t stream;
+  ASSERT_EQ(0, cudaStreamCreate(&stream));
+#if IS_TRT_VERSION_GE(8600)
+  ASSERT_EQ(engine.getNbIOTensors(), 2);
+  void* buffers[2];
+  for (int i = 0; i < 2; ++i) {
+    ASSERT_EQ(0, cudaMalloc(&buffers[i], sizeof(float)));
+    auto tensor_name = engine.getIOTensorName(i);
+    context->setTensorAddress(tensor_name, buffers[i]);
+  }
+  ASSERT_EQ(
+      0,
+      cudaMemcpyAsync(
+          buffers[0], input, sizeof(float), cudaMemcpyHostToDevice, stream));
+  context->enqueueV3(stream);
+  ASSERT_EQ(
+      0,
+      cudaMemcpyAsync(
+          output, buffers[1], sizeof(float), cudaMemcpyDeviceToHost, stream));
+  cudaStreamSynchronize(stream);
+  cudaStreamDestroy(stream);
+  ASSERT_EQ(0, cudaFree(buffers[0]));
+  ASSERT_EQ(0, cudaFree(buffers[1]));
+#else
+  ASSERT_EQ(engine.getNbBindings(), 2);
+  const int input_index = engine.getBindingIndex(kInputTensor);
+  const int output_index = engine.getBindingIndex(kOutputTensor);
+  // Create GPU buffers and a stream
+  std::vector<void*> buffers(2);
+  ASSERT_EQ(0, cudaMalloc(&buffers[input_index], sizeof(float)));
+  ASSERT_EQ(0, cudaMalloc(&buffers[output_index], sizeof(float)));
+  ASSERT_EQ(0, cudaStreamCreate(&stream));
+  // Copy the input to the GPU, execute the network, and copy the output back.
+  ASSERT_EQ(0,
+            cudaMemcpyAsync(buffers[input_index],
+                            input,
+                            sizeof(float),
+                            cudaMemcpyHostToDevice,
+                            stream));
+  context->enqueue(1, buffers.data(), stream, nullptr);
+  ASSERT_EQ(0,
+            cudaMemcpyAsync(output,
+                            buffers[output_index],
+                            sizeof(float),
+                            cudaMemcpyDeviceToHost,
+                            stream));
+  cudaStreamSynchronize(stream);
+
+  // Release the stream and the buffers
+  cudaStreamDestroy(stream);
+  ASSERT_EQ(0, cudaFree(buffers[input_index]));
+  ASSERT_EQ(0, cudaFree(buffers[output_index]));
+#endif
+}
+
+TEST(TensorrtTest, BasicFunction) {
+  // Create the network serialized model.
+  nvinfer1::IHostMemory* model = CreateNetwork();
+
+  // Use the model to create an engine and an execution context.
+  Logger logger;
+  nvinfer1::IRuntime* runtime = createInferRuntime(&logger);
+  nvinfer1::ICudaEngine* engine =
+      runtime->deserializeCudaEngine(model->data(), model->size());
+  delete model;
+  nvinfer1::IExecutionContext* context = engine->createExecutionContext();
+
+  // Execute the network.
+  float input = 1234;
+  float output;
+  Execute(context, &input, &output);
+  EXPECT_EQ(output, input * 2 + 3);
+
+  // Destroy the engine.
+  delete context;
+  delete engine;
+  delete runtime;
+}