samples/variables_and_state/variables_and_state.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "variables_and_state.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "FH3IRpYTta2v"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FH3IRpYTta2v"
      },
      "source": [
        "##### Copyright 2021 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mWGa71_Ct2ug",
        "cellView": "form"
      },
      "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"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h5s6ncerSpc5"
      },
      "source": [
        "# Variables and State\n",
        "\n",
        "This notebook\n",
        "\n",
        "1. Creates a TensorFlow program with basic tf.Variable use\n",
        "2. Imports that program into IREE's compiler\n",
        "3. Compiles the imported program to an IREE VM bytecode module\n",
        "4. Tests running the compiled VM module using IREE's runtime\n",
        "5. Downloads compilation artifacts for use with the native (C API) sample application"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s2bScbYkP6VZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "af779558-f4b9-4aaf-d7fa-8ad3b8ae60d0"
      },
      "source": [
        "#@title General setup\n",
        "\n",
        "import os\n",
        "import tempfile\n",
        "\n",
        "ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n",
        "os.makedirs(ARTIFACTS_DIR, exist_ok=True)\n",
        "print(f\"Using artifacts directory '{ARTIFACTS_DIR}'\")"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using artifacts directory '/tmp/iree/colab_artifacts'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!python -m pip install --upgrade tensorflow"
      ],
      "metadata": {
        "id": "xqd9Ka9Mq_5j"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "print(\"TensorFlow version: \", tf.__version__)"
      ],
      "metadata": {
        "id": "ELeIPAUprA8b",
        "outputId": "54b234dd-e6d0-41eb-88b3-27e271cac077",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "TensorFlow version:  2.16.1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dBHgjTjGPOJ7"
      },
      "source": [
        "## Create a program using TensorFlow and import it into IREE\n",
        "\n",
        "This program uses `tf.Variable` to track state internal to the program then exports functions which can be used to interact with that variable.\n",
        "\n",
        "Note that each function we want to be callable from our compiled program needs\n",
        "to use `@tf.function` with an `input_signature` specified.\n",
        "\n",
        "References:\n",
        "\n",
        "* [\"Introduction to Variables\" Guide](https://www.tensorflow.org/guide/variable)\n",
        "* [`tf.Variable` reference](https://www.tensorflow.org/api_docs/python/tf/Variable)\n",
        "* [`tf.function` reference](https://www.tensorflow.org/api_docs/python/tf/function)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hwApbPstraWZ"
      },
      "source": [
        "#@title Define a simple \"counter\" TensorFlow module\n",
        "\n",
        "class CounterModule(tf.Module):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.counter = tf.Variable(0)\n",
        "\n",
        "  @tf.function(input_signature=[])\n",
        "  def get_value(self):\n",
        "    return self.counter\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])\n",
        "  def set_value(self, new_value):\n",
        "    self.counter.assign(new_value)\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])\n",
        "  def add_to_value(self, x):\n",
        "    self.counter.assign(self.counter + x)\n",
        "\n",
        "  @tf.function(input_signature=[])\n",
        "  def reset_value(self):\n",
        "    self.counter.assign(0)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k4aMPI2C7btB"
      },
      "source": [
        "%%capture\n",
        "!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://iree.dev/pip-release-links.html"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Print version information for future notebook users to reference.\n",
        "!iree-compile --version"
      ],
      "metadata": {
        "id": "uU2IkylVrEkk",
        "outputId": "2571c1a2-fb8a-4e2c-d2f0-ac1753a71240",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "IREE (https://iree.dev):\n",
            "  IREE compiler version 20240511.890 @ a3b7e12f1ae3b4d0da9cc5dfa5fb7865b178ec4b\n",
            "  LLVM version 19.0.0git\n",
            "  Optimized build\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3nSXZiZ_X8-P",
        "outputId": "f0ee8ed7-3601-4a03-cf42-4ced767b36ae"
      },
      "source": [
        "#@title Import the TensorFlow program into IREE as MLIR\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "from iree.compiler import tf as tfc\n",
        "\n",
        "compiler_module = tfc.compile_module(\n",
        "    CounterModule(), import_only=True,\n",
        "    output_mlir_debuginfo=False)\n",
        "clear_output()  # Skip over TensorFlow's output.\n",
        "\n",
        "# Save the imported MLIR to disk.\n",
        "imported_mlirbc_path = os.path.join(ARTIFACTS_DIR, \"counter.mlirbc\")\n",
        "with open(imported_mlirbc_path, \"wb\") as output_file:\n",
        "  output_file.write(compiler_module)\n",
        "print(f\"Wrote MLIR to path '{imported_mlirbc_path}'\")\n",
        "\n",
        "# Copy MLIR bytecode to MLIR text and see how the compiler views this program.\n",
        "# Note the 'stablehlo' and 'ml_program' ops and the public (exported) functions.\n",
        "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"counter.mlir\")\n",
        "!iree-ir-tool copy {imported_mlirbc_path} -o {imported_mlir_path}\n",
        "print(\"Counter MLIR:\")\n",
        "!cat {imported_mlir_path}"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wrote MLIR to path '/tmp/iree/colab_artifacts/counter.mlirbc'\n",
            "Counter MLIR:\n",
            "module {\n",
            "  ml_program.global public mutable @vars.__sm_node1__counter(dense<0> : tensor<i32>) : tensor<i32>\n",
            "  func.func @add_to_value(%arg0: tensor<i32>) {\n",
            "    %vars.__sm_node1__counter = ml_program.global_load @vars.__sm_node1__counter : tensor<i32>\n",
            "    %0 = stablehlo.add %vars.__sm_node1__counter, %arg0 : tensor<i32>\n",
            "    ml_program.global_store @vars.__sm_node1__counter = %0 : tensor<i32>\n",
            "    return\n",
            "  }\n",
            "  func.func @get_value() -> tensor<i32> {\n",
            "    %vars.__sm_node1__counter = ml_program.global_load @vars.__sm_node1__counter : tensor<i32>\n",
            "    return %vars.__sm_node1__counter : tensor<i32>\n",
            "  }\n",
            "  func.func @reset_value() {\n",
            "    %c = stablehlo.constant dense<0> : tensor<i32>\n",
            "    ml_program.global_store @vars.__sm_node1__counter = %c : tensor<i32>\n",
            "    return\n",
            "  }\n",
            "  func.func @set_value(%arg0: tensor<i32>) {\n",
            "    ml_program.global_store @vars.__sm_node1__counter = %arg0 : tensor<i32>\n",
            "    return\n",
            "  }\n",
            "}"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WCiRV6KRh3iA"
      },
      "source": [
        "## Test the imported program\n",
        "\n",
        "_Note: you can stop after each step and use intermediate outputs with other tools outside of Colab._\n",
        "\n",
        "_See the [README](https://github.com/iree-org/iree/tree/main/samples/variables_and_state#changing-compilation-options) for more details and example command line instructions._\n",
        "\n",
        "* _The \"imported MLIR\" can be used by IREE's generic compiler tools_\n",
        "* _The \"flatbuffer blob\" can be saved and used by runtime applications_\n",
        "\n",
        "_The specific point at which you switch from Python to native tools will depend on your project._"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GF0dzDsbaP2w",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b8ffb4b6-52f6-4f44-fed6-c04f1ce01fe2"
      },
      "source": [
        "#@title Compile the imported MLIR further into an IREE VM bytecode module\n",
        "\n",
        "from iree.compiler import compile_str\n",
        "\n",
        "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"], input_type=\"stablehlo\")\n",
        "\n",
        "# Save the compiled program to disk.\n",
        "flatbuffer_path = os.path.join(ARTIFACTS_DIR, \"counter_vmvx.vmfb\")\n",
        "with open(flatbuffer_path, \"wb\") as output_file:\n",
        "  output_file.write(flatbuffer_blob)\n",
        "print(f\"Wrote .vmfb to path '{flatbuffer_path}'\")"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wrote .vmfb to path '/tmp/iree/colab_artifacts/counter_vmvx.vmfb'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h8cmF6nAfza0",
        "outputId": "754925ca-229a-4924-b578-88c1615888b8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "#@title Test running the compiled VM module using IREE's runtime\n",
        "\n",
        "from iree import runtime as ireert\n",
        "\n",
        "config = ireert.Config(\"local-task\")\n",
        "ctx = ireert.SystemContext(config=config)\n",
        "vm_module = ireert.VmModule.from_flatbuffer(ctx.instance, flatbuffer_blob)\n",
        "ctx.add_vm_module(vm_module)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-10-e57c61828074>:7: UserWarning: Making copy of unaligned VmModule buffer. It is recommended to make this deterministic by calling `copy_buffer` to always make a copy or `mmap` to efficiently load from a file. This warning can be silenced by adding `warn_if_copy=False` to `from_buffer`\n",
            "  vm_module = ireert.VmModule.from_flatbuffer(ctx.instance, flatbuffer_blob)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CQffg1iQatkb"
      },
      "source": [
        "# Our @tf.functions are accessible by name on the module named 'module'\n",
        "counter = ctx.modules.module\n",
        "\n",
        "# These are buggy in Python but should still work from C\n",
        "# TODO(scotttodd): figure out why and fix\n",
        "\n",
        "# print(counter.get_value().to_host())\n",
        "# counter.set_value(101)\n",
        "# print(counter.get_value().to_host())\n",
        "\n",
        "# counter.add_to_value(20)\n",
        "# print(counter.get_value().to_host())\n",
        "# counter.add_to_value(-50)\n",
        "# print(counter.get_value().to_host())\n",
        "\n",
        "# counter.reset_value()\n",
        "# print(counter.get_value().to_host())"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCvwX1IEokm6"
      },
      "source": [
        "## Download compilation artifacts"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bUaNUkS2ohRj",
        "outputId": "c0971809-fb09-49dc-b92c-acfd672912be",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 104
        }
      },
      "source": [
        "ARTIFACTS_ZIP = \"/tmp/variables_and_state_colab_artifacts.zip\"\n",
        "\n",
        "print(f\"Zipping '{ARTIFACTS_DIR}' to '{ARTIFACTS_ZIP}' for download...\")\n",
        "!cd {ARTIFACTS_DIR} && zip -r {ARTIFACTS_ZIP} .\n",
        "\n",
        "# Note: you can also download files using Colab's file explorer\n",
        "try:\n",
        "  from google.colab import files\n",
        "  print(\"Downloading the artifacts zip file...\")\n",
        "  files.download(ARTIFACTS_ZIP)\n",
        "except ImportError:\n",
        "  print(\"Missing google_colab Python package, can't download files\")"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Zipping '/tmp/iree/colab_artifacts' to '/tmp/variables_and_state_colab_artifacts.zip' for download...\n",
            "  adding: counter.mlir (deflated 74%)\n",
            "  adding: counter.mlirbc (deflated 29%)\n",
            "  adding: counter_vmvx.vmfb (deflated 67%)\n",
            "Downloading the artifacts zip file...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "download(\"download_566d6155-54d4-4ee2-b8f5-aa9110fb9358\", \"variables_and_state_colab_artifacts.zip\", 4679)"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ]
}