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

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

@@ -57,3 +57,4 @@ Submodules
    logging
    ptq
    ts
+   fx

docsrc/index.rst

@@ -1,42 +1,29 @@
-.. user_guide:
-Torch-TensorRT (FX Path) User Guide
-========================
-Torch-TensorRT (FX Path) is a tool that can convert a PyTorch model through torch.FX to an TensorRT engine optimized targeting running on Nvidia GPUs. TensorRT is the inference engine developed by Nvidia which composed of various kinds of optimization including kernel fusion, graph optimization, low precision, etc..
-This tool is developed in Python environment providing most usability to researchers and engineers. There are a few stages that a user want to use this tool and we will introduce them here.
-
-
-Installation
-------------
-Torch-TensorRT (FX Path) is in ``Beta`` phase and always recommended to work with PyTorch nightly.  
+.. _getting_started_with_fx:
 
+Torch-TensorRT (FX Frontend) User Guide
+========================
+Torch-TensorRT (FX Frontend) is a tool that can convert a PyTorch model through ``torch.fx`` to an
+TensorRT engine optimized targeting running on Nvidia GPUs. TensorRT is the inference engine
+developed by NVIDIA which composed of various kinds of optimization including kernel fusion,
+graph optimization, low precision, etc.. This tool is developed in Python environment which allows this
+workflow to be very accessible to researchers and engineers. There are a few stages that a
+user want to use this tool and we will introduce them here.
 
-* Method 1. Follow the instrucions for Torch-TensorRT
-* Method 2. To install FX path only (Python path) and avoid the C++ build for torchscript path
+> Torch-TensorRT (FX Frontend) is in ``Beta`` and currently it is recommended to work with PyTorch nightly.
 
 .. code-block:: shell
 
-    $ conda create --name python_env python=3.8
-    $ conda activate python_env
-
-    # Recommend to install PyTorch 1.12 and later
-    $ conda install pytorch torchvision torchtext cudatoolkit=11.3 -c pytorch-nightly
-
-    # Install TensorRT python package
-    $ pip3 install nvidia-pyindex
-    $ pip3 install nvidia-tensorrt==8.2.4.2
-    $ git clone https://github.com/pytorch/TensorRT.git
-    $ cd TensorRT/py && python setup.py install --fx-only && cd ..
-
-    $ pyton -c "import torch_tensorrt.fx"
     # Test an example by
     $ python py/torch_tensorrt/fx/example/lower_example.py
 
 
 Converting a PyTorch Model to TensorRT Engine
 ---------------------------------------------
-In general, users are welcome to use the ``compile()`` to finish the conversion from a model to tensorRT engine. It is a wrapper API that consists of the major steps needed to finish this converison. Please refer to ``lower_example.py`` file in ``examples/fx``.
+In general, users are welcome to use the ``compile()`` to finish the conversion from a model to tensorRT engine. It is a
+wrapper API that consists of the major steps needed to finish this converison. Please refer to ``lower_example.py`` file in ``examples/fx``.
 
-In this section, we will go through an example to illustrate the major steps that FX path uses. Users can refer to ``fx2trt_example.py`` file in ``examples/fx``.
+In this section, we will go through an example to illustrate the major steps that fx path uses.
+Users can refer to ``fx2trt_example.py`` file in ``examples/fx``.
 
 * **Step 1: Trace the model with acc_tracer**
 Acc_tracer is a tracer inheritated from FX tracer. It comes with args normalizer to convert all args to kwargs and pass to TRT converters.
@@ -276,7 +263,7 @@ In the custom mapper function, we construct an acc op node and return it. The no
 
 The last step would be *adding unit test* for the new acc op or mapper function we added. The place to add the unit test is here `test_acc_tracer.py <https://github.com/pytorch/TensorRT/blob/master/py/torch_tensorrt/fx/test/tracer/test_acc_tracer.py>`_.
 
-* **Step 2. Add a new fx2trt converter**
+* **Step 2. Add a new converter**
 
 All the developed converters for acc ops are all in `acc_op_converter.py <https://github.com/pytorch/TensorRT/blob/master/py/torch_tensorrt/fx/converters/acc_ops_converters.py>`_. It could give you a good example of how the converter is added.