Skip to content

Commit d657e45

Browse files
zewenli98zewenli98
committed
feat: support activation dynamo converters
1 parent 65c6494 commit d657e45

File tree

9 files changed

+774
-79
lines changed

9 files changed

+774
-79
lines changed

‎py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

+142-15
Original file line numberDiff line numberDiff line change
@@ -169,55 +169,182 @@ def aten_ops_gelu(
169169
)
170170

171171

172-
@dynamo_tensorrt_converter(torch.ops.aten.matmul) # type: ignore[misc]
173-
@dynamo_tensorrt_converter(torch.ops.aten.mm.default) # type: ignore[misc]
174-
def aten_ops_matmul(
172+
@dynamo_tensorrt_converter(torch.ops.aten.relu.default)
173+
def aten_ops_relu(
175174
network: TRTNetwork,
176175
target: Target,
177176
args: Tuple[Argument, ...],
178177
kwargs: Dict[str, Argument],
179178
name: str,
180179
) -> Union[TRTTensor, Sequence[TRTTensor]]:
181-
return impl.matmul.matrix_multiply(
182-
network, target, SourceIR.ATEN, name, args[0], args[1]
180+
return impl.activation.relu(
181+
network,
182+
target,
183+
SourceIR.ATEN,
184+
name,
185+
args[0],
183186
)
184187

185188

186-
@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default) # type: ignore[misc]
187-
def aten_ops_layernorm(
189+
@dynamo_tensorrt_converter(torch.ops.aten.sigmoid.default)
190+
def aten_ops_sigmoid(
188191
network: TRTNetwork,
189192
target: Target,
190193
args: Tuple[Argument, ...],
191194
kwargs: Dict[str, Argument],
192195
name: str,
193196
) -> Union[TRTTensor, Sequence[TRTTensor]]:
194-
return impl.normalization.layer_norm(
197+
return impl.activation.sigmoid(
195198
network,
196199
target,
197200
SourceIR.ATEN,
198201
name,
199202
args[0],
200-
args[1],
201-
args[2],
202-
args[3],
203-
args[4],
204203
)
205204

206205

207-
@dynamo_tensorrt_converter(torch.ops.aten.relu.default) # type: ignore[misc]
208-
def aten_ops_relu(
206+
@dynamo_tensorrt_converter(torch.ops.aten.tanh.default)
207+
def aten_ops_tanh(
209208
network: TRTNetwork,
210209
target: Target,
211210
args: Tuple[Argument, ...],
212211
kwargs: Dict[str, Argument],
213212
name: str,
214213
) -> Union[TRTTensor, Sequence[TRTTensor]]:
215-
return impl.activation.relu(
214+
return impl.activation.tanh(
215+
network,
216+
target,
217+
SourceIR.ATEN,
218+
name,
219+
args[0],
220+
)
221+
222+
223+
@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
224+
def aten_ops_leaky_relu(
225+
network: TRTNetwork,
226+
target: Target,
227+
args: Tuple[Argument, ...],
228+
kwargs: Dict[str, Argument],
229+
name: str,
230+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
231+
return impl.activation.leaky_relu(
232+
network,
233+
target,
234+
SourceIR.ATEN,
235+
name,
236+
args[0],
237+
args_bounds_check(args, 1, 0.01),
238+
)
239+
240+
241+
@dynamo_tensorrt_converter(torch.ops.aten.elu.default)
242+
def aten_ops_elu(
243+
network: TRTNetwork,
244+
target: Target,
245+
args: Tuple[Argument, ...],
246+
kwargs: Dict[str, Argument],
247+
name: str,
248+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
249+
return impl.activation.elu(
250+
network,
251+
target,
252+
SourceIR.ATEN,
253+
name,
254+
args[0],
255+
alpha=args_bounds_check(args, 1, 1.0),
256+
beta=args_bounds_check(args, 2, None),
257+
)
258+
259+
260+
@dynamo_tensorrt_converter(torch.ops.aten.softplus.default)
261+
def aten_ops_softplus(
262+
network: TRTNetwork,
263+
target: Target,
264+
args: Tuple[Argument, ...],
265+
kwargs: Dict[str, Argument],
266+
name: str,
267+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
268+
return impl.activation.softplus(
216269
network,
217270
target,
218271
SourceIR.ATEN,
219272
name,
220273
args[0],
274+
beta=args_bounds_check(args, 1, 1),
275+
)
276+
277+
278+
@dynamo_tensorrt_converter(torch.ops.aten.clip.default)
279+
def aten_ops_clip(
280+
network: TRTNetwork,
281+
target: Target,
282+
args: Tuple[Argument, ...],
283+
kwargs: Dict[str, Argument],
284+
name: str,
285+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
286+
return impl.activation.clip(
287+
network,
288+
target,
289+
SourceIR.ATEN,
290+
name,
291+
args[0],
292+
alpha=args_bounds_check(args, 1),
293+
beta=args_bounds_check(args, 2),
294+
)
295+
296+
297+
@dynamo_tensorrt_converter(torch.ops.aten.hardsigmoid.default)
298+
def aten_ops_hard_sigmoid(
299+
network: TRTNetwork,
300+
target: Target,
301+
args: Tuple[Argument, ...],
302+
kwargs: Dict[str, Argument],
303+
name: str,
304+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
305+
return impl.activation.hard_sigmoid(
306+
network,
307+
target,
308+
SourceIR.ATEN,
309+
name,
310+
args[0],
311+
alpha=args_bounds_check(args, 1, 1 / 6),
312+
beta=args_bounds_check(args, 2, 1 / 2),
313+
)
314+
315+
316+
@dynamo_tensorrt_converter(torch.ops.aten.matmul) # type: ignore[misc]
317+
@dynamo_tensorrt_converter(torch.ops.aten.mm.default) # type: ignore[misc]
318+
def aten_ops_matmul(
319+
network: TRTNetwork,
320+
target: Target,
321+
args: Tuple[Argument, ...],
322+
kwargs: Dict[str, Argument],
323+
name: str,
324+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
325+
return impl.matmul.matrix_multiply(
326+
network, target, SourceIR.ATEN, name, args[0], args[1]
327+
)
328+
329+
330+
@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default) # type: ignore[misc]
331+
def aten_ops_layernorm(
332+
network: TRTNetwork,
333+
target: Target,
334+
args: Tuple[Argument, ...],
335+
kwargs: Dict[str, Argument],
336+
name: str,
337+
) -> Union[TRTTensor, Sequence[TRTTensor]]:
338+
return impl.normalization.layer_norm(
339+
network,
340+
target,
341+
SourceIR.ATEN,
342+
name,
343+
args[0],
344+
args[1],
345+
args[2],
346+
args[3],
347+
args[4],
221348
)
222349

223350

‎py/torch_tensorrt/dynamo/conversion/impl/activation.py

-63
This file was deleted.

‎py/torch_tensorrt/dynamo/conversion/impl/activation/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from .ops import *

‎py/torch_tensorrt/dynamo/conversion/impl/activation/base.py

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
from typing import Any, Callable, Optional
2+
3+
import tensorrt as trt
4+
from torch.fx.node import Target
5+
from torch_tensorrt.dynamo._SourceIR import SourceIR
6+
from torch_tensorrt.fx.converters.converter_utils import (
7+
mark_as_int8_layer,
8+
set_layer_name,
9+
)
10+
from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
11+
12+
13+
def convert_activation(
14+
network: TRTNetwork,
15+
target: Target,
16+
source_ir: Optional[SourceIR],
17+
name: str,
18+
operation_type: trt.ActivationType,
19+
input_val: TRTTensor,
20+
alpha: Optional[Any] = None,
21+
beta: Optional[Any] = None,
22+
dyn_range_fn: Optional[Callable[[float, float], Any]] = None,
23+
) -> TRTTensor:
24+
"""
25+
Add a TensorRT Activation layer to `network`.
26+
"""
27+
if not isinstance(input_val, TRTTensor):
28+
raise RuntimeError(
29+
f"{operation_type} received input {input_val} that is not part "
30+
"of the TensorRT region!"
31+
)
32+
layer = network.add_activation(input_val, operation_type)
33+
if alpha is not None:
34+
layer.alpha = alpha
35+
if beta is not None:
36+
layer.beta = beta
37+
set_layer_name(layer, target, name, source_ir)
38+
39+
if input_val.dynamic_range is not None:
40+
dyn_range = dyn_range_fn(input_val.dynamic_range)
41+
mark_as_int8_layer(layer, dyn_range)
42+
return layer.get_output(0)

0 commit comments

Comments
 (0)