samples/colab/pytorch_jit.ipynb

{
  "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 Just-in-time (JIT) 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 eager execution within a PyTorch session using [IREE](https://github.com/iree-org/iree) and [torch-mlir](https://github.com/llvm/torch-mlir) under the covers."
      ]
    },
    {
      "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.8+ (supports Python 3.12)\n",
        "!python -m pip install torch --index-url https://download.pytorch.org/whl/cpu"
      ]
    },
    {
      "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": [
        "## Sample JIT workflow\n",
        "\n",
        "1. **(Optional)** Set `TURBINE_LOG_LEVEL=debug` to see verbose compilation output including intermediate MLIR IR\n",
        "2. Define a program using `torch.nn.Module`\n",
        "3. Run `torch.compile(module, backend=\"turbine_cpu\")`\n",
        "4. Use the resulting `OptimizedModule` as you would a regular `nn.Module`\n",
        "\n",
        "Useful documentation:\n",
        "\n",
        "* [PyTorch Modules](https://pytorch.org/docs/stable/notes/modules.html) (`nn.Module`) as building blocks for stateful computation\n",
        "* [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) as an interface to TorchDynamo and optimizing using backend compilers like Turbine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Enable debug logging to see MLIR output\n",
        "%env TURBINE_LOG_LEVEL=debug"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "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": [
        "opt_linear_module = torch.compile(linear_module, backend=\"turbine_cpu\")\n",
        "print(\"Compiled module using Turbine. New module type is\", type(opt_linear_module))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "args = torch.randn(4)\n",
        "turbine_output = opt_linear_module(args)\n",
        "\n",
        "print(\"Weight:\", linear_module.weight)\n",
        "print(\"Bias:\", linear_module.bias)\n",
        "print(\"Args:\", args)\n",
        "print(\"Output:\", turbine_output)"
      ]
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 0
}