| { |
| "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 --pre iree-base-compiler iree-base-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": {} |
| } |
| ] |
| } |
| ] |
| } |
