{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UUXnh11hA75x"
   },
   "source": [
    "##### Copyright 2023 The IREE Authors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "id": "FqsvmKpjBJO2"
   },
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
    "# See https://llvm.org/LICENSE.txt for license information.\n",
    "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "38UDc27KBPD1"
   },
   "source": [
    "# <img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/1/10/PyTorch_logo_icon.svg/640px-PyTorch_logo_icon.svg.png\" height=\"20px\"> PyTorch Ahead-of-time (AOT) export workflows using <img src=\"https://raw.githubusercontent.com/iree-org/iree/main/docs/website/docs/assets/images/IREE_Logo_Icon_Color.svg\" height=\"20px\"> IREE\n",
    "\n",
    "This notebook shows how to use [iree-turbine](https://github.com/iree-org/iree-turbine) for export from a PyTorch session to [IREE](https://github.com/iree-org/iree), leveraging [torch-mlir](https://github.com/llvm/torch-mlir) under the covers.\n",
    "\n",
    "iree-turbine contains both a \"simple\" AOT exporter and an underlying advanced\n",
    "API for complicated models and full feature availability. This notebook shows\n",
    "some of the features available in the \"advanced\" toolkit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jbcW5jMLK8gK"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "cellView": "form",
    "id": "KsPubQSvCbXd"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "#@title Uninstall existing packages\n",
    "#   This avoids some warnings when installing specific PyTorch packages below.\n",
    "!python -m pip uninstall -y fastai torchaudio torchdata torchtext torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oO1tirq2ggmO",
    "outputId": "2f36a84a-ac8d-453c-c8ed-3d19161b8866"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://download.pytorch.org/whl/test/cpu\n",
      "Collecting torch==2.3.0\n",
      "  Downloading https://download.pytorch.org/whl/test/cpu/torch-2.3.0%2Bcpu-cp310-cp310-linux_x86_64.whl (190.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.4/190.4 MB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (4.12.2)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (1.13.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (2024.10.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch==2.3.0) (3.0.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch==2.3.0) (1.3.0)\n",
      "Installing collected packages: torch\n",
      "  Attempting uninstall: torch\n",
      "    Found existing installation: torch 2.5.0+cu121\n",
      "    Uninstalling torch-2.5.0+cu121:\n",
      "      Successfully uninstalled torch-2.5.0+cu121\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "timm 1.0.11 requires torchvision, which is not installed.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed torch-2.3.0+cpu\n"
     ]
    }
   ],
   "source": [
    "#@title Install Pytorch 2.3.0 (for CPU)\n",
    "!python -m pip install --index-url https://download.pytorch.org/whl/test/cpu --upgrade torch==2.3.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4iJFDHbsAzo4",
    "outputId": "0234c0d1-94b0-4a4d-8876-2feb320c4ae5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting iree-turbine\n",
      "  Downloading iree_turbine-2.5.0-py3-none-any.whl.metadata (5.7 kB)\n",
      "Requirement already satisfied: numpy>=1.26.3 in /usr/local/lib/python3.10/dist-packages (from iree-turbine) (1.26.4)\n",
      "Collecting iree-compiler (from iree-turbine)\n",
      "  Downloading iree_compiler-20241104.1068-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (615 bytes)\n",
      "Collecting iree-runtime (from iree-turbine)\n",
      "  Downloading iree_runtime-20241104.1068-cp310-cp310-manylinux_2_28_x86_64.whl.metadata (798 bytes)\n",
      "Requirement already satisfied: torch>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from iree-turbine) (2.3.0+cpu)\n",
      "Requirement already satisfied: Jinja2>=3.1.3 in /usr/local/lib/python3.10/dist-packages (from iree-turbine) (3.1.4)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from Jinja2>=3.1.3->iree-turbine) (3.0.2)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=2.3.0->iree-turbine) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch>=2.3.0->iree-turbine) (4.12.2)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=2.3.0->iree-turbine) (1.13.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=2.3.0->iree-turbine) (3.4.2)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=2.3.0->iree-turbine) (2024.10.0)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=2.3.0->iree-turbine) (1.3.0)\n",
      "Downloading iree_turbine-2.5.0-py3-none-any.whl (271 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m271.3/271.3 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading iree_compiler-20241104.1068-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (70.7 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m70.7/70.7 MB\u001b[0m \u001b[31m13.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading iree_runtime-20241104.1068-cp310-cp310-manylinux_2_28_x86_64.whl (8.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.0/8.0 MB\u001b[0m \u001b[31m59.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: iree-runtime, iree-compiler, iree-turbine\n",
      "Successfully installed iree-compiler-20241104.1068 iree-runtime-20241104.1068 iree-turbine-2.5.0\n"
     ]
    }
   ],
   "source": [
    "#@title Install iree-turbine\n",
    "!python -m pip install --pre iree-turbine -f https://iree.dev/pip-release-links.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nkVLzRpcDnVL",
    "outputId": "f989db15-1644-4c9e-f307-ea5e2abbbc82"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installed iree-turbine, Version: 2.5.0\n",
      "\n",
      "Installed IREE, compiler version information:\n",
      "IREE (https://iree.dev):\n",
      "  IREE compiler version 20241104.1068 @ 9c85e30df30d6efcf68a7a1b594e89322bd6085d\n",
      "  LLVM version 20.0.0git\n",
      "  Optimized build\n",
      "\n",
      "Installed PyTorch, version: 2.3.0+cpu\n"
     ]
    }
   ],
   "source": [
    "#@title Report version information\n",
    "!echo \"Installed iree-turbine, $(python -m pip show iree_turbine | grep Version)\"\n",
    "\n",
    "!echo -e \"\\nInstalled IREE, compiler version information:\"\n",
    "!iree-compile --version\n",
    "\n",
    "import torch\n",
    "print(\"\\nInstalled PyTorch, version:\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1Mi3YR75LBxl"
   },
   "source": [
    "## Advanced AOT toolkit examples\n",
    "\n",
    "1. Define a PyTorch program using `torch.nn.Module`\n",
    "2. Define the API and properties of that program by using `aot.CompiledModule`\n",
    "3. Export the program using `aot.export()`\n",
    "4. Compile to a deployable artifact\n",
    "  * a: By staying within a Python session\n",
    "  * b: By outputting MLIR and continuing using native tools\n",
    "\n",
    "Useful documentation:\n",
    "\n",
    "* [IREE PyTorch guide](https://iree.dev/guides/ml-frameworks/pytorch/)\n",
    "* [PyTorch Modules](https://pytorch.org/docs/stable/notes/modules.html) (`nn.Module`) as building blocks for stateful computation\n",
    "* IREE compiler and runtime [Python bindings](https://www.iree.dev/reference/bindings/python/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "oPdjrmPZMNz6"
   },
   "outputs": [],
   "source": [
    "#@title 1. Define a program using `torch.nn.Module`\n",
    "torch.manual_seed(0)\n",
    "\n",
    "class LinearModule(torch.nn.Module):\n",
    "  def __init__(self, in_features, out_features):\n",
    "    super().__init__()\n",
    "    self.weight = torch.nn.Parameter(torch.randn(in_features, out_features))\n",
    "    self.bias = torch.nn.Parameter(torch.randn(out_features))\n",
    "\n",
    "  def forward(self, input):\n",
    "    return (input @ self.weight) + self.bias\n",
    "\n",
    "linear_module = LinearModule(4, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "Ua3tNtUIozoa"
   },
   "outputs": [],
   "source": [
    "#@title 2. Define the API and properties of that program by using aot.CompiledModule\n",
    "\n",
    "import iree.turbine.aot as aot\n",
    "\n",
    "example_weight = torch.randn(4, 3)\n",
    "example_bias = torch.randn(3)\n",
    "\n",
    "class CompiledLinearModule(aot.CompiledModule):\n",
    "  params = aot.export_parameters(linear_module, mutable=True)\n",
    "  compute = aot.jittable(linear_module.forward)\n",
    "\n",
    "  def main(self, x=aot.AbstractTensor(4)):\n",
    "    return self.compute(x)\n",
    "\n",
    "  def get_weight(self):\n",
    "    return self.params[\"weight\"]\n",
    "\n",
    "  def set_weight(self, weight=aot.abstractify(example_weight)):\n",
    "    self.params[\"weight\"] = weight\n",
    "\n",
    "  def get_bias(self):\n",
    "    return self.params[\"bias\"]\n",
    "\n",
    "  def set_bias(self, bias=aot.abstractify(example_bias)):\n",
    "    self.params[\"bias\"] = bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "eK2fWVfiSQ8f"
   },
   "outputs": [],
   "source": [
    "#@title 3. Export the program using `aot.export()`\n",
    "export_output = aot.export(CompiledLinearModule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eMRNdFdos900",
    "outputId": "fae696f2-8dbe-4873-f392-953673b6094f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.4178505 -1.2343317 -7.4767942]\n"
     ]
    }
   ],
   "source": [
    "#@title 4a. Compile fully to a deployable artifact, in our existing Python session\n",
    "\n",
    "# Staying in Python gives the API a chance to reuse memory, improving\n",
    "# performance when compiling large programs.\n",
    "\n",
    "compiled_binary = export_output.compile(save_to=None)\n",
    "\n",
    "# Use the IREE runtime API to test the compiled program.\n",
    "import numpy as np\n",
    "import iree.runtime as ireert\n",
    "\n",
    "config = ireert.Config(\"local-task\")\n",
    "vm_module = ireert.load_vm_module(\n",
    "    ireert.VmModule.wrap_buffer(config.vm_instance, compiled_binary.map_memory()),\n",
    "    config,\n",
    ")\n",
    "\n",
    "input = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)\n",
    "result = vm_module.main(input)\n",
    "print(result.to_host())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0AdkXY8VNL2-",
    "outputId": "66335f65-2e9a-4a3a-b9ae-638490063a6e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "module @compiled_linear {\n",
      "  util.global private mutable @_params.weight = dense_resource<_params.weight> : tensor<4x3xf32>\n",
      "  util.global private mutable @_params.bias = dense_resource<_params.bias> : tensor<3xf32>\n",
      "  func.func @main(%arg0: tensor<4xf32>) -> tensor<3xf32> attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\", torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
      "    %0 = torch_c.from_builtin_tensor %arg0 : tensor<4xf32> -> !torch.vtensor<[4],f32>\n",
      "    %1 = call @forward(%0) : (!torch.vtensor<[4],f32>) -> !torch.vtensor<[3],f32>\n",
      "    %2 = torch_c.to_builtin_tensor %1 : !torch.vtensor<[3],f32> -> tensor<3xf32>\n",
      "    return %2 : tensor<3xf32>\n",
      "  }\n",
      "  func.func private @forward(%arg0: !torch.vtensor<[4],f32>) -> !torch.vtensor<[3],f32> {\n",
      "    %int0 = torch.constant.int 0\n",
      "    %0 = torch.aten.unsqueeze %arg0, %int0 : !torch.vtensor<[4],f32>, !torch.int -> !torch.vtensor<[1,4],f32>\n",
      "    %_params.weight = util.global.load @_params.weight : tensor<4x3xf32>\n",
      "    %1 = torch_c.from_builtin_tensor %_params.weight : tensor<4x3xf32> -> !torch.vtensor<[4,3],f32>\n",
      "    %2 = torch.aten.mm %0, %1 : !torch.vtensor<[1,4],f32>, !torch.vtensor<[4,3],f32> -> !torch.vtensor<[1,3],f32>\n",
      "    %int0_0 = torch.constant.int 0\n",
      "    %3 = torch.aten.squeeze.dim %2, %int0_0 : !torch.vtensor<[1,3],f32>, !torch.int -> !torch.vtensor<[3],f32>\n",
      "    %_params.bias = util.global.load @_params.bias : tensor<3xf32>\n",
      "    %4 = torch_c.from_builtin_tensor %_params.bias : tensor<3xf32> -> !torch.vtensor<[3],f32>\n",
      "    %int1 = torch.constant.int 1\n",
      "    %5 = torch.aten.add.Tensor %3, %4, %int1 : !torch.vtensor<[3],f32>, !torch.vtensor<[3],f32>, !torch.int -> !torch.vtensor<[3],f32>\n",
      "    return %5 : !torch.vtensor<[3],f32>\n",
      "  }\n",
      "  func.func @get_weight() -> tensor<4x3xf32> attributes {torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
      "    %_params.weight = util.global.load @_params.weight : tensor<4x3xf32>\n",
      "    return %_params.weight : tensor<4x3xf32>\n",
      "  }\n",
      "  func.func @set_weight(%arg0: tensor<4x3xf32>) attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\"} {\n",
      "    util.global.store %arg0, @_params.weight : tensor<4x3xf32>\n",
      "    return\n",
      "  }\n",
      "  func.func @get_bias() -> tensor<3xf32> attributes {torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
      "    %_params.bias = util.global.load @_params.bias : tensor<3xf32>\n",
      "    return %_params.bias : tensor<3xf32>\n",
      "  }\n",
      "  func.func @set_bias(%arg0: tensor<3xf32>) attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\"} {\n",
      "    util.global.store %arg0, @_params.bias : tensor<3xf32>\n",
      "    return\n",
      "  }\n",
      "}\n",
      "\n",
      "{-#\n",
      "  dialect_resources: {\n",
      "    builtin: {\n",
      "      _params.weight: \"0x040000005C3FC53F503C96BE49710BC0B684113FA1D18ABF2D05B3BF7A83CE3EE588563F442138BF0B83CEBE18BD18BFC6673A3E\",\n",
      "      _params.bias: \"0x04000000074F5BBF99E08C3FAB1C89BF\"\n",
      "    }\n",
      "  }\n",
      "#-}\n",
      "EXEC @main\n",
      "result[0]: hal.buffer_view\n",
      "3xf32=1.41785 -1.23433 -7.47679\n"
     ]
    }
   ],
   "source": [
    "#@title 4b. Output MLIR then continue from Python or native tools later\n",
    "\n",
    "# Leaving Python allows for file system checkpointing and grants access to\n",
    "# native development workflows.\n",
    "\n",
    "mlir_file_path = \"/tmp/linear_module_pytorch.mlirbc\"\n",
    "vmfb_file_path = \"/tmp/linear_module_pytorch_llvmcpu.vmfb\"\n",
    "\n",
    "export_output.print_readable()\n",
    "export_output.save_mlir(mlir_file_path)\n",
    "\n",
    "!iree-compile --iree-input-type=torch --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu --iree-llvmcpu-target-cpu=host {mlir_file_path} -o {vmfb_file_path}\n",
    "!iree-run-module --module={vmfb_file_path} --device=local-task --function=main --input=\"4xf32=[1.0, 2.0, 3.0, 4.0]\""
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "UUXnh11hA75x"
   ],
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}