Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit e1c7ed2

Browse files
peri044gs-olive
peri044
and
gs-olive
committedMay 4, 2023
feat: Refactor FX APIs under dynamo namespace for parity with TS APIs (#1807)
Signed-off-by: Dheeraj Peri <[email protected]> Co-authored-by: gs-olive <[email protected]>
1 parent 76d879c commit e1c7ed2

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+3834
-159
lines changed
 

‎.circleci/config.yml

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -521,6 +521,7 @@ commands:
521521
- store_artifacts:
522522
path: /tmp/testlogs
523523

524+
# =================== FX tests start ======================== #
524525
test-fx_core:
525526
description: "Test the fx core"
526527
steps:
@@ -720,6 +721,61 @@ commands:
720721
- store_artifacts:
721722
path: /tmp/testlogs
722723

724+
# =================== FX tests end ======================== #
725+
726+
# =================== Dynamo tests start ======================== #
727+
test-dynamo-fx_ts:
728+
description: "Test the Dynamo fx_ts_compat path"
729+
steps:
730+
- run:
731+
name: Run Dynamo fx_ts_compat core tests
732+
command: |
733+
cd py/torch_tensorrt/dynamo/fx_ts_compat/test
734+
pushd core/
735+
pytest --junitxml=/tmp/artifacts/test_results/dynamo/fx_ts_compat/test_results.xml
736+
popd
737+
738+
- store_test_results:
739+
path: /tmp/artifacts
740+
- store_artifacts:
741+
path: /tmp/testlogs
742+
743+
test-dynamo-torch_compile-core:
744+
description: "Test the Dynamo torch_compile path"
745+
steps:
746+
- run:
747+
name: Run Dynamo torch_compile core tests
748+
command: |
749+
cd py/torch_tensorrt/dynamo/torch_compile
750+
pushd test/
751+
pytest --junitxml=/tmp/artifacts/test_results/dynamo/torch_compile/test_results.xml
752+
popd
753+
754+
- store_test_results:
755+
path: /tmp/artifacts
756+
- store_artifacts:
757+
path: /tmp/testlogs
758+
759+
test-dynamo-torch_compile:
760+
description: "Test the Dynamo torch_compile path"
761+
steps:
762+
- run:
763+
name: Run Dynamo torch_compile E2E tests
764+
command: |
765+
cd py/torch_tensorrt/dynamo/
766+
pushd test/
767+
pip3 install timm
768+
pip3 install transformers
769+
pytest --junitxml=/tmp/artifacts/test_results/dynamo/torch_compile/test_results.xml --ir torch_compile
770+
popd
771+
772+
- store_test_results:
773+
path: /tmp/artifacts
774+
- store_artifacts:
775+
path: /tmp/testlogs
776+
777+
# =================== Dynamo tests end ======================== #
778+
723779
# Define a job to be invoked later in a workflow.
724780
# See: https://circleci.com/docs/2.0/configuration-reference/#jobs
725781
jobs:
@@ -911,6 +967,39 @@ jobs:
911967
- dump-test-env
912968
- test-fx-no-aten
913969

970+
test-py-dynamo-x86_64-linux:
971+
parameters:
972+
torch-build:
973+
type: string
974+
torch-build-index:
975+
type: string
976+
trt-version-long:
977+
type: string
978+
machine:
979+
image: ubuntu-2004-cuda-11.4:202110-01
980+
resource_class: gpu.nvidia.large
981+
steps:
982+
- checkout
983+
- attach_workspace:
984+
at: /tmp/dist/
985+
- install-torch-from-index:
986+
torch-build: << parameters.torch-build >>
987+
torch-build-index: << parameters.torch-build-index >>
988+
- create-py-env:
989+
trt-version-long: << parameters.trt-version-long >>
990+
- install-cudnn
991+
# - run:
992+
# name: "Set LD_LIBRARY_PATH path to include the installed CUDNN"
993+
# command: export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
994+
- run:
995+
name: "Install torch-tensorrt"
996+
command: pip3 install --pre /tmp/dist/x86_64-linux/*cp39-cp39*.whl
997+
# We install torch after torch-trt because pip automatically enforces the version constraint otherwise
998+
- dump-test-env
999+
- test-dynamo-torch_compile
1000+
- test-dynamo-torch_compile-core
1001+
- test-dynamo-fx_ts
1002+
9141003
package-x86_64-linux:
9151004
parameters:
9161005
enabled:
@@ -1300,6 +1389,13 @@ workflows:
13001389
requires:
13011390
- build-x86_64-linux
13021391

1392+
- test-py-dynamo-x86_64-linux:
1393+
torch-build: << pipeline.parameters.torch-build >>
1394+
torch-build-index: << pipeline.parameters.torch-build-index >>
1395+
trt-version-long: << pipeline.parameters.trt-version-long >>
1396+
requires:
1397+
- build-x86_64-linux
1398+
13031399
- build-x86_64-linux:
13041400
name: build-x86_64-linux-legacy
13051401
torch-build: << pipeline.parameters.torch-build-legacy >>
@@ -1374,6 +1470,13 @@ workflows:
13741470
requires:
13751471
- package-x86_64-linux
13761472

1473+
- test-py-dynamo-x86_64-linux:
1474+
torch-build: << pipeline.parameters.torch-build >>
1475+
torch-build-index: << pipeline.parameters.torch-build-index >>
1476+
trt-version-long: << pipeline.parameters.trt-version-long >>
1477+
requires:
1478+
- package-x86_64-linux
1479+
13771480
on-push:
13781481
jobs:
13791482
- build-x86_64-linux:
@@ -1407,6 +1510,13 @@ workflows:
14071510
requires:
14081511
- build-x86_64-linux
14091512

1513+
- test-py-dynamo-x86_64-linux:
1514+
torch-build: << pipeline.parameters.torch-build >>
1515+
torch-build-index: << pipeline.parameters.torch-build-index >>
1516+
trt-version-long: << pipeline.parameters.trt-version-long >>
1517+
requires:
1518+
- build-x86_64-linux
1519+
14101520
- build-x86_64-linux-cmake:
14111521
torch-build: << pipeline.parameters.torch-build >>
14121522
torch-build-index: << pipeline.parameters.torch-build-index >>

‎py/setup.py

Lines changed: 41 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,10 @@ def run(self):
362362
"torch_tensorrt.fx.tools",
363363
"torch_tensorrt.fx.tracer.acc_tracer",
364364
"torch_tensorrt.fx.tracer.dispatch_tracer",
365+
"torch_tensorrt.dynamo",
366+
"torch_tensorrt.dynamo.fx_ts_compat",
367+
"torch_tensorrt.dynamo.fx_ts_compat.passes",
368+
"torch_tensorrt.dynamo.fx_ts_compat.tools",
365369
]
366370
package_dir = {
367371
"torch_tensorrt.fx": "torch_tensorrt/fx",
@@ -370,11 +374,47 @@ def run(self):
370374
"torch_tensorrt.fx.tools": "torch_tensorrt/fx/tools",
371375
"torch_tensorrt.fx.tracer.acc_tracer": "torch_tensorrt/fx/tracer/acc_tracer",
372376
"torch_tensorrt.fx.tracer.dispatch_tracer": "torch_tensorrt/fx/tracer/dispatch_tracer",
377+
"torch_tensorrt.dynamo": "torch_tensorrt/dynamo",
378+
"torch_tensorrt.dynamo.fx_ts_compat": "torch_tensorrt/dynamo/fx_ts_compat",
379+
"torch_tensorrt.dynamo.fx_ts_compat.passes": "torch_tensorrt/dynamo/fx_ts_compat/passes",
380+
"torch_tensorrt.dynamo.fx_ts_compat.tools": "torch_tensorrt/dynamo/fx_ts_compat/tools",
373381
}
374382

375383
with open("README.md", "r", encoding="utf-8") as fh:
376384
long_description = fh.read()
377385

386+
if FX_ONLY:
387+
package_data_list = [
388+
"_Input.py",
389+
]
390+
else:
391+
package_data_list = [
392+
"lib/*",
393+
"include/torch_tensorrt/*.h",
394+
"include/torch_tensorrt/core/*.h",
395+
"include/torch_tensorrt/core/conversion/*.h",
396+
"include/torch_tensorrt/core/conversion/conversionctx/*.h",
397+
"include/torch_tensorrt/core/conversion/converters/*.h",
398+
"include/torch_tensorrt/core/conversion/evaluators/*.h",
399+
"include/torch_tensorrt/core/conversion/tensorcontainer/*.h",
400+
"include/torch_tensorrt/core/conversion/var/*.h",
401+
"include/torch_tensorrt/core/ir/*.h",
402+
"include/torch_tensorrt/core/lowering/*.h",
403+
"include/torch_tensorrt/core/lowering/passes/*.h",
404+
"include/torch_tensorrt/core/partitioning/*.h",
405+
"include/torch_tensorrt/core/partitioning/segmentedblock/*.h",
406+
"include/torch_tensorrt/core/partitioning/partitioninginfo/*.h",
407+
"include/torch_tensorrt/core/partitioning/partitioningctx/*.h",
408+
"include/torch_tensorrt/core/plugins/*.h",
409+
"include/torch_tensorrt/core/plugins/impl/*.h",
410+
"include/torch_tensorrt/core/runtime/*.h",
411+
"include/torch_tensorrt/core/util/*.h",
412+
"include/torch_tensorrt/core/util/logging/*.h",
413+
"bin/*",
414+
"BUILD",
415+
"WORKSPACE",
416+
]
417+
378418
setup(
379419
name="torch_tensorrt",
380420
version=__version__,
@@ -418,32 +458,7 @@ def run(self):
418458
python_requires=">=3.8",
419459
include_package_data=True,
420460
package_data={
421-
"torch_tensorrt": [
422-
"lib/*",
423-
"include/torch_tensorrt/*.h",
424-
"include/torch_tensorrt/core/*.h",
425-
"include/torch_tensorrt/core/conversion/*.h",
426-
"include/torch_tensorrt/core/conversion/conversionctx/*.h",
427-
"include/torch_tensorrt/core/conversion/converters/*.h",
428-
"include/torch_tensorrt/core/conversion/evaluators/*.h",
429-
"include/torch_tensorrt/core/conversion/tensorcontainer/*.h",
430-
"include/torch_tensorrt/core/conversion/var/*.h",
431-
"include/torch_tensorrt/core/ir/*.h",
432-
"include/torch_tensorrt/core/lowering/*.h",
433-
"include/torch_tensorrt/core/lowering/passes/*.h",
434-
"include/torch_tensorrt/core/partitioning/*.h",
435-
"include/torch_tensorrt/core/partitioning/segmentedblock/*.h",
436-
"include/torch_tensorrt/core/partitioning/partitioninginfo/*.h",
437-
"include/torch_tensorrt/core/partitioning/partitioningctx/*.h",
438-
"include/torch_tensorrt/core/plugins/*.h",
439-
"include/torch_tensorrt/core/plugins/impl/*.h",
440-
"include/torch_tensorrt/core/runtime/*.h",
441-
"include/torch_tensorrt/core/util/*.h",
442-
"include/torch_tensorrt/core/util/logging/*.h",
443-
"bin/*",
444-
"BUILD",
445-
"WORKSPACE",
446-
],
461+
"torch_tensorrt": package_data_list,
447462
},
448463
exclude_package_data={
449464
"": ["*.cpp"],

0 commit comments

Comments
 (0)