wehub-resource-sync
682 lines
28 KiB
Python
682 lines
28 KiB
Python
# 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 ctypes
|
|
import hashlib
|
|
import logging
|
|
|
|
import paddle
|
|
|
|
paddle.base.core.register_paddle_plugin()
|
|
import tensorrt as trt
|
|
|
|
import paddle
|
|
from paddle import pir
|
|
from paddle.base.core import clear_shape_info, get_value_shape_range_info
|
|
from paddle.base.log_helper import get_logger
|
|
|
|
from .impls.activation import * # noqa: F403
|
|
from .impls.attribute import * # noqa: F403
|
|
from .impls.common import * # noqa: F403
|
|
from .impls.conv import * # noqa: F403
|
|
from .impls.creation import * # noqa: F403
|
|
from .impls.einsum import * # noqa: F403
|
|
from .impls.input import * # noqa: F403
|
|
from .impls.linalg import * # noqa: F403
|
|
from .impls.logic import * # noqa: F403
|
|
from .impls.manipulation import * # noqa: F403
|
|
from .impls.math import * # noqa: F403
|
|
from .impls.norm import * # noqa: F403
|
|
from .impls.ops import * # noqa: F403
|
|
from .impls.others import * # noqa: F403
|
|
from .impls.pooling import * # noqa: F403
|
|
from .impls.search import * # noqa: F403
|
|
from .impls.stat import * # noqa: F403
|
|
from .impls.vision import * # noqa: F403
|
|
from .register import converter_registry
|
|
from .util import (
|
|
RefitManager,
|
|
RefitRole,
|
|
TensorRTConfigManager,
|
|
TensorRTConstantManager,
|
|
all_ops_into_trt,
|
|
get_cache_path,
|
|
get_trt_version,
|
|
get_trt_version_list,
|
|
is_shape_tensor,
|
|
map_dtype,
|
|
remove_duplicate_value,
|
|
set_dynamic_range,
|
|
weight_to_tensor,
|
|
zero_dims_to_one_dims,
|
|
)
|
|
|
|
version_list = get_trt_version_list()
|
|
|
|
_logger = get_logger(
|
|
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
|
)
|
|
|
|
|
|
class PaddleToTensorRTConverter:
|
|
def __init__(self, paddle_program, scope, trt_config=None):
|
|
self.scope = scope
|
|
self.program = paddle_program
|
|
self.trt_config = trt_config
|
|
self.constant_manager = TensorRTConstantManager()
|
|
self.refit_manager = RefitManager()
|
|
params = paddle_program.global_block().all_parameters()
|
|
param_dict = {}
|
|
# save parameters
|
|
for v in params:
|
|
name = v.get_defining_op().attrs()["parameter_name"]
|
|
weight_tensor = self.scope.var(name).get_tensor()
|
|
self.constant_manager.set_constant_value(name, weight_tensor, v)
|
|
|
|
self.input_info = {}
|
|
self.trt_output_value_map = {}
|
|
self.engine_num = 0
|
|
# init tensorrt plugin
|
|
trt_plugin_lib = ctypes.CDLL('libnvinfer_plugin.so')
|
|
trt_plugin_lib.initLibNvInferPlugins(None, "")
|
|
|
|
def find_graph_inputs_outputs(self, group_op):
|
|
operations = next(iter(group_op.blocks())).ops
|
|
all_values = {}
|
|
output_values = {}
|
|
|
|
graph_output_values = []
|
|
|
|
def __is_output_value(value):
|
|
for op in value.all_used_ops():
|
|
if op.name() == "cf.yield":
|
|
return True
|
|
return False
|
|
|
|
# Collect all output values from all operations
|
|
for op in operations:
|
|
for result in op.results():
|
|
output_values[result.id] = result
|
|
all_values[result.id] = result
|
|
if __is_output_value(result):
|
|
graph_output_values.append(result)
|
|
for operand in op.operands():
|
|
source = operand.source()
|
|
if not source.initialized():
|
|
_logger.warning(f"Skipping uninitialized source: {source}")
|
|
continue
|
|
else:
|
|
all_values[source.id] = source
|
|
|
|
# Input values are those that are in all_values but not in output_values
|
|
input_values = [
|
|
value
|
|
for value_id, value in all_values.items()
|
|
if value_id not in output_values
|
|
]
|
|
|
|
return input_values, graph_output_values
|
|
|
|
def convert_subgraph_to_trt(self, program, group_op):
|
|
from .export import PrecisionMode
|
|
|
|
trt_manager = TensorRTConfigManager(self.trt_config)
|
|
if self.trt_config is not None and self.trt_config.ops_run_float:
|
|
_logger.info(f"force_fp32_ops: {trt_manager.get_force_fp32_ops()}")
|
|
|
|
if not self.trt_config.disable_logging:
|
|
_logger.info(f"start process {group_op}")
|
|
|
|
operations = next(iter(group_op.blocks())).ops
|
|
input_values, output_values = self.find_graph_inputs_outputs(group_op)
|
|
builder = trt.Builder(trt.Logger(trt.Logger.ERROR))
|
|
network = builder.create_network(
|
|
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
|
)
|
|
profile = builder.create_optimization_profile()
|
|
|
|
# Mapping from Value id to TensorRT ITensor
|
|
value_to_trt_tensor = {}
|
|
min_shape_map = {}
|
|
opt_shape_map = {}
|
|
max_shape_map = {}
|
|
min_value_map = {}
|
|
opt_value_map = {}
|
|
max_value_map = {}
|
|
input_names = []
|
|
new_input_values = []
|
|
refit_param_name = []
|
|
precision_mode = PrecisionMode.FP32
|
|
if self.trt_config is not None:
|
|
precision_mode = self.trt_config.precision_mode
|
|
# Because one of the inputs to pd_op.concat is builtin.combine,
|
|
# during the conversion process using the converter,
|
|
# it is necessary to obtain the input of builtin.combine.
|
|
origin_input_value = []
|
|
for value in input_values:
|
|
defining_op = value.get_defining_op()
|
|
if defining_op.name() == "builtin.combine":
|
|
for operand in defining_op.operands():
|
|
source = operand.source()
|
|
origin_input_value.append(source)
|
|
else:
|
|
origin_input_value.append(value)
|
|
|
|
origin_input_value = remove_duplicate_value(origin_input_value)
|
|
# create TRT Weight and TRT Input
|
|
for value in origin_input_value:
|
|
defining_op = value.get_defining_op()
|
|
if defining_op.name() == "builtin.parameter":
|
|
param_name = defining_op.attrs()["parameter_name"]
|
|
refit_param_name.append(param_name)
|
|
weight = trt.Weights(
|
|
self.constant_manager.get_constant_value(param_name)
|
|
)
|
|
if self.trt_config.refit_params_path:
|
|
paddle_shape = value.shape
|
|
trt_shape = trt.Dims(paddle_shape)
|
|
constant_layer = network.add_constant(trt_shape, weight)
|
|
constant_layer.name = param_name
|
|
value_to_trt_tensor[value.id] = constant_layer.get_output(0)
|
|
self.refit_manager.set_trt_weight_tensor(
|
|
constant_layer.get_output(0).name, weight
|
|
)
|
|
self.refit_manager.set_mapping(
|
|
param_name, param_name, RefitRole.CONSTANT
|
|
)
|
|
else:
|
|
value_to_trt_tensor[value.id] = weight
|
|
elif defining_op.name() == "builtin.constant":
|
|
constant_value_name = defining_op.attrs()["value"]
|
|
constant_tensor = self.scope.var(
|
|
constant_value_name
|
|
).get_tensor()
|
|
self.constant_manager.set_constant_value(
|
|
constant_value_name, constant_tensor, value
|
|
)
|
|
constant_tensor = trt.Weights(
|
|
self.constant_manager.get_constant_value(
|
|
constant_value_name
|
|
)
|
|
)
|
|
if self.trt_config.refit_params_path:
|
|
paddle_shape = value.shape
|
|
trt_shape = trt.Dims(paddle_shape)
|
|
constant_layer = network.add_constant(
|
|
trt_shape, constant_tensor
|
|
)
|
|
constant_layer.name = constant_value_name
|
|
value_to_trt_tensor[value.id] = constant_layer.get_output(0)
|
|
self.refit_manager.set_trt_weight_tensor(
|
|
constant_layer.get_output(0).name, constant_tensor
|
|
)
|
|
else:
|
|
value_to_trt_tensor[value.id] = constant_tensor
|
|
else:
|
|
shape = value.shape
|
|
dtype = map_dtype(value.dtype.name)
|
|
input_name = f"input_{value.id}"
|
|
# 0-dims -> 1-dims
|
|
if len(shape) == 0:
|
|
shape = [1]
|
|
input_tensor = network.add_input(
|
|
name=input_name, dtype=dtype, shape=shape
|
|
)
|
|
input_names.append(input_name)
|
|
new_input_values.append(value)
|
|
value_to_trt_tensor[value.id] = input_tensor
|
|
|
|
for op in operations:
|
|
# Adding marker labels to builtin ops facilitates convert processing, but they ultimately do not enter the TensorRT subgraph.
|
|
if op.name() == "builtin.split" or op.name() == "builtin.combine":
|
|
continue
|
|
operands = []
|
|
for operand in op.operands():
|
|
source = operand.source()
|
|
if not source.initialized():
|
|
operands.append(None)
|
|
continue
|
|
vec_type = source.type().as_vec_type()
|
|
if vec_type is not None and len(vec_type.as_list()) == 0:
|
|
continue
|
|
define_op_name = source.get_defining_op().name()
|
|
if define_op_name == "builtin.combine":
|
|
operand_list = []
|
|
for combined_operand in source.get_defining_op().operands():
|
|
combined_source = combined_operand.source()
|
|
combined_source_id = combined_source.id
|
|
if combined_source_id in value_to_trt_tensor:
|
|
trt_input_tensor = weight_to_tensor(
|
|
network,
|
|
combined_source,
|
|
value_to_trt_tensor[combined_source_id],
|
|
op.name(),
|
|
)
|
|
trt_input_tensor = zero_dims_to_one_dims(
|
|
network, trt_input_tensor
|
|
)
|
|
operand_list.append(trt_input_tensor)
|
|
else:
|
|
raise RuntimeError(
|
|
f'{combined_source_id} not found in value_to_trt_tensor'
|
|
)
|
|
operands.append(operand_list)
|
|
else:
|
|
source_id = source.id
|
|
if source_id in value_to_trt_tensor:
|
|
trt_input_tensor = weight_to_tensor(
|
|
network,
|
|
source,
|
|
value_to_trt_tensor[source_id],
|
|
op.name(),
|
|
)
|
|
trt_input_tensor = zero_dims_to_one_dims(
|
|
network, trt_input_tensor
|
|
)
|
|
operands.append(trt_input_tensor)
|
|
else:
|
|
raise RuntimeError(
|
|
f'{source_id} not found in value_to_trt_tensor'
|
|
)
|
|
|
|
if precision_mode.value == PrecisionMode.INT8.value:
|
|
set_dynamic_range(op, operands)
|
|
trt_outs = self.convert(network, op, operands)
|
|
|
|
results = []
|
|
|
|
for idx, result in enumerate(op.results()):
|
|
if result.is_combine():
|
|
# empty vec value condition
|
|
if len(result.type().as_vec_type().as_list()) == 0:
|
|
results.append(result)
|
|
continue
|
|
used_ops = result.all_used_ops()
|
|
for use_op in used_ops:
|
|
if use_op.name() == "builtin.split":
|
|
split_outputs = use_op.results()
|
|
results.extend(split_outputs)
|
|
else:
|
|
results.append(result)
|
|
|
|
for idx, result in enumerate(results):
|
|
if idx < len(trt_outs):
|
|
value_to_trt_tensor[result.id] = trt_outs[idx]
|
|
else:
|
|
value_to_trt_tensor[result.id] = None
|
|
|
|
# Set TRT min/opt/max input shape and the value of shape tensor
|
|
for i, value in enumerate(origin_input_value):
|
|
trt_input = value_to_trt_tensor[value.id]
|
|
defining_op_name = value.get_defining_op().name()
|
|
if (
|
|
defining_op_name == "builtin.parameter"
|
|
or defining_op_name == "builtin.constant"
|
|
):
|
|
# constant/parameter condition, needn't get min/opt/max shape
|
|
continue
|
|
input_name = trt_input.name
|
|
if not self.trt_config.disable_logging:
|
|
_logger.info(
|
|
f"set shape of {value}, op is: {value.get_defining_op()}"
|
|
)
|
|
min_shape = []
|
|
opt_shape = []
|
|
max_shape = []
|
|
min_value = []
|
|
opt_value = []
|
|
max_value = []
|
|
|
|
value_define_op = value.get_defining_op()
|
|
# if the input value is generated by the other trt_engine_op, so the shape is searched by origin value
|
|
if (
|
|
value_define_op.name() == "builtin.split"
|
|
and value_define_op.operand_source(0).get_defining_op().name()
|
|
== "pd_op.tensorrt_engine"
|
|
):
|
|
min_shape = self.input_info[value.id]["min_shape"]
|
|
opt_shape = self.input_info[value.id]["opt_shape"]
|
|
max_shape = self.input_info[value.id]["max_shape"]
|
|
if trt_input.is_shape_tensor:
|
|
min_value = self.input_info[value.id]["min_value"]
|
|
opt_value = self.input_info[value.id]["opt_value"]
|
|
max_value = self.input_info[value.id]["max_value"]
|
|
else:
|
|
min_shape = get_value_shape_range_info(
|
|
value, False, paddle.base.core.ShapeMode.kMIN
|
|
)
|
|
opt_shape = get_value_shape_range_info(
|
|
value, False, paddle.base.core.ShapeMode.kOPT
|
|
)
|
|
max_shape = get_value_shape_range_info(
|
|
value, False, paddle.base.core.ShapeMode.kMAX
|
|
)
|
|
|
|
if trt_input.is_shape_tensor:
|
|
min_value = get_value_shape_range_info(
|
|
value, True, paddle.base.core.ShapeMode.kMIN
|
|
)
|
|
opt_value = get_value_shape_range_info(
|
|
value, True, paddle.base.core.ShapeMode.kOPT
|
|
)
|
|
max_value = get_value_shape_range_info(
|
|
value, True, paddle.base.core.ShapeMode.kMAX
|
|
)
|
|
if not trt_input.is_shape_tensor:
|
|
if not self.trt_config.disable_logging:
|
|
_logger.info(f"set min_shape of {value} as {min_shape}")
|
|
_logger.info(f"set opt_shape of {value} as {opt_shape}")
|
|
_logger.info(f"set max_shape of {value} as {max_shape}")
|
|
profile.set_shape(
|
|
input_name, min=min_shape, opt=opt_shape, max=max_shape
|
|
)
|
|
else:
|
|
if not self.trt_config.disable_logging:
|
|
_logger.info(
|
|
f"set min_value of shape input: {value} as {min_value}"
|
|
)
|
|
_logger.info(
|
|
f"set opt_value of shape input: {value} as {opt_value}"
|
|
)
|
|
_logger.info(
|
|
f"set max_value of shape input: {value} as {max_value}"
|
|
)
|
|
profile.set_shape_input(
|
|
input_name, min=min_value, opt=opt_value, max=max_value
|
|
)
|
|
|
|
min_shape_map[input_name] = min_shape
|
|
opt_shape_map[input_name] = opt_shape
|
|
max_shape_map[input_name] = max_shape
|
|
min_value_map[input_name] = min_value
|
|
opt_value_map[input_name] = opt_value
|
|
max_value_map[input_name] = max_value
|
|
|
|
out_shapes = []
|
|
out_names = []
|
|
out_types = []
|
|
for out_index in range(len(output_values)):
|
|
result_value = output_values[out_index]
|
|
output_tensor = value_to_trt_tensor[result_value.id]
|
|
if output_tensor is None:
|
|
out_names.append("")
|
|
out_shapes.append([])
|
|
out_types.append(None)
|
|
continue
|
|
network.mark_output(output_tensor)
|
|
out_names.append(output_tensor.name)
|
|
out_shapes.append(result_value.shape)
|
|
out_types.append(result_value.dtype)
|
|
if group_op.result(out_index).use_empty():
|
|
# if result value is not used, it doesn't need get shape, continue
|
|
continue
|
|
min_shape = []
|
|
opt_shape = []
|
|
max_shape = []
|
|
if len(result_value.shape) != 0:
|
|
min_shape = get_value_shape_range_info(
|
|
result_value, False, paddle.base.core.ShapeMode.kMIN
|
|
)
|
|
opt_shape = get_value_shape_range_info(
|
|
result_value, False, paddle.base.core.ShapeMode.kOPT
|
|
)
|
|
max_shape = get_value_shape_range_info(
|
|
result_value, False, paddle.base.core.ShapeMode.kMAX
|
|
)
|
|
|
|
min_value = []
|
|
opt_value = []
|
|
max_value = []
|
|
if is_shape_tensor(result_value):
|
|
min_value = get_value_shape_range_info(
|
|
result_value, True, paddle.base.core.ShapeMode.kMIN
|
|
)
|
|
opt_value = get_value_shape_range_info(
|
|
result_value, True, paddle.base.core.ShapeMode.kOPT
|
|
)
|
|
max_value = get_value_shape_range_info(
|
|
result_value, True, paddle.base.core.ShapeMode.kMAX
|
|
)
|
|
|
|
self.input_info[result_value.id] = {
|
|
"min_shape": min_shape,
|
|
"opt_shape": opt_shape,
|
|
"max_shape": max_shape,
|
|
"min_value": min_value,
|
|
"opt_value": opt_value,
|
|
"max_value": max_value,
|
|
}
|
|
|
|
config = builder.create_builder_config()
|
|
if self.trt_config and self.trt_config.refit_params_path:
|
|
config.set_flag(trt.BuilderFlag.REFIT)
|
|
config.add_optimization_profile(profile)
|
|
if version_list[0] > 8 or (
|
|
version_list[0] == 8 and version_list[1] >= 6
|
|
): # trt version >= 8.6
|
|
config.builder_optimization_level = (
|
|
self.trt_config.optimization_level
|
|
)
|
|
config.set_memory_pool_limit(
|
|
trt.MemoryPoolType.WORKSPACE, self.trt_config.workspace_size
|
|
)
|
|
|
|
if precision_mode.value == PrecisionMode.FP16.value:
|
|
if builder.platform_has_fast_fp16:
|
|
config.set_flag(trt.BuilderFlag.FP16)
|
|
_logger.info("Run Paddle-TRT FP16 mode")
|
|
else:
|
|
_logger.warning(
|
|
"Hardware does not support FP16. Continuing in FP32 mode."
|
|
)
|
|
elif precision_mode.value == PrecisionMode.BF16.value:
|
|
if version_list[0] >= 9:
|
|
if builder.platform_has_fast_bfp16 and hasattr(
|
|
builder, 'platform_has_fast_bf16'
|
|
):
|
|
config.set_flag(trt.BuilderFlag.BF16)
|
|
_logger.info("Run Paddle-TRT BF16 mode")
|
|
else:
|
|
_logger.warning(
|
|
"Hardware does not support BF16. Continuing in FP32 mode."
|
|
)
|
|
else:
|
|
if builder.platform_has_fast_fp16:
|
|
config.set_flag(trt.BuilderFlag.FP16)
|
|
_logger.warning(
|
|
"Because the version of TensorRT is less than 9.0, run Paddle-TRT FP16 mode"
|
|
)
|
|
else:
|
|
_logger.warning(
|
|
"Hardware does not support FP16. Continuing in FP32 mode."
|
|
)
|
|
elif precision_mode.value == PrecisionMode.INT8.value:
|
|
config.set_flag(trt.BuilderFlag.INT8)
|
|
_logger.info("Run Paddle-TRT INT8 mode")
|
|
elif self.trt_config is not None:
|
|
_logger.info(
|
|
f"Default precision mode {self.trt_config.precision_mode}"
|
|
)
|
|
|
|
if (
|
|
version_list[0] > 8
|
|
or version_list[0] == 8
|
|
and version_list[1] >= 2
|
|
and version_list[2] >= 1
|
|
):
|
|
if self.trt_config is not None and self.trt_config.ops_run_float:
|
|
config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
|
|
|
|
trt_engine = builder.build_serialized_network(network, config)
|
|
assert trt_engine is not None, (
|
|
'Failed to build engine. please see ERROR log from trt.Logger'
|
|
)
|
|
trt_params = paddle.base.libpaddle.TRTEngineParams()
|
|
trt_params.min_input_shape = min_shape_map
|
|
trt_params.max_input_shape = max_shape_map
|
|
trt_params.optim_input_shape = opt_shape_map
|
|
trt_params.min_shape_tensor = min_value_map
|
|
trt_params.max_shape_tensor = max_value_map
|
|
trt_params.optim_shape_tensor = opt_value_map
|
|
trt_params.use_cuda_graph = self.trt_config.use_cuda_graph
|
|
all_nodes_offload_to_trt = all_ops_into_trt(self.program)
|
|
if self.trt_config.use_cuda_graph and not all_nodes_offload_to_trt:
|
|
_logger.info(
|
|
"You have enabled CudaGraph, but not the entire graph offload to "
|
|
"trt, now return to normal mode."
|
|
)
|
|
trt_params.use_cuda_graph = False
|
|
if self.trt_config.refit_params_path:
|
|
trt_params.refit_params_path = self.trt_config.refit_params_path
|
|
trt_params.refit_param_name = refit_param_name
|
|
trt_params.refit_param_names2trt_names = (
|
|
self.refit_manager.get_all_mappings()
|
|
)
|
|
group_str = str(group_op)
|
|
engine_name = (
|
|
int(hashlib.sha256(group_str.encode('utf-8')).hexdigest(), 16)
|
|
% 10**8
|
|
)
|
|
CACHE_ROOT = get_cache_path(self.trt_config.save_model_dir)
|
|
CACHE_FILE = f"{CACHE_ROOT}/engine_{engine_name}_{self.engine_num}.trt"
|
|
with open(CACHE_FILE, "wb") as f:
|
|
f.write(trt_engine)
|
|
PIR_DUMP_FILE = (
|
|
f"{CACHE_ROOT}/engine_{engine_name}_{self.engine_num}.pir"
|
|
)
|
|
with open(PIR_DUMP_FILE, "w") as f:
|
|
f.write(group_str)
|
|
trt_params.engine_serialized_data = CACHE_FILE
|
|
|
|
with paddle.pir_utils.IrGuard(), paddle.pir.core.program_guard(program):
|
|
pir.set_insertion_point(group_op)
|
|
out = paddle._C_ops.tensorrt_engine(
|
|
new_input_values,
|
|
trt_params,
|
|
input_names,
|
|
out_names,
|
|
out_shapes,
|
|
out_types,
|
|
"",
|
|
)
|
|
|
|
for out_index in range(len(out)):
|
|
if group_op.result(out_index).use_empty():
|
|
# if result value is not been used, it doesn't need get shape, continue
|
|
continue
|
|
ori_value = output_values[out_index]
|
|
current_value = out[out_index]
|
|
orin_min_shape = self.input_info[ori_value.id]["min_shape"]
|
|
orin_opt_shape = self.input_info[ori_value.id]["opt_shape"]
|
|
orin_max_shape = self.input_info[ori_value.id]["max_shape"]
|
|
orin_min_value = self.input_info[ori_value.id]["min_value"]
|
|
orin_opt_value = self.input_info[ori_value.id]["opt_value"]
|
|
orin_max_value = self.input_info[ori_value.id]["max_value"]
|
|
self.input_info[current_value.id] = {
|
|
"min_shape": orin_min_shape,
|
|
"opt_shape": orin_opt_shape,
|
|
"max_shape": orin_max_shape,
|
|
"min_value": orin_min_value,
|
|
"opt_value": orin_opt_value,
|
|
"max_value": orin_max_value,
|
|
}
|
|
|
|
return out
|
|
|
|
def convert(self, network, paddle_op, inputs):
|
|
trt_version = get_trt_version()
|
|
op_name = paddle_op.name()
|
|
if op_name in ["cf.yield"]:
|
|
return
|
|
else:
|
|
converter_func = converter_registry.get(op_name, trt_version)
|
|
if converter_func is None:
|
|
raise NotImplementedError(
|
|
f"Converter for {op_name} not implemented."
|
|
)
|
|
outs = converter_func(network, paddle_op, inputs)
|
|
if isinstance(outs, trt.ITensor):
|
|
return (outs,)
|
|
else:
|
|
return outs
|
|
|
|
def convert_program_to_trt(self):
|
|
for op in self.program.global_block().ops:
|
|
if op.name() == "cinn_op.group" or op.name() == "builtin.group":
|
|
if not self.trt_config.disable_logging:
|
|
_logger.info(f"start process {op.name()}")
|
|
self.engine_num += 1
|
|
new_out = self.convert_subgraph_to_trt(self.program, op)
|
|
orin_out_values = op.results()
|
|
for o_i in range(len(orin_out_values)):
|
|
orin_out_values[o_i].replace_all_uses_with(new_out[o_i])
|
|
|
|
self.program.global_block().remove_op(op)
|
|
|
|
save_one_parameter = (
|
|
False # We need to keep at least one parameter for save
|
|
)
|
|
for op in self.program.global_block().ops:
|
|
if op.name() == "builtin.parameter":
|
|
parameter_name = op.attrs()["parameter_name"]
|
|
if (
|
|
not save_one_parameter
|
|
and "constant_folding" not in parameter_name
|
|
):
|
|
save_one_parameter = True
|
|
continue
|
|
if op.results()[0].use_empty():
|
|
self.program.global_block().remove_op(op)
|
|
if op.name() == "builtin.constant":
|
|
# builtin.constant can't be saved/loaded, we need del it
|
|
if op.results()[0].use_empty():
|
|
self.program.global_block().remove_op(op)
|
|
else:
|
|
constant_result = op.results()[0]
|
|
constant_value_name = op.attrs()["value"]
|
|
out_dtype = np.dtype(
|
|
paddle.pir.core.datatype_to_str[constant_result.dtype]
|
|
)
|
|
tensor_data = self.scope.var(
|
|
constant_value_name
|
|
).get_tensor()
|
|
constant_array = np.array(
|
|
tensor_data, dtype=out_dtype
|
|
).tolist()
|
|
|
|
if isinstance(constant_array, (int, float)):
|
|
constant_array = [constant_array]
|
|
|
|
# convert builtin.constant to pd_op.full_int_array/full and then delete it
|
|
with paddle.pir.core.program_guard(self.program):
|
|
paddle.base.libpaddle.pir.reset_insertion_point_to_start()
|
|
if len(constant_array) == 1:
|
|
full_value = paddle._C_ops.full(
|
|
[1],
|
|
constant_array[0],
|
|
constant_result.dtype,
|
|
paddle.CUDAPlace(0),
|
|
)
|
|
else:
|
|
full_value = paddle._C_ops.full_int_array(
|
|
constant_array,
|
|
constant_result.dtype,
|
|
paddle.CUDAPlace(0),
|
|
)
|
|
op.replace_all_uses_with([full_value])
|
|
self.program.global_block().remove_op(op)
|
|
|
|
# Call clear_shape_info to clear the previous shape information
|
|
clear_shape_info()
|