{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2023 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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": {},
      "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": {},
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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": null,
      "metadata": {},
      "outputs": [],
      "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.11.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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": null,
      "metadata": {},
      "outputs": [],
      "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": {},
      "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": null,
      "metadata": {},
      "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": null,
      "metadata": {},
      "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": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title 3. Export the program using `aot.export()`\n",
        "export_output = aot.export(CompiledLinearModule)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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": {},
      "outputs": [],
      "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": {},
  "nbformat": 4,
  "nbformat_minor": 0
}