Skip to content

Commit 3fe5d13

Browse files
committed
Addressing review comment and removing FX portion
1 parent 5231d36 commit 3fe5d13

File tree

1 file changed

+29
-91
lines changed

1 file changed

+29
-91
lines changed

docsrc/contributors/fx_converters.rst

Lines changed: 29 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,27 @@
1-
.. _conversion:
1+
.. _dynamo_conversion:
22

3-
FX Converters
3+
Dynamo Converters
44
==================
5-
The FX converter library in Torch-TensorRT is located in ``TensorRT/py/torch_tensorrt/fx/converters`` (Converters present in FX will soon be deprecated) and ``TensorRT/py/torch_tensorrt/dynamo/conversion``.
6-
FX converters are categorized into - ``aten_ops_converters``, ``acc_ops_converters`` and ``nn_ops_converters``, while dynamo converters only cover ``aten_ops_converters``
7-
The individual converters present in the above folders are useful for the quantization workflow.
5+
The dynamo converter library in Torch-TensorRT is located in ``TensorRT/py/torch_tensorrt/dynamo/conversion``.
86

9-
The dynamo converters are registered using the ``dynamo_tensorrt_converter`` and the FX converters are registered using the ``tensorrt_converter``.
10-
Since FX converters will be deprecated soon, this document will focus more on the dynamo converters.
117

128

139
Steps
1410
==================
1511

16-
Operation Sets
12+
Operation Set
1713
-------------------
18-
There are three different converter sets for FX in torch_tensorrt. Depending on whether the operation is generated using acc_trace, aten_trace or fx_trace, the converters are categorized to one of the three operation sets -
19-
``aten_ops_converters``, ``acc_ops_converters`` or ``nn_ops_converters``. The converters are registered using ``tensorrt_converter`` decorator for FX and ``dynamo_tensorrt_converter`` for dynamo. The function decorated
14+
The converters in dynamo are produced by ``aten_trace`` and falls under ``aten_ops_converters`` ( FX earlier had ``acc_ops_converters``, ``aten_ops_converters`` or ``nn_ops_converters`` depending on the trace through which it was produced). The converters are registered using ``dynamo_tensorrt_converter`` for dynamo. The function decorated
2015
has the arguments - ``network, target, args, kwargs, name``, which is common across all the operators schema.
2116
These functions are mapped in the ``aten`` converter registry dictionary (at present a compilation of FX and dynamo converters, FX will be deprecated soon), with key as the function target name.
2217

23-
* acc_ops_converters
24-
* acc_trace is produced by ``torch_tensorrt.fx.tracer.acc_tracer.acc_tracer.trace``.
25-
* aten_ops
26-
There are two options at present for this
27-
#. Dynamo: aten_trace is produced by ``torch_tensorrt.dynamo.backend.compile``. The second round of trace is produced by ``aot_torch_tensorrt_aten_backend`` by invoking ``aot_module_simplified`` from ``torch._functorch.aot_autograd``
28-
#. FX: aten_trace is produced by ``torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer.trace``. This flow is more common currently, but this will soon be deprecated in torch_tensorrt.
29-
* nn_ops
30-
* symbolic_trace is produced by ``torch.fx._symbolic_trace``.
18+
* aten_trace is produced by ``torch_tensorrt.dynamo.backend.compile`` or ``torch_tensorrt.dynamo.backend.export``.
19+
The second round of trace in compile produced by ``aot_torch_tensorrt_aten_backend`` by invoking ``aot_module_simplified`` from ``torch._functorch.aot_autograd``,
20+
Both these simplify the torch operators to reduced set of Aten operations.
21+
3122

3223
As mentioned above, if you would like to add a new converter, its implementation will be included in ``TensorRT/py/torch_tensorrt/dynamo/conversion/impl``
33-
Although there is a corresponding implementation of the converters included in the common implementation library present in ``TensorRT/py/torch_tensorrt/fx/impl`` for FX converters, this documentation focuses on the implementation of the ``aten_ops`` converters in dynamo. There might be some steps involved in reorganizing files for ``acc_ops`` converters. This is discussed in more detail in the next section.
24+
Although there is a corresponding implementation of the converters included in the common implementation library present in ``TensorRT/py/torch_tensorrt/fx/impl`` for FX converters, this documentation focuses on the implementation of the ``aten_ops`` converters in dynamo.
3425

3526

3627
Converter implementation
@@ -40,60 +31,14 @@ Each of them is detailed with the help of an example
4031

4132
* Registration
4233

43-
The converter needs to be registered with the appropriate op code in the ``tensorrt_converter`` and ``dynamo_tensorrt_converter``.
34+
The converter needs to be registered with the appropriate op code in the ``dynamo_tensorrt_converter``.
4435

4536
* Activation type
4637

4738
Example: ``leaky_relu``
4839

49-
* acc_ops_converters
50-
51-
* FX_converters (soon to be deprecated)
52-
53-
Define in ``py/torch_tensorrt/fx/converters/acc_ops_converters``. One needs to register the opcode generated in the trace, with ``tensorrt_converter`` decorator. Op code to be used for the registration or the converter registry key in this case is ``acc_ops.leaky_relu``
54-
55-
.. code-block:: python
56-
57-
@tensorrt_converter(acc_ops.leaky_relu)
58-
def acc_ops_leaky_relu(
59-
network: TRTNetwork,
60-
target: Target,
61-
args: Tuple[Argument, ...],
62-
kwargs: Dict[str, Argument],
63-
name: str,
64-
) -> Union[TRTTensor, Sequence[TRTTensor]]:
65-
input_val = kwargs["input"]
66-
negative_slope = kwargs["negative_slope"]
67-
operation_type = trt.ActivationType.LEAKY_RELU
68-
return activation.leaky_relu(
69-
network, target, SourceIR.ACC, name, kwargs["input"], kwargs["negative_slope"]
70-
)
71-
72-
Note since the above is deprecated, you may need to revisit these files only for file reorganization.
73-
74-
* Dynamo_converters
75-
76-
The ``acc_ops`` are not present in dynamo converters
7740

78-
* aten_ops_converters
79-
80-
* FX_converters (soon to be deprecated)
81-
82-
Define in ``py/torch_tensorrt/fx/converters/aten_ops_converters``. One needs to register the opcode generated in the trace with ``tensorrt_converter`` decorator. Op code to be used for the registration or the converter registry key in this case is ``torch.ops.aten.leaky_relu.default``
83-
84-
.. code-block:: python
85-
86-
@tensorrt_converter(torch.ops.aten.leaky_relu.default)
87-
def aten_ops_leaky_relu(
88-
network: TRTNetwork,
89-
target: Target,
90-
args: Tuple[Argument, ...],
91-
kwargs: Dict[str, Argument],
92-
name: str,
93-
) -> Union[TRTTensor, Sequence[TRTTensor]]:
94-
return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
95-
96-
* Dynamo_converters
41+
* aten_ops_converters: Dynamo_converters
9742

9843
Define in ``py/torch_tensorrt/dynamo/conversion/aten_ops_converters``. One needs to register the opcode generated in the trace with ``dynamo_tensorrt_converter`` decorator. Op code to be used for the registration or the converter registry key in this case is ``torch.ops.aten.leaky_relu.default``
9944

@@ -109,7 +54,7 @@ Each of them is detailed with the help of an example
10954
) -> Union[TRTTensor, Sequence[TRTTensor]]:
11055
return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
11156
112-
The ``tensorrt_converter`` and ``dynamo_tensorrt_converter`` are similar decorator functions with some differences.
57+
The ``tensorrt_converter`` (used for FX registration) and ``dynamo_tensorrt_converter`` are similar decorator functions with some differences.
11358

11459
#. Both register the converters in the registeries (python dictionaries) - ``CONVERTERS`` and ``DYNAMO_CONVERTERS`` respectively. These are two dictioneries which are concatenated to form the overall converter registry
11560
#. The dictionary is keyed on the ``OpOverLoad`` which is mentioned in more detail below with examples
@@ -136,21 +81,19 @@ Each of them is detailed with the help of an example
13681
The function decorated by ``tensorrt_converter`` and ``dynamo_tensorrt_converter`` has the following arguments which are automatically generated by the trace functions mentioned above.
13782

13883
#. network : Node in the form of ``call_module`` or ``call_function`` having the target as the key
139-
#. target: Target key in the ``call_module`` or ``call_function`` above. eg: ``torch.ops.aten_.leaky_relu.default``. Note that ``torch.ops.aten._leaky_relu`` is the ``OpOverloadPacket`` while ``torch.ops.aten_.leaky_relu.default`` is ``OpOverload``. The
84+
#. target: Target key in the ``call_module`` or ``call_function`` above. eg: ``torch.ops.aten_.leaky_relu.default``. Note that ``torch.ops.aten._leaky_relu`` is the ``OpOverloadPacket`` while ``torch.ops.aten_.leaky_relu.default`` is ``OpOverload``.
14085
#. args: The arguments passed in the ``call_module`` or ``call_function`` above
14186
#. kwargs: The kwargs passed in the ``call_module`` or ``call_function`` above
14287
#. name: String containing the name of the target
14388

144-
As a user writing new converters, one just needs to take care that the approriate arguments are extracted from the trace generated to the implementation function in the implementation lib function ``activation.leaky_relu`` (which we will discuss below in detail). As one can see in the example above, the trace for ``acc_op`` and ``aten_op`` is different.
145-
``Acc_ops`` has arguments in the ``args`` whereas ``aten_ops`` has arguments in the ``kwargs`` in the trace.
146-
89+
As a user writing new converters, one just needs to take care that the approriate arguments are extracted from the trace generated to the implementation function in the implementation lib function ``activation.leaky_relu`` (which we will discuss below in detail).
14790

14891
* Operation type
14992

15093
Example: ``fmod``
15194

15295
It follows the same steps as the above converter. In this case the opcode is ``torch.ops.aten.fmod.Scalar`` or ``torch.ops.aten.fmod.Tensor``.
153-
Hence both the opcodes are registered in ``py/torch_tensorrt/fx/converters/aten_ops_converters`` and ``py/torch_tensorrt/dynamo/conversion/aten_ops_converters``. The opcode is ``acc_ops.fmod`` in ``py/torch_tensorrt/fx/converters/acc_ops_converters``.
96+
Hence both the opcodes are registered in ``py/torch_tensorrt/dynamo/conversion/aten_ops_converters``.
15497
Note that ``torch.ops.aten.fmod`` is the ``OpOverLoadPacket`` while the registry is keyed on ``torch.ops.aten.fmod.Scalar`` or ``torch.ops.aten.fmod.Tensor``, which is ``OpOverLoad``
15598

15699
Example: ``embedding``
@@ -165,7 +108,7 @@ Each of them is detailed with the help of an example
165108
torch.ops.aten.embedding.default, capability_validator=embedding_param_validator
166109
)
167110
168-
So if there is a new converted in which certain special cases are not to be supported then they can be specified in the ``capability_validator``.
111+
So if there is a new converter in which certain special cases are not to be supported then they can be specified in the ``capability_validator``.
169112

170113
* Evaluator type
171114

@@ -177,14 +120,13 @@ Each of them is detailed with the help of an example
177120

178121
* Implementation Library
179122

180-
The converters across all the above three opsets have the common implementation library. FX converters would be ``py/torch_tensorrt/fx/converters/impl`` and dynamo converters would be ``py/torch_tensorrt/dynamo/conversion/impl``
181-
Again as mentioned above, one should focus on the dynamo converters which are implemented in ``py/torch_tensorrt/dynamo/conversion/impl``
123+
The dynamo converters would be located in ``py/torch_tensorrt/dynamo/conversion/impl``
182124

183125
* Activation
184126

185127
Example: ``leaky_relu``
186128

187-
The implementation is to be placed in present in ``py/torch_tensorrt/fx/impl/activation.py``. This is where all the activation functions are defined and implemented.
129+
The implementation is to be placed in present in ``py/torch_tensorrt/dynamo/conversion/impl/activation.py``. This is where all the activation functions are defined and implemented.
188130

189131
.. code-block:: python
190132
@@ -202,22 +144,22 @@ Each of them is detailed with the help of an example
202144

203145
#. network : ``network`` passed from the decorated function registration
204146
#. target: ``target`` passed from the decorated function registration
205-
#. source_ir: Enum attribute. ``SourceIR`` enum is defined in ``py/torch_tensorrt/fx/converters/impl/converter_utils``
147+
#. source_ir: Enum attribute. ``SourceIR`` enum is defined in ``py/torch_tensorrt/dynamo/conversion/impl/converter_utils``
206148
#. name: ``name`` passed from the decorated function registration
207149
#. input_val: Approriate arguments extracted from the decorated function registration from args or kwargs
208150
#. alpha: Approriate arguments extracted from the decorated function registration from args or kwargs. If not None, it will set the alpha attribute of the created TensorRT activation layer eg: Used in leaky_relu, elu, hardtanh
209151
#. beta: Approriate arguments extracted from the decorated function registration from args or kwargs. If not None, it will set the beta attribute of the created TensorRT activation layer eg: Used in hardtanh
210152
#. dyn_range_fn: A optional function which takes the dynamic range of a TensorRT Tensor and returns the output dynamic range
211153

212-
The implementation functions call the ``convert_activation`` function in ``py/torch_tensorrt/fx/impl/activation.py``. This function will add the approriate activation layer via ``network.add_activation``.
154+
The implementation functions call the ``convert_activation`` function in ``py/torch_tensorrt/dynamo/conversion/impl/activation.py``. This function will add the approriate activation layer via ``network.add_activation``.
213155

214156
* Operator
215157

216-
The implementation is to be placed in ``py/torch_tensorrt/fx/impl/elementwise/ops.py`` for FX and ``py/torch_tensorrt/dynamo/conversion/impl`` for dynamo. This is where all the elementwise functions are defined and implemented.
158+
The implementation is to be placed in ``py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py`` for dynamo. This is where all the elementwise functions are defined and implemented.
217159
For a new operator, one should identify the category to which it belongs. Following are some examples
218160

219-
#. Elementwise operators like ``fmod`` is present in ``py/torch_tensorrt/dynamo/conversion/impl/elementwise``. The ``py/torch_tensorrt/fx/impl/elementwise/base`` contains base functions for elementwise operator.
220-
#. Unary operators like ``sqrt`` will be present in ``py/torch_tensorrt/dynamo/conversion/impl/unary``. The ``py/torch_tensorrt/fx/impl/unary/base`` contains base functions for unary operator.
161+
#. Elementwise operators like ``fmod`` is present in ``py/torch_tensorrt/dynamo/conversion/impl/elementwise``. The ``py/torch_tensorrt/dynamo/conversion/impl/elementwise/base`` contains base functions for elementwise operator.
162+
#. Unary operators like ``sqrt`` will be present in ``py/torch_tensorrt/dynamo/conversion/impl/unary``. The ``py/torch_tensorrt/dynamo/conversion/impl/unary/base`` contains base functions for unary operator.
221163
#. Normalization operators like ``softmax``, ``layer_norm``, ``batch_norm`` will be present in ``py/torch_tensorrt/dynamo/conversion/impl/normalization``. Since there are no base operations common to all, there is no base file. But one can choose to implement a base file, if there are common functions across all normalization operations
222164
#. Individual operators like ``slice``, ``select``, ``where``, ``embedding`` will be present in ``py/torch_tensorrt/dynamo/conversion/impl/*.py``. They will have individual operator implementation with the same API structure as above but with different individual arguments
223165

@@ -244,21 +186,17 @@ Each of them is detailed with the help of an example
244186
return torch.add(
245187
torch.mul(input_, beta), torch.mul(torch.matmul(mat1, mat2), alpha)
246188
)
189+
190+
Note that there are some pre-existing dynamo decompositions in torch directory, in which case they should be used,
191+
In that case please enable the decompositions in ``py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py`` in ``torch_enabled_decompositions``.
192+
Similarly you can choose to disable any in ``torch_disabled_decompositions``. Please note that the ones already defined in the lowering will take precedence over torch lowering ops.
193+
247194

248195

249196

250197
Tests
251198
-----
252199

253-
* FX testing:
254-
255-
Implement the fx tests in ``py/torch_tensorrt/fx/test/converters/aten_op/test_<operator_name>_aten.py``. Derive the test class from ``DispatchTestCase``, with parameterized testing to implement different test cases. Check for the following two conditions
256-
257-
#. Compare the results for ``dispatch_tracer.aten_trace`` and torch.
258-
#. Test the ``expected_op``. You can find examples in the above tests. This op will be called by the model and needs to be specified in the test so that the test checks that the approriate converter is invoked
259-
260-
The tests should fail if any of the above two conditions fail
261-
262200
* Dynamo testing:
263201

264202
Dynamo tests are present for the lowering ops in ``py/torch_tensorrt/dynamo/backend/test/test_decompositions.py``. The above converters will soon be ported to dynamo tests

0 commit comments

Comments
 (0)