Skip to content

Commit 4dbf214

Browse files
apboseapbose
committed
FX converter doc
1 parent 65c6494 commit 4dbf214

File tree

1 file changed

+196
-0
lines changed

1 file changed

+196
-0
lines changed

‎docsrc/contributors/fx_converters.rst

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
.. _conversion:
2+
3+
FX Converters
4+
==================
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.
8+
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+
11+
Steps
12+
==================
13+
14+
Operation Sets
15+
-------------------
16+
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
18+
has the arguments - ``network, target, args, kwargs, name``, which is common across all the operators schema.
19+
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.
20+
21+
* acc_ops_converters
22+
* acc_trace is produced by ``torch_tensorrt.fx.tracer.acc_tracer.acc_tracer.trace``.
23+
* aten_ops
24+
There are two options at present for this
25+
#. 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``
26+
#. 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.
27+
* nn_ops
28+
* symbolic_trace is produced by ``torch.fx._symbolic_trace``.
29+
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+
33+
Converter implementation
34+
------------------------
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.
36+
Each of them is detailed with the help of an example
37+
38+
* Registration
39+
40+
The converter needs to be registered with the appropriate op code in the ``tensorrt_converter``.
41+
42+
* Activation
43+
44+
Example: ``leaky_relu``
45+
46+
* acc_ops_converters
47+
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``
49+
50+
.. code-block:: python
51+
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+
)
66+
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``
70+
71+
.. code-block:: python
72+
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.
84+
85+
#. 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``
87+
#. args: The arguments passed in the ``call_module`` or ``call_function`` above
88+
#. kwargs: The kwargs passed in the ``call_module`` or ``call_function`` above
89+
#. name: String containing the name of the target
90+
91+
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.
92+
``Acc_ops`` has arguments in the ``args`` whereas ``aten_ops`` has arguments in the ``kwargs`` in the trace.
93+
94+
95+
* Operation type
96+
97+
Example: ``fmod``
98+
99+
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``.
101+
102+
103+
* Implementation Library
104+
105+
The converters across all the above three opsets have the common implementation library ``py/torch_tensorrt/fx/converters/impl``
106+
107+
* Activation
108+
109+
Example: ``leaky_relu``
110+
111+
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.
112+
113+
.. code-block:: python
114+
115+
def leaky_relu(
116+
network: TRTNetwork,
117+
target: Target,
118+
source_ir: Optional[SourceIR],
119+
name: str,
120+
input_val: TRTTensor,
121+
alpha: Optional[Any],
122+
):
123+
#implementation
124+
125+
The implementation function has the following arguments.
126+
127+
#. network : ``network`` passed from the decorated function registration
128+
#. target: ``target`` passed from the decorated function registration
129+
#. source_ir: Enum attribute. ``SourceIR`` enum is defined in ``py/torch_tensorrt/fx/converters/impl/converter_utils``
130+
#. name: ``name`` passed from the decorated function registration
131+
#. input_val: Approriate arguments extracted from the decorated function registration from args or kwargs
132+
#. 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
133+
#. 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
134+
#. dyn_range_fn: A optional function which takes the dynamic range of a TensorRT Tensor and returns the output dynamic range
135+
136+
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``.
137+
138+
* Operator
139+
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.
141+
For a new operator, one should identify the category to which it belongs. Following are some examples
142+
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
147+
148+
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``
150+
151+
152+
* Lowering type
153+
154+
There are some converters which can be decomposed into suboperations and need not have seperate converter registration.
155+
Such converters can be implemented via ``lowering passes``
156+
157+
Example: ``addmm``
158+
159+
The decompositions are registered via ``register_decomposition`` in ``py/torch_tensorrt/dynamo/backend/lowering/_decompositions.py``
160+
We define ``addmm_replacement`` and replace it with the torch ops, which will have their corresponding converters called.
161+
162+
.. code-block:: python
163+
164+
@register_decomposition(torch.ops.aten.addmm, registry=DECOMPOSITIONS)
165+
def addmm_replacement(
166+
input_: torch.Tensor, mat1: torch.Tensor, mat2: torch.Tensor, *, beta=1, alpha=1
167+
) -> torch.Tensor:
168+
return torch.add(
169+
torch.mul(input_, beta), torch.mul(torch.matmul(mat1, mat2), alpha)
170+
)
171+
172+
173+
174+
Tests
175+
-----
176+
177+
* FX testing:
178+
179+
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
180+
181+
#. Compare the results for ``dispatch_tracer.aten_trace`` and torch.
182+
#. 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
183+
184+
The tests should fail if any of the above two conditions fail
185+
186+
* Dynamo testing:
187+
188+
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
189+
190+
#. Compare the results for ``fx.symbolic_trace`` and ``torch_tensorrt.dynamo.compile``.
191+
#. Test for the ``expected_op`` and the ``unexpected_op``.
192+
193+
#. ``expected_op``: Operations the operations are lowered to. eg: ``mul`` and ``add`` for ``addmm``
194+
#. ``unexpected_op``: Original operation. eg: ``addmm`` for ``addmm``
195+
196+
The tests should fail if any of the above two conditions fail

0 commit comments

Comments
 (0)