| { |
| "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": "cd13fe22-d885-42d0-8c2f-169bbdc43e58" |
| }, |
| "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": "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", |
| "colab": { |
| "base_uri": "https://localhost:8080/" |
| }, |
| "outputId": "5a26c9c9-1ed2-4181-affb-4d9f986c3da2" |
| }, |
| "source": [ |
| "#@title Define a simple \"counter\" TensorFlow module\n", |
| "\n", |
| "import tensorflow as tf\n", |
| "\n", |
| "# Print version information for future notebook users to reference.\n", |
| "print(\"TensorFlow version: \", tf.__version__)\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": 3, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "name": "stdout", |
| "text": [ |
| "TensorFlow version: 2.12.0\n" |
| ] |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "k4aMPI2C7btB" |
| }, |
| "source": [ |
| "%%capture\n", |
| "!python -m pip install iree-compiler iree-tools-tf -f https://openxla.github.io/iree/pip-release-links.html" |
| ], |
| "execution_count": 4, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "colab": { |
| "base_uri": "https://localhost:8080/" |
| }, |
| "id": "3nSXZiZ_X8-P", |
| "outputId": "7319a86c-a764-499e-ec46-edffc26b7929" |
| }, |
| "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", |
| "# Print the imported MLIR to see how the compiler views this TensorFlow program.\n", |
| "# Note the 'stablehlo' and 'ml_program' ops and the public (exported) functions.\n", |
| "print(\"Counter MLIR:\\n```\\n%s```\\n\" % compiler_module)\n", |
| "\n", |
| "# Save the imported MLIR to disk.\n", |
| "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"counter.mlir\")\n", |
| "with open(imported_mlir_path, \"wt\") as output_file:\n", |
| " output_file.write(compiler_module)\n", |
| "print(f\"Wrote MLIR to path '{imported_mlir_path}'\")" |
| ], |
| "execution_count": 5, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "name": "stdout", |
| "text": [ |
| "Counter MLIR:\n", |
| "```\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", |
| " %0 = ml_program.global_load @vars.__sm_node1__counter : tensor<i32>\n", |
| " %1 = stablehlo.add %0, %arg0 : tensor<i32>\n", |
| " ml_program.global_store @vars.__sm_node1__counter = %1 : tensor<i32>\n", |
| " return\n", |
| " }\n", |
| " func.func @get_value() -> tensor<i32> {\n", |
| " %0 = ml_program.global_load @vars.__sm_node1__counter : tensor<i32>\n", |
| " return %0 : tensor<i32>\n", |
| " }\n", |
| " func.func @reset_value() {\n", |
| " %0 = stablehlo.constant dense<0> : tensor<i32>\n", |
| " ml_program.global_store @vars.__sm_node1__counter = %0 : 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", |
| "}\n", |
| "```\n", |
| "\n", |
| "Wrote MLIR to path '/tmp/iree/colab_artifacts/counter.mlir'\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/openxla/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": "6TV6_Hdu6Xlf" |
| }, |
| "source": [ |
| "%%capture\n", |
| "!python -m pip install iree-compiler -f https://openxla.github.io/iree/pip-release-links.html" |
| ], |
| "execution_count": 6, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "GF0dzDsbaP2w", |
| "colab": { |
| "base_uri": "https://localhost:8080/" |
| }, |
| "outputId": "c527d4ae-7cb8-4e3e-867e-c93592b87151" |
| }, |
| "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=\"mhlo\")\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": 7, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "name": "stdout", |
| "text": [ |
| "Wrote .vmfb to path '/tmp/iree/colab_artifacts/counter_vmvx.vmfb'\n" |
| ] |
| } |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "G7g5eXYL6hWb" |
| }, |
| "source": [ |
| "%%capture\n", |
| "!python -m pip install iree-runtime -f https://openxla.github.io/iree/pip-release-links.html" |
| ], |
| "execution_count": 8, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "h8cmF6nAfza0" |
| }, |
| "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": 9, |
| "outputs": [] |
| }, |
| { |
| "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": 10, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "wCvwX1IEokm6" |
| }, |
| "source": [ |
| "## Download compilation artifacts" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "bUaNUkS2ohRj", |
| "outputId": "c687a21d-9337-4426-be58-e3bf6a08adf8", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 86 |
| } |
| }, |
| "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": 11, |
| "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 71%)\n", |
| " adding: counter_vmvx.vmfb (deflated 65%)\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_4e66c11f-2a2c-4501-82e1-31e95ba67607\", \"variables_and_state_colab_artifacts.zip\", 3941)" |
| ] |
| }, |
| "metadata": {} |
| } |
| ] |
| } |
| ] |
| } |
