Skip to content

Commit ea47d19

Browse files
committed
fix: Add truncate_long_and_double to Dynamo
- Add default, setting, and function arguments for `truncate_long_and_double` in Dynamo
1 parent 4ed3e3f commit ea47d19

File tree

3 files changed

+9
-2
lines changed

3 files changed

+9
-2
lines changed

py/torch_tensorrt/dynamo/backend/__init__.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
VERSION_COMPATIBLE,
2222
OPTIMIZATION_LEVEL,
2323
USE_EXPERIMENTAL_RT,
24+
TRUNCATE_LONG_AND_DOUBLE,
2425
)
2526

2627

@@ -44,7 +45,7 @@ def compile(
4445
dla_local_dram_size=1073741824,
4546
dla_global_dram_size=536870912,
4647
calibrator=None,
47-
truncate_long_and_double=False,
48+
truncate_long_and_double=TRUNCATE_LONG_AND_DOUBLE,
4849
require_full_compilation=False,
4950
min_block_size=MIN_BLOCK_SIZE,
5051
torch_executed_ops=[],
@@ -63,7 +64,7 @@ def compile(
6364
"The Dynamo backend is an experimental feature, for which only the "
6465
+ "following arguments are supported: "
6566
+ "{enabled_precisions, debug, workspace_size, min_block_size, "
66-
+ "torch_executed_ops, pass_through_build_failures}"
67+
+ "truncate_long_and_double, torch_executed_ops, pass_through_build_failures}"
6768
)
6869

6970
if not isinstance(inputs, collections.abc.Sequence):
@@ -100,6 +101,7 @@ def compile(
100101
version_compatible=version_compatible,
101102
optimization_level=optimization_level,
102103
use_experimental_rt=use_experimental_rt,
104+
truncate_long_and_double=truncate_long_and_double,
103105
**kwargs,
104106
)
105107

@@ -127,6 +129,7 @@ def create_backend(
127129
version_compatible: bool = VERSION_COMPATIBLE,
128130
optimization_level: Optional[int] = OPTIMIZATION_LEVEL,
129131
use_experimental_rt: bool = USE_EXPERIMENTAL_RT,
132+
truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE,
130133
**kwargs,
131134
):
132135
"""Create torch.compile backend given specified arguments
@@ -160,6 +163,7 @@ def create_backend(
160163
version_compatible=version_compatible,
161164
optimization_level=optimization_level,
162165
use_experimental_rt=use_experimental_rt,
166+
truncate_long_and_double=truncate_long_and_double,
163167
)
164168

165169
return partial(

py/torch_tensorrt/dynamo/common/_defaults.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,4 @@
1010
VERSION_COMPATIBLE = False
1111
OPTIMIZATION_LEVEL = None
1212
USE_EXPERIMENTAL_RT = False
13+
TRUNCATE_LONG_AND_DOUBLE = False

py/torch_tensorrt/dynamo/common/_settings.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
VERSION_COMPATIBLE,
1313
OPTIMIZATION_LEVEL,
1414
USE_EXPERIMENTAL_RT,
15+
TRUNCATE_LONG_AND_DOUBLE,
1516
)
1617

1718

@@ -27,3 +28,4 @@ class CompilationSettings:
2728
version_compatible: bool = VERSION_COMPATIBLE
2829
optimization_level: Optional[int] = OPTIMIZATION_LEVEL
2930
use_experimental_rt: bool = USE_EXPERIMENTAL_RT
31+
truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE

0 commit comments

Comments
 (0)