| import contextlib |
| import threading |
|
|
| try: |
| from types import NoneType |
| except ImportError: |
| NoneType = type(None) |
| from typing import ByteString, Iterable, MutableMapping |
|
|
| import tensorrt as trt |
| import torch |
|
|
| from tensorrt_llm._utils import get_extra_attr, np_dtype_to_trt, set_extra_attr |
| from tensorrt_llm.logger import logger |
| from tensorrt_llm.network import PluginInfo, get_plugin_info |
|
|
| LAYER_TYPE_2_CLASS = { |
| trt.LayerType.ACTIVATION: trt.IActivationLayer, |
| trt.LayerType.CONCATENATION: trt.IConcatenationLayer, |
| trt.LayerType.CONSTANT: trt.IConstantLayer, |
| trt.LayerType.ELEMENTWISE: trt.IElementWiseLayer, |
| trt.LayerType.FILL: trt.IFillLayer, |
| trt.LayerType.GATHER: trt.IGatherLayer, |
| trt.LayerType.MATRIX_MULTIPLY: trt.IMatrixMultiplyLayer, |
| trt.LayerType.REDUCE: trt.IReduceLayer, |
| trt.LayerType.SELECT: trt.ISelectLayer, |
| trt.LayerType.SHUFFLE: trt.IShuffleLayer, |
| trt.LayerType.SLICE: trt.ISliceLayer, |
| trt.LayerType.SOFTMAX: trt.ISoftMaxLayer, |
| trt.LayerType.UNARY: trt.IUnaryLayer, |
| trt.LayerType.SHAPE: trt.IShapeLayer, |
| trt.LayerType.ASSERTION: trt.IAssertionLayer, |
| trt.LayerType.CAST: trt.ICastLayer, |
| trt.LayerType.NORMALIZATION: trt.INormalizationLayer, |
| trt.LayerType.IDENTITY: trt.IIdentityLayer, |
| trt.LayerType.PLUGIN_V2: trt.IPluginV2Layer, |
| } |
|
|
|
|
| def to_subclass_layer(trt_layer): |
| trt_layer.__class__ = LAYER_TYPE_2_CLASS[trt_layer.type] |
|
|
|
|
| def to_base_class_layer(trt_layer): |
| trt_layer.__class__ = trt.ILayer |
|
|
|
|
| def to_trt_weights(ndarray): |
| weight = trt.Weights( |
| np_dtype_to_trt(ndarray.dtype), |
| ndarray.ctypes.data, |
| ndarray.size, |
| ) |
| |
| set_extra_attr(weight, "numpy", ndarray) |
| return weight |
|
|
|
|
| @contextlib.contextmanager |
| def silent_trt_logger(): |
| min_severity = logger.trt_logger.min_severity |
| logger.trt_logger.min_severity = trt.Logger.ERROR |
| yield |
| logger.trt_logger.min_severity = min_severity |
|
|
|
|
| def compare_tensor(trt_tensor, new_trt_tensor): |
| assert trt_tensor.name == new_trt_tensor.name |
| assert trt_tensor.dtype == new_trt_tensor.dtype |
| assert tuple(trt_tensor.shape) == tuple(new_trt_tensor.shape) |
| assert trt_tensor.broadcast_across_batch == new_trt_tensor.broadcast_across_batch |
| assert trt_tensor.location == new_trt_tensor.location |
| assert trt_tensor.is_network_input == new_trt_tensor.is_network_input |
| assert trt_tensor.is_network_output == new_trt_tensor.is_network_output |
| assert trt_tensor.dynamic_range == new_trt_tensor.dynamic_range |
| assert trt_tensor.is_shape_tensor == new_trt_tensor.is_shape_tensor |
| assert trt_tensor.is_execution_tensor == new_trt_tensor.is_execution_tensor |
| assert trt_tensor.allowed_formats == new_trt_tensor.allowed_formats |
|
|
|
|
| def compare_network(trt_network, new_trt_network): |
| assert trt_network.num_inputs == new_trt_network.num_inputs |
| for i in range(trt_network.num_inputs): |
| input = trt_network.get_input(i) |
| new_input = new_trt_network.get_input(i) |
| compare_tensor(input, new_input) |
| assert trt_network.num_outputs == new_trt_network.num_outputs |
| for i in range(trt_network.num_outputs): |
| output = trt_network.get_output(i) |
| new_output = new_trt_network.get_output(i) |
| compare_tensor(output, new_output) |
| assert trt_network.num_layers == new_trt_network.num_layers |
| for index, new_index in zip(get_sorted_layer_ids(trt_network), |
| get_sorted_layer_ids(new_trt_network)): |
| layer = trt_network.get_layer(index) |
| new_layer = new_trt_network.get_layer(new_index) |
| assert layer.name == new_layer.name |
| assert layer.type == new_layer.type |
| assert layer.precision_is_set == new_layer.precision_is_set |
| assert layer.precision == new_layer.precision |
| assert layer.num_inputs == new_layer.num_inputs |
| for j in range(layer.num_inputs): |
| input = layer.get_input(j) |
| new_input = new_layer.get_input(j) |
| if input is None: |
| assert new_input is None |
| else: |
| assert new_input is not None |
| compare_tensor(input, new_input) |
| assert layer.num_outputs == new_layer.num_outputs |
| for j in range(layer.num_outputs): |
| output = layer.get_output(j) |
| new_output = new_layer.get_output(j) |
| compare_tensor(output, new_output) |
| assert layer.output_type_is_set(j) == new_layer.output_type_is_set( |
| j) |
| if layer.output_type_is_set(j): |
| assert layer.get_output_type(j) == new_layer.get_output_type(j) |
|
|
|
|
| def get_sorted_layer_ids(trt_network): |
| inputs = set() |
| for i in range(trt_network.num_inputs): |
| inputs.add(trt_network.get_input(i).name) |
| layer_ids = [*range(trt_network.num_layers)] |
| sorted_layer_ids = [] |
| walked_tensors = set(inputs) |
| while len(layer_ids) > 0: |
| layer_id = layer_ids.pop(0) |
| layer = trt_network.get_layer(layer_id) |
| no_dependencies = True |
| for j in range(layer.num_inputs): |
| input = layer.get_input(j) |
| if input is None: |
| continue |
| if input.name in walked_tensors: |
| continue |
| else: |
| no_dependencies = False |
| break |
| if no_dependencies: |
| sorted_layer_ids.append(layer_id) |
| for j in range(layer.num_outputs): |
| output = layer.get_output(j) |
| if output is None: |
| continue |
| walked_tensors.add(output.name) |
| else: |
| layer_ids.append(layer_id) |
| assert len(sorted_layer_ids) == trt_network.num_layers |
| return sorted_layer_ids |
|
|
|
|
| def to_tuple(values): |
| if isinstance(values, (int, float, str, bool, NoneType, ByteString)): |
| return values |
| elif isinstance(values, (trt.Dims, trt.Permutation)): |
| if values.__len__() < 0: |
| return None |
| else: |
| return tuple(values) |
| elif isinstance(values, Iterable): |
| return tuple(to_tuple(v) for v in values) |
| elif isinstance(values, MutableMapping): |
| return tuple((k, to_tuple(v)) for k, v in values.items()) |
| else: |
| return values |
|
|
|
|
| _base_layer_attr_names = set(dir(trt.ILayer)) |
|
|
|
|
| def get_cache_key(layer, shapes, values, dtypes=None, updated_attrs=None): |
| updated_attrs = updated_attrs or {} |
| layer_type = layer.type |
| to_subclass_layer(layer) |
| attr_names = set(dir(layer)) - _base_layer_attr_names |
| if layer_type == trt.LayerType.CONSTANT: |
| attr_names.remove("weights") |
| elif layer_type == trt.LayerType.SHUFFLE: |
| if layer.num_inputs >= 2: |
| attr_names.remove("reshape_dims") |
| elif layer_type == trt.LayerType.SLICE: |
| if layer.num_inputs >= 2 and layer.get_input(1) is not None: |
| attr_names.remove("start") |
| if layer.num_inputs >= 3 and layer.get_input(2) is not None: |
| attr_names.remove("shape") |
| if layer.num_inputs >= 4 and layer.get_input(3) is not None: |
| attr_names.remove("stride") |
| elif layer_type == trt.LayerType.FILL: |
| attr_names.remove("is_alpha_beta_int64") |
| if layer.num_inputs >= 1 and layer.get_input(0) is not None: |
| attr_names.remove("shape") |
| if layer.num_inputs >= 2 and layer.get_input(1) is not None: |
| attr_names.remove("alpha") |
| if layer.num_inputs >= 3 and layer.get_input(2) is not None: |
| attr_names.remove("beta") |
| if layer_type != trt.LayerType.PLUGIN_V2: |
| attr_key = tuple( |
| (name, to_tuple(updated_attrs.get(name) or getattr(layer, name))) |
| for name in sorted(attr_names)) |
| else: |
| network = get_trt_network(layer) |
| plugin_info = get_plugin_info(network, layer.name) |
| assert plugin_info is not None, f"layer {layer.name} does not register plugin info" |
| attr_key = tuple( |
| (name, tuple(updated_attrs.get(name) or data)) |
| for name, data in sorted(plugin_info.pfc_as_list.items())) |
| to_base_class_layer(layer) |
| shape_key = () |
| value_key = () |
| dtype_key = () |
| for i in range(layer.num_inputs): |
| input = layer.get_input(i) |
| if input is not None: |
| shape_key += (tuple(shapes[input.name]), ) |
| if input.name in values: |
| value = values[input.name] |
| |
| |
| if isinstance(value, torch.Tensor): |
| value = None |
| else: |
| value = tuple(value) |
| value_key += (value, ) |
| else: |
| value_key += (None, ) |
| if dtypes is not None: |
| dtype_key += (dtypes[input.name], ) |
| else: |
| shape_key += (None, ) |
| value_key += (None, ) |
| dtype_key += (None, ) |
| if dtypes is not None: |
| for i in range(layer.num_outputs): |
| output = layer.get_output(i) |
| dtype_key += (dtypes[output.name], ) |
| cache_key = (layer.type, attr_key, shape_key, value_key) |
| if dtypes is not None: |
| cache_key += (dtype_key, ) |
| return cache_key |
|
|
|
|
| def get_trt_network(layer: trt.ILayer): |
| network = get_extra_attr(layer, "network") |
| assert network is not None |
| return network |
|
|
|
|
| def set_trt_network(layer: trt.ILayer, network: trt.INetworkDefinition): |
| set_extra_attr(layer, "network", network) |
|
|
|
|
| def get_updated_plugin(plugin_info: PluginInfo, updated_attrs): |
| fields = [] |
| for field in plugin_info.pfc: |
| name = field.name |
| if name in updated_attrs: |
| field = trt.PluginField(name, updated_attrs[name], field.type) |
| else: |
| field = trt.PluginField(name, plugin_info.pfc_as_ndarray[name], |
| field.type) |
| fields.append(field) |
| pfc = trt.PluginFieldCollection(fields) |
| plugin = plugin_info.plugin_creator.create_plugin(plugin_info.plugin_name, |
| pfc) |
| new_plugin_info = PluginInfo(plugin_info.plugin_creator, |
| plugin_info.plugin_name, pfc) |
| return plugin, new_plugin_info |
|
|
|
|
| _builder_flags = threading.local() |
| _strongly_typed = threading.local() |
|
|
|
|
| def get_builder_flags(): |
| return getattr(_builder_flags, 'value', 0) |
|
|
|
|
| def get_strongly_typed(): |
| return getattr(_strongly_typed, 'value', False) |
|
|
|
|
| @contextlib.contextmanager |
| def current_flags(builder_flags, strongly_typed): |
| previous_builder_flags = get_builder_flags() |
| _builder_flags.value = builder_flags |
| previous_strongly_typed = get_strongly_typed() |
| _strongly_typed.value = strongly_typed |
| yield |
| _builder_flags.value = previous_builder_flags |
| _strongly_typed.value = previous_strongly_typed |
|
|
|
|
| def get_engine_information(engine_file) -> str: |
| with open(engine_file, "rb") as f: |
| engine_buffer = f.read() |
| runtime = trt.Runtime(logger.trt_logger) |
| engine = runtime.deserialize_cuda_engine(engine_buffer) |
| inspector = engine.create_engine_inspector() |
| return inspector.get_engine_information(trt.LayerInformationFormat.JSON) |
|
|
|
|
| def print_engine_info(engine_file) -> dict: |
| with open(engine_file, "rb") as f: |
| engine_buffer = f.read() |
| from tensorrt_llm.runtime.session import Session |
| Session.from_serialized_engine(engine_buffer)._print_engine_info() |
|
|
