# Copyright (c) 2024 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. import copy import unittest import numpy as np import paddle from paddle.base import core from paddle.tensorrt.converter import PaddleToTensorRTConverter from paddle.tensorrt.export import ( Input, PrecisionMode, TensorRTConfig, ) from paddle.tensorrt.util import ( mark_builtin_op, run_pir_pass, run_trt_partition, warmup_shape_infer, ) class TensorRTBaseTest(unittest.TestCase): def __init__(self, methodName='runTest'): super().__init__(methodName) self.python_api = None self.api_args = None self.program_config = None self.min_shape = None self.opt_shape = None self.max_shape = None self.target_marker_op = "" self.dynamic_shape_data = {} self.disable_passes = [ "constant_folding_pass", "dead_code_elimination_pass", ] def create_fake_program(self): if self.python_api is None: raise ValueError( "The unittest must specify a python api that will be used for building pir program." ) main_program = paddle.static.Program() startup_program = paddle.static.Program() with paddle.static.program_guard(main_program, startup_program): api_args = copy.deepcopy(self.api_args) for feed_name in self.program_config["feed_list"]: if self.api_args[feed_name] is None: continue if isinstance(self.api_args[feed_name], dict): new_list_args = [] for sub_arg_name, sub_arg_value in self.api_args[ feed_name ].items(): if ( feed_name in self.min_shape.keys() and feed_name in self.opt_shape.keys() and feed_name in self.max_shape.keys() ): input_shape_without_dynamic_dim = ( sub_arg_value.shape[1:] ) input_dynamic_shape = [-1] input_dynamic_shape.extend( input_shape_without_dynamic_dim ) input_shape = input_dynamic_shape else: input_shape = [] input_shape_without_dynamic_dim = ( sub_arg_value.shape[0:] ) input_shape.extend(input_shape_without_dynamic_dim) input_dtype = sub_arg_value.dtype input_data = paddle.static.data( name=sub_arg_name, shape=input_shape, dtype=input_dtype, ) new_list_args.append(input_data) api_args[feed_name] = new_list_args else: empty_min_max_shape = ( self.min_shape is None or self.max_shape is None or self.opt_shape is None ) if ( not empty_min_max_shape and feed_name in self.min_shape.keys() and feed_name in self.opt_shape.keys() and feed_name in self.max_shape.keys() ): # dynamic shape condition input_shape_without_dynamic_dim = self.api_args[ feed_name ].shape[1:] input_shape = [-1] input_shape.extend(input_shape_without_dynamic_dim) else: input_shape = self.api_args[feed_name].shape input_dtype = self.api_args[feed_name].dtype input_data = paddle.static.data( name=feed_name, shape=input_shape, dtype=input_dtype, ) api_args[feed_name] = input_data actual_args = [] for name, value in api_args.items(): actual_args.append(value) output = self.python_api(*actual_args) fetch_list = [] if isinstance(output, tuple): fetch_list = [out for out in list(output) if out is not None] else: fetch_list.append(output) return main_program, startup_program, fetch_list def run_program(self, main_program, fetch_list): place = ( paddle.CUDAPlace(0) if core.is_compiled_with_cuda() else paddle.CPUPlace() ) exe = paddle.static.Executor(place) feed_data = dict() # noqa: C408 for feed_name in self.program_config["feed_list"]: if self.api_args[feed_name] is None: continue if isinstance(self.api_args[feed_name], dict): for sub_arg_name, sub_arg_value in self.api_args[ feed_name ].items(): feed_data[sub_arg_name] = sub_arg_value else: feed_data[feed_name] = self.api_args[feed_name] ret = exe.run(main_program, feed=feed_data, fetch_list=fetch_list) return ret def prepare_feed(self): for arg_name, arg_value in self.api_args.items(): # deal with condition that input is a list tensor if ( isinstance(self.api_args[arg_name], list) and arg_name in self.program_config["feed_list"] ): new_list_args = dict() # noqa: C408 for i in range(len(self.api_args[arg_name])): sub_arg_name = arg_name + str(i) new_list_args[sub_arg_name] = self.api_args[arg_name][i] self.api_args[arg_name] = new_list_args def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode="fp32"): paddle.framework.set_flags({"FLAGS_trt_min_group_size": 1}) with paddle.pir_utils.IrGuard(): self.prepare_feed() main_program, startup_program, fetch_list = ( self.create_fake_program() ) place = ( paddle.CUDAPlace(0) if core.is_compiled_with_cuda() else paddle.CPUPlace() ) exe = paddle.static.Executor(place) # init all parameter exe.run(startup_program) fetch_num = len(fetch_list) if isinstance(fetch_list[0], list): fetch_index = [i for i, v in enumerate(fetch_list)] else: fetch_index = [v.index() for v in fetch_list] output_expected = self.run_program(main_program, fetch_list) min_shape_data = dict() # noqa: C408 opt_shape_data = dict() # noqa: C408 max_shape_data = dict() # noqa: C408 for feed_name in self.program_config["feed_list"]: if self.api_args[feed_name] is None: continue if isinstance(self.api_args[feed_name], dict): # shape_tensor if ( feed_name not in self.min_shape.keys() and feed_name not in self.max_shape.keys() and feed_name not in self.opt_shape.keys() ): for sub_feed_name, sub_feed_value in self.api_args[ feed_name ].items(): min_shape_data[sub_feed_name] = sub_feed_value opt_shape_data[sub_feed_name] = sub_feed_value max_shape_data[sub_feed_name] = sub_feed_value continue else: # not shape_tensor for i in range(len(self.min_shape[feed_name])): sub_feed_name = feed_name + str(i) min_shape_data[sub_feed_name] = np.random.randn( *self.min_shape[feed_name][i] ).astype( self.api_args[feed_name][sub_feed_name].dtype ) opt_shape_data[sub_feed_name] = np.random.randn( *self.opt_shape[feed_name][i] ).astype( self.api_args[feed_name][sub_feed_name].dtype ) max_shape_data[sub_feed_name] = np.random.randn( *self.max_shape[feed_name][i] ).astype( self.api_args[feed_name][sub_feed_name].dtype ) else: # shape_tensor is list if ( feed_name not in self.min_shape.keys() and feed_name not in self.max_shape.keys() and feed_name not in self.opt_shape.keys() ): min_shape_data[feed_name] = self.api_args[feed_name] opt_shape_data[feed_name] = self.api_args[feed_name] max_shape_data[feed_name] = self.api_args[feed_name] continue else: if self.dynamic_shape_data: min_shape_data[feed_name] = self.dynamic_shape_data[ feed_name ](self.min_shape[feed_name]) opt_shape_data[feed_name] = self.dynamic_shape_data[ feed_name ](self.opt_shape[feed_name]) max_shape_data[feed_name] = self.dynamic_shape_data[ feed_name ](self.max_shape[feed_name]) else: min_shape_data[feed_name] = np.random.randn( *self.min_shape[feed_name] ).astype(self.api_args[feed_name].dtype) opt_shape_data[feed_name] = np.random.randn( *self.opt_shape[feed_name] ).astype(self.api_args[feed_name].dtype) max_shape_data[feed_name] = np.random.randn( *self.max_shape[feed_name] ).astype(self.api_args[feed_name].dtype) # run pir pass(including some constant fold pass, dead code elimination pass, fusion pass and trt_op_marker_pass) main_program = run_pir_pass( main_program, disable_passes=self.disable_passes, ) # delete unused op for op in main_program.global_block().ops: if ( op.name() == "builtin.constant" or op.name() == "builtin.parameter" ): if op.results()[0].use_empty(): main_program.global_block().remove_op(op) scope = paddle.static.global_scope() main_program = warmup_shape_infer( main_program, feeds=[min_shape_data, opt_shape_data, max_shape_data], scope=scope, ) for op in main_program.global_block().ops[::-1]: # Remove all invalid fetch op if op.name() == "pd_op.fetch": main_program.global_block().remove_op(op) # Adding marker labels to builtin ops facilitates convert processing, but they ultimately do not enter the TensorRT subgraph. mark_builtin_op(main_program) # run trt_sub_graph_extract_pass() program_with_trt = run_trt_partition(main_program) # run TRTConverter(would lower group_op into tensorrt_engine_op) trt_config = None input = Input( min_input_shape=self.min_shape, optim_input_shape=self.opt_shape, max_input_shape=self.max_shape, ) trt_config = TensorRTConfig(inputs=[input]) trt_config.disable_logging = False if precision_mode == "fp16": trt_config.precision_mode = PrecisionMode.FP16 converter = PaddleToTensorRTConverter( program_with_trt, scope, trt_config ) converter.convert_program_to_trt() # check whether has trt op has_trt_op = False for op in program_with_trt.global_block().ops: if op.name() == "pd_op.tensorrt_engine": has_trt_op = True self.assertEqual(has_trt_op, True) trt_fetch_list = [] split_op = program_with_trt.global_block().ops[-1] if split_op.name() == "builtin.split": trt_fetch_list = [ split_op.result(index) for index in fetch_index ] else: raise ValueError( "The last op of convert pir Program in test must be split op that is the next op of pd_op.engine." ) output_trt = self.run_program(program_with_trt, trt_fetch_list) # Check that the results are close to each other within a tolerance of 1e-3 for i in range(fetch_num): np.testing.assert_allclose( output_expected[i], output_trt[i], rtol=rtol, atol=atol, ) def check_marker(self, expected_result): paddle.framework.set_flags({"FLAGS_trt_min_group_size": 1}) with paddle.pir_utils.IrGuard(): main_program, startup_program, fetch_list = ( self.create_fake_program() ) main_program = run_pir_pass( main_program, disable_passes=self.disable_passes, ) marker_result = False for op in main_program.global_block().ops: if op.name() == self.target_marker_op: marker_result = op.attrs().get("__l_trt__", False) self.assertEqual(marker_result, expected_result)