| { |
| "nbformat": 4, |
| "nbformat_minor": 0, |
| "metadata": { |
| "colab": { |
| "name": "dynamic_shapes.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": [ |
| "# Dynamic Shapes\n", |
| "\n", |
| "This notebook\n", |
| "\n", |
| "1. Creates a TensorFlow program with dynamic shapes\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": "bfab242f-9bb2-44bd-de62-24fd6df58be4" |
| }, |
| "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", |
| "NOTE: as in other domains, providing more information to a compiler allows it\n", |
| "to generate more efficient code. As a general rule, the slowest varying\n", |
| "dimensions of program data like batch index or timestep are safer to treat as\n", |
| "dynamic than faster varying dimensions like image x/y/channel. See\n", |
| "[this paper](https://arxiv.org/pdf/2006.03031.pdf) for a discussion of the\n", |
| "challenges imposed by dynamic shapes and one project's approach to addressing\n", |
| "them." |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "hwApbPstraWZ" |
| }, |
| "source": [ |
| "#@title Define a sample TensorFlow module using dynamic shapes\n", |
| "\n", |
| "import tensorflow as tf\n", |
| "\n", |
| "class DynamicShapesModule(tf.Module):\n", |
| " # reduce_sum_1d (dynamic input size, static output size)\n", |
| " # e.g. [1, 2, 3] -> 6\n", |
| " @tf.function(input_signature=[tf.TensorSpec([None], tf.int32)])\n", |
| " def reduce_sum_1d(self, values):\n", |
| " return tf.math.reduce_sum(values)\n", |
| " \n", |
| " # reduce_sum_2d (partially dynamic input size, static output size)\n", |
| " # e.g. [[1, 2, 3], [10, 20, 30]] -> [11, 22, 33]\n", |
| " @tf.function(input_signature=[tf.TensorSpec([None, 3], tf.int32)])\n", |
| " def reduce_sum_2d(self, values):\n", |
| " return tf.math.reduce_sum(values, 0)\n", |
| "\n", |
| " # add_one (dynamic input size, dynamic output size)\n", |
| " # e.g. [1, 2, 3] -> [2, 3, 4]\n", |
| " @tf.function(input_signature=[tf.TensorSpec([None], tf.int32)])\n", |
| " def add_one(self, values):\n", |
| " return tf.math.add(values, tf.constant(1, dtype=tf.int32))" |
| ], |
| "execution_count": 3, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "k4aMPI2C7btB" |
| }, |
| "source": [ |
| "%%capture\n", |
| "!python -m pip install iree-compiler iree-tools-tf -f https://github.com/google/iree/releases" |
| ], |
| "execution_count": 4, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "colab": { |
| "base_uri": "https://localhost:8080/" |
| }, |
| "id": "3nSXZiZ_X8-P", |
| "outputId": "f193b051-1935-4fcc-ba4a-8a5d7a2f3c51" |
| }, |
| "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", |
| " DynamicShapesModule(), 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 program.\n", |
| "print(\"Dynamic Shapes MLIR:\\n```\\n%s```\\n\" % compiler_module.decode(\"utf-8\"))\n", |
| "\n", |
| "# Save the imported MLIR to disk.\n", |
| "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"dynamic_shapes.mlir\")\n", |
| "with open(imported_mlir_path, \"wt\") as output_file:\n", |
| " output_file.write(compiler_module.decode(\"utf-8\"))\n", |
| "print(f\"Wrote MLIR to path '{imported_mlir_path}'\")" |
| ], |
| "execution_count": 5, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "name": "stdout", |
| "text": [ |
| "Dynamic Shapes MLIR:\n", |
| "```\n", |
| "\"builtin.module\"() ( {\n", |
| " \"builtin.func\"() ( {\n", |
| " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n", |
| " %0 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?xi32>\n", |
| " %1 = \"std.call\"(%0) {callee = @__inference_add_one_70} : (tensor<?xi32>) -> tensor<?xi32>\n", |
| " %2 = \"iree_input.cast.tensor_to_buffer_view\"(%1) : (tensor<?xi32>) -> !iree_input.buffer_view\n", |
| " \"std.return\"(%2) : (!iree_input.buffer_view) -> ()\n", |
| " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22v\\22:1}\", sym_name = \"add_one\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n", |
| " \"builtin.func\"() ( {\n", |
| " ^bb0(%arg0: tensor<?xi32>): // no predecessors\n", |
| " %0 = \"mhlo.constant\"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>\n", |
| " %1 = \"chlo.broadcast_add\"(%arg0, %0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>\n", |
| " \"std.return\"(%1) : (tensor<?xi32>) -> ()\n", |
| " }) {arg_attrs = [{tf._user_specified_name = \"values\"}], sym_name = \"__inference_add_one_70\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<?>], type = (tensor<?xi32>) -> tensor<?xi32>} : () -> ()\n", |
| " \"builtin.func\"() ( {\n", |
| " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n", |
| " %0 = \"mhlo.constant\"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>\n", |
| " %1 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?xi32>\n", |
| " %2 = \"mhlo.reduce\"(%1, %0) ( {\n", |
| " ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>): // no predecessors\n", |
| " %4 = \"mhlo.add\"(%arg1, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i32>\n", |
| " \"mhlo.return\"(%4) : (tensor<i32>) -> ()\n", |
| " }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<?xi32>, tensor<i32>) -> tensor<i32>\n", |
| " %3 = \"iree_input.cast.tensor_to_buffer_view\"(%2) : (tensor<i32>) -> !iree_input.buffer_view\n", |
| " \"std.return\"(%3) : (!iree_input.buffer_view) -> ()\n", |
| " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22v\\22:1}\", sym_name = \"reduce_sum_1d\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n", |
| " \"builtin.func\"() ( {\n", |
| " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n", |
| " %0 = \"mhlo.constant\"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>\n", |
| " %1 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?x3xi32>\n", |
| " %2 = \"mhlo.reduce\"(%1, %0) ( {\n", |
| " ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>): // no predecessors\n", |
| " %4 = \"mhlo.add\"(%arg1, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i32>\n", |
| " \"mhlo.return\"(%4) : (tensor<i32>) -> ()\n", |
| " }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<?x3xi32>, tensor<i32>) -> tensor<3xi32>\n", |
| " %3 = \"iree_input.cast.tensor_to_buffer_view\"(%2) : (tensor<3xi32>) -> !iree_input.buffer_view\n", |
| " \"std.return\"(%3) : (!iree_input.buffer_view) -> ()\n", |
| " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,2,null,3]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,3]],\\22v\\22:1}\", sym_name = \"reduce_sum_2d\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n", |
| "}) : () -> ()\n", |
| "\n", |
| "```\n", |
| "\n", |
| "Wrote MLIR to path '/tmp/iree/colab_artifacts/dynamic_shapes.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/google/iree/tree/main/iree/samples/dynamic_shapes#instructions) 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://github.com/google/iree/releases" |
| ], |
| "execution_count": 6, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "GF0dzDsbaP2w" |
| }, |
| "source": [ |
| "#@title Compile the imported MLIR further into an IREE VM bytecode module\n", |
| "\n", |
| "from iree.compiler import compile_str\n", |
| "\n", |
| "# Note: we'll use the dylib-llvm-aot backend since it has the best support\n", |
| "# for dynamic shapes among our compiler targets.\n", |
| "\n", |
| "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"dylib-llvm-aot\"], input_type=\"mhlo\")\n", |
| "\n", |
| "# Note: the dylib-llvm-aot target produces platform-specific code. Since you\n", |
| "# may need to recompile it yourself using the appropriate\n", |
| "# `-iree-llvm-target-triple` flag, we skip saving it to disk and downloading it." |
| ], |
| "execution_count": 7, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "G7g5eXYL6hWb" |
| }, |
| "source": [ |
| "%%capture\n", |
| "!python -m pip install iree-runtime -f https://github.com/google/iree/releases" |
| ], |
| "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", |
| "vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)\n", |
| "config = ireert.Config(\"dylib\")\n", |
| "ctx = ireert.SystemContext(config=config)\n", |
| "ctx.add_vm_module(vm_module)" |
| ], |
| "execution_count": 9, |
| "outputs": [] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "colab": { |
| "base_uri": "https://localhost:8080/" |
| }, |
| "id": "CQffg1iQatkb", |
| "outputId": "dab42bd6-0291-41fa-d664-3739fbfa1908" |
| }, |
| "source": [ |
| "import numpy as np\n", |
| "\n", |
| "# Our @tf.functions are accessible by name on the module named 'module'\n", |
| "dynamic_shapes_program = ctx.modules.module\n", |
| "\n", |
| "print(dynamic_shapes_program.reduce_sum_1d(np.array([1, 10, 100], dtype=np.int32)))\n", |
| "print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30]], dtype=np.int32)))\n", |
| "print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30], [100, 200, 300]], dtype=np.int32)))\n", |
| "print(dynamic_shapes_program.add_one(np.array([1, 10, 100], dtype=np.int32)))" |
| ], |
| "execution_count": 10, |
| "outputs": [ |
| { |
| "output_type": "stream", |
| "name": "stdout", |
| "text": [ |
| "111\n", |
| "[11 22 33]\n", |
| "[111 222 333]\n", |
| "[ 2 11 101]\n" |
| ] |
| } |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "id": "wCvwX1IEokm6" |
| }, |
| "source": [ |
| "## Download compilation artifacts" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "metadata": { |
| "id": "bUaNUkS2ohRj", |
| "outputId": "0382c67c-72a9-4310-a3a4-550708c515aa", |
| "colab": { |
| "base_uri": "https://localhost:8080/", |
| "height": 69 |
| } |
| }, |
| "source": [ |
| "ARTIFACTS_ZIP = \"/tmp/dynamic_shapes_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/dynamic_shapes_colab_artifacts.zip' for download...\n", |
| " adding: dynamic_shapes.mlir (deflated 80%)\n", |
| "Downloading the artifacts zip file...\n" |
| ] |
| }, |
| { |
| "output_type": "display_data", |
| "data": { |
| "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", |
| " " |
| ], |
| "text/plain": [ |
| "<IPython.core.display.Javascript object>" |
| ] |
| }, |
| "metadata": {} |
| }, |
| { |
| "output_type": "display_data", |
| "data": { |
| "application/javascript": [ |
| "download(\"download_f673a666-f490-4fc5-865b-7d6da53b46a9\", \"dynamic_shapes_colab_artifacts.zip\", 831)" |
| ], |
| "text/plain": [ |
| "<IPython.core.display.Javascript object>" |
| ] |
| }, |
| "metadata": {} |
| } |
| ] |
| } |
| ] |
| } |
