<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2020 The IREE Authors
Step1:  Image edge detection module
Setup
Step2:  High Level Compilation With IREE
Step3:  Low-Level Compilation
Overview
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
Explanation: Copyright 2020 The IREE Authors
End of explanation
!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://github.com/google/iree/releases
#@title Imports
import os
import tempfile
from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf
from iree import runtime as ireert
from iree.tf.support import module_utils
from iree.compiler import compile_str
from iree.compiler import tf as tfc
#@title Setup Artifacts Directory
# Used in the low-level compilation section.
ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), "iree", "colab_artifacts")
os.makedirs(ARTIFACTS_DIR, exist_ok=True)
#@title Define the EdgeDetectionModule
class EdgeDetectionModule(tf.Module):
  @tf.function(input_signature=[tf.TensorSpec([1, 128, 128, 1], tf.float32)])
  def edge_detect_sobel_operator(self, image):
    # https://en.wikipedia.org/wiki/Sobel_operator
    sobel_x = tf.constant([[-1.0, 0.0, 1.0],
                           [-2.0, 0.0, 2.0],
                           [-1.0, 0.0, 1.0]],
                          dtype=tf.float32, shape=[3, 3, 1, 1])    
    sobel_y = tf.constant([[ 1.0,  2.0,  1.0],
                           [ 0.0,  0.0,  0.0],
                           [-1.0, -2.0, -1.0]],
                          dtype=tf.float32, shape=[3, 3, 1, 1])
    gx = tf.nn.conv2d(image, sobel_x, 1, "SAME")
    gy = tf.nn.conv2d(image, sobel_y, 1, "SAME")
    return tf.math.sqrt(gx * gx + gy * gy)
tf_module = EdgeDetectionModule()
#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg) and run the module with TF
def load_image(path_to_image):
  image = tf.io.read_file(path_to_image)
  image = tf.image.decode_image(image, channels=1)
  image = tf.image.convert_image_dtype(image, tf.float32)
  image = tf.image.resize(image, (128, 128))
  image = image[tf.newaxis, :]
  return image
content_path = tf.keras.utils.get_file(
    'YellowLabradorLooking_new.jpg',
    'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')
image = load_image(content_path).numpy()
def show_images(image, edges):
  fig, axs = plt.subplots(1, 2)
  axs[0].imshow(image.reshape(128, 128), cmap="gray")
  axs[0].set_title("Input image")
  axs[1].imshow(edges.reshape(128, 128), cmap="gray")
  axs[1].set_title("Output image")
  axs[0].axis("off")
  axs[1].axis("off")
  fig.tight_layout()
  fig.show()
# Invoke the function with the image as an argument
tf_edges = tf_module.edge_detect_sobel_operator(image).numpy()
# Plot the input and output images
show_images(image, tf_edges)
Explanation: Image edge detection module
Setup
End of explanation
#@markdown ### Backend Configuration
backend_choice = "iree_vmvx (CPU)" #@param [ "iree_vmvx (CPU)", "iree_llvmaot (CPU)", "iree_vulkan (GPU/SwiftShader)" ]
backend_choice = backend_choice.split(" ")[0]
backend = module_utils.BackendInfo(backend_choice)
#@title Compile and Run the EdgeDetectionModule with IREE.
module = backend.compile_from_class(EdgeDetectionModule)
# Compute the edges using the compiled module and display the result.
iree_edges = module.edge_detect_sobel_operator(image)
show_images(image, iree_edges)
Explanation: High Level Compilation With IREE
End of explanation
#@title Construct a module containing the edge detection function
# Do *not* further compile to a bytecode module for a particular backend.
# 
# By stopping at mhlo in text format, we can more easily take advantage of
# future compiler improvements within IREE and can use iree_bytecode_module to
# compile and bundle the module into a sample application. For a production
# application, we would probably want to freeze the version of IREE used and
# compile as completely as possible ahead of time, then use some other scheme
# to load the module into the application at runtime.
compiler_module = tfc.compile_module(EdgeDetectionModule(), import_only=True)
print("Edge Detection MLIR: ", compiler_module.decode('utf-8'))
edge_detection_mlir_path = os.path.join(ARTIFACTS_DIR, "edge_detection.mlir")
with open(edge_detection_mlir_path, "wt") as output_file:
  output_file.write(compiler_module.decode('utf-8'))
print(f"Wrote MLIR to path '{edge_detection_mlir_path}'")
#@title Compile and prepare to test the edge detection module
flatbuffer_blob = compile_str(compiler_module, target_backends=["vmvx"], input_type="mhlo")
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
# Register the module with a runtime context.
config = ireert.Config(backend.driver)
ctx = ireert.SystemContext(config=config)
ctx.add_vm_module(vm_module)
edge_detect_sobel_operator_f = ctx.modules.module["edge_detect_sobel_operator"]
low_level_iree_edges = edge_detect_sobel_operator_f(image)
show_images(image, low_level_iree_edges)
Explanation: Low-Level Compilation
Overview:
Convert the tf.Module into an IREE compiler module (using mhlo)
Save the MLIR assembly from the module into a file (can stop here to use it from another application)
Compile the mhlo MLIR into a VM module for IREE to execute
Run the VM module through IREE's runtime to test the edge detection function
End of explanation
<END_TASK>