{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.SPDX-License-Identifier: Apache-2.0\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use\n",
    "\n",
    "this file except in compliance with the License. You may obtain a copy of the License at\n",
    "\n",
    "\n",
    "\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "\n",
    "See the License for the specific language governing permissions and\n",
    "\n",
    "limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Constructing a Network with TensorRT Layer APIs\n",
    "\n",
    "In this notebook, you'll learn how to move beyond pre-built model formats and directly construct a neural network using TensorRT's versatile Layer APIs. This approach offers fine-grained control over your network architecture and optimizations.\n",
    "\n",
    "Specifically, we will cover:\n",
    "\n",
    "1.  **Building a Recurrent Network (LSTM) from Scratch:** Understand how to define each layer of a Long Short-Term Memory (LSTM) cell and then use these components to construct an entire recurrent LSTM layer. This involves using various Layer API functionalities like `add_constant`, `add_matrix_multiply`, `add_elementwise`, `add_slice`, and `add_activation`.\n",
    "2.  **Implementing Loops for Recurrence:** Utilize TensorRT's `add_loop` functionality to efficiently handle the recurrent nature of the LSTM, processing an input sequence step-by-step.\n",
    "3.  **Monitoring Build Progress:** Implement an `IProgressMonitor` to track the engine creation process in real-time, providing visibility into potentially long build times.\n",
    "4.  **Creating Version-Compatible Engines:** Learn to save TensorRT engines with the `BuilderFlag.VERSION_COMPATIBLE` flag, enhancing their portability across different TensorRT patch versions and compatible hardware.\n",
    "\n",
    "This sample uses a small, single-layer LSTM to keep the focus on these core TensorRT API features. We'll also verify its output against an equivalent NumPy implementation to ensure correctness."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "While importing models via ONNX offers convenience, constructing networks directly with TensorRT APIs provides fine-grained control over the network definition. The **[TensorRT Layer API](https://docs.nvidia.com/deeplearning/tensorrt/latest/python_api/infer/Graph/Layers.html)** enables users to define each layer explicitly, offering flexibility and optimization opportunities.\n",
    "\n",
    "To facilitate understanding and verification, this demonstration employs small tensors, allowing for direct comparison with an equivalent NumPy implementation.\n",
    "\n",
    "> **Note: This sample assumes familiarity with the basic concepts of Long Short-Term Memory (LSTM) networks. If you're new to LSTMs, you might find it helpful to review their structure and operation before proceeding.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Prerequisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install numpy tensorrt polygraphy --extra-index-url https://pypi.ngc.nvidia.com\n",
    "import tensorrt as trt\n",
    "import numpy as np\n",
    "from typing import Tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simplify, network parameters and weight initializations use small, illustrative values and dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === Network Parameters & Weights Initialization ===\n",
    "batch_size = 1\n",
    "seq_len = 5      # Length of the sequence\n",
    "input_size = 1   # Dimension of input vector at each time step\n",
    "hidden_size = 2  # Dimension of hidden state and cell state\n",
    "num_units = 1\n",
    "\n",
    "# --- Create Fixed Dummy Weights and Biases (NumPy arrays with dummy values) ---\n",
    "# These will be used by both the TensorRT build and the NumPy verification\n",
    "w_val, u_val, b_val = 0.01, 0.05, 0.3\n",
    "initial_h_val = 0.1\n",
    "initial_c_val = 0.2\n",
    "\n",
    "# Define shapes\n",
    "w_shape = (input_size, 4 * hidden_size) # e.g., [1, 8] for layer 0\n",
    "u_shape = (hidden_size, 4 * hidden_size)       # e.g., [2, 8]\n",
    "b_shape = (4 * hidden_size,)                   # e.g., [8]\n",
    "initial_h_shape = (batch_size, hidden_size)\n",
    "initial_c_shape = (batch_size, hidden_size)\n",
    "\n",
    "# Create NumPy arrays\n",
    "np_weight_W = np.full(w_shape, w_val, dtype=np.float32)\n",
    "np_weight_U = np.full(u_shape, u_val, dtype=np.float32)\n",
    "np_bias = np.full(b_shape, b_val, dtype=np.float32)\n",
    "np_initial_h = np.full(initial_h_shape, initial_h_val, dtype=np.float32)\n",
    "np_initial_c = np.full(initial_c_shape, initial_c_val, dtype=np.float32)\n",
    "\n",
    "# Create inputs for the network\n",
    "np_inputs = np.ones((seq_len, batch_size, input_size), dtype=np.float32)\n",
    "\n",
    "print(\"NumPy Weights Initialized:\")\n",
    "print(f\"  W shape : {np_weight_W.shape}\")\n",
    "print(f\"  U shape : {np_weight_U.shape}\")\n",
    "print(f\"  Bias shape : {np_bias.shape}\")\n",
    "print(f\"  Initial H shape : {np_initial_h.shape}\")\n",
    "print(f\"  Initial C shape : {np_initial_c.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Defining LSTM Operations with the Layer API\n",
    "\n",
    "This step involves defining the LSTM operations by adding layers to the TensorRT `INetworkDefinition`.\n",
    "\n",
    "### Typical Usage Pattern for TensorRT Layer APIs\n",
    "\n",
    "When adding layers to a TensorRT network using the Layer API, the common pattern is:\n",
    "\n",
    "1.  **Add the layer:** Use a `network.add_*` method (e.g., `network.add_matrix_multiply`) to add the desired layer. This method takes input tensors and layer-specific parameters, returning an `ILayer` object representing the newly added layer.\n",
    "2.  **Configure the layer:** Access the returned `ILayer` object to configure its properties. This is optional but useful for naming layer's name and output tensors name for easier debugging and more helpful logs. \n",
    "\n",
    "```python\n",
    "# Example: Adding and configuring a generic layer\n",
    "\n",
    "# 1. Add the layer (replace with a specific layer like add_matrix_multiply)\n",
    "layer = network.add_some_layer(input_tensor, ...)\n",
    "\n",
    "# 2. Configure the layer (optional)\n",
    "output_tensor = layer.get_output(0)\n",
    "output_tensor.name = 'my_layer_output'  # Name the output\n",
    "# ... other configurations ...\n",
    "```\n",
    "\n",
    "For a comprehensive list of available layer types and their specific methods and properties, consult the official [TensorRT Layer API documentation](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Layers.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRT_LOGGER = trt.Logger(trt.Logger.INFO)\n",
    "\n",
    "def add_lstm_unit(network: trt.INetworkDefinition,\n",
    "                  input_x: trt.ITensor,      # Shape: [batch_size, input_size]\n",
    "                  prev_h: trt.ITensor,       # Shape: [batch_size, hidden_size]\n",
    "                  prev_c: trt.ITensor,       # Shape: [batch_size, hidden_size]\n",
    "                  W: np.ndarray,            # Shape: [input_size, 4 * hidden_size]\n",
    "                  U: np.ndarray,            # Shape: [hidden_size, 4 * hidden_size]\n",
    "                  bias: np.ndarray,         # Shape: [4 * hidden_size]\n",
    "                  hidden_size: int,\n",
    "                  input_size: int\n",
    "                  ) -> Tuple[trt.ITensor, trt.ITensor]:\n",
    "    \"\"\"\n",
    "    Adds the computations for a single LSTM time step.\n",
    "    Assumes input tensors have a leading batch dimension.\n",
    "    \"\"\"\n",
    "    batch_size = input_x.shape[0] # Get batch size from input\n",
    "\n",
    "    # Create constant layers for weights and biases\n",
    "    W_layer = network.add_constant(W.shape, trt.Weights(W))\n",
    "    W_layer.get_output(0).name = \"W_const\"\n",
    "    U_layer = network.add_constant(U.shape, U)\n",
    "    U_layer.get_output(0).name = \"U_const\"\n",
    "    # Reshape bias for broadcasting: [4*hidden] -> [1, 4*hidden]\n",
    "    bias_reshaped_np = np.expand_dims(bias.copy(), axis=0)\n",
    "    bias_layer = network.add_constant(bias_reshaped_np.shape, bias_reshaped_np)\n",
    "    bias_layer.get_output(0).name = \"Bias_const\"\n",
    "\n",
    "\n",
    "    # Linear transformations: Wx = input_x * W ; Uh = prev_h * U\n",
    "    # Wx = [batch, input] * [input, 4*hidden] = [batch, 4*hidden]\n",
    "    mm_wx = network.add_matrix_multiply(input_x, trt.MatrixOperation.NONE,\n",
    "                                        W_layer.get_output(0), trt.MatrixOperation.NONE)\n",
    "    mm_wx.get_output(0).name = \"Wx\"\n",
    "\n",
    "    # Uh = [batch, hidden] * [hidden, 4*hidden] = [batch, 4*hidden]\n",
    "    mm_uh = network.add_matrix_multiply(prev_h, trt.MatrixOperation.NONE,\n",
    "                                        U_layer.get_output(0), trt.MatrixOperation.NONE)\n",
    "    mm_uh.get_output(0).name = \"Uh\"\n",
    "\n",
    "\n",
    "    # Combined gates = Wx + Uh + Bias\n",
    "    gates_wx_uh = network.add_elementwise(mm_wx.get_output(0), mm_uh.get_output(0),\n",
    "                                         trt.ElementWiseOperation.SUM)\n",
    "    gates_wx_uh.get_output(0).name = \"Wx_plus_Uh\"\n",
    "\n",
    "    gates = network.add_elementwise(gates_wx_uh.get_output(0), bias_layer.get_output(0),\n",
    "                                    trt.ElementWiseOperation.SUM)\n",
    "\n",
    "    gates_output = gates.get_output(0) # Shape [batch, 4*hidden]\n",
    "    gates_output.name = \"Gates_Combined\"\n",
    "\n",
    "    # Split the combined gates tensor [batch, 4*hidden] -> four [batch, hidden] gate tensors (Input, Forget, Candidate, Output)\n",
    "    def add_gate_slice(index):\n",
    "        gate_slice_layer = network.add_slice(input=gates_output,\n",
    "                                       start=(0, index * hidden_size), # Start [batch_idx=0, col_idx]\n",
    "                                       shape=(batch_size, hidden_size), # Slice shape\n",
    "                                       stride=(1, 1))                   # Stride\n",
    "        return gate_slice_layer.get_output(0)\n",
    "\n",
    "    slice_i = add_gate_slice(0)\n",
    "    slice_i.name = \"Slice_I\"\n",
    "    slice_f = add_gate_slice(1)\n",
    "    slice_f.name = \"Slice_F\"\n",
    "    slice_c = add_gate_slice(2)\n",
    "    slice_c.name = \"Slice_C_candidate\" # Cell candidate\n",
    "    slice_o = add_gate_slice(3)\n",
    "    slice_o.name = \"Slice_O\"\n",
    "\n",
    "    # Apply activations\n",
    "    act_i_layer = network.add_activation(slice_i, trt.ActivationType.SIGMOID)\n",
    "    act_i = act_i_layer.get_output(0)\n",
    "    act_i.name = \"Gate_I\"\n",
    "    act_f_layer = network.add_activation(slice_f, trt.ActivationType.SIGMOID)\n",
    "    act_f = act_f_layer.get_output(0)\n",
    "    act_f.name = \"Gate_F\"\n",
    "    act_c_layer = network.add_activation(slice_c, trt.ActivationType.TANH)\n",
    "    act_c = act_c_layer.get_output(0)\n",
    "    act_c.name = \"Gate_C_candidate\"\n",
    "    act_o_layer = network.add_activation(slice_o, trt.ActivationType.SIGMOID)\n",
    "    act_o = act_o_layer.get_output(0)\n",
    "    act_o.name = \"Gate_O\"\n",
    "\n",
    "    # Cell state update: c_t = f_t * c_{t-1} + i_t * g_t\n",
    "    term1_c = network.add_elementwise(act_f, prev_c, trt.ElementWiseOperation.PROD)\n",
    "    term2_c = network.add_elementwise(act_i, act_c, trt.ElementWiseOperation.PROD)\n",
    "    next_c_layer = network.add_elementwise(term1_c.get_output(0), term2_c.get_output(0), trt.ElementWiseOperation.SUM)\n",
    "\n",
    "    next_c = next_c_layer.get_output(0)\n",
    "    next_c.name = \"next_c\" # Shape [batch, hidden]\n",
    "\n",
    "    # Hidden state update: h_t = o_t * tanh(c_t)\n",
    "    tanh_c_layer = network.add_activation(next_c, trt.ActivationType.TANH)\n",
    "    tanh_c = tanh_c_layer.get_output(0)\n",
    "    next_h_layer = network.add_elementwise(act_o, tanh_c, trt.ElementWiseOperation.PROD)\n",
    "\n",
    "    next_h = next_h_layer.get_output(0)\n",
    "    next_h.name = \"next_h\" # Shape [batch, hidden]\n",
    "\n",
    "    return next_h, next_c\n",
    "\n",
    "\n",
    "def add_lstm_layer(network: trt.INetworkDefinition,\n",
    "                   input_sequence: trt.ITensor, # Shape: [seq_len, batch_size, input_size]\n",
    "                   hidden_size: int,\n",
    "                   seq_len: int,\n",
    "                   weight_W: np.ndarray, # [input_size, 4*hidden] or [hidden, 4*hidden]\n",
    "                   weight_U: np.ndarray, # [hidden, 4*hidden]\n",
    "                   bias: np.ndarray    # [4*hidden]\n",
    "                   ) -> trt.ITensor:\n",
    "    \"\"\"\n",
    "    Adds a LSTM to the network by adding one lstm_unit, and run multiple times with loops.\n",
    "    \"\"\"\n",
    "    # Infer batch_size and input_size from the input tensor shape\n",
    "    assert len(input_sequence.shape) == 3, f\"Input sequence tensor must have 3 dimensions [seq, batch, input]. Got shape {input_sequence.shape}\"\n",
    "    input_size = input_sequence.shape[2]\n",
    "\n",
    "    # Shape: [batch_size, hidden_size]\n",
    "    initial_h = network.add_constant(np_initial_h.shape, np_initial_h).get_output(0)\n",
    "    initial_h.name = \"Initial_H\"\n",
    "    initial_c = network.add_constant(np_initial_c.shape, np_initial_c).get_output(0)\n",
    "    initial_c.name = \"Initial_C\"\n",
    "\n",
    "    loop = network.add_loop()\n",
    "    loop.name = \"Time_Loop_Layer\"\n",
    "\n",
    "    # add_trip_limit determines when the loop should stop. For here we want the loop to run seq_len times.\n",
    "    trip_limit = network.add_constant((), np.array([seq_len], dtype=np.int32)).get_output(0)\n",
    "    loop.add_trip_limit(trip_limit, trt.TripLimit.COUNT)\n",
    "\n",
    "    # Recurrences for hidden and cell states\n",
    "    h_recurrence = loop.add_recurrence(initial_h)\n",
    "    c_recurrence = loop.add_recurrence(initial_c)\n",
    "    prev_h_tensor = h_recurrence.get_output(0)\n",
    "    prev_h_tensor.name = \"Prev_H\"\n",
    "    prev_c_tensor = c_recurrence.get_output(0)\n",
    "    prev_c_tensor.name = \"Prev_C\"\n",
    "\n",
    "    # add_iterator iterates through slices of the input sequence along the specified axis, providing one slice per iteration.\n",
    "    x_t_iterator = loop.add_iterator(input_sequence, axis=0)\n",
    "    x_t = x_t_iterator.get_output(0)\n",
    "    x_t.name = \"x_t\"\n",
    "\n",
    "\n",
    "    # Call the LSTM unit function\n",
    "    next_h, next_c = add_lstm_unit(network=network,\n",
    "                                    input_x=x_t,\n",
    "                                    prev_h=prev_h_tensor,\n",
    "                                    prev_c=prev_c_tensor,\n",
    "                                    W=weight_W,\n",
    "                                    U=weight_U,\n",
    "                                    bias=bias,\n",
    "                                    hidden_size=hidden_size,\n",
    "                                    input_size=input_size)\n",
    "\n",
    "    # Feed the computed states back into the recurrence inputs\n",
    "    h_recurrence.set_input(1, next_h)\n",
    "    c_recurrence.set_input(1, next_c)\n",
    "\n",
    "    # add_loop_output() collects the values in the loop and outputs them. For this example, we concatenate the values along the first axis.\n",
    "    loop_output_h = loop.add_loop_output(next_h, trt.LoopOutput.CONCATENATE, axis=0)\n",
    "\n",
    "    # when using CONCATENATE, the second input must be the trip limit.\n",
    "    loop_output_h.set_input(1, trip_limit)\n",
    "    loop_output_h.get_output(0).name = \"Hidden_Sequence\"\n",
    "\n",
    "    # --- End of time step loop definition ---\n",
    "\n",
    "    layer_output_sequence = loop_output_h.get_output(0)\n",
    "\n",
    "    # The final output sequence is the sequence from the last layer\n",
    "    if layer_output_sequence is None:\n",
    "         raise RuntimeError(\"LSTM layer output was not generated (num_layers may be 0)\")\n",
    "    layer_output_sequence.name = \"Final_LSTM_Output_Sequence\"\n",
    "    return layer_output_sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Build the Network\n",
    "\n",
    "Now that we have the LSTM layer implementation (`add_lstm_layer`), let's proceed to build the TensorRT `INetworkDefinition`.\n",
    "This involves defining the network structure by:\n",
    "1. Adding the input tensor using `network.add_input`.\n",
    "2. Adding the LSTM layer using our custom `add_lstm_layer` function.\n",
    "3. Marking the LSTM layer's output tensor as the network's final output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = trt.Builder(TRT_LOGGER)\n",
    "network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))\n",
    "\n",
    "# === Network Definition ===\n",
    "# Shape: [seq_len, batch_size, input_size] -> e.g., [5, 1, 1]\n",
    "input_tensor = network.add_input(name='input', dtype=trt.float32, shape=(seq_len, batch_size, input_size))\n",
    "\n",
    "# --- Add SINGLE LSTM Layer ---\n",
    "lstm_output = add_lstm_layer(network=network,\n",
    "                                input_sequence=input_tensor,\n",
    "                                hidden_size=hidden_size,\n",
    "                                seq_len=seq_len,\n",
    "                                weight_W=np_weight_W, \n",
    "                                weight_U=np_weight_U, \n",
    "                                bias=np_bias) \n",
    "# lstm_output shape: [seq_len, batch_size, hidden_size] -> e.g., [5, 1, 2]\n",
    "\n",
    "# --- Mark Output ---\n",
    "lstm_output.name = 'hidden_state_sequence'\n",
    "network.mark_output(lstm_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Build the Engine\n",
    "\n",
    "Now that we have defined the network (`INetworkDefinition`), the next step is to build the optimized TensorRT engine. This process involves using the `trt.Builder` along with an `trt.BuilderConfig` object to specify how the engine should be built.\n",
    "\n",
    "The `IBuilderConfig` allows you to control various aspects of the build process, such as:\n",
    "*   Setting memory constraints (e.g., workspace size using `set_memory_pool_limit`).\n",
    "*   Setting builder flags to control optimization strategies and compatibility.\n",
    "\n",
    "Once the network and configuration are ready, the `builder.build_serialized_network(network, config)` method is called to produce the serialized engine, which can then be saved to a file or used directly.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Defining a Progress Monitor\n",
    "Building a TensorRT engine can sometimes take a while, especially for complex models. Don't worry if the build seems long! TensorRT offers a helpful tool called `IProgressMonitor`. This interface lets you track the build process step-by-step, making it easier to monitor progress and even debug if needed. \n",
    "\n",
    "### Implementing `IProgressMonitor`\n",
    "\n",
    "To use the progress monitor, inherit from `trt.IProgressMonitor` and override its key methods:\n",
    "\n",
    "*   `phase_start(self, phase_name, parent_phase, num_steps)`: TensorRT calls this method when it begins a significant phase of the build process (e.g., \"Parsing ONNX Model\", \"Building Engine\"). \n",
    "    *   `phase_name`: Name of the phase starting.\n",
    "    *   `parent_phase`: Name of the parent phase, if this is a sub-phase (can be `None`).\n",
    "    *   `num_steps`: The total number of steps expected for this phase.\n",
    "*   `step_complete(self, phase_name, step)`: Called after each incremental step within a phase is completed.\n",
    "    *   `phase_name`: Name of the current phase.\n",
    "    *   `step`: The index of the step that just finished (0-based).\n",
    "    *   *Your implementation* usually updates the corresponding progress indicator.\n",
    "    *   **Crucially, this method must return `True` to allow the build to continue.** Returning `False` or `None` will signal TensorRT to cancel the build.\n",
    "*   `phase_finish(self, phase_name)`: Called when a phase (and all its steps) is completed.\n",
    "    *   `phase_name`: Name of the phase that finished.\n",
    "    *   *Your implementation* typically finalizes and removes the progress indicator for this phase.\n",
    "\n",
    "After that, hook it with `IBuilderConfig` by setting `config.progress_monitor = MyProgressMonitor()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleProgressMonitor(trt.IProgressMonitor):\n",
    "    def __init__(self):\n",
    "        trt.IProgressMonitor.__init__(self)\n",
    "        self._active_phases = 0\n",
    "\n",
    "    def phase_start(self, phase_name, parent_phase, num_steps):\n",
    "        print(f\"[ProgressMonitor] Phase Start: {phase_name} ({num_steps} steps)\")\n",
    "        self._active_phases += 1\n",
    "\n",
    "    def phase_finish(self, phase_name):\n",
    "        print(f\"[ProgressMonitor] Phase Finish: {phase_name}\")\n",
    "        self._active_phases -= 1\n",
    "\n",
    "    def step_complete(self, phase_name, step):\n",
    "        print(f\"[ProgressMonitor] Step Complete: {phase_name}, Step {step}\")\n",
    "        return True\n",
    "\n",
    "    @property\n",
    "    def active_phases(self):\n",
    "        return self._active_phases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) Version Compatible Engine\n",
    "TensorRT engines are typically optimized for the specific GPU and TensorRT version they are built on. This maximizes performance but can cause incompatibility if the deployment environment differs.\n",
    "\n",
    "The `trt.BuilderFlag.VERSION_COMPATIBLE` flag addresses this by creating a more portable engine. This engine is less sensitive to minor variations in TensorRT versions or GPU models (within a compatible family), potentially at the cost of some performance compared to a non-compatible engine optimized for the exact target. It also reduces the need to rebuild the engine for every minor TensorRT update. Version compatibility is supported from TensorRT 8.6 onwards; the plan must be built with a version at least 8.6 or higher, and the runtime must also be 8.6 or higher.\n",
    "\n",
    "### Use Cases\n",
    "*   Deploying across diverse hardware fleets with compatible GPUs/TRT versions.\n",
    "*   Distributing applications where end-user system configurations vary.\n",
    "*   Simplifying maintenance by avoiding frequent rebuilds for minor updates.\n",
    "\n",
    "### How it Works\n",
    "Enabling `trt.BuilderFlag.VERSION_COMPATIBLE` instructs TensorRT to use more generic optimizations. By default, this flag also causes a copy of a \"lean runtime\" (a version-specific, stripped-down runtime component) to be packaged within the engine plan file. When you deserialize this engine plan on a compatible system, TensorRT recognizes the embedded lean runtime, loads it, and uses this runtime to deserialize and execute the rest of the plan. \n",
    "\n",
    "Because this process involves loading and executing code (the lean runtime) directly from the engine plan file, you must explicitly indicate that you trust the origin and integrity of the plan. This is done by setting `runtime.engine_host_code_allowed = True` on your `trt.Runtime` instance before attempting to deserialize the engine.\n",
    "\n",
    "> **Considerations for Multiple Version-Compatible Engines:**\n",
    "If deploying many version-compatible engines, the embedded lean runtime in each plan can lead to large overall application sizes. An alternative is to exclude the runtime from the engine plan (using `trt.BuilderFlag.EXCLUDE_LEAN_RUNTIME`) and load it manually. This approach can significantly reduce the total deployment footprint. For detailed instructions, refer to the NVIDIA TensorRT documentation on [Manually Loading the Runtime](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html#manually-loading-the-runtime)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENGINE_FILE_PATH = './lstm_network.plan'\n",
    "\n",
    "config = builder.create_builder_config()\n",
    "config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 28) # 256MB\n",
    "config.progress_monitor = SimpleProgressMonitor()\n",
    "config.set_flag(trt.BuilderFlag.VERSION_COMPATIBLE)\n",
    "\n",
    "print(\"Building engine...\")\n",
    "serialized_engine = builder.build_serialized_network(network, config)\n",
    "\n",
    "print(\"Engine build completed.\")\n",
    "with open(ENGINE_FILE_PATH, 'wb') as f:\n",
    "    f.write(serialized_engine)\n",
    "print(f\"Engine saved to {ENGINE_FILE_PATH}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "Once the TensorRT engine is built, the next step is typically to run inference to verify its functionality and performance. The standard process involves creating a runtime and execution context, managing GPU memory for inputs and outputs, transferring data between host and device, and executing the engine etc. While this process provides fine-grained control, it involves boilerplate code. This standard procedure was demonstrated in detail in Sample 1.\n",
    "\n",
    "In this sample, we'll simplify the inference process by using **[Polygraphy](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy)**, a versatile toolkit included with TensorRT that automates many underlying details, such as:\n",
    "*   Context creation\n",
    "*   Buffer management\n",
    "*   Data transfers\n",
    "\n",
    "> **Important Note:** While Polygraphy is excellent for debugging and testing due to its ease of use, it may introduce overhead.\n",
    "> For optimal performance in deployment scenarios, consider handcrafting the inference code as demonstrated in the `1_run_onnx_with_tensorrt` sample.\n",
    "\n",
    "For more examples, please refer to [polygraphy examples](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from polygraphy.backend.common import BytesFromPath\n",
    "from polygraphy.backend.trt import EngineFromBytes, TrtRunner\n",
    "\n",
    "def run_inference_with_polygraphy(h_input: np.ndarray) -> np.ndarray:\n",
    "    input_name = 'input'\n",
    "    output_name = 'hidden_state_sequence'\n",
    "\n",
    "    # Prepare the feed dictionary required by Polygraphy\n",
    "    # Ensure input is contiguous C-style array, which Polygraphy prefers.\n",
    "    h_input_contiguous = np.ascontiguousarray(h_input)\n",
    "    feed_dict = {input_name: h_input_contiguous}\n",
    "\n",
    "    print(f\"Loading engine from: {ENGINE_FILE_PATH}\")\n",
    "    outputs = None\n",
    "    load_engine = EngineFromBytes(BytesFromPath(ENGINE_FILE_PATH))\n",
    "    with TrtRunner(load_engine) as runner:\n",
    "        outputs = runner.infer(feed_dict=feed_dict)\n",
    "        # Polygraphy automatically synchronizes, so no explicit stream sync needed here\n",
    "\n",
    "    output_sequence = outputs[output_name]\n",
    "    print(f\"Output '{output_name}' shape: {output_sequence.shape}, dtype: {output_sequence.dtype}\")\n",
    "    return output_sequence\n",
    "\n",
    "output_sequence = run_inference_with_polygraphy(np_inputs) \n",
    "\n",
    "if output_sequence is not None:\n",
    "    print(f\"\\nInput Sequence (shape {np_inputs.shape}):\\n{np_inputs}\")\n",
    "    print(f\"\\nOutput Hidden State Sequence (shape {output_sequence.shape}):\\n{output_sequence}\")\n",
    "else:\n",
    "    print(\"Inference failed.\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verifying the Output (Comparison with Equivalent Operations in NumPy)\n",
    "\n",
    "To ensure our TensorRT LSTM implementation is correct, we'll compare its output with a reference implementation in NumPy. This is a common practice to validate custom layer logic.\n",
    "\n",
    "The NumPy version will mimic the same LSTM cell computations and unroll the loop over the time sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid_np(x):\n",
    "    x_clipped = np.clip(x, -500, 500)  # avoid overflow\n",
    "    return 1.0 / (1.0 + np.exp(-x_clipped))\n",
    "\n",
    "\n",
    "def tanh_np(x):\n",
    "    x_clipped = np.clip(x, -100, 100)  # avoid overflow\n",
    "    return np.tanh(x_clipped)\n",
    "\n",
    "\n",
    "def lstm_step_numpy(x_t, prev_h, prev_c, W, U, bias):\n",
    "    # W: shape [input_size, 4*hidden_size]\n",
    "    # U: shape [hidden_size, 4*hidden_size]\n",
    "    # bias: shape [4*hidden_size]\n",
    "    # x_t: shape [batch_size, input_size]\n",
    "    # prev_h, prev_c: shape [batch_size, hidden_size]\n",
    "\n",
    "    hidden_size_ = prev_h.shape[1]\n",
    "\n",
    "    Wx = x_t @ W  # Shape [batch_size, 4*hidden_size]\n",
    "    Uh = prev_h @ U  # Shape [batch_size, 4*hidden_size]\n",
    "    gates = Wx + Uh + bias\n",
    "\n",
    "    # Split gates\n",
    "    i = gates[:, 0 * hidden_size_ : 1 * hidden_size_]\n",
    "    f = gates[:, 1 * hidden_size_ : 2 * hidden_size_]\n",
    "    c = gates[:, 2 * hidden_size_ : 3 * hidden_size_]  # Cell candidate\n",
    "    o = gates[:, 3 * hidden_size_ : 4 * hidden_size_]\n",
    "\n",
    "    i_act = sigmoid_np(i)\n",
    "    f_act = sigmoid_np(f)\n",
    "    c_act = tanh_np(c)\n",
    "    o_act = sigmoid_np(o)\n",
    "\n",
    "    next_c = f_act * prev_c + i_act * c_act\n",
    "    next_h = o_act * tanh_np(next_c)\n",
    "\n",
    "    return next_h, next_c\n",
    "\n",
    "\n",
    "def lstm_layer_numpy(input_sequence_np, np_W, np_U, np_bias):\n",
    "    seq_len_ = input_sequence_np.shape[0]\n",
    "    final_output_sequence_np = None\n",
    "    h = np_initial_h.copy()\n",
    "    c = np_initial_c.copy()\n",
    "\n",
    "    layer_output_sequence_list = []\n",
    "\n",
    "    for t in range(seq_len_):\n",
    "        # Slice the input sequence for this time step\n",
    "        x_t = input_sequence_np[t, :, :]\n",
    "\n",
    "        h, c = lstm_step_numpy(x_t, h, c, np_W, np_U, np_bias)\n",
    "        layer_output_sequence_list.append(h)\n",
    "        layer_output_sequence_np = np.stack(layer_output_sequence_list, axis=0)\n",
    "        final_output_sequence_np = layer_output_sequence_np\n",
    "\n",
    "    return final_output_sequence_np\n",
    "\n",
    "\n",
    "numpy_output_sequence = lstm_layer_numpy(np_inputs, np_weight_W, np_weight_U, np_bias)\n",
    "print(\"\\n--- NumPy LSTM Calculation Results ---\")\n",
    "print(f\"Input Sequence (all ones, shape {np_inputs.shape}):\\n{np_inputs}\")\n",
    "print(f\"\\nNumPy Output Hidden State Sequence (shape {numpy_output_sequence.shape}):\\n{numpy_output_sequence}\")\n",
    "print(\"\\n--- Comparison ---\")\n",
    "diff = np.abs(output_sequence - numpy_output_sequence)\n",
    "max_diff = np.max(diff) if diff.size > 0 else 0.0\n",
    "print(f\"Max absolute difference: {max_diff}\")\n",
    "assert np.allclose(\n",
    "    output_sequence, numpy_output_sequence, atol=1e-5\n",
    "), f\"Output sequence mismatch between TensorRT and NumPy, max diff: {max_diff}\"\n",
    "print(\"Notebook executed successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion and Next Steps\n",
    "\n",
    "Congratulations! You have successfully:\n",
    "- Defined an LSTM cell and layer using TensorRT's Layer APIs.\n",
    "- Implemented a recurrent loop with `add_loop`.\n",
    "- Monitored the engine build process using `IProgressMonitor`.\n",
    "- Built a version-compatible TensorRT engine.\n",
    "- Performed inference using the built engine via Polygraphy.\n",
    "- Verified the results against a NumPy implementation.\n",
    "\n",
    "This sample demonstrates the fundamental building blocks for creating custom network architectures in TensorRT. From here, you can explore:\n",
    "- More complex network structures.\n",
    "- Different types of layers available in the TensorRT API.\n",
    "- Advanced loop constructs and conditional logic.\n",
    "- Further optimization techniques if performance is critical (though for this sample, we focused on API usage).\n",
    "\n",
    "By mastering the Layer API, you gain the power to optimize virtually any deep learning model for inference on NVIDIA GPUs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
