Update aoti tutorial (#3224)

Co-authored-by: Svetlana Karslioglu <svekars@meta.com>

Author: angelayi

recipes_source/torch_export_aoti_python.py

@@ -3,7 +3,7 @@
 """
 .. meta::
    :description: An end-to-end example of how to use AOTInductor for Python runtime.
-   :keywords: torch.export, AOTInductor, torch._inductor.aot_compile, torch._export.aot_load
+   :keywords: torch.export, AOTInductor, torch._inductor.aoti_compile_and_package, aot_compile, torch._export.aoti_load_package
 
 ``torch.export`` AOTInductor Tutorial for Python runtime (Beta)
 ===============================================================
@@ -14,19 +14,18 @@
 #
 # .. warning::
 #
-#     ``torch._inductor.aot_compile`` and ``torch._export.aot_load`` are in Beta status and are subject to backwards compatibility
-#     breaking changes. This tutorial provides an example of how to use these APIs for model deployment using Python runtime.
+#     ``torch._inductor.aoti_compile_and_package`` and
+#     ``torch._inductor.aoti_load_package`` are in Beta status and are subject
+#     to backwards compatibility breaking changes. This tutorial provides an
+#     example of how to use these APIs for model deployment using Python
+#     runtime.
 #
-# It has been shown `previously <https://pytorch.org/docs/stable/torch.compiler_aot_inductor.html#>`__ how AOTInductor can be used 
-# to do Ahead-of-Time compilation of PyTorch exported models by creating
-# a shared library that can be run in a non-Python environment.
-#
-#
-# In this tutorial, you will learn an end-to-end example of how to use AOTInductor for Python runtime.
-# We will look at how  to use :func:`torch._inductor.aot_compile` along with :func:`torch.export.export` to generate a 
-# shared library. Additionally, we will examine how to execute the shared library in Python runtime using :func:`torch._export.aot_load`.
-# You will learn about the speed up seen in the first inference time using AOTInductor, especially when using 
-# ``max-autotune`` mode which can take some time to execute.
+# It has been shown `previously
+# <https://pytorch.org/docs/stable/torch.compiler_aot_inductor.html#>`__ how
+# AOTInductor can be used to do Ahead-of-Time compilation of PyTorch exported
+# models by creating an artifact that can be run in a non-Python environment.
+# In this tutorial, you will learn an end-to-end example of how to use
+# AOTInductor for Python runtime.
 #
 # **Contents**
 #
@@ -36,115 +35,169 @@
 ######################################################################
 # Prerequisites
 # -------------
-# * PyTorch 2.4 or later
+# * PyTorch 2.6 or later
 # * Basic understanding of ``torch.export`` and AOTInductor
 # * Complete the `AOTInductor: Ahead-Of-Time Compilation for Torch.Export-ed Models <https://pytorch.org/docs/stable/torch.compiler_aot_inductor.html#>`_ tutorial
 
 ######################################################################
 # What you will learn
 # ----------------------
-# * How to use AOTInductor for python runtime.
-# * How  to use :func:`torch._inductor.aot_compile` along with :func:`torch.export.export` to generate a shared library
-# * How to run a shared library in Python runtime using :func:`torch._export.aot_load`.
-# * When do you use AOTInductor for python runtime
+# * How to use AOTInductor for Python runtime.
+# * How to use :func:`torch._inductor.aoti_compile_and_package` along with :func:`torch.export.export` to generate a compiled artifact
+# * How to load and run the artifact in a Python runtime using :func:`torch._export.aot_load`.
+# * When to you use AOTInductor with a Python runtime
 
 ######################################################################
 # Model Compilation
 # -----------------
 #
-# We will use the TorchVision pretrained `ResNet18` model and TorchInductor on the 
-# exported PyTorch program using :func:`torch._inductor.aot_compile`.
+# We will use the TorchVision pretrained ``ResNet18`` model as an example.
 #
-# .. note::
+# The first step is to export the model to a graph representation using
+# :func:`torch.export.export`. To learn more about using this function, you can
+# check out the `docs <https://pytorch.org/docs/main/export.html>`_ or the
+# `tutorial <https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html>`_.
 #
-#       This API also supports :func:`torch.compile` options like ``mode``
-#       This means that if used on a CUDA enabled device, you can, for example, set ``"max_autotune": True``
-#       which leverages Triton based matrix multiplications & convolutions, and enables CUDA graphs by default.
+# Once we have exported the PyTorch model and obtained an ``ExportedProgram``,
+# we can apply :func:`torch._inductor.aoti_compile_and_package` to AOTInductor
+# to compile the program to a specified device, and save the generated contents
+# into a ".pt2" artifact.
 #
-# We also specify ``dynamic_shapes`` for the batch dimension. In this example, ``min=2`` is not a bug and is 
-# explained in `The 0/1 Specialization Problem <https://docs.google.com/document/d/16VPOa3d-Liikf48teAOmxLc92rgvJdfosIy-yoT38Io/edit?fbclid=IwAR3HNwmmexcitV0pbZm_x1a4ykdXZ9th_eJWK-3hBtVgKnrkmemz6Pm5jRQ#heading=h.ez923tomjvyk>`__
-
+# .. note::
+#
+#       This API supports the same available options that :func:`torch.compile`
+#       has, such as ``mode`` and ``max_autotune`` (for those who want to enable
+#       CUDA graphs and leverage Triton based matrix multiplications and
+#       convolutions)
 
 import os
 import torch
+import torch._inductor
 from torchvision.models import ResNet18_Weights, resnet18
 
 model = resnet18(weights=ResNet18_Weights.DEFAULT)
 model.eval()
 
 with torch.inference_mode():
+    inductor_configs = {}
 
-    # Specify the generated shared library path
-    aot_compile_options = {
-            "aot_inductor.output_path": os.path.join(os.getcwd(), "resnet18_pt2.so"),
-    }
     if torch.cuda.is_available():
         device = "cuda"
-        aot_compile_options.update({"max_autotune": True})
+        inductor_configs["max_autotune"] = True
     else:
         device = "cpu"
 
     model = model.to(device=device)
     example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
 
-    # min=2 is not a bug and is explained in the 0/1 Specialization Problem
-    batch_dim = torch.export.Dim("batch", min=2, max=32)
     exported_program = torch.export.export(
         model,
         example_inputs,
-        # Specify the first dimension of the input x as dynamic
-        dynamic_shapes={"x": {0: batch_dim}},
     )
-    so_path = torch._inductor.aot_compile(
-        exported_program.module(),
-        example_inputs,
-        # Specify the generated shared library path
-        options=aot_compile_options
+    path = torch._inductor.aoti_compile_and_package(
+        exported_program,
+        package_path=os.path.join(os.getcwd(), "resnet18.pt2"),
+        inductor_configs=inductor_configs
     )
 
+######################################################################
+# The result of :func:`aoti_compile_and_package` is an artifact "resnet18.pt2"
+# which can be loaded and executed in Python and C++.
+# 
+# The artifact itself contains a bunch of AOTInductor generated code, such as
+# a generated C++ runner file, a shared library compiled from the C++ file, and
+# CUDA binary files, aka cubin files, if optimizing for CUDA.
+# 
+# Structure-wise, the artifact is a structured ``.zip`` file, with the following 
+# specification:
+#
+# .. code::
+#    .
+#    ├── archive_format
+#    ├── version
+#    ├── data
+#    │   ├── aotinductor
+#    │   │   └── model
+#    │   │       ├── xxx.cpp            # AOTInductor generated cpp file
+#    │   │       ├── xxx.so             # AOTInductor generated shared library
+#    │   │       ├── xxx.cubin          # Cubin files (if running on CUDA)
+#    │   │       └── xxx_metadata.json  # Additional metadata to save
+#    │   ├── weights
+#    │   │  └── TBD
+#    │   └── constants
+#    │      └── TBD
+#    └── extra
+#        └── metadata.json
+#
+# We can use the following command to inspect the artifact contents:
+#
+# .. code:: bash
+#
+#    $ unzip -l resnet18.pt2
+#
+# .. code::
+#
+#    Archive:  resnet18.pt2
+#      Length      Date    Time    Name
+#    ---------  ---------- -----   ----
+#            1  01-08-2025 16:40   version
+#            3  01-08-2025 16:40   archive_format
+#        10088  01-08-2025 16:40   data/aotinductor/model/cagzt6akdaczvxwtbvqe34otfe5jlorktbqlojbzqjqvbfsjlge4.cubin
+#        17160  01-08-2025 16:40   data/aotinductor/model/c6oytfjmt5w4c7onvtm6fray7clirxt7q5xjbwx3hdydclmwoujz.cubin
+#        16616  01-08-2025 16:40   data/aotinductor/model/c7ydp7nocyz323hij4tmlf2kcedmwlyg6r57gaqzcsy3huneamu6.cubin
+#        17776  01-08-2025 16:40   data/aotinductor/model/cyqdf46ordevqhiddvpdpp3uzwatfbzdpl3auj2nx23uxvplnne2.cubin
+#        10856  01-08-2025 16:40   data/aotinductor/model/cpzfebfgrusqslui7fxsuoo4tvwulmrxirc5tmrpa4mvrbdno7kn.cubin
+#        14608  01-08-2025 16:40   data/aotinductor/model/c5ukeoz5wmaszd7vczdz2qhtt6n7tdbl3b6wuy4rb2se24fjwfoy.cubin
+#        11376  01-08-2025 16:40   data/aotinductor/model/csu3nstcp56tsjfycygaqsewpu64l5s6zavvz7537cm4s4cv2k3r.cubin
+#        10984  01-08-2025 16:40   data/aotinductor/model/cp76lez4glmgq7gedf2u25zvvv6rksv5lav4q22dibd2zicbgwj3.cubin
+#        14736  01-08-2025 16:40   data/aotinductor/model/c2bb5p6tnwz4elgujqelsrp3unvkgsyiv7xqxmpvuxcm4jfl7pc2.cubin
+#        11376  01-08-2025 16:40   data/aotinductor/model/c6eopmb2b4ngodwsayae4r5q6ni3jlfogfbdk3ypg56tgpzhubfy.cubin
+#        11624  01-08-2025 16:40   data/aotinductor/model/chmwe6lvoekzfowdbiizitm3haiiuad5kdm6sd2m6mv6dkn2zk32.cubin
+#        15632  01-08-2025 16:40   data/aotinductor/model/c3jop5g344hj3ztsu4qm6ibxyaaerlhkzh2e6emak23rxfje6jam.cubin
+#        25472  01-08-2025 16:40   data/aotinductor/model/chaiixybeiuuitm2nmqnxzijzwgnn2n7uuss4qmsupgblfh3h5hk.cubin
+#       139389  01-08-2025 16:40   data/aotinductor/model/cvk6qzuybruhwxtfblzxiov3rlrziv5fkqc4mdhbmantfu3lmd6t.cpp
+#           27  01-08-2025 16:40   data/aotinductor/model/cvk6qzuybruhwxtfblzxiov3rlrziv5fkqc4mdhbmantfu3lmd6t_metadata.json
+#     47195424  01-08-2025 16:40   data/aotinductor/model/cvk6qzuybruhwxtfblzxiov3rlrziv5fkqc4mdhbmantfu3lmd6t.so
+#    ---------                     -------
+#     47523148                     18 files
+
 
 ######################################################################
 # Model Inference in Python
 # -------------------------
 #
-# Typically, the shared object generated above is used in a non-Python environment. In PyTorch 2.3, 
-# we added a new API called :func:`torch._export.aot_load` to load the shared library in the Python runtime.
-# The API follows a structure similar to the :func:`torch.jit.load` API . You need to specify the path 
-# of the shared library and the device where it should be loaded.
+# To load and run the artifact in Python, we can use :func:`torch._inductor.aoti_load_package`.
 #
-# .. note::
-#      In the example above, we specified ``batch_size=1`` for inference and  it still functions correctly even though we specified ``min=2`` in 
-#      :func:`torch.export.export`.
-
 
 import os
 import torch
+import torch._inductor
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_so_path = os.path.join(os.getcwd(), "resnet18_pt2.so")
+model_path = os.path.join(os.getcwd(), "resnet18.pt2")
 
-model = torch._export.aot_load(model_so_path, device)
-example_inputs = (torch.randn(1, 3, 224, 224, device=device),)
+compiled_model = torch._inductor.aoti_load_package(model_path)
+example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
 
 with torch.inference_mode():
-    output = model(example_inputs)
+    output = compiled_model(example_inputs)
+
 
 ######################################################################
-# When to use AOTInductor for Python Runtime
-# ------------------------------------------
+# When to use AOTInductor with a Python Runtime
+# ---------------------------------------------
 #
-# One of the requirements for using AOTInductor is that the model shouldn't have any graph breaks.
-# Once this requirement is met, the primary use case for using AOTInductor Python Runtime is for
-# model deployment using Python.
-# There are mainly two reasons why you would use AOTInductor Python Runtime:
+# There are mainly two reasons why one would use AOTInductor with a Python Runtime:
 #
-# -  ``torch._inductor.aot_compile`` generates a shared library. This is useful for model
-#    versioning for deployments and tracking model performance over time.
+# -  ``torch._inductor.aoti_compile_and_package`` generates a singular
+#    serialized artifact. This is useful for model versioning for deployments
+#    and tracking model performance over time.
 # -  With :func:`torch.compile` being a JIT compiler, there is a warmup
-#    cost associated with the first compilation. Your deployment needs to account for the
-#    compilation time taken for the first inference. With AOTInductor, the compilation is
-#    done offline using ``torch.export.export`` & ``torch._indutor.aot_compile``. The deployment
-#    would only load the shared library using ``torch._export.aot_load`` and run inference.
+#    cost associated with the first compilation. Your deployment needs to
+#    account for the compilation time taken for the first inference. With
+#    AOTInductor, the compilation is done ahead of time using
+#    ``torch.export.export`` and ``torch._inductor.aoti_compile_and_package``.
+#    At deployment time, after loading the model, running inference does not
+#    have any additional cost.
 #
 #
 # The section below shows the speedup achieved with AOTInductor for first inference
@@ -185,7 +238,7 @@ def timed(fn):
 
 torch._dynamo.reset()
 
-model = torch._export.aot_load(model_so_path, device)
+model = torch._inductor.aoti_load_package(model_path)
 example_inputs = (torch.randn(1, 3, 224, 224, device=device),)
 
 with torch.inference_mode():
@@ -217,8 +270,7 @@ def timed(fn):
 # ----------
 #
 # In this recipe, we have learned how to effectively use the AOTInductor for Python runtime by 
-# compiling and loading a pretrained ``ResNet18`` model using the ``torch._inductor.aot_compile``
-# and ``torch._export.aot_load`` APIs. This process demonstrates the practical application of 
-# generating a shared library and running it within a Python environment, even with dynamic shape
-# considerations and device-specific optimizations. We also looked at the advantage of using 
+# compiling and loading a pretrained ``ResNet18`` model. This process
+# demonstrates the practical application of generating a compiled artifact and
+# running it within a Python environment. We also looked at the advantage of using 
 # AOTInductor in model deployments, with regards to speed up in first inference time.