Add basic Export docs

Summary: Title

Differential Revision: D47808548

fbshipit-source-id: ea620536fb23e90c69d0080316924ca2260b55b9

Author: tugsbayasgalan

README.md

@@ -40,6 +40,7 @@ Compared to the legacy Lite Interpreter, there are some major benefits:
 - [Setting up ExecuTorch from GitHub](/docs/website/docs/tutorials/setting_up_executorch.md)
 - [Exporting to Executorch](/docs/website/docs/tutorials/exporting_to_executorch.md)
     - [EXIR Spec](/docs/website/docs/ir_spec/00_exir.md)
+    - [Exporting manual](/docs/website/docs/export/00_export_manual.md)
     - [Delegate to a backend](/docs/website/docs/tutorials/backend_delegate.md)
 - [Executorch Google Colab](https://colab.research.google.com/drive/1oJBt3fj_Tr3FE7L9RdUgSKK9XzJfUv4F#scrollTo=fC4CB3kFhHPJ)
 

docs/website/docs/export/00_export_manual.md

@@ -0,0 +1,27 @@
+PT2.0 Export Manual
+
+# Context
+At a high level, the goal of PT2 Export is to enable the execution of entire PyTorch programs by other means than the “eager” PyTorch runtime, with a representation that is amenable to meet the optimization and targeting goals of specialized use cases. Specifically, we want to enable users to convert their PyTorch models to a standardized IR, decoupled from its execution, that various domain-specific runtimes can transform and execute independently. This conversion is powered by Dynamo’s technique for sound whole-graph capture—capturing a graph without any “breaks” that would require the eager runtime to fall back to Python. The rest of this wiki documents a snapshot of PT2 Export as of early May 2023—what we consider an "MVP" release. Please note that this project is under active and heavy development: while this snapshot should give a fairly accurate picture of the final state, some some details might change in the coming weeks / months based on feedback. If you have any issues, please file an issue on Github and tag "export".
+
+
+Modules and Entrypoints
+Constraints API
+Control Flow Operators
+Custom Operators
+Compiler Passes on Exported Artifact
+Non-strict Mode
+# Documentation
+- [Overview](./overview.md)
+  - [Background](./background.md)
+  - [Overall Workflow](./overall_workflow.md)
+  - [Soundness](./soundness.md)
+  - [Errors](./errors.md)
+- [Export API Reference](./export_api_reference.md)
+  - [Modules and Entrypoints](./modules_and_entrypoints.md)
+  - [Constraints API](./constraint_apis.md)
+  - [Control Flow Operators](../ir_spec/control_flow.md)
+  - [Custom Operators](./custom_operators.md)
+- [Exported Programs](../ir_spec/00_exir.md#exportedprogram)
+- [ExportDB](./export_db.md)
+
+ir_spec/control_flow

docs/website/docs/export/background.md

@@ -0,0 +1,28 @@
+<h1> Background </h1>
+
+# Setup
+Let's say you have a function that you want to export. You export it by passing example inputs (tensors) as arguments to `torch._export.export`. The exported program can then be called with other inputs.
+```python
+from torch._export import export
+
+def foo(x):  # expect x to be a tensor
+  ...
+
+t = torch.rand(4, 8, 32)  # example input tensor
+exported_foo = export(foo, t)
+
+# expect that exported_foo can now be called with other input tensor
+```
+More generally, the function to be exported can take multiple inputs, and the function itself could be a torch.nn.Module (with a forward method). See [Export API Reference] (./export_api_reference.md).
+
+# Graph Breaks
+The PT2 compiler is a "tracing" compiler, which means that it compiles the execution path—or "trace"—of your function on your example inputs. The intermediate representation of such a trace is a graph. In eager mode it is usual to have graph breaks, where the compiler can fail to trace some parts of the code; this is fine because it can always fall back to the Python interpreter to fill these gaps. However in export mode we do not want any graph breaks: we want the compiler to capture the entire execution in a single graph.
+
+## Rewriting Code
+Graph breaks can arise either because of missing support for Python features, or because the compiler cannot decide which control flow path to continue tracing on. In most cases, it is possible to rewrite code to avoid such graph breaks and complete the export.When the compiler's tracing mechanism does not support some Python feature, we strive to provide a workaround as part of the error message. Over time, we expect to fill in such gaps. On the other hand, not being able to decide which control flow path to continue tracing on is a necessary limitation of the compiler. You are required to use special operators to unblock such cases. See [Control Flow Operators](../ir_spec/control_flow.md).
+
+# Shapes
+Recall that while we need example inputs for export, we must generalize the exported program to be callable with other inputs. The main mechanism for this generalization is through reuse of shapes (of tensors). Next, let us dive deeper into shapes.
+
+## Static and Dynamic Dimensions
+The shape of a tensor is a tuple of dimensions. Roughly speaking, the exported program can be called with inputs of the same shape as the example inputs. By default, we assume that dimensions are static: the compiler assumes they are going to be the same, and specializes the exported program to those dimensions.However, some dimensions, such as a batch dimension, are expected to not be the same—the example inputs passed to export may have a different batch size as inputs to the exported program. Such dimensions must be marked dynamic. See [Soundness](./soundness.md) to learn how to specify properties of such dimensions.

docs/website/docs/export/constraint_apis.md

@@ -0,0 +1,73 @@
+<h1> Constraint APIs </h1>
+
+To enable the export of input shape-dependent models, it is necessary for users to impose constraints on the tracing inputs, ensuring the safe traversal of the model during the export process. For a comprehensive understanding of how to utilize this feature, please refer to [this overview](./soundness.md). To express these constraints, we have developed the following set of APIs.
+
+# dynamic_dim
+```python
+def dynamic_dim(x: torch.Tensor, dim: int):
+    """
+    Marks the dimension `dim` of input `x` as unbounded dynamic
+    """
+    pass
+```
+
+It is possible to impose specific bounds on the marked dynamic dimension. For example:
+```python
+constraints = [dynamic_dim(x, 0) >= 3, dynamic_dim(x, 0) <= 6]
+```
+
+By passing the above `constraints` to export, we effectively establish the range of acceptable values for the 0th dimension of tensor x, constraining it between 3 and 6. Consequently, the PT2 Export functionality can safely trace through the following program:
+```python
+def f(x):
+    if x.shape[0] > 3:
+       return x.cos()
+    return x.sin()
+```
+
+Moreover, it is possible to impose specific equalities between marked dynamic dimensions. For example:
+```python
+constraints = [dynamic_dim(x, 0) == dynamic_dim(y, 0)]
+```
+
+This means that whatever the 0th dimensions of tensors x and y may be, they must be the same. This is useful to export the following program, which implicitly requires this condition because of the semantics of `torch.add`
+
+```python
+def f(x, y):
+    return x + y
+```
+
+# constrain_as_value
+```python
+def constrain_as_value(symbol, min: Optional[int], max: Optional[int]):
+    """
+    Adds a minimum and/or maximum constraint on the intermediate symbol during the tracing phase.
+    """
+```
+
+The `constrain_as_value` function informs us that the specified symbol is guaranteed to fall within the provided minimum and maximum values. If no minimum or maximum values are provided, the symbol is assumed to be unbounded. Here's a concrete example of its usage within a model:
+```python
+def f(x, y):
+    b = y.item()
+    constrain_as_value(b, 3, 5)
+    if b > 3:
+       return x.cos()
+    return x.sin()
+```
+
+# constrain_as_size
+```python
+def constrain_as_size(symbol, min: Optional[int] = 2, max: Optional[int]):
+    """
+    Adds a minimum and/or maximum constraint on the intermediate symbol during the tracing phase,
+    with additional checks to ensure the constrained value can be used to construct a tensor shape.
+    """
+```
+
+The `constrain_as_size` API is similar to constrain_as_value but includes additional verifications to ensure that the constrained value can be used to construct a tensor shape. For instance, our tracer specializes in handling shape sizes of 0 or 1, so this API explicitly raises an error if the constrain_as_size is used with a minimum value less than 2. Here's an example of its usage:
+```python
+def f(x, y):
+    b = y.item()
+    constrain_as_size(b, 3, 5)
+    z = torch.ones(b, 4)
+    return x.sum() + z.sum()
+```

docs/website/docs/export/custom_operators.md

@@ -0,0 +1,7 @@
+<h1> Custom Operators </h1>
+
+To ensure a successful export of your model, it is necessary to provide a META implementation for any custom operators used. Custom operators refer to operators that are not part of the "aten" or "prim" namespaces. You have the flexibility to implement the META functionality in either Python or C++.
+
+Note that the official API for registering custom meta kernels is currently undergoing intensive development. While the final API is being refined, you can refer to the documentation [here](https://docs.google.com/document/d/1GgvOe7C8_NVOMLOCwDaYV1mXXyHMXY7ExoewHqooxrs/edit#heading=h.64r4npvq0w0). In this document, you can find detailed instructions on how to write a Python meta function by searching for the "Out-of-tree" section within the "How to write a Python meta function" section.
+
+By following the guidelines outlined in the documentation, you can ensure that your custom operators are properly registered and integrated into the export process. We recommend staying updated with our latest announcements and documentation for any updates or improvements to the official API.

docs/website/docs/export/errors.md

@@ -0,0 +1,46 @@
+<h1> Errors </h1>
+
+In this section we discuss errors that can commonly arise during export.
+
+# Expected Graph Breaks
+## Unsupported Features
+In PT2 Export, we are primarily reusing the same tracing mechanism—Dynamo—that we use in eager mode. Recall that in eager mode, graph breaks are expected—we always have a fallback option. A consequence of this design is that Dynamo has incomplete coverage of PyTorch and Python features. (That said, the fewer graph breaks there are, generally speaking, the better performance we can expect—even in eager mode—because it enables optimizations to apply over larger regions of code. Thus we are actively working on filling in coverage gaps to avoid graph breaks where possible.)Unfortunately, this means that you may encounter graph breaks during export due to Dynamo coverage gaps. In such cases, you should expect to get an error that includes a link to [ExportDB](./export_db.md). The corresponding entry should show a minimal negative example (failure) and a minimal positive example (success) that should help you understand the limitation and the workaround, i.e., how to fix the error by rewriting code.
+
+# Constraint Violations
+Recall that you can specify constraints on dynamic dimensions, which encode the soundness conditions for export. It is possible that these constraints are not valid.## Various Cases
+Specifically, the compiler may find that:
+- A dynamic dimension must be equal to a constant.
+  - In this case, this dimension must be static: you cannot mark it dynamic.
+- A dynamic dimension must be in a range that does not follow the specified range, i.e., is not entirely included between the specified lower and upper bounds.
+  - In this case, you need to adjust the specified bounds.
+  - Note that when bounds are not specified, they are implicitly assumed to be [2, infinity).
+  - For technical reasons that are difficult to explain here, they are assumed to be not 0 or 1. This is not a bug, and does not necessarily mean that your exported program will not work for dimensions 0 or 1. It does mean, though, that you should test for these cases.
+- A dynamic dimension must be equal to another dynamic dimension that it is not specified equal to.
+   - In this case, you need to add the missing equality.
+   - By default, all dynamic dimensions are assumed to be independent.
+   - For legacy reasons that are difficult to explain here, you might find spurious implicitly assumed equalities when dimensions in your example inputs happen to be equal. If you ever encounter such a case, please report it as a bug.
+
+## Using the Compiler as a Guide
+See [this overview](./soundness.md/#Constraint Violations and How to Fix Them) of how to fix such errors. Briefly:
+* You should see generated functions specializations and specify_constraints on the console that respectively summarize which dimensions are assumed static and what the necessary constraints on the remaining dynamic dimensions are.
+* If you agree with this information, you can copy-paste and call specify_constraints with your example inputs to specify constraints, and you can copy-paste and call specializations on your example inputs to assert their constant values.
+* If you do not agree and would like to provide tighter constraints, feel free to modify specify_constraints; the compiler will be happy to accept.
+* If you do not agree and would like looser constraints, please use TORCH_LOGs=dynamic to enable INFO-level dynamic-shape logging, which will guide you to where the inferred constraints come from. You can also try TORCH_LOGs=+dynamic to enable (further, verbose) DEBUG-level logging.
+   * Note that you might have to change your code or your expectations based on this information. If you are absolutely convinced that the compiler has a bug, please report it! For example, there are tricky cases where the constraints may come from non-user code, like a fast path in the compiler itself. We encourage you to try different example inputs to avoid such constraints.
+
+# Missing META Kernels for Operators
+## ATen Operators
+In the unfortunate case where your model uses an ATen operator that is not supported yet, you may get an obscure error of the form:
+```python
+Unable to find op(FakeTensor, FakeTensor, ...)
+```
+Please report a bug if you encounter this error
+
+## Custom Operators
+In this case you should follow the instructions at [Custom Operators](./custom_operators.md). Note that the current mechanism is not ideal, but will be updated soon to make it easy for you to register custom operators.
+
+# Validation Errors
+Note that we do not do any validation of the exported program yet; this is planned for the near future.I n these cases you should report a bug since the issue is likely in PyTorch.
+## Correctness
+The export workflow should complain when the exported program behaves differently than the eager program by running the example inputs through both. ## Serialization roundtrip failure
+The export workflow should serialize and deserialize the exported program, and then run the correctness test again.

docs/website/docs/export/export_api_reference.md

@@ -0,0 +1,6 @@
+<h1> Export API Reference </h1>
+
+- [Modules and Entrypoints] (./modules_and_entrypoints.md)
+- [Constraints API] (./constraint_apis.md)
+- [Control Flow Operators] (../ir_spec/control_flow.md)
+- [Custom Operators] (./custom_operators.md)

docs/website/docs/export/exportdb.md

@@ -0,0 +1,5 @@
+<h1> ExportDB </h1>
+
+ExportDB is a [centralized] (https://pytorch.org/docs/main/generated/exportdb/index.html) dataset for recording working and non working export usage examples.
+
+We try to cover a lot of corner cases from PyTorch and Python so that users can have a better understanding of the capabilities and caveats of export, by browsing the example database. We also expect that errors that arise during export will have actionable messages that point to examples in this database.

docs/website/docs/export/modules_and_entrypoints.md

@@ -0,0 +1,84 @@
+<h1> Modules and Entrypoints </h1>
+
+<b> Disclaimer: Please note that at present, we do not offer any backward compatibility guarantees for the following APIs. While we are committed to minimizing significant API changes, it is important to understand that we are currently in an intensive development phase, and as such, we reserve the right to modify implementation details and top-level API parameters.  We are constantly striving to enhance our offerings and deliver the best possible experience to our users. However, during this phase, it is essential to remain aware that certain adjustments may be necessary to improve functionality, stability, or meet evolving requirements. </b>
+
+# Export API
+
+At the top level, the export API is defined as follows:
+
+```python
+def export(
+    m: Union[torch.nn.Module, Callable[..., Any]],
+    args: Union[Dict[str, Tuple[Value, ...]], Tuple[Value, ...]],
+    constraints: Optional[List[Constraint]] = None,
+) -> ExportedProgram:
+    """
+    Traces either an nn.Module's forward function or just a callable with PyTorch
+    operations inside and produce a ExportedProgram.
+
+    Args:
+        m: the `nn.Module` or callable to trace.
+
+        args: Tracing example inputs.
+
+        constraints: A list of constraints on the dynamic arguments specifying
+            their possible range of their shapes
+
+    Returns:
+        An ExportedProgram containing the traced method.
+    """
+```
+
+# Exported Artifact
+The export call returns a custom export artifact called ExportedProgram:
+
+```python
+class ExportedProgram:
+    graph_module: torch.fx.GraphModule
+    graph_signature: ExportGraphSignature
+    call_spec: CallSpec
+    state_dict: Dict[str, Any]
+    symbol_to_range: Dict[sympy.Symbol, Tuple[int, int]]
+
+    @property
+    def graph(self):
+        return self.graph_module.graph
+
+    def transform(self, *passes: PassType) -> "ExportedProgram":
+        # Runs graph based transformations on the given ExportedProgram
+        # and returns a new transformed ExportedProgram
+        ...
+
+    def add_runtime_assertions(self) -> "ExportedProgram":
+        # Adds runtime assertions based on the constraints
+        ...
+
+# Information to maintain user calling/returning specs
+@dataclasses.dataclass
+class CallSpec:
+    in_spec: Optional[pytree.TreeSpec] = None
+    out_spec: Optional[pytree.TreeSpec] = None
+
+
+# Extra information for joint graphs
+@dataclasses.dataclass
+class ExportBackwardSignature:
+    gradients_to_parameters: Dict[str, str]
+    gradients_to_user_inputs: Dict[str, str]
+    loss_output: str
+
+
+@dataclasses.dataclass
+class ExportGraphSignature:
+    parameters: List[str]
+    buffers: List[str]
+
+    user_inputs: List[str]
+    user_outputs: List[str]
+    inputs_to_parameters: Dict[str, str]
+    inputs_to_buffers: Dict[str, str]
+
+    buffers_to_mutate: Dict[str, str]
+
+    backward_signature: Optional[ExportBackwardSignature]
+```