# ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- # Modified from TensorRT demo diffusion, which has the following license: # # SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # 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 gc import os import pathlib from collections import OrderedDict import numpy as np import onnx import onnx_graphsurgeon as gs import tensorrt as trt import torch from cuda import cudart from diffusion_models import PipelineInfo from engine_builder import EngineBuilder, EngineType from polygraphy.backend.common import bytes_from_path from polygraphy.backend.trt import ( CreateConfig, ModifyNetworkOutputs, Profile, engine_from_bytes, engine_from_network, network_from_onnx_path, save_engine, ) from trt_utilities import TRT_LOGGER # Map of numpy dtype -> torch dtype numpy_to_torch_dtype_dict = { np.int32: torch.int32, np.int64: torch.int64, np.float16: torch.float16, np.float32: torch.float32, } def _cuda_assert(cuda_ret): err = cuda_ret[0] if err != cudart.cudaError_t.cudaSuccess: raise RuntimeError( f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" ) if len(cuda_ret) > 1: return cuda_ret[1] return None class TensorrtEngine: def __init__( self, engine_path, ): self.engine_path = engine_path self.engine = None self.context = None self.buffers = OrderedDict() self.tensors = OrderedDict() self.cuda_graph_instance = None def __del__(self): del self.engine del self.context del self.buffers del self.tensors def refit(self, onnx_path, onnx_refit_path): def convert_int64(arr): if len(arr.shape) == 0: return np.int32(arr) return arr def add_to_map(refit_dict, name, values): if name in refit_dict: assert refit_dict[name] is None if values.dtype == np.int64: values = convert_int64(values) refit_dict[name] = values print(f"Refitting TensorRT engine with {onnx_refit_path} weights") refit_nodes = gs.import_onnx(onnx.load(onnx_refit_path)).toposort().nodes # Construct mapping from weight names in refit model -> original model name_map = {} for n, node in enumerate(gs.import_onnx(onnx.load(onnx_path)).toposort().nodes): refit_node = refit_nodes[n] assert node.op == refit_node.op # Constant nodes in ONNX do not have inputs but have a constant output if node.op == "Constant": name_map[refit_node.outputs[0].name] = node.outputs[0].name # Handle scale and bias weights elif node.op == "Conv": if node.inputs[1].__class__ == gs.Constant: name_map[refit_node.name + "_TRTKERNEL"] = node.name + "_TRTKERNEL" if node.inputs[2].__class__ == gs.Constant: name_map[refit_node.name + "_TRTBIAS"] = node.name + "_TRTBIAS" # For all other nodes: find node inputs that are initializers (gs.Constant) else: for i, inp in enumerate(node.inputs): if inp.__class__ == gs.Constant: name_map[refit_node.inputs[i].name] = inp.name def map_name(name): if name in name_map: return name_map[name] return name # Construct refit dictionary refit_dict = {} refitter = trt.Refitter(self.engine, TRT_LOGGER) all_weights = refitter.get_all() for layer_name, role in zip(all_weights[0], all_weights[1]): # for specialized roles, use a unique name in the map: if role == trt.WeightsRole.KERNEL: name = layer_name + "_TRTKERNEL" elif role == trt.WeightsRole.BIAS: name = layer_name + "_TRTBIAS" else: name = layer_name assert name not in refit_dict, "Found duplicate layer: " + name refit_dict[name] = None for n in refit_nodes: # Constant nodes in ONNX do not have inputs but have a constant output if n.op == "Constant": name = map_name(n.outputs[0].name) print(f"Add Constant {name}\n") add_to_map(refit_dict, name, n.outputs[0].values) # Handle scale and bias weights elif n.op == "Conv": if n.inputs[1].__class__ == gs.Constant: name = map_name(n.name + "_TRTKERNEL") add_to_map(refit_dict, name, n.inputs[1].values) if n.inputs[2].__class__ == gs.Constant: name = map_name(n.name + "_TRTBIAS") add_to_map(refit_dict, name, n.inputs[2].values) # For all other nodes: find node inputs that are initializers (AKA gs.Constant) else: for inp in n.inputs: name = map_name(inp.name) if inp.__class__ == gs.Constant: add_to_map(refit_dict, name, inp.values) for layer_name, weights_role in zip(all_weights[0], all_weights[1]): if weights_role == trt.WeightsRole.KERNEL: custom_name = layer_name + "_TRTKERNEL" elif weights_role == trt.WeightsRole.BIAS: custom_name = layer_name + "_TRTBIAS" else: custom_name = layer_name # Skip refitting Trilu for now; scalar weights of type int64 value 1 - for clip model if layer_name.startswith("onnx::Trilu"): continue if refit_dict[custom_name] is not None: refitter.set_weights(layer_name, weights_role, refit_dict[custom_name]) else: print(f"[W] No refit weights for layer: {layer_name}") if not refitter.refit_cuda_engine(): print("Failed to refit!") exit(0) def build( self, onnx_path, fp16, input_profile=None, enable_refit=False, enable_preview=False, enable_all_tactics=False, timing_cache=None, update_output_names=None, ): print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") p = Profile() if input_profile: for name, dims in input_profile.items(): assert len(dims) == 3 p.add(name, min=dims[0], opt=dims[1], max=dims[2]) config_kwargs = {} if not enable_all_tactics: config_kwargs["tactic_sources"] = [] network = network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]) if update_output_names: print(f"Updating network outputs to {update_output_names}") network = ModifyNetworkOutputs(network, update_output_names) engine = engine_from_network( network, config=CreateConfig( fp16=fp16, refittable=enable_refit, profiles=[p], load_timing_cache=timing_cache, **config_kwargs ), save_timing_cache=timing_cache, ) save_engine(engine, path=self.engine_path) def load(self): print(f"Loading TensorRT engine: {self.engine_path}") self.engine = engine_from_bytes(bytes_from_path(self.engine_path)) def activate(self, reuse_device_memory=None): if reuse_device_memory: self.context = self.engine.create_execution_context_without_device_memory() self.context.device_memory = reuse_device_memory else: self.context = self.engine.create_execution_context() def allocate_buffers(self, shape_dict=None, device="cuda"): for idx in range(self.engine.num_io_tensors): binding = self.engine[idx] if shape_dict and binding in shape_dict: shape = shape_dict[binding] else: shape = self.engine.get_binding_shape(binding) dtype = trt.nptype(self.engine.get_binding_dtype(binding)) if self.engine.binding_is_input(binding): self.context.set_binding_shape(idx, shape) tensor = torch.empty(tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]).to(device=device) self.tensors[binding] = tensor def infer(self, feed_dict, stream, use_cuda_graph=False): for name, buf in feed_dict.items(): self.tensors[name].copy_(buf) for name, tensor in self.tensors.items(): self.context.set_tensor_address(name, tensor.data_ptr()) if use_cuda_graph: if self.cuda_graph_instance is not None: _cuda_assert(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream)) _cuda_assert(cudart.cudaStreamSynchronize(stream)) else: # do inference before CUDA graph capture noerror = self.context.execute_async_v3(stream) if not noerror: raise ValueError("ERROR: inference failed.") # capture cuda graph _cuda_assert( cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) ) self.context.execute_async_v3(stream) self.graph = _cuda_assert(cudart.cudaStreamEndCapture(stream)) from cuda import nvrtc result, major, minor = nvrtc.nvrtcVersion() assert result == nvrtc.nvrtcResult(0) if major < 12: self.cuda_graph_instance = _cuda_assert( cudart.cudaGraphInstantiate(self.graph, b"", 0) ) # cuda < 12 else: self.cuda_graph_instance = _cuda_assert(cudart.cudaGraphInstantiate(self.graph, 0)) # cuda >= 12 else: noerror = self.context.execute_async_v3(stream) if not noerror: raise ValueError("ERROR: inference failed.") return self.tensors class TensorrtEngineBuilder(EngineBuilder): """ Helper class to hide the detail of TensorRT Engine from pipeline. """ def __init__( self, pipeline_info: PipelineInfo, max_batch_size=16, hf_token=None, device="cuda", use_cuda_graph=False, ): """ Initializes the ONNX Runtime TensorRT ExecutionProvider Engine Builder. Args: pipeline_info (PipelineInfo): Version and Type of pipeline. max_batch_size (int): Maximum batch size for dynamic batch engine. hf_token (str): HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints. device (str): device to run. use_cuda_graph (bool): Use CUDA graph to capture engine execution and then launch inference """ super().__init__( EngineType.TRT, pipeline_info, max_batch_size=max_batch_size, hf_token=hf_token, device=device, use_cuda_graph=use_cuda_graph, ) self.stream = None self.shared_device_memory = None def load_resources(self, image_height, image_width, batch_size): super().load_resources(image_height, image_width, batch_size) self.stream = _cuda_assert(cudart.cudaStreamCreate()) def teardown(self): super().teardown() if self.shared_device_memory: cudart.cudaFree(self.shared_device_memory) cudart.cudaStreamDestroy(self.stream) del self.stream def load_engines( self, engine_dir, framework_model_dir, onnx_dir, onnx_opset, opt_batch_size, opt_image_height, opt_image_width, force_export=False, force_optimize=False, force_build=False, static_batch=False, static_shape=True, enable_refit=False, enable_preview=False, enable_all_tactics=False, timing_cache=None, onnx_refit_dir=None, ): """ Build and load engines for TensorRT accelerated inference. Export ONNX models first, if applicable. Args: engine_dir (str): Directory to write the TensorRT engines. framework_model_dir (str): Directory to write the framework model ckpt. onnx_dir (str): Directory to write the ONNX models. onnx_opset (int): ONNX opset version to export the models. opt_batch_size (int): Batch size to optimize for during engine building. opt_image_height (int): Image height to optimize for during engine building. Must be a multiple of 8. opt_image_width (int): Image width to optimize for during engine building. Must be a multiple of 8. force_export (bool): Force re-exporting the ONNX models. force_optimize (bool): Force re-optimizing the ONNX models. force_build (bool): Force re-building the TensorRT engine. static_batch (bool): Build engine only for specified opt_batch_size. static_shape (bool): Build engine only for specified opt_image_height & opt_image_width. Default = True. enable_refit (bool): Build engines with refit option enabled. enable_preview (bool): Enable TensorRT preview features. enable_all_tactics (bool): Enable all tactic sources during TensorRT engine builds. timing_cache (str): Path to the timing cache to accelerate build or None onnx_refit_dir (str): Directory containing refit ONNX models. """ # Create directory for directory in [engine_dir, onnx_dir]: if not os.path.exists(directory): print(f"[I] Create directory: {directory}") pathlib.Path(directory).mkdir(parents=True) self.load_models(framework_model_dir) # Export models to ONNX for model_name, obj in self.models.items(): if model_name == "vae" and self.vae_torch_fallback: continue profile_id = obj.get_profile_id( opt_batch_size, opt_image_height, opt_image_width, static_batch, static_shape ) engine_path = self.get_engine_path(engine_dir, model_name, profile_id) if force_export or force_build or not os.path.exists(engine_path): onnx_path = self.get_onnx_path(model_name, onnx_dir, opt=False) onnx_opt_path = self.get_onnx_path(model_name, onnx_dir, opt=True) if force_export or not os.path.exists(onnx_opt_path): if force_export or not os.path.exists(onnx_path): print(f"Exporting model: {onnx_path}") model = obj.load_model(framework_model_dir, self.hf_token) with torch.inference_mode(), torch.autocast("cuda"): inputs = obj.get_sample_input(1, opt_image_height, opt_image_width) torch.onnx.export( model, inputs, onnx_path, export_params=True, opset_version=onnx_opset, do_constant_folding=True, input_names=obj.get_input_names(), output_names=obj.get_output_names(), dynamic_axes=obj.get_dynamic_axes(), ) del model torch.cuda.empty_cache() gc.collect() else: print(f"Found cached model: {onnx_path}") # Optimize onnx if force_optimize or not os.path.exists(onnx_opt_path): print(f"Generating optimizing model: {onnx_opt_path}") obj.optimize_trt(onnx_path, onnx_opt_path) else: print(f"Found cached optimized model: {onnx_opt_path} ") # Build TensorRT engines for model_name, obj in self.models.items(): if model_name == "vae" and self.vae_torch_fallback: continue profile_id = obj.get_profile_id( opt_batch_size, opt_image_height, opt_image_width, static_batch, static_shape ) engine_path = self.get_engine_path(engine_dir, model_name, profile_id) engine = TensorrtEngine(engine_path) onnx_opt_path = self.get_onnx_path(model_name, onnx_dir, opt=True) if force_build or not os.path.exists(engine.engine_path): engine.build( onnx_opt_path, fp16=True, input_profile=obj.get_input_profile( opt_batch_size, opt_image_height, opt_image_width, static_batch, static_shape, ), enable_refit=enable_refit, enable_preview=enable_preview, enable_all_tactics=enable_all_tactics, timing_cache=timing_cache, update_output_names=None, ) self.engines[model_name] = engine # Load TensorRT engines for model_name in self.models: if model_name == "vae" and self.vae_torch_fallback: continue self.engines[model_name].load() if onnx_refit_dir: onnx_refit_path = self.get_onnx_path(model_name, onnx_refit_dir, opt=True) if os.path.exists(onnx_refit_path): self.engines[model_name].refit(onnx_opt_path, onnx_refit_path) def max_device_memory(self): max_device_memory = 0 for engine in self.engines.values(): max_device_memory = max(max_device_memory, engine.engine.device_memory_size) return max_device_memory def activate_engines(self, shared_device_memory=None): if shared_device_memory is None: max_device_memory = self.max_device_memory() _, shared_device_memory = cudart.cudaMalloc(max_device_memory) self.shared_device_memory = shared_device_memory # Load and activate TensorRT engines for engine in self.engines.values(): engine.activate(reuse_device_memory=self.shared_device_memory) def run_engine(self, model_name, feed_dict): return self.engines[model_name].infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph)