{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {},
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2020 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Image edge detection module"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!python -m pip install --upgrade tensorflow"
      ],
      "metadata": {},
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "!python -m pip install --pre iree-base-compiler iree-base-runtime iree-tools-tf -f https://iree.dev/pip-release-links.html"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Imports\n",
        "\n",
        "import os\n",
        "import tempfile\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from iree import runtime as ireert\n",
        "from iree.tf.support import module_utils\n",
        "from iree.compiler import compile_str\n",
        "from iree.compiler import tf as tfc\n",
        "\n",
        "# Print version information for future notebook users to reference.\n",
        "print(\"TensorFlow version: \", tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Setup Artifacts Directory\n",
        "\n",
        "# Used in the low-level compilation section.\n",
        "ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n",
        "os.makedirs(ARTIFACTS_DIR, exist_ok=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Define the EdgeDetectionModule\n",
        "class EdgeDetectionModule(tf.Module):\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec([1, 128, 128, 1], tf.float32)])\n",
        "  def edge_detect_sobel_operator(self, image):\n",
        "    # https://en.wikipedia.org/wiki/Sobel_operator\n",
        "    sobel_x = tf.constant([[-1.0, 0.0, 1.0],\n",
        "                           [-2.0, 0.0, 2.0],\n",
        "                           [-1.0, 0.0, 1.0]],\n",
        "                          dtype=tf.float32, shape=[3, 3, 1, 1])\n",
        "    sobel_y = tf.constant([[ 1.0,  2.0,  1.0],\n",
        "                           [ 0.0,  0.0,  0.0],\n",
        "                           [-1.0, -2.0, -1.0]],\n",
        "                          dtype=tf.float32, shape=[3, 3, 1, 1])\n",
        "    gx = tf.nn.conv2d(image, sobel_x, 1, \"SAME\")\n",
        "    gy = tf.nn.conv2d(image, sobel_y, 1, \"SAME\")\n",
        "    return tf.math.sqrt(gx * gx + gy * gy)\n",
        "\n",
        "tf_module = EdgeDetectionModule()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg) and run the module with TF\n",
        "def load_image(path_to_image):\n",
        "  image = tf.io.read_file(path_to_image)\n",
        "  image = tf.image.decode_image(image, channels=1)\n",
        "  image = tf.image.convert_image_dtype(image, tf.float32)\n",
        "  image = tf.image.resize(image, (128, 128))\n",
        "  image = image[tf.newaxis, :]\n",
        "  return image\n",
        "\n",
        "content_path = tf.keras.utils.get_file(\n",
        "    'YellowLabradorLooking_new.jpg',\n",
        "    'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')\n",
        "image = load_image(content_path).numpy()\n",
        "\n",
        "def show_images(image, edges):\n",
        "  fig, axs = plt.subplots(1, 2)\n",
        "\n",
        "  axs[0].imshow(image.reshape(128, 128), cmap=\"gray\")\n",
        "  axs[0].set_title(\"Input image\")\n",
        "  axs[1].imshow(edges.reshape(128, 128), cmap=\"gray\")\n",
        "  axs[1].set_title(\"Output image\")\n",
        "\n",
        "  axs[0].axis(\"off\")\n",
        "  axs[1].axis(\"off\")\n",
        "  fig.tight_layout()\n",
        "  fig.show()\n",
        "\n",
        "# Invoke the function with the image as an argument\n",
        "tf_edges = tf_module.edge_detect_sobel_operator(image).numpy()\n",
        "\n",
        "# Plot the input and output images\n",
        "show_images(image, tf_edges)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## High Level Compilation With IREE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@markdown ### Backend Configuration\n",
        "\n",
        "backend_choice = \"iree_vmvx (CPU)\" #@param [ \"iree_vmvx (CPU)\", \"iree_llvmcpu (CPU)\", \"iree_vulkan (GPU/SwiftShader)\" ]\n",
        "backend_choice = backend_choice.split(\" \")[0]\n",
        "backend = module_utils.BackendInfo(backend_choice)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Compile and Run the EdgeDetectionModule with IREE.\n",
        "module = backend.compile_from_class(EdgeDetectionModule)\n",
        "\n",
        "# Compute the edges using the compiled module and display the result.\n",
        "iree_edges = module.edge_detect_sobel_operator(image)\n",
        "show_images(image, iree_edges)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Low-Level Compilation\n",
        "\n",
        "Overview:\n",
        "\n",
        "1.  Convert the `tf.Module` into an IREE compiler module (using `stablehlo`)\n",
        "2.  Save the MLIR assembly from the module into a file (can stop here to use it from another application)\n",
        "3.  Compile the `stablehlo` MLIR into a VM module for IREE to execute\n",
        "4.  Run the VM module through IREE's runtime to test the edge detection function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Construct a module containing the edge detection function\n",
        "\n",
        "# Do *not* further compile to a bytecode module for a particular backend.\n",
        "#\n",
        "# By stopping at stablehlo, we can more easily take advantage of future\n",
        "# compiler improvements within IREE. For a production application, we would\n",
        "# probably want to freeze the version of IREE used and compile as completely as\n",
        "# possible ahead of time, then use some other scheme to load the module into\n",
        "# the application at runtime.\n",
        "compiler_module = tfc.compile_module(\n",
        "    EdgeDetectionModule(), import_only=True)\n",
        "\n",
        "# Save the imported MLIR to disk.\n",
        "imported_mlirbc_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.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",
        "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"edge_detection.mlir\")\n",
        "!iree-ir-tool copy {imported_mlirbc_path} -o {imported_mlir_path}\n",
        "print(\"Edge Detection MLIR:\")\n",
        "!cat {imported_mlir_path}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "#@title Compile and prepare to test the edge detection module\n",
        "\n",
        "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"], input_type=\"stablehlo\")\n",
        "\n",
        "# Register the module with a runtime context.\n",
        "config = ireert.Config(backend.driver)\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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "edge_detect_sobel_operator_f = ctx.modules.module[\"edge_detect_sobel_operator\"]\n",
        "\n",
        "low_level_iree_edges = edge_detect_sobel_operator_f(image)\n",
        "\n",
        "show_images(image, low_level_iree_edges)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}