[dynamo, docs] cherry pick torch.compile programming model docs into 2.8 (pytorch#159373)

* [dynamo, docs] cherry pick torch.compile programming model docs into 2.8

* revert requirements-docs.txt

* add remaining docs, update conf.py with myst_nb

Author: williamwen42

docs/source/compile/_static/dynamo_summary_diagram.png

@@ -0,0 +1,15 @@
+import functools
+import os
+
+import torch
+
+
+# to lower notebook execution time while hiding backend="eager"
+torch.compile = functools.partial(torch.compile, backend="eager")
+
+# to clear torch logs format
+os.environ["TORCH_LOGS_FORMAT"] = ""
+torch._logging._internal.DEFAULT_FORMATTER = (
+    torch._logging._internal._default_formatter()
+)
+torch._logging._internal._init_logs()

docs/source/compile/header_code.py

@@ -0,0 +1,142 @@
+---
+file_format: mystnb
+kernelspec:
+  name: python3
+mystnb:
+  execution_timeout: 30
+  execution_show_tb: True
+  merge_streams: True
+---
+
+```{code-cell}
+:tags: [remove-cell]
+import torch
+
+import header_code
+
+torch._logging.set_logs(graph_breaks=True)
+```
+
+# Common Graph Breaks
+
+Below are some common graph breaks and some workarounds.
+
+## Incorrect Code
+Your code might contain errors (meaning it doesn't execute even without `torch.compile`). In the example below, there's a typo in the `torch.sin` call due to an extra argument. **Always disable `torch.compile` to check if the code runs correctly.**
+
+
+```{code-cell}
+@torch.compile
+def fn(x):
+    y = torch.sin(x, x)
+    return y
+
+try:
+    fn(torch.ones(3, 3))
+except Exception as e:
+    pass
+```
+
+Dynamo makes a best-effort attempt to hint if a graph break is caused by your code.
+But it can still sometimes be difficult to tell from the logs if the graph break is caused by an error in your code,
+is a more complicated graph break, or is a `torch.compile` bug. In order to differentiate, we recommend trying to run your code without `torch.compile` to see if you still get the error reported by the graph break.
+
+## Data-dependent operations
+
+`torch.compile` graph breaks on data-dependent operations such as data-dependent control flow (if-statements, loops with tensors) and direct tensor data accesses (`.item`, `.data_ptr`).
+
+```{code-cell}
+@torch.compile
+def fn(x):
+    y = x.sum()
+    if y > 0:
+        return x + y.item()
+    return x - y.item()
+
+print(fn(torch.ones(3, 3)))
+```
+
+The general workaround for these graph breaks is to avoid doing data-dependent operations. Some specific workarounds are:
+
+- If your control flow doesn't actually depend on data values, consider modifying your code to perform control flow on constants.
+
+
+```{code-cell}
+# old
+x = torch.randn(3, 3)
+@torch.compile
+def fn(y):
+    if x.sum() > 0:
+        return y + x
+    else:
+        return y - x
+
+print(fn(torch.ones(3, 3)))
+```
+
+```{code-cell}
+# new
+x = torch.randn(3, 3)
+cond = (x.sum() > 0).item()
+@torch.compile
+def fn(y):
+    if cond:
+        return y + x
+    else:
+        return y - x
+
+print(fn(torch.ones(3, 3)))
+```
+
+- Use higher-order ops like {ref}`cond` in place of data-dependent control flow
+
+
+```{code-cell}
+# old
+@torch.compile
+def fn(x):
+    if x.sum() > 0:
+        return x + 1
+    return x - 1
+
+print(fn(torch.ones(3, 3)))
+```
+
+```{code-cell}
+# new
+@torch.compile
+def fn(x):
+    return torch.cond(
+        x.sum() > 0,
+        lambda x: x + 1,
+        lambda x: x - 1,
+        (x,),
+    )
+
+print(fn(torch.ones(3, 3)))
+```
+
+- If you have a `.item()` call, try `torch._dynamo.config.capture_scalar_outputs = True`
+or `TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1`.
+- Wrap problematic parts of the function in a custom operator
+
+## Printing and logging
+
+Printing/logging/issuing warnings will result in a graph break.
+You can try working around this by using `torch._dynamo.config.reorderable_logging_functions`.
+This config is used to reorder logging functions so that they are called at the end of the
+traced function, thus avoiding a graph break.
+However, the logged contents may differ if, for example, a mutation occurs.
+
+
+```{code-cell}
+torch._dynamo.config.reorderable_logging_functions.add(print)
+
+@torch.compile
+def fn(x):
+    x += 1
+    print("log!")
+    return torch.sin(x)
+
+print(fn(torch.ones(3, 3)))
+```

docs/source/compile/programming_model.common_graph_breaks.md

@@ -0,0 +1,75 @@
+---
+file_format: mystnb
+kernelspec:
+  name: python3
+mystnb:
+  execution_timeout: 30
+  execution_show_tb: True
+  merge_streams: True
+---
+
+```{code-cell}
+:tags: [remove-cell]
+import torch
+
+import header_code
+
+torch._logging.set_logs(graph_breaks=True, graph_code=True)
+```
+
+# Disabling and Suppressing Errors
+For some model architectures, there are portions of the model which are particularly difficult to compile -
+either there are many graph breaks, or there are crashes.
+You may want to explicitly disable these portions of the model which are problematic so that you can apply
+`torch.compile` to the parts that work. You can do this by using the `@torch.compiler.disable` decorator.
+When `torch.compile` attempts to call a disabled function, it breaks the graph and skips tracing the disabled function,
+resuming tracing after the call. By default, all recursive calls made from a disabled function are also disabled.
+Use the `recursive=False` option to allow compilation for recursive calls.
+
+```{code-cell}
+def inner1(x):
+    torch._dynamo.graph_break()  # not traced
+    return x + 1  # not traced
+
+@torch.compiler.disable
+def outer1(x):
+    x = x + 2  # not traced
+    torch._dynamo.graph_break()  # not traced
+    return inner1(x)
+
+@torch.compile
+def f(x):
+    x = outer1(x)
+    return x + 4  # traced
+
+print(f(torch.ones(3)))
+```
+
+```{code-cell}
+def inner2(x):
+    torch._dynamo.graph_break()  # traced
+    return x + 1  # traced
+
+@torch.compiler.disable(recursive=False)
+def outer2(x):
+    x = x + 2  # not traced
+    torch._dynamo.graph_break()  # not traced
+    return inner2(x)
+
+@torch.compile
+def g(x):
+    x = outer2(x)
+    return x + 4  # traced
+
+print(g(torch.ones(3)))
+```
+
+For example, one can use `torch.compiler.disable` to disable `torch.compile` on sparse architecture in
+recommendation models, as the sparse arch is difficult to compile.
+Preprocessing and logging functions are other examples of functions that typically cause
+a lot of graph breaks and do not get value from being compiled.
+
+If you are experiencing compiler crashes and you want to continue regardless,
+you can set `torch._dynamo.config.suppress_errors = True`.
+When the compiler crashes, we will just skip tracing the function and try again later.
+**This is not best practice** - it is better to eventually manually add `disable` annotations as necessary.

docs/source/compile/programming_model.compiler_disable.md

@@ -0,0 +1,12 @@
+# Custom Operators
+
+**Summary:**
+- Use custom operators to have `torch.compile` treat a function as opaque. `torch.compile` will never trace into the function and Inductor (the backend) will run the function as-is.
+
+You may wish to use a custom operator in any of the following situations:
+- Your code calls some C/C++/CUDA code. Dynamo is a Python bytecode interpreter and generally does not know how to handle calls to C/C++/CUDA functions that are bound to Python.
+- Dynamo and non-strict tracing have trouble tracing through a function and you want it to be ignored by `torch.compile`.
+
+Please see [the Python custom ops tutorial](https://pytorch.org/tutorials/advanced/python_custom_ops.html#python-custom-ops-tutorial)for more details on how to wrap a Python function into a `torch.compile`-understood custom operator.
+
+For more advanced use cases, you may wish to use our C++ Custom Operator API; please see [here](https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html) for more information.