{
  "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": {}
        }
      ]
    }
  ]
}