Skip to content

Commit 65d1113

Browse files
apboseapbose
committed
adding dynamo documentation
1 parent 4dbf214 commit 65d1113

File tree

1 file changed

+129
-52
lines changed

1 file changed

+129
-52
lines changed

‎docsrc/contributors/fx_converters.rst

Lines changed: 129 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,21 @@
22

33
FX Converters
44
==================
5-
The FX converter library in Torch-TensorRT is located in ``TensorRT/py/torch_tensorrt/dynamo/converters`` and ``TensorRT/py/torch_tensorrt/fx/converters`` (soon to be deprecated)
6-
They are categorized into - ``aten_ops_converters``, ``acc_ops_converters`` and ``nn_ops_converters``.
7-
The individual converters present are useful for the quantization workflow.
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.
8+
89
The dynamo converters are registered using the ``dynamo_tensorrt_converter`` and the FX converters are registered using the ``tensorrt_converter``.
9-
Since FX converters will be deprecated soon, this document will focus on the dynamo converters.
10+
Since FX converters will be deprecated soon, this document will focus more on the dynamo converters.
11+
1012

1113
Steps
1214
==================
1315

1416
Operation Sets
1517
-------------------
1618
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 -
17-
``aten_ops_converters``, ``acc_ops_converters`` or ``nn_ops_converters``. The converters are registered using ``tensorrt_converter`` decorator. The function decorated
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
1820
has the arguments - ``network, target, args, kwargs, name``, which is common across all the operators schema.
1921
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.
2022

@@ -27,83 +29,157 @@ These functions are mapped in the ``aten`` converter registry dictionary (at pre
2729
* nn_ops
2830
* symbolic_trace is produced by ``torch.fx._symbolic_trace``.
2931

30-
The implementation of the above converter set is to be included in the common implementation library present in ``TensorRT/py/torch_tensorrt/fx/impl``.
31-
This documentation focuses on the implementation of the ``aten_ops`` converters. There might be some steps involved in reorganizing files for ``acc_ops`` converters. This is discussed in more detail in the next section.
32+
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.
34+
3235

3336
Converter implementation
3437
------------------------
35-
In this section, we illustrate the steps to be implemented for writing a converter. We divide them according to activation, operator or lowering pass implementation.
38+
In this section, we illustrate the steps to be implemented for writing a converter. We divide them according to activation, operator, lowering pass implementation or evaluator.
3639
Each of them is detailed with the help of an example
3740

3841
* Registration
3942

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

42-
* Activation
45+
* Activation type
4346

4447
Example: ``leaky_relu``
4548

4649
* acc_ops_converters
50+
51+
* FX_converters (soon to be deprecated)
4752

48-
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``
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
77+
78+
* aten_ops_converters
4979

50-
.. code-block:: python
80+
* FX_converters (soon to be deprecated)
5181

52-
@tensorrt_converter(acc_ops.leaky_relu)
53-
def acc_ops_leaky_relu(
54-
network: TRTNetwork,
55-
target: Target,
56-
args: Tuple[Argument, ...],
57-
kwargs: Dict[str, Argument],
58-
name: str,
59-
) -> Union[TRTTensor, Sequence[TRTTensor]]:
60-
input_val = kwargs["input"]
61-
negative_slope = kwargs["negative_slope"]
62-
operation_type = trt.ActivationType.LEAKY_RELU
63-
return activation.leaky_relu(
64-
network, target, SourceIR.ACC, name, kwargs["input"], kwargs["negative_slope"]
65-
)
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
6685
67-
* aten_ops_converters
68-
69-
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``
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])
7095
71-
.. code-block:: python
96+
* Dynamo_converters
97+
98+
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``
99+
100+
.. code-block:: python
72101
73-
@tensorrt_converter(torch.ops.aten.leaky_relu.default)
74-
def aten_ops_leaky_relu(
75-
network: TRTNetwork,
76-
target: Target,
77-
args: Tuple[Argument, ...],
78-
kwargs: Dict[str, Argument],
79-
name: str,
80-
) -> Union[TRTTensor, Sequence[TRTTensor]]:
81-
return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
82-
83-
The function decorated by ``tensorrt_converter`` has the following arguments which are automatically generated by the trace functions mentioned above.
102+
@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
103+
def aten_ops_leaky_relu(
104+
network: TRTNetwork,
105+
target: Target,
106+
args: Tuple[Argument, ...],
107+
kwargs: Dict[str, Argument],
108+
name: str,
109+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
110+
return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
111+
112+
The ``tensorrt_converter`` and ``dynamo_tensorrt_converter`` are similar decorator functions with some differences.
113+
114+
#. 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
115+
#. The dictionary is keyed on the ``OpOverLoad`` which is mentioned in more detail below with examples
116+
#. Both return the decorated converter implementation
117+
#. The ``CONVERTERS`` directly registers the decorated ``converter_implementation`` function, while ``DYNAMO_CONVERTERS`` has additionational arguments and registers the ``ConverterSupport`` object
118+
#. The additional arguments are:
119+
120+
.. code-block:: python
121+
def dynamo_tensorrt_converter(
122+
key: Target,
123+
enabled: bool = True,
124+
capability_validator: Optional[Callable[[Node], bool]] = None,
125+
priority: ConverterPriority = ConverterPriority.STANDARD,
126+
) -> Callable[[Any], Union[TRTTensor, Sequence[TRTTensor]]]:
127+
128+
#. key: Node target for which the converter is implemented for (for example, torch.ops.aten.leaky_relu.Tensor)
129+
#. enabled: Whether the converter should be enabled/cached or not
130+
#. capability_validator: Function which evaluates whether a node is valid for conversion by the decorated converter. It defaults to None, implying the capability_validator function is always true. This means all nodes of "key" kind can be supported by this converter by default. See ``embedding`` example for more details
131+
#. priority: Converter's level of priority relative to other converters with the same target
132+
133+
#. The ``ConverterSupport`` is a compilation of ``converter_implementation`` and ``capability_validator``.
134+
135+
136+
The function decorated by ``tensorrt_converter`` and ``dynamo_tensorrt_converter`` has the following arguments which are automatically generated by the trace functions mentioned above.
84137

85138
#. network : Node in the form of ``call_module`` or ``call_function`` having the target as the key
86-
#. target: Target key in the ``call_module`` or ``call_function`` above. eg: ``torch.ops.aten_.leaky_relu.default``
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
87140
#. args: The arguments passed in the ``call_module`` or ``call_function`` above
88141
#. kwargs: The kwargs passed in the ``call_module`` or ``call_function`` above
89142
#. name: String containing the name of the target
90143

91144
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.
92145
``Acc_ops`` has arguments in the ``args`` whereas ``aten_ops`` has arguments in the ``kwargs`` in the trace.
93146

94-
147+
95148
* Operation type
96149

97150
Example: ``fmod``
98151

99152
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``.
100-
Hence both the opcodes are registered in ``py/torch_tensorrt/fx/converters/aten_ops_converters``. The opcode is ``acc_ops.fmod`` in ``py/torch_tensorrt/fx/converters/acc_ops_converters``.
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``.
154+
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``
155+
156+
Example: ``embedding``
157+
158+
It follows the same steps as the above converter. In this case the opcode is ``torch.ops.aten.embedding.default``.
159+
There are some converters which have special cases to be accounted for. In those cases, one should use ``capability_validators`` to register the converter using ``@dynamo_tensorrt_converter``
160+
We illustrate this through ``torch.ops.aten.embedding.default``. It has parameters - ``scale_grad_by_freq`` and ``sparse`` which are not currently supported by the implementation.
161+
In such cases we can write validator ``embedding_param_validator`` which implements that given those paramters the converter is not supported and register the converter by
162+
163+
.. code-block:: python
164+
@dynamo_tensorrt_converter(
165+
torch.ops.aten.embedding.default, capability_validator=embedding_param_validator
166+
)
167+
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``.
169+
170+
* Evaluator type
171+
172+
Example: ``operator.getitem``
173+
174+
Evaluators are categorized as so since they do not make any modification to the graph. This is implemented in ``py/torch_tensorrt/dynamo/conversion/op_evaluators.py``, with the corresponding ``capbility_validator``.
175+
The opcode is ``operator.getitem``.
101176

102177

103178
* Implementation Library
104179

105-
The converters across all the above three opsets have the common implementation library ``py/torch_tensorrt/fx/converters/impl``
106-
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``
182+
107183
* Activation
108184

109185
Example: ``leaky_relu``
@@ -137,16 +213,16 @@ Each of them is detailed with the help of an example
137213

138214
* Operator
139215

140-
The implementation is to be placed in ``py/torch_tensorrt/fx/impl/elementwise/ops.py``. This is where all the elementwise functions are defined and implemented.
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.
141217
For a new operator, one should identify the category to which it belongs. Following are some examples
142218

143-
#. Elementwise operators like ``fmod`` is present in ``py/torch_tensorrt/fx/impl/elementwise``. The ``py/torch_tensorrt/fx/impl/elementwise/base`` contains base functions for elementwise operator.
144-
#. Unary operators like ``sqrt`` will be present in ``py/torch_tensorrt/fx/impl/unary``. The ``py/torch_tensorrt/fx/impl/unary/base`` contains base functions for unary operator.
145-
#. Normalization operators like ``softmax``, ``layer_norm``, ``batch_norm`` will be present in ``py/torch_tensorrt/fx/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
146-
#. Individual operators like ``slice``, ``select``, ``where``, ``embedding`` will be present in ``py/torch_tensorrt/fx/impl/*.py``. They will have individual operator implementation with the same API structure as above but with different individual arguments
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.
221+
#. 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
222+
#. 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
147223

148224
Please note that the above operators would have common functions to be implemented which should be placed in
149-
``py/torch_tensorrt/fx/impl/converter_utils.py``
225+
``py/torch_tensorrt/dynamo/conversion/impl/converter_utils.py``
150226

151227

152228
* Lowering type
@@ -191,6 +267,7 @@ The tests should fail if any of the above two conditions fail
191267
#. Test for the ``expected_op`` and the ``unexpected_op``.
192268

193269
#. ``expected_op``: Operations the operations are lowered to. eg: ``mul`` and ``add`` for ``addmm``
270+
#. Note that specify that ``disable_passes= True`` for cases where you would not want lowering passes (which should be the default when testing converters)
194271
#. ``unexpected_op``: Original operation. eg: ``addmm`` for ``addmm``
195272

196273
The tests should fail if any of the above two conditions fail

0 commit comments

Comments
 (0)