Represents a TensorRT algorithm variant, which can be uniquely represented by an implementation ID, tactic ID, and I/O tensor information.
| 123 | |
| 124 | @mod.export() |
| 125 | class Algorithm: |
| 126 | """ |
| 127 | Represents a TensorRT algorithm variant, which can be uniquely represented |
| 128 | by an implementation ID, tactic ID, and I/O tensor information. |
| 129 | """ |
| 130 | |
| 131 | @staticmethod |
| 132 | def from_trt(context, algorithm): |
| 133 | """ |
| 134 | Creates a Polygraphy ``Algorithm`` instance from a TensorRT |
| 135 | ``IAlgorithmContext`` and ``IAlgorithm``. |
| 136 | |
| 137 | Args: |
| 138 | context (trt.IAlgorithmContext): |
| 139 | The algorithm context corresponding to the layer. |
| 140 | algorithm (trt.IAlgorithm): |
| 141 | The algorithm variant provided by TensorRT. |
| 142 | |
| 143 | Returns: |
| 144 | Algorithm |
| 145 | """ |
| 146 | |
| 147 | implementation = algorithm.algorithm_variant.implementation |
| 148 | tactic = algorithm.algorithm_variant.tactic |
| 149 | inputs = tuple(TensorInfo.from_trt(algorithm.get_algorithm_io_info(i)) for i in range(context.num_inputs)) |
| 150 | outputs = tuple( |
| 151 | TensorInfo.from_trt(algorithm.get_algorithm_io_info(i)) |
| 152 | for i in range(context.num_inputs, context.num_inputs + context.num_outputs) |
| 153 | ) |
| 154 | return Algorithm(implementation, tactic, inputs, outputs) |
| 155 | |
| 156 | def __init__(self, implementation, tactic, inputs, outputs): |
| 157 | """ |
| 158 | Args: |
| 159 | implementation (int): |
| 160 | The implementation for this Algorithm. |
| 161 | tactic (int): |
| 162 | The tactic for this Algorithm. |
| 163 | inputs (Sequence[TensorInfo]): |
| 164 | A sequence of TensorInfos for each input. |
| 165 | outputs (Sequence[TensorInfo]): |
| 166 | A sequence of TensorInfos for each output. |
| 167 | """ |
| 168 | self.implementation = implementation |
| 169 | self.tactic = tactic |
| 170 | |
| 171 | def check_io(lst, name): |
| 172 | for index, io in enumerate(lst): |
| 173 | check_is_instance(io, TensorInfo, f"{name}[{index}]") |
| 174 | |
| 175 | check_io(inputs, "inputs") |
| 176 | check_io(outputs, "outputs") |
| 177 | |
| 178 | # Use tuples here so the class is hashable. |
| 179 | self.inputs = tuple(inputs) |
| 180 | self.outputs = tuple(outputs) |
| 181 | |
| 182 | def __str__(self): |