diff --git a/test/cpp/inference/tensorrt/test_tensorrt_engine_instruction.cc b/test/cpp/inference/tensorrt/test_tensorrt_engine_instruction.cc index 7bd29a9f1adbc..37f03996609c6 100644 --- a/test/cpp/inference/tensorrt/test_tensorrt_engine_instruction.cc +++ b/test/cpp/inference/tensorrt/test_tensorrt_engine_instruction.cc @@ -85,20 +85,79 @@ TEST(TensorRTEngineInstructionTest, test_tensorrt_engine_instruction) { nvinfer1::DataType::kFLOAT, raw_bias, size); auto *x = engine->DeclareInput( "x", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{-1, 1, 1, 1}); - auto *fc_layer = TRT_ENGINE_ADD_LAYER( - engine, FullyConnected, *x, size, weight.get(), bias.get()); - PADDLE_ENFORCE_NOT_NULL(fc_layer, - common::errors::InvalidArgument( - "TRT fully connected layer building failed.")); + auto *flatten_layer = engine->network()->addShuffle(*x); + PADDLE_ENFORCE_NOT_NULL( + flatten_layer, + common::errors::InvalidArgument( + "Unable to build the TensorRT shuffle layer for the input tensor " + "'x'. " + "This usually indicates the TensorRT network failed to allocate the " + "intermediate reshape layer.")); + flatten_layer->setReshapeDimensions(nvinfer1::Dims2{-1, 1}); + + auto *weight_layer = TRT_ENGINE_ADD_LAYER( + engine, Constant, nvinfer1::Dims2{1, 1}, weight.get()); + PADDLE_ENFORCE_NOT_NULL( + weight_layer, + common::errors::InvalidArgument("TensorRT failed to create the constant " + "layer for parameter 'weight'. " + "Please confirm the TensorRT builder " + "supports constant initialisation " + "for the provided weight shape.")); + + auto *bias_layer = + TRT_ENGINE_ADD_LAYER(engine, Constant, nvinfer1::Dims2{1, 1}, bias.get()); + PADDLE_ENFORCE_NOT_NULL( + bias_layer, + common::errors::InvalidArgument( + "TensorRT failed to create the constant layer for parameter 'bias'. " + "Check whether the provided bias data matches the expected shape.")); + + auto *matmul_layer = TRT_ENGINE_ADD_LAYER(engine, + MatrixMultiply, + *flatten_layer->getOutput(0), + nvinfer1::MatrixOperation::kNONE, + *weight_layer->getOutput(0), + nvinfer1::MatrixOperation::kNONE); + PADDLE_ENFORCE_NOT_NULL( + matmul_layer, + common::errors::InvalidArgument( + "TensorRT returned a null matrix-multiply layer while fusing the " + "fully-connected op. Verify the network input ranks and TensorRT " + "version.")); + + 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( + "TensorRT could not construct the elementwise-add layer for bias " + "fusion. Ensure the bias tensor uses broadcastable dimensions.")); - engine->DeclareOutput(fc_layer, 0, "y"); + auto *reshape_layer = engine->network()->addShuffle(*add_layer->getOutput(0)); + PADDLE_ENFORCE_NOT_NULL( + reshape_layer, + common::errors::InvalidArgument( + "TensorRT could not emit the final shuffle layer to restore the " + "output shape. Confirm the shape tensor and inferred dimensions are " + "valid.")); + reshape_layer->setReshapeDimensions(nvinfer1::Dims4{-1, 1, 1, 1}); + + engine->DeclareOutput(reshape_layer, 0, "y"); std::vector input_names = {"x", ""}; std::vector output_names = {"y"}; std::vector> outputs_shape = {{1}}; std::vector outputs_dtype = {phi::DataType::FLOAT32}; 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 nvinfer1::IHostMemory *serialized_engine_data = engine->Serialize(); std::ofstream outFile("engine_serialized_data.bin", std::ios::binary); @@ -220,7 +279,10 @@ TEST(TensorRTEngineInstructionTest, test_tensorrt_engine_instruction_dynamic) { layer->setInput(1, *shape); PADDLE_ENFORCE_NOT_NULL( layer, - common::errors::InvalidArgument("TRT shuffle layer building failed.")); + common::errors::InvalidArgument( + "TensorRT failed to construct the dynamic shuffle layer that " + "consumes the runtime shape tensor. Please check the provided " + "shape binding.")); engine->DeclareOutput(layer, 0, "y"); engine->FreezeNetwork(); @@ -401,14 +463,19 @@ TEST(PluginTest, test_generic_plugin) { creator->createPlugin("pir_generic_plugin", plugin_collection.get()); PADDLE_ENFORCE_NOT_NULL( generic_plugin, - common::errors::InvalidArgument("TRT create generic plugin failed.")); + common::errors::InvalidArgument( + "TensorRT plugin registry returned nullptr while creating " + "'pir_generic_plugin'. Verify the plugin has been registered before " + "building the engine.")); std::vector plugin_inputs; plugin_inputs.emplace_back(x); auto plugin_layer = engine->network()->addPluginV2( plugin_inputs.data(), plugin_inputs.size(), *generic_plugin); - PADDLE_ENFORCE_NOT_NULL(plugin_layer, - common::errors::InvalidArgument( - "TRT generic plugin layer building failed.")); + PADDLE_ENFORCE_NOT_NULL( + plugin_layer, + common::errors::InvalidArgument( + "TensorRT failed to add the generic plugin layer to the network. " + "Ensure the plugin inputs match the expected TensorRT types.")); engine->DeclareOutput(plugin_layer, 0, "y"); std::vector input_names = {"x"}; @@ -417,7 +484,11 @@ TEST(PluginTest, test_generic_plugin) { std::vector outputs_dtype = {phi::DataType::FLOAT32}; 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 nvinfer1::IHostMemory *serialized_engine_data = engine->Serialize(); std::ofstream outFile("engine_serialized_data.bin", std::ios::binary); outFile.write(static_cast(serialized_engine_data->data()),