apache · tqchen · Aug 22, 2024 · Aug 20, 2024 · Aug 21, 2024 · Aug 21, 2024
diff --git a/docs/arch/index.rst b/docs/arch/index.rst
@@ -408,15 +408,6 @@ Frontends ingest models from different frameworks into the TVM stack.
 
    frontend/tensorflow
 
-
-Security
----------
-.. toctree::
-   :maxdepth: 1
-
-   security
-
-
 microTVM
 --------
 .. toctree::

diff --git a/docs/conf.py b/docs/conf.py
@@ -39,6 +39,7 @@
 import re
 import sys
 from textwrap import dedent, indent
+from typing import List
 from unittest.mock import patch
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -712,10 +713,50 @@ def update_alias_docstring(name, obj, lines):
             lines.append(".. rubric:: Alias of %s:`%s.%s`" % (obj_type, amod, target_name))
 
 
+tvm_class_name_rewrite_map = {
+    "tvm.tir": ["Var", "Call"],
+    "tvm.relax": ["Var", "Call"],
+    "tvm.relax.frontend.nn": ["Module"],
+}
+
+
+def distinguish_class_name(name: str, lines: List[str]):
+    """Distinguish the docstring of type annotations.
+
+    In the whole TVM, there are many classes with the same name but in different modules,
+    e.g. ``tir.Var``, ``relax.Var``. This function is used to distinguish them in the docstring,
+    by adding the module name as prefix.
+
+    To be specific, this function will check the current object name, and if it in the specific
+    module with specific name, it will add the module name as prefix to the class name to prevent
+    the confusion. Further, we only add the prefix to those standalone class name, but skip
+    the pattern of `xx.Var`, `Var.xx` and `xx.Var.xx`.
+
+    Parameters
+    ----------
+    name : str
+        The full name of the object in the doc.
+
+    lines : list
+        The docstring lines, need to be modified inplace.
+    """
+    remap = {}
+    for module_name in tvm_class_name_rewrite_map:
+        if name.startswith(module_name):
+            short_name = module_name[4:] if module_name.startswith("tvm.") else module_name
+            for class_name in tvm_class_name_rewrite_map[module_name]:
+                remap.update({class_name: f"{short_name}.{class_name}"})
+
+    for k, v in remap.items():
+        for i in range(len(lines)):
+            lines[i] = re.sub(rf"(?<!\.)\b{k}\b(?!\.)", v, lines[i])
+
+
 def process_docstring(app, what, name, obj, options, lines):
     """Sphinx callback to process docstring"""
     if callable(obj) or inspect.isclass(obj):
         update_alias_docstring(name, obj, lines)
+        distinguish_class_name(name, lines)
 
 
 from legacy_redirect import build_legacy_redirect

diff --git a/docs/dev/how_to/relay_add_op.rst b/docs/dev/how_to/relay_add_op.rst
@@ -196,14 +196,12 @@ on how to do this, we recommend looking up the tutorials on :ref:`tensor
 expressions <tutorial-tensor-expr-get-started>`, :ref:`TVM's operator inventory
 (topi) <tutorial-topi>` and looking at the example cumulative sum and product
 implementations found in `python/tvm/topi/scan.py`_ and the gpu versions in
-`python/tvm/topi/cuda/scan.py`_. In the case of our cumulative sum and product
-operations we write things directly in :ref:`TIR <api-python-tir>` which is the
-representation where tensor expressions and topi will lower into.
+`python/tvm/topi/cuda/scan.py`_.
 
 .. _python/tvm/topi/scan.py: https://github.com/apache/tvm/blob/main/python/tvm/topi/scan.py
 .. _python/tvm/topi/cuda/scan.py: https://github.com/apache/tvm/blob/main/python/tvm/topi/cuda/scan.py
 
-5. Hooking up Compute and Strategy with Relay
+1. Hooking up Compute and Strategy with Relay
 ---------------------------------------------
 
 After you have implemented your compute function we now need to glue it to our

diff --git a/docs/index.rst b/docs/index.rst
@@ -51,23 +51,29 @@ driving its costs down.
 
 .. toctree::
    :maxdepth: 1
-   :caption: Architecture  Guide
+   :caption: API Reference
 
-   arch/index
+   reference/api/python/index
+   reference/api/links
 
 .. toctree::
    :maxdepth: 1
-   :caption: Topic Guides
+   :caption: Legacy
 
+   reference/langref/index
+   arch/index
    topic/microtvm/index
    topic/vta/index
 
 .. toctree::
    :maxdepth: 1
-   :caption: Reference Guide
+   :caption: About
 
-   reference/langref/index
-   reference/api/python/index
-   reference/api/links
    reference/publications
+   reference/security
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Index
+
    genindex
diff --git a/docs/reference/api/python/dlight.rst b/docs/reference/api/python/dlight.rst
@@ -0,0 +1,22 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+tvm.dlight
+----------
+.. automodule:: tvm.dlight
+   :members:
+   :imported-members:
diff --git a/docs/reference/api/python/index.rst b/docs/reference/api/python/index.rst
@@ -18,34 +18,89 @@
 Python API
 ==========
 
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm
+
+    error
+    ir
+    instrument
+    transform
+    target
+    driver
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.runtime
+
+    runtime/runtime
+    runtime/ndarray
+    runtime/relax_vm
+    runtime/disco
+    runtime/profiling
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.relax
+
+    relax/relax
+    relax/analysis
+    relax/block_builder
+    relax/frontend
+    relax/op
+    relax/transform
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.tir
+
+    tir/tir
+    tir/analysis
+    tir/schedule
+    tir/stmt_functor
+    tir/transform
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.te
+
+    te
+    topi
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.meta_schedule
+
+    meta_schedule
+
+.. toctree::
+    :maxdepth: 1
+    :caption: tvm.dlight
+
+    dlight
+
+.. toctree::
+    :maxdepth: 1
+    :caption: Misc
+
+    rpc
+    contrib
 
 .. toctree::
-   :maxdepth: 2
-
-   runtime
-   ndarray
-   error
-   ir
-   target
-   tir
-   te
-   driver
-   relay/index
-   relay/frontend
-   relay/nn
-   relay/vision
-   relay/image
-   relay/transform
-   relay/analysis
-   relay/backend
-   relay/dataflow_pattern
-   relay/testing
-   autotvm
-   auto_scheduler
-   meta_schedule
-   rpc
-   micro
-   contrib
-   graph_executor
-   topi
-   vta/index
+    :maxdepth: 1
+    :caption: Legacy
+
+    relay/index
+    relay/frontend
+    relay/nn
+    relay/vision
+    relay/image
+    relay/transform
+    relay/analysis
+    relay/backend
+    relay/dataflow_pattern
+    relay/testing
+    autotvm
+    auto_scheduler
+    micro
+    graph_executor
diff --git a/docs/reference/api/python/instrument.rst b/docs/reference/api/python/instrument.rst
@@ -0,0 +1,22 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+tvm.instrument
+--------------
+.. automodule:: tvm.instrument
+   :members:
+   :imported-members:
diff --git a/docs/reference/api/python/ir.rst b/docs/reference/api/python/ir.rst
@@ -21,19 +21,3 @@ tvm.ir
    :members:
    :imported-members:
    :autosummary:
-
-
-tvm.instrument
---------------
-.. automodule:: tvm.instrument
-   :members:
-   :imported-members:
-   :autosummary:
-
-
-tvm.transform
--------------
-.. automodule:: tvm.transform
-   :members:
-   :imported-members:
-   :autosummary:
diff --git a/docs/reference/api/python/relax/analysis.rst b/docs/reference/api/python/relax/analysis.rst
@@ -0,0 +1,22 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+tvm.relax.analysis
+------------------
+.. automodule:: tvm.relax.analysis
+   :members:
+   :imported-members:
diff --git a/docs/reference/api/python/relax/block_builder.rst b/docs/reference/api/python/relax/block_builder.rst
@@ -0,0 +1,21 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+tvm.relax.block_builder
+-----------------------
+.. automodule:: tvm.relax.block_builder
+   :members: