Selectively enable different frontends (pytorch#2693)

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

.github/workflows/build-test.yml

@@ -227,3 +227,30 @@ jobs:
         ${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/
         ${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/
         popd
+
+  tests-py-core:
+    name: Test core [Python]
+    needs: [generate-matrix, build]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/tensorrt
+            package-name: torch_tensorrt
+            pre-script: packaging/pre_build_script.sh
+    uses: pytorch/tensorrt/.github/workflows/linux-test.yml@main
+    with:
+      job-name: tests-py-core
+      repository: "pytorch/tensorrt"
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      script: |
+        export USE_HOST_DEPS=1
+        pushd .
+        cd tests/py/core
+        ${CONDA_RUN} python -m pip install --pre pytest-xdist timm transformers parameterized expecttest==0.1.6 --use-deprecated=legacy-resolver
+        ${CONDA_RUN} python -m pytest -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml .
+        popd

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 exclude: ^.github/actions/assigner/dist
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: check-yaml
       - id: trailing-whitespace
@@ -16,38 +16,38 @@ repos:
           - --fix=lf
         exclude: ^docs
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v16.0.6
+    rev: v18.1.1
     hooks:
       - id: clang-format
         types_or: [c++, c, cuda]
   - repo: https://github.com/keith/pre-commit-buildifier
-    rev: 6.1.0.2
+    rev: 6.4.0
     hooks:
     -   id: buildifier
         args:
           - --warnings=all
     -   id: buildifier-lint
   - repo: https://github.com/abravalheri/validate-pyproject
-    rev: v0.13
+    rev: v0.16
     hooks:
       - id: validate-pyproject
   - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         name: isort (python)
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: 'v1.4.1'
+    rev: 'v1.9.0'
     hooks:
     -   id: mypy
         exclude: "^py/torch_tensorrt/fx|^examples|^tests|^py/torch_tensorrt/dynamo/_experimental|^tools|^docs|noxfile.py|setup.py|versions.py"
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.278
+    rev: v0.3.3
     hooks:
       - id: ruff
   - repo: https://github.com/psf/black
-    rev: 24.1.1
+    rev: 24.3.0
     hooks:
       - id: black
         exclude: ^examples/custom_converters/elu_converter/setup.py|^docs

BUILD

@@ -33,6 +33,14 @@ pkg_tar(
     ],
 )
 
+pkg_tar(
+    name = "include_rt",
+    package_dir = "include/torch_tensorrt",
+    deps = [
+        "//core/runtime:include",
+    ],
+)
+
 pkg_tar(
     name = "include",
     srcs = [
@@ -55,6 +63,18 @@ pkg_tar(
     package_dir = "lib/",
 )
 
+pkg_tar(
+    name = "lib_rt",
+    srcs = select({
+        ":windows": ["//cpp/lib:torch_tensorrt_runtime.dll"],
+        "//conditions:default": [
+            "//cpp/lib:libtorchtrt_runtime.so",
+        ],
+    }),
+    mode = "0755",
+    package_dir = "lib/",
+)
+
 pkg_tar(
     name = "bin",
     srcs = [
@@ -82,3 +102,18 @@ pkg_tar(
         "//conditions:default": [":bin"],
     }),
 )
+
+pkg_tar(
+    name = "libtorchtrt_runtime",
+    srcs = [
+        "//:LICENSE",
+        "//bzl_def:BUILD",
+        "//bzl_def:WORKSPACE",
+    ],
+    extension = "tar.gz",
+    package_dir = "torch_tensorrt_runtime",
+    deps = [
+        ":include_rt",
+        ":lib_rt",
+    ],
+)

core/runtime/RTDevice.cpp

@@ -7,8 +7,6 @@ namespace torch_tensorrt {
 namespace core {
 namespace runtime {
 
-const std::string DEVICE_INFO_DELIM = "%";
-
 typedef enum { ID_IDX = 0, SM_MAJOR_IDX, SM_MINOR_IDX, DEVICE_TYPE_IDX, DEVICE_NAME_IDX } SerializedDeviceInfoIndex;
 
 RTDevice::RTDevice() : id{-1}, major{-1}, minor{-1}, device_type{nvinfer1::DeviceType::kGPU} {}

core/runtime/RTDevice.h

@@ -6,6 +6,8 @@ namespace torch_tensorrt {
 namespace core {
 namespace runtime {
 
+const std::string DEVICE_INFO_DELIM = "%";
+
 struct RTDevice {
   int64_t id; // CUDA device id
   int64_t major; // CUDA compute major version

core/runtime/register_jit_hooks.cpp

@@ -116,6 +116,7 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
 TORCH_LIBRARY(tensorrt, m) {
   m.def("execute_engine", execute_engine);
   m.def("SERIALIZED_ENGINE_BINDING_DELIM", []() -> std::string { return std::string(1, TRTEngine::BINDING_DELIM); });
+  m.def("SERIALIZED_RT_DEVICE_DELIM", []() -> std::string { return DEVICE_INFO_DELIM; });
   m.def("ABI_VERSION", []() -> std::string { return ABI_VERSION; });
   m.def("get_multi_device_safe_mode", []() -> bool { return MULTI_DEVICE_SAFE_MODE; });
   m.def("set_multi_device_safe_mode", [](bool multi_device_safe_mode) -> void {

docsrc/py_api/torch_tensorrt.rst

@@ -37,10 +37,6 @@ Classes
    :members:
    :special-members: __init__
 
-.. autoclass:: TRTModuleNext
-   :members:
-   :special-members: __init__
-
 Enums
 -------
 
@@ -50,7 +46,7 @@ Enums
 
 .. autoclass:: EngineCapability
 
-.. autoclass:: TensorFormat
+.. autoclass:: memory_format
 
 Submodules
 ----------

examples/dynamo/torch_compile_stable_diffusion.py

@@ -18,9 +18,8 @@
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 import torch
-from diffusers import DiffusionPipeline
-
 import torch_tensorrt
+from diffusers import DiffusionPipeline
 
 model_id = "CompVis/stable-diffusion-v1-4"
 device = "cuda:0"
@@ -39,7 +38,7 @@
     backend=backend,
     options={
         "truncate_long_and_double": True,
-        "precision": torch.float16,
+        "enabled_precisions": {torch.float32, torch.float16},
     },
     dynamic=False,
 )