{
  "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": "4fcb2888-584f-4817-ef30-03f0c4912d14"
      },
      "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",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ae93a2b8-69a3-4fef-ce5c-22e425836589"
      },
      "source": [
        "#@title Define a sample TensorFlow module using dynamic shapes\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 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": [
        {
          "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": "95b23080-0164-4104-9230-1bec7f537bc3"
      },
      "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)\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)\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",
            "module {\n",
            "  func.func @add_one(%arg0: tensor<?xi32>) -> tensor<?xi32> {\n",
            "    %0 = stablehlo.constant dense<1> : tensor<i32>\n",
            "    %1 = shape.shape_of %arg0 : tensor<?xi32> -> tensor<1xindex>\n",
            "    %2 = stablehlo.dynamic_broadcast_in_dim %0, %1, dims = [] : (tensor<i32>, tensor<1xindex>) -> tensor<?xi32>\n",
            "    %3 = stablehlo.add %arg0, %2 : tensor<?xi32>\n",
            "    return %3 : tensor<?xi32>\n",
            "  }\n",
            "  func.func @reduce_sum_1d(%arg0: tensor<?xi32>) -> tensor<i32> {\n",
            "    %0 = stablehlo.constant dense<0> : tensor<i32>\n",
            "    %1 = stablehlo.reduce(%arg0 init: %0) across dimensions = [0] : (tensor<?xi32>, tensor<i32>) -> tensor<i32>\n",
            "     reducer(%arg1: tensor<i32>, %arg2: tensor<i32>)  {\n",
            "      %2 = stablehlo.add %arg1, %arg2 : tensor<i32>\n",
            "      stablehlo.return %2 : tensor<i32>\n",
            "    }\n",
            "    return %1 : tensor<i32>\n",
            "  }\n",
            "  func.func @reduce_sum_2d(%arg0: tensor<?x3xi32>) -> tensor<3xi32> {\n",
            "    %0 = stablehlo.constant dense<0> : tensor<i32>\n",
            "    %1 = stablehlo.reduce(%arg0 init: %0) across dimensions = [0] : (tensor<?x3xi32>, tensor<i32>) -> tensor<3xi32>\n",
            "     reducer(%arg1: tensor<i32>, %arg2: tensor<i32>)  {\n",
            "      %2 = stablehlo.add %arg1, %arg2 : tensor<i32>\n",
            "      stablehlo.return %2 : tensor<i32>\n",
            "    }\n",
            "    return %1 : tensor<3xi32>\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/openxla/iree/tree/main/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://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": "e134a866-4110-43f6-8651-e42a16f7685e"
      },
      "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 LLVM CPU 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=[\"llvm-cpu\"], input_type=\"stablehlo\")\n",
        "\n",
        "# Save the compiled program to disk.\n",
        "flatbuffer_path = os.path.join(ARTIFACTS_DIR, \"dynamic_shapes_cpu.vmfb\")\n",
        "with open(flatbuffer_path, \"wb\") as output_file:\n",
        "  output_file.write(flatbuffer_blob)\n",
        "print(f\"Wrote compiled program to path '{flatbuffer_path}'\")"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wrote compiled program to path '/tmp/iree/colab_artifacts/dynamic_shapes_cpu.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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CQffg1iQatkb",
        "outputId": "c8c2e32c-ed6d-4fc6-cb1c-84725624670c"
      },
      "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)).to_host())\n",
        "print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30]], dtype=np.int32)).to_host())\n",
        "print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30], [100, 200, 300]], dtype=np.int32)).to_host())\n",
        "print(dynamic_shapes_program.add_one(np.array([1, 10, 100], dtype=np.int32)).to_host())"
      ],
      "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": "c54d811d-297b-4ae3-901b-ff6327a74c6a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 86
        }
      },
      "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_cpu.vmfb (deflated 63%)\n",
            "  adding: dynamic_shapes.mlir (deflated 74%)\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_c17b6919-479c-422b-a1c6-d44b38e21698\", \"dynamic_shapes_colab_artifacts.zip\", 5389)"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ]
}