[fbgemm_gpu][docs] Import overview docs for Jagged Tensor Ops

- Add overview documentation for Jagged Tensor Ops

- Add more docstrings for quantize ops

.github/scripts/fbgemm_gpu_build.bash

@@ -420,6 +420,7 @@ build_fbgemm_gpu_install () {
   # fbgemm_gpu/ subdirectory present
   cd - || return 1
   (test_python_import_package "${env_name}" fbgemm_gpu) || return 1
+  cd - || return 1
 
   echo "[BUILD] FBGEMM-GPU build + install completed"
 }

fbgemm_gpu/README.md

@@ -13,7 +13,7 @@ packages (2.1+) that are built against those CUDA versions.
 
 Only Intel/AMD CPUs with AVX2 extensions are currently supported.
 
-See our [Documentation](docs/README.md) for more information.
+See our [Documentation](https://pytorch.org/FBGEMM) for more information.
 
 
 ## Installation

fbgemm_gpu/docs/requirements.txt

@@ -10,6 +10,7 @@ breathe
 bs4
 docutils
 lxml
+myst-parser
 sphinx-lint
 sphinx-serve
 six

fbgemm_gpu/docs/src/conf.py

@@ -49,9 +49,11 @@
 # ones.
 extensions = [
     "breathe",
+    "myst_parser",
     "sphinx.ext.autodoc",
     "sphinx.ext.autosectionlabel",
     "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
 ]
 

fbgemm_gpu/docs/src/general/ContactInfo.rst

@@ -0,0 +1,67 @@
+Testing FBGEMM_GPU
+------------------
+
+The tests (in the ``fbgemm_gpu/test/`` directory) and benchmarks (in the
+``fbgemm_gpu/bench/`` directory) provide good examples on how to use FBGEMM_GPU.
+
+FBGEMM_GPU Tests
+~~~~~~~~~~~~~~~~
+
+To run the tests after building / installing the FBGEMM_GPU package:
+
+.. code:: sh
+
+  # From the /fbgemm_gpu/ directory
+  cd test
+
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning quantize_ops_test.py
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning sparse_ops_test.py
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_embedding_inference_converter_test.py
+
+Testing with the CUDA Variant
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For the FBGEMM_GPU CUDA package, GPUs will be automatically detected and
+used for testing. To run the tests and benchmarks on a GPU-capable
+device in CPU-only mode, ``CUDA_VISIBLE_DEVICES=-1`` must be set in the
+environment:
+
+.. code:: sh
+
+  # Enable for running in CPU-only mode (when on a GPU-capable machine)
+  export CUDA_VISIBLE_DEVICES=-1
+
+  # Enable for debugging failed kernel executions
+  export CUDA_LAUNCH_BLOCKING=1
+
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py
+
+Testing with the ROCm Variant
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+For ROCm machines, testing against a ROCm GPU needs to be enabled with
+``FBGEMM_TEST_WITH_ROCM=1`` set in the environment:
+
+.. code:: sh
+
+  # From the /fbgemm_gpu/ directory
+  cd test
+
+  export FBGEMM_TEST_WITH_ROCM=1
+  # Enable for debugging failed kernel executions
+  export HIP_LAUNCH_BLOCKING=1
+
+  python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py
+
+FBGEMM_GPU Benchmarks
+~~~~~~~~~~~~~~~~~~~~~
+
+To run the benchmarks:
+
+.. code:: sh
+
+  # From the /fbgemm_gpu/ directory
+  cd bench
+
+  python split_table_batched_embeddings_benchmark.py uvm

fbgemm_gpu/docs/src/general/DocsInstructions.rst

@@ -199,6 +199,7 @@ description:
       /// @param param2 Description of param #2
       ///
       /// @return Description of the method's return value.
+      ///
       /// @throw fbgemm_gpu::my_error if an error occurs
       ///
       /// @note This is an example note.
@@ -233,3 +234,39 @@ description:
 
 #.  Verify the changes by building the docs locally or submitting a PR for a
     Netlify preview.
+
+
+Sphinx Documentation Pointers
+-----------------------------
+
+Adding References to Other Sections
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To reference other sections in the documentation, an anchor must first be
+created above the target section:
+
+.. code:: rst
+
+  .. _fbgemm-gpu.docs.example.reference:
+
+  Example Section Header
+  ----------------------
+
+  NOTES:
+
+  #.  The reference anchor must start with an underscore, i.e. ``_``.
+
+  #.  There must be an empty line between the anchor and its target.
+
+The anchor can then be referenced elsewhere in the docs:
+
+.. code:: rst
+
+  Referencing the section :ref:`fbgemm-gpu.docs.example.reference` from
+  another page in the docs.
+
+  Referencing the section with
+  :ref:`custom text <fbgemm-gpu.docs.example.reference>` from another page
+  in the docs.
+
+  Note that the prefix underscore is not needed when referencing the anchor.

fbgemm_gpu/docs/src/general/TestInstructions.rst

@@ -1,8 +1,8 @@
 Testing FBGEMM_GPU
 ------------------
 
-The tests (in the ``test/`` directoy) and benchmarks (in the ``bench/``
-directory) provide good examples on how to use FBGEMM_GPU.
+The tests (in the ``fbgemm_gpu/test/`` directory) and benchmarks (in the
+``fbgemm_gpu/bench/`` directory) provide good examples on how to use FBGEMM_GPU.
 
 FBGEMM_GPU Tests
 ~~~~~~~~~~~~~~~~
@@ -29,7 +29,9 @@ environment:
 
 .. code:: sh
 
+  # Enable for running in CPU-only mode (when on a GPU-capable machine)
   export CUDA_VISIBLE_DEVICES=-1
+
   # Enable for debugging failed kernel executions
   export CUDA_LAUNCH_BLOCKING=1
 

fbgemm_gpu/docs/src/index.rst

@@ -20,7 +20,15 @@ library.
    general/TestInstructions.rst
    general/DocsInstructions.rst
 
-.. _fbgemm-gpu.docs.toc.python:
+.. _fbgemm-gpu.docs.toc.overview:
+
+.. toctree::
+   :maxdepth: 2
+   :caption: FBGEMM_GPU Overview
+
+   overview/jagged-tensor-ops/JaggedTensorOps.rst
+
+.. _fbgemm-gpu.docs.toc.api.python:
 
 .. toctree::
    :maxdepth: 2
@@ -29,7 +37,7 @@ library.
    python-api/table_batched_embedding_ops.rst
    python-api/jagged_tensor_ops.rst
 
-.. _fbgemm-gpu.docs.toc.cpp:
+.. _fbgemm-gpu.docs.toc.api.cpp:
 
 .. toctree::
    :maxdepth: 2