writing doc

docs/_static/css/style.css

@@ -1,6 +1,155 @@
+/* css style file taken from https://github.com/thu-ml/tianshou/blob/master/docs/_static/css/style.css */
+
 .wy-side-nav-search {
     background-color: #fff;
 }
 a.icon.icon-home, a.icon.icon-home::before {
     color: #2980b9;
+}
+
+pre {
+    white-space: pre-wrap !important;       /* Since CSS 2.1 */
+    white-space: -moz-pre-wrap !important;  /* Mozilla, since 1999 */
+    white-space: -pre-wrap !important;      /* Opera 4-6 */
+    white-space: -o-pre-wrap !important;    /* Opera 7 */
+    word-wrap: break-word !important;       /* Internet Explorer 5.5+ */
+}
+
+body {
+    font-family: "Lato","proxima-nova","Helvetica Neue",Arial,sans-serif;
+}
+
+/* Default header fonts are ugly */
+h1, h2, .rst-content .toctree-wrapper p.caption, h3, h4, h5, h6, legend, p.caption {
+    font-family: "Lato","proxima-nova","Helvetica Neue",Arial,sans-serif;
+}
+
+/* Use white for docs background */
+.wy-side-nav-search {
+    background-color: #fff;
+}
+
+.wy-nav-content {
+    max-width: 1200px !important;
+}
+
+.wy-nav-content-wrap, .wy-menu li.current > a  {
+    background-color: #fff;
+}
+
+.wy-side-nav-search>a img.logo {
+    width: 80%;
+    margin-top: 10px;
+}
+
+@media screen and (min-width: 1400px) {
+    .wy-nav-content-wrap {
+        background-color: #fff;
+    }
+
+    .wy-nav-content {
+        background-color: #fff;
+    }
+}
+
+/* Fixes for mobile */
+.wy-nav-top {
+    background-color: #fff;
+    background-repeat: no-repeat;
+    background-position: center;
+    padding: 0;
+    margin: 0.4045em 0.809em;
+    color: #333;
+}
+
+.wy-nav-top > a {
+    display: none;
+}
+
+@media screen and (max-width: 768px) {
+    .wy-side-nav-search>a img.logo {
+        height: 60px;
+    }
+}
+
+/* This is needed to ensure that logo above search scales properly */
+.wy-side-nav-search a {
+    display: block;
+}
+
+/* This ensures that multiple constructors will remain in separate lines. */
+.rst-content dl:not(.docutils) dt {
+    display: table;
+}
+
+/* Use our red for literals (it's very similar to the original color) */
+.rst-content tt.literal, .rst-content tt.literal, .rst-content code.literal {
+    color: #4692BC;
+}
+
+.rst-content tt.xref, a .rst-content tt, .rst-content tt.xref,
+.rst-content code.xref, a .rst-content tt, a .rst-content code {
+    color: #404040;
+}
+
+/* Change link colors (except for the menu) */
+
+a {
+    color: #4692BC;
+}
+
+a:hover {
+    color: #4692BC;
+}
+
+
+a:visited {
+    color: #4692BC;
+}
+
+.wy-menu a {
+    color: #b3b3b3;
+}
+
+.wy-menu a:hover {
+    color: #b3b3b3;
+}
+
+/* Default footer text is quite big */
+footer {
+    font-size: 80%;
+}
+
+footer .rst-footer-buttons {
+    font-size: 125%; /* revert footer settings - 1/80% = 125% */
+}
+
+footer p {
+    font-size: 100%;
+}
+
+.ethical-rtd {
+    display: none;
+}
+
+.ethical-fixedfooter {
+    display: none;
+}
+
+.ethical-content {
+    display: none;
+}
+
+/* For hidden headers that appear in TOC tree */
+/* see http://stackoverflow.com/a/32363545/3343043 */
+.rst-content .hidden-section {
+    display: none;
+}
+
+nav .hidden-section {
+    display: inherit;
+}
+
+.wy-side-nav-search>div.version {
+    color: #000;
 }

docs/walk_through.rst

@@ -27,6 +27,7 @@ In this tutorial, we will learn how does PyTorch compiler work for the following
 The code tries to implement a strange sigmoid function :math:`\frac{1}{e^{-x} + 5}`, and scales the output according to its activation value, then multiplies the output with another tensor ``y``.
 
 The tutorial intends to cover the following aspects of PyTorch compiler:
+
 - Dynamo (graph capture)
 - AOTAutograd (forward graph and backward graph)
 - Inductor (compile graph to kernel)
@@ -69,4 +70,84 @@ We write this funny function ``f``, that contains a module call, and a ``torch.l
         elif isinstance(mod, B):
             return -x
 
-That's it! You can call it your fist Just-In-Time compiler, although it is `compiled` by your brain rather than an automated program.
+That's it! We can call it our fist Just-In-Time compiler, although it is `compiled` by our brain rather than an automated program.
+
+The basic workflow of a Just-In-Time compiler is: right before the function is executed, it analyzes if the execution can be optimized, and what is the condition under which the function execution can be optimized. Hopefully, the condition is general enough for new inputs, so that the benfit outweights the cost of Just-In-Time compilation.
+
+This leads to two basic concepts in Just-In-Time compilers: guards, and compiled code. Guards are conditions when the functions can be optimized, and compiled code is the optimized version of functions. In the above simple Just-In-Time compiler example, ``isinstance(mod, A)`` is a guard, and ``return 2 * x`` is the corresponding compiled code that is equivalent to the original code under the guarding condition, but is significantly faster.
+
+And if we want to be rigorous, our Just-In-Time example should be updated as follows:
+
+.. code-block:: python
+
+    def f(x, mod):
+        if isinstance(x, torch.Tensor) and isinstance(mod, A):
+            return 2 * x
+        elif isinstance(x, torch.Tensor) and isinstance(mod, B):
+            return -x
+        else:
+            y = mod(x)
+            z = torch.log(y)
+            return z
+
+We have to check each parameter so that our guards are sound, and also fallback to the original code if we fail to optimize the code.
+
+Going more rigorous, the above example is actually an Ahead-of-time compiler: we inspect all the available source code, and before running any function, we write the optimized function in terms of guards and compiled code. A real Just-In-Time procedure should be:
+
+.. code-block:: python
+
+    def f(x, mod):
+        for guard, compiled_code in f.compiled_entries:
+            if guard(x, mod):
+                return compiled_code(x, mod)
+        try:
+            guard, compiled_code = compile_and_optimize(x, mod)
+            f.compiled_entries.append([guard, compiled_code])
+            return compiled_code(x, mod)
+        except FailToCompileError:
+            y = mod(x)
+            z = torch.log(y)
+            return z
+
+A Just-In-Time compiler just optimizes for what it has seen. Everytime it sees a new input that does not satisfy any guarding condition, it compiles a new guard and compiled code for the new input.
+
+Let's explain it step-by-step:
+
+.. code-block:: python
+
+    import torch
+
+    class A(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x):
+            return torch.exp(2 * x)
+
+    class B(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x):
+            return torch.exp(-x)
+
+    @just_in_time_compile
+    def f(x, mod):
+        y = mod(x)
+        z = torch.log(y)
+        return z
+
+    a = A()
+    b = B()
+    x = torch.randn((5, 5, 5))
+    # before executing f(x, a), f.compiled_entries == []
+    # after executing f(x, a), f.compiled_entries == [Guard("isinstance(x, torch.Tensor) and isinstance(mod, A)"), CompiledCode("return 2 * x")]
+    f(x, a)
+    # the second call of f(x, a) hit a condition, so we can just execute the compiled code
+    f(x, a)
+    # f(x, b) will trigger compilation and add a new compiled entry
+    # before executing f(x, b), f.compiled_entries == [Guard("isinstance(x, torch.Tensor) and isinstance(mod, A)"), CompiledCode("return 2 * x")]
+    # after executing f(x, b), f.compiled_entries == [Guard("isinstance(x, torch.Tensor) and isinstance(mod, A)"), CompiledCode("return 2 * x"), Guard("isinstance(x, torch.Tensor) and isinstance(mod, B)"), CompiledCode("return -x")]
+    f(x, b)
+    # the second call of f(x, b) hit a condition, so we can just execute the compiled code
+    f(x, b)