[Docs] Add more documentation (#401)

.gitignore

@@ -204,3 +204,12 @@ build-release
 
 # intermediate files
 /gallery/**/*.json
+
+# hidet model files
+*.hidet
+
+# lock files
+*.lock
+
+# experiments folder
+/experiments

docs/source/how-to-guides/add-new-operator/index.rst

@@ -1,6 +1,8 @@
 Add New Operator
 ================
 
+
+
 Hidet is designed to be extensible. It is easy to add new operators to Hidet. There are two ways to add and schedule
 an operator.
 

docs/source/index.rst

@@ -36,6 +36,7 @@ Hidet is an open-source DNN inference framework, it features
   :maxdepth: 1
   :caption: Developer Guide
 
+  gallery/developer-guides/add-torch-operator-mapping
   how-to-guides/add-new-operator/index
   gallery/developer-guides/add-operator-resolve-rule
   gallery/developer-guides/add-subgraph-rewrite-rule

docs/source/python_api/data_types.rst

@@ -1,6 +1,9 @@
 hidet.dtypes
 ============
 
+Hidet supports the following primitive data types, which can be used as the ``dtype`` parameter of functions like
+:func:`hidet.zeros` and :func:`hidet.ones`:.
+
 .. data:: hidet.uint8
 .. data:: hidet.uint16
 .. data:: hidet.uint32
@@ -14,3 +17,4 @@ hidet.dtypes
 .. data:: hidet.float64
 .. data:: hidet.bfloat16
 .. data:: hidet.tfloat32
+.. data:: hidet.boolean

docs/source/python_api/driver.rst → docs/source/python_api/drivers.rst

@@ -2,5 +2,6 @@ hidet.drivers
 -------------
 
 .. automodule:: hidet.drivers
-  :members:
-  :autosummary:
+   :members:
+   :imported-members:
+   :autosummary:

docs/source/python_api/ffi/index.rst

@@ -0,0 +1,8 @@
+hidet.ffi
+---------
+
+.. automodule:: hidet.ffi
+   :members:
+   :imported-members:
+   :autosummary:
+

docs/source/python_api/index.rst

@@ -15,8 +15,10 @@ Python API
    cuda
    tensor
    data_types
+   drivers
    ops/index
    graph/index
    runtime/index
+   ffi/index
    utils/index
    testing/index

docs/source/python_api/option.rst

@@ -4,3 +4,4 @@ hidet.option
 .. automodule:: hidet.option
   :members:
   :autosummary:
+  :member-order: groupwise

gallery/developer-guides/add-torch-operator-mapping.py

@@ -0,0 +1,204 @@
+"""
+Add PyTorch Operator Mapping
+============================
+
+This guide describes how to add an operator mapping for PyTorch.
+
+.. graphviz::
+  :caption: The workflow of hidet backend of :code:`torch.compile(..., backend='hidet')`.
+
+  digraph {
+      // rankdir=LR;
+      splines=curved;
+      node [
+          shape=box, style="rounded, filled",
+          height=0.4, width=0.6, margin="0.2,0.10",
+          fillcolor="#EEF0E5",
+          color="#163020",
+          fontcolor="#163020",
+      ];
+      edge [
+          color="#163020",
+          fontcolor="#163020",
+      ];
+
+
+      graph [style="rounded, dashed"]
+          a [label="PyTorch nn.Module"];
+          b [label="torch.fx.Graph"];
+          c [label="hidet.FlowGraph"];
+          d [label="hidet.runtime.CompiledGraph"];
+
+          a -> b [label="   Step 1: PyTorch Dynamo"];
+          b -> c [label="   Step 2: Operator mapping"];
+          c -> d [label="   Step 3: FlowGraph building"];
+  }
+
+During step 2, we convert each pytorch operator to a hidet operator. In a `torch.fx.Graph`, there are three kinds of
+operators that need to be converted:
+
+1. functions (e.g., :code:`torch.nn.functional.relu`, :code:`torch.relu`, :code:`operator.add`, etc.)
+2. modules (e.g., :code:`torch.nn.ReLU`, :code:`torch.nn.Linear`, etc.)
+3. tensor methods (e.g., :code:`torch.Tensor.squeeze`, :code:`torch.Tensor.to`, etc.)
+
+In this guide, we will show how to add the operator mapping for all the three kinds of operators.
+
+1. Prepare Environment
+----------------------
+First, we remove some existing operator mapping (i.e., conversion) rules for demonstration purpose, and define an
+example model.
+"""
+import operator
+import torch
+from torch import nn
+
+# hidet employs an interpreter to convert a fx.Graph to FlowGraph
+from hidet.graph.frontend.torch.interpreter import Registry
+
+# the following three modules register the conversion rules
+import hidet.graph.frontend.torch.register_functions
+import hidet.graph.frontend.torch.register_modules
+import hidet.graph.frontend.torch.register_methods
+
+# we remove the rules for the following operators for demonstration purpose
+# we will add them back later
+del Registry.registered_functions[torch.nn.functional.relu]
+del Registry.registered_functions[operator.add]
+del Registry.registered_modules[torch.nn.Linear]
+del Registry.registered_methods[torch.Tensor.flatten]
+
+
+class Model(nn.Module):
+    """a model used nn.Linear, nn.functional.relu, operator.add and Tensor.flatten"""
+
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(10, 10)
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = torch.nn.functional.relu(x)
+        x = x + x
+        return x.flatten()
+
+
+# %%
+# 2. Compile and Run the Model
+# ----------------------------
+# If we compile and run the model, we will get an error that complains about the missing conversion rules for
+# :code:`torch.nn.Linear`, :code:`torch.nn.functional.relu` and :code:`operator.add`.
+
+
+def run_model():
+    model = Model().cuda()
+    model_opt = torch.compile(model, backend='hidet')
+
+    x = torch.randn(10, 10, device='cuda')
+    y1 = model_opt(x)
+    y2 = model(x)
+    torch.testing.assert_close(actual=y1, expected=y2)
+    print('success!')
+
+
+try:
+    run_model()
+except Exception as e:
+    print(e)
+
+# %%
+# 3. Add Operator Mappings
+# ------------------------
+#
+from typing import Optional
+from hidet import ops
+from hidet import Tensor
+from hidet.graph.frontend.torch.interpreter import (
+    register_function,
+    register_module,
+    register_method,
+    HidetModule,
+)
+
+
+# register the conversion rule for torch.nn.functional.relu
+@register_function(torch.nn.functional.relu)
+def torch_relu(x: Tensor, inplace: bool = False):  # the signature must match the original function
+    # the parameter `x` is hidet.Tensor instead of torch.Tensor
+    # we also need to return a hidet.Tensor instead of torch.Tensor
+    _ = inplace  # ignore inplace
+    return ops.relu(x)
+
+
+@register_function(operator.add)
+def operator_add(x: Tensor, y: Tensor):
+    return ops.add(x, y)
+
+
+@register_module(torch.nn.Linear)
+class HidetLinear(
+    HidetModule
+):  # HidetModule is a tool class that helps us to convert a torch.nn.Module
+    def __init__(self, torch_module: torch.nn.Module):
+        super().__init__(torch_module)
+        # inside the class, we can access the parameter of the torch module via
+        # `self.param(name: str, optional: bool = False) -> Tensor`
+        # and the returned tensor is a hidet.Tensor
+        self.transposed_weight: Tensor = ops.transpose(self.param('weight'), [1, 0])
+        self.bias: Optional[Tensor] = self.param('bias', optional=True)
+
+    def __call__(self, x: Tensor) -> Tensor:
+        # similarly, the parameter `x` is hidet.Tensor instead of torch.Tensor
+        y = ops.matmul(x, self.transposed_weight)
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+
+
+# %%
+# If we run the model again, it will complain about the missing conversion rule for :code:`torch.Tensor.flatten`.
+# It does not complain about missing conversion rule for :code:`torch.Tensor.flatten` before because we can not
+# know the type of the method's class (i.e., :code:`torch.Tensor`) before we actually run the model.
+#
+try:
+    run_model()
+except Exception as e:
+    print(e)
+
+
+# %%
+# Thus, we need to add the conversion rule for :code:`torch.Tensor.flatten` later as well.
+
+
+@register_method(torch.Tensor.flatten)
+def tensor_flatten(self: Tensor, start_dim=0, end_dim=-1):
+    return ops.flatten(self, start_dim=start_dim, end_dim=end_dim)
+
+
+run_model()
+
+# %%
+# We put all the registration code in the following three modules:
+#
+# 1. :code:`hidet.graph.frontend.torch.register_functions` (all the functions in `torch.nn.functional.*` and
+#    `operator.*`)
+# 2. :code:`hidet.graph.frontend.torch.register_modules` (all the modules in `torch.nn.*`)
+# 3. :code:`hidet.graph.frontend.torch.register_methods` (all the methods in `torch.Tensor.*`)
+#
+# Lots of operators have already been registered in the above three modules, and they are also good examples for us
+# to learn how to add operator mapping.
+#
+# Usually, we will use the existing operators in hidet (defined in `hidet.ops.*`) to implement the pytorch operators.
+# If there are no corresponding operators in hidet, we can add the missing operators to `hidet.ops.*` by following the
+# guide :doc:`/how-to-guides/add-new-operator/index`.
+#
+# .. note::
+#    The operator mapping rules are registered in the global registry. Thus, if we register the same operator mapping
+#    rules multiple times, only the last registration will take effect.
+
+# %%
+# 4. Summary
+# ----------
+# In this guide, we show how to add operator mapping for PyTorch. We first remove some existing operator mapping rules
+# for demonstration purpose, and then add them back. We also show how to add operator mapping for functions, modules
+# and tensor methods.
+#

python/hidet/__init__.py

@@ -22,6 +22,7 @@
 from . import logging
 from . import cuda
 from . import distributed
+from . import ffi
 
 from .version import __version__
 

python/hidet/ffi/__init__.py

@@ -11,6 +11,7 @@
 # limitations under the License.
 from .ffi import _LIB
 from .runtime_api import runtime_api
+from .shared_lib import SharedLibrary
 
 from . import callbacks
 from . import crt

python/hidet/ffi/shared_lib.py

@@ -22,30 +22,39 @@ class SharedLibrary:
     """
     Manage the loaded dynamic libraries.
 
-    Why we need this module?
-    ------------------------
+    **Why we need this module?**
 
     The ctypes.CDLL class does not provide a way to unload the loaded library. When a library is loaded, it will never
     be unloaded until the program exits. However, when we tune an operator, we need to generate hundreds of kernels, and
     each kernel will be compiled into a shared library. If we do not unload the shared library, we would load tens of
-    thousands of shared libraries, which will trigger memory error like:
-      "cannot apply additional memory protection after relocation"
-    (I also see other error messages).
+    thousands of shared libraries, which will trigger memory error like (I also see other error messages):
+
+    | "cannot apply additional memory protection after relocation"
+
     To solve this problem, we need to unload the shared library after we use it. Thus, whenever we need to load a shared
     library, we should use this module instead of the ctypes.CDLL class. The SharedLibrary class will keep track of the
     loaded libraries, and unload them when no one references them.
 
-    The current implementation only supports *nix systems. Will add support for Windows when we plan to support Windows
+    The current implementation only supports Linux systems. Will add support for Windows when we plan to support Windows
     in the project-level.
 
-    Usage
-    -----
+    **Usage**
+
+    .. code-block::
+
+      lib = SharedLibrary('./libhidet.so')
+      func = lib['func_name']
+      del lib
+      # at this point, the lib will not be unloaded because we still
+      # hold a reference to a function in the library.
+      del func
+      # the 'libhidet.so' will be unloaded via dlclose after the last
+      # reference to it or its functions are deleted.
 
-    >>> lib = SharedLibrary('./libhidet.so')
-    >>> func = lib['func_name']
-    >>> del func
-    >>> del lib
-    >>> # the 'libhidet.so' will be unloaded via dlclose after the last reference to it is deleted.
+    Parameters
+    ----------
+    lib_path: str
+        The path to the shared library (e.g., './lib.so')
     """
 
     loaded_cdll_libraries: Dict[str, ctypes.CDLL] = {}

python/hidet/graph/frontend/torch/dynamo_backends.py

@@ -30,7 +30,7 @@
 logger = logging.getLogger(__name__)
 
 
-def get_flow_graph(interpreter, example_inputs) -> FlowGraph:
+def get_flow_graph(interpreter, example_inputs):
     # prepare dummy and symbolic inputs for correctness and flow graph construction
     inputs: List[Union[Tensor, SymbolVar, int, bool, float]] = []  # for flow graph construction
     for example_input in example_inputs:
@@ -63,7 +63,7 @@ def get_flow_graph(interpreter, example_inputs) -> FlowGraph:
     output_format, output_tensors = serialize_output(output)
     input_tensors = [x for x in inputs if isinstance(x, hidet.Tensor)]
 
-    return (hidet.trace_from(output_tensors, inputs=input_tensors), inputs, output_format)
+    return hidet.trace_from(output_tensors, inputs=input_tensors), inputs, output_format
 
 
 def get_compiled_graph(flow_graph: FlowGraph):