Merge pull request #1294 from pytorch/docgen_update

docs: Update docgen task

Author: narendasan

.github/workflows/docgen.yml

@@ -10,9 +10,9 @@ on:
 
 jobs:
   build-docs:
-    runs-on: ubuntu-18.04
+    runs-on: ubuntu-20.04
     container:
-      image: docker.pkg.github.com/pytorch/tensorrt/docgen:latest
+      image: nvidia/cuda:11.3.1-devel-ubuntu20.04
       credentials:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
@@ -22,23 +22,64 @@ jobs:
           rm -rf /usr/share/dotnet
           rm -rf /opt/ghc
           rm -rf "/usr/local/share/boost"
+          rm -rf /usr/local/cuda/cuda-*
+      - name: Install Python
+        run:
+          apt update
+          apt install -y gcc git curl wget make zlib1g-dev bzip2 libbz2-dev lzma lzma-dev libreadline-dev libsqlite3-dev libssl-dev libffi-dev doxygen pandoc
+          mkdir -p /opt/circleci
+          git clone https://github.com/pyenv/pyenv.git /opt/circleci/.pyenv
+          export PYENV_ROOT="/opt/circleci/.pyenv"
+          export PATH="$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH"
+          pyenv install 3.9.4
+          pyenv global 3.9.4
+          python3 -m pip install --upgrade pip
+          python3 -m pip install wheel
       - uses: actions/checkout@v2
         with:
           ref: ${{github.head_ref}}
       - name: Get HEAD SHA
         id: vars
         run: echo "::set-output name=sha::$(git rev-parse --short HEAD)"
+      - name: Get Bazel version
+        id: bazel_info
+        run: echo "::set-output name=version::$(cat .bazelversion)"
+      - name: Install Bazel
+        run: |
+          wget -q https://github.com/bazelbuild/bazel/releases/download/${{ steps.bazel_info.outputs.version }}/bazel-${{ steps.bazel_info.outputs.version }}-linux-x86_64 -O /usr/bin/bazel
+          chmod a+x /usr/bin/bazel
+      - name: Install cudnn + tensorrt
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
+          mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
+          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
+          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 536F8F1DE80F6A35
+          apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
+          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+          apt-get update
+          apt-get install -y libcudnn8 libcudnn8-dev
+
+          apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
+          add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+          apt-get update
+
+          apt-get install -y libnvinfer8 libnvinfer-plugin8 libnvinfer-dev libnvinfer-plugin-dev
+      - name: Install Torch
+        run: |
+          python3 -m pip install -r py/requirements.txt
       - name: Build Python Package
         run: |
-          cp docker/WORKSPACE.docker WORKSPACE
+          cp toolchains/ci_workspaces/WORKSPACE.x86_64 WORKSPACE
           cd py
-          python3 setup.py install
+          pip install -e .
+          cd ..
       - name: Generate New Docs
         run: |
           cd docsrc
-          pip3 install -r requirements.txt
+          python3 -m pip install -r requirements.txt
           python3 -c "import torch_tensorrt; print(torch_tensorrt.__version__)"
           make html
+          cd ..
       - uses: stefanzweifel/git-auto-commit-action@v4
         with:
         # Required

core/compiler.cpp

@@ -426,7 +426,8 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
       auto outputIsCollection = conversion::OutputIsCollection(g->block());
       if (cfg.partition_info.enabled &&
           (cfg.lower_info.forced_fallback_modules.size() == 0 &&
-           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) && !outputIsCollection) {
+           cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) &&
+          !outputIsCollection) {
         LOG_INFO("Skipping partitioning since model is fully supported");
       }
 

docsrc/Makefile

@@ -37,8 +37,8 @@ endif
 	rm -rf $(SOURCEDIR)/_tmp
 
 html:
-	mkdir -p $(SOURCEDIR)/_notebooks
-	cp -r $(SOURCEDIR)/../notebooks/*.ipynb $(SOURCEDIR)/_notebooks
+#	mkdir -p $(SOURCEDIR)/_notebooks
+#	cp -r $(SOURCEDIR)/../notebooks/*.ipynb $(SOURCEDIR)/_notebooks
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 	mkdir -p $(DESTDIR)
 	cp -r $(BUILDDIR)/html/* $(DESTDIR)

docsrc/WORKSPACE.docs

@@ -1,7 +1,7 @@
-.. _getting_started:
+.. _getting_started_cpp:
 
-Getting Started with C++
-========================
+Using Torch-TensorRT in  C++
+==============================
 
 If you haven't already, acquire a tarball of the library by following the instructions in :ref:`Installation`
 

docsrc/tutorials/torchtrtc.rst renamed to docsrc/cli/torchtrtc.rst

@@ -3,8 +3,16 @@
 Using Torch-TensorRT in Python
 *******************************
 
-Torch-TensorRT Python API accepts a ```torch.nn.Module`` as an input. Under the hood, it uses ``torch.jit.script`` to convert the input module into a
-TorchScript module. To compile your input ```torch.nn.Module`` with Torch-TensorRT, all you need to do is provide the module and inputs
+The Torch-TensorRT Python API supports a number of unique usecases compared to the CLI and C++ APIs which solely support TorchScript compilation.
+
+Torch-TensorRT Python API can accept a ``torch.nn.Module``, ``torch.jit.ScriptModule``, or ``torch.fx.GraphModule`` as an input.
+Depending on what is provided one of the two frontends (TorchScript or FX) will be selected to compile the module. Provided the
+module type is supported, users may explicitly set which frontend they would like to use using the ``ir`` flag for ``compile``.
+If given a ``torch.nn.Module`` and the ``ir`` flag is set to either ``default`` or ``torchscript`` the module will be run through
+``torch.jit.script`` to convert the input module into a  TorchScript module.
+
+
+To compile your input ``torch.nn.Module`` with Torch-TensorRT, all you need to do is provide the module and inputs
 to Torch-TensorRT and you will be returned an optimized TorchScript module to run or add into another PyTorch module. Inputs
 is a list of ``torch_tensorrt.Input`` classes which define input's shape, datatype and memory format. You can also specify settings such as
 operating precision for the engine or target device. After compilation you can save the module just like any other module
@@ -46,6 +54,5 @@ to load in a deployment application. In order to load a TensorRT/TorchScript mod
     input_data = input_data.to("cuda").half()
     result = trt_ts_module(input_data)
 
-Torch-TensorRT python API also provides ``torch_tensorrt.ts.compile`` which accepts a TorchScript module as input.
-The torchscript module can be obtained via scripting or tracing (refer to :ref:`creating_torchscript_module_in_python`). ``torch_tensorrt.ts.compile`` accepts a Torchscript module
-and a list of ``torch_tensorrt.Input`` classes.
+Torch-TensorRT Python API also provides ``torch_tensorrt.ts.compile`` which accepts a TorchScript module as input and ``torch_tensorrt.fx.compile`` which accepts a FX GraphModule as input.
+

docsrc/tutorials/getting_started_with_cpp_api.rst renamed to docsrc/getting_started/getting_started_with_cpp_api.rst

@@ -22,53 +22,51 @@ More Information / System Architecture:
 Getting Started
 ----------------
 * :ref:`installation`
-* :ref:`getting_started`
+* :ref:`getting_started_with_python_api`
+* :ref:`getting_started_cpp`
+
+.. toctree::
+   :caption: Getting Started
+   :maxdepth: 1
+   :hidden:
+
+   getting_started/installation
+   getting_started/getting_started_with_python_api
+   getting_started/getting_started_with_cpp_api
+
+
+Tutorials
+------------
+* :ref:`creating_a_ts_mod`
+* :ref:`getting_started_with_fx`
 * :ref:`ptq`
-* :ref:`torchtrtc`
-* :ref:`use_from_pytorch`
 * :ref:`runtime`
-* :ref:`using_dla`
 * :ref:`serving_torch_tensorrt_with_triton`
-* :ref:`user_guide`
+* :ref:`use_from_pytorch`
+* :ref:`using_dla`
+* :ref:`notebooks`
 
 .. toctree::
-   :caption: Getting Started
+   :caption: Tutorials
    :maxdepth: 1
    :hidden:
 
-   tutorials/installation
-   tutorials/getting_started_with_cpp_api
-   tutorials/getting_started_with_python_api
    tutorials/creating_torchscript_module_in_python
+   tutorials/getting_started_with_fx_path
    tutorials/ptq
-   tutorials/torchtrtc
-   tutorials/use_from_pytorch
    tutorials/runtime
-   tutorials/using_dla
    tutorials/serving_torch_tensorrt_with_triton
-   tutorials/getting_started_with_fx_path
-
-.. toctree::
-   :caption: Notebooks
-   :maxdepth: 1
-   :hidden:
-
-   _notebooks/CitriNet-example
-   _notebooks/dynamic-shapes
-   _notebooks/EfficientNet-example
-   _notebooks/Hugging-Face-BERT
-   _notebooks/lenet-getting-started
-   _notebooks/Resnet50-example
-   _notebooks/ssd-object-detection-demo
-   _notebooks/vgg-qat
-
+   tutorials/use_from_pytorch
+   tutorials/using_dla
+   tutorials/notebooks
 
 Python API Documenation
 ------------------------
 * :ref:`torch_tensorrt_py`
 * :ref:`torch_tensorrt_logging_py`
 * :ref:`torch_tensorrt_ptq_py`
 * :ref:`torch_tensorrt_ts_py`
+* :ref:`torch_tensorrt_fx_py`
 
 .. toctree::
    :caption: Python API Documenation
@@ -79,6 +77,7 @@ Python API Documenation
    py_api/logging
    py_api/ptq
    py_api/ts
+   py_api/fx
 
 C++ API Documenation
 ----------------------
@@ -99,6 +98,18 @@ C++ API Documenation
    _cpp_api/namespace_torch_tensorrt__torchscript
    _cpp_api/namespace_torch_tensorrt__ptq
 
+CLI Documentation
+---------------------
+* :ref:`torchtrtc`
+
+.. toctree::
+   :caption: CLI Documenation
+   :maxdepth: 0
+   :hidden:
+
+   cli/torchtrtc
+
+
 Contributor Documentation
 --------------------------------
 * :ref:`system_overview`

docsrc/tutorials/getting_started_with_python_api.rst renamed to docsrc/getting_started/getting_started_with_python_api.rst

@@ -0,0 +1,31 @@
+.. _torch_tensorrt_fx_py:
+
+torch_tensorrt.fx
+===================
+
+.. currentmodule:: torch_tensorrt.fx
+
+.. automodule torch_tensorrt.ts
+   :undoc-members:
+
+.. automodule:: torch_tensorrt.fx
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Functions
+------------
+
+.. autofunction:: compile
+
+
+Classes
+--------
+
+.. autoclass:: TRTModule
+
+.. autoclass:: InputTensorSpec
+
+.. autoclass:: TRTInterpreter
+
+.. autoclass:: TRTInterpreterResult