<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 2021 The IREE Authors
Step1:  Dynamic Shapes
This notebook
Creates a TensorFlow program with dynamic shapes
Imports that program into IREE's compiler
Compiles the imported program to an IREE VM bytecode module
Tests running the compiled VM module using IREE's runtime
Downloads compilation artifacts for use with the native (C API) sample application
Step2:  Create a program using TensorFlow and import it into IREE
NOTE
Step3:  Test the imported program
Note
Step4:  Download compilation artifacts
<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 2021 The IREE Authors
End of explanation
#@title General setup
import os
import tempfile
ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), "iree", "colab_artifacts")
os.makedirs(ARTIFACTS_DIR, exist_ok=True)
print(f"Using artifacts directory '{ARTIFACTS_DIR}'")
Explanation: Dynamic Shapes
This notebook
Creates a TensorFlow program with dynamic shapes
Imports that program into IREE's compiler
Compiles the imported program to an IREE VM bytecode module
Tests running the compiled VM module using IREE's runtime
Downloads compilation artifacts for use with the native (C API) sample application
End of explanation
#@title Define a sample TensorFlow module using dynamic shapes
import tensorflow as tf
class DynamicShapesModule(tf.Module):
  # reduce_sum_1d (dynamic input size, static output size)
  #   e.g. [1, 2, 3] -> 6
  @tf.function(input_signature=[tf.TensorSpec([None], tf.int32)])
  def reduce_sum_1d(self, values):
    return tf.math.reduce_sum(values)
    
  # reduce_sum_2d (partially dynamic input size, static output size)
  #   e.g. [[1, 2, 3], [10, 20, 30]] -> [11, 22, 33]
  @tf.function(input_signature=[tf.TensorSpec([None, 3], tf.int32)])
  def reduce_sum_2d(self, values):
    return tf.math.reduce_sum(values, 0)
  # add_one (dynamic input size, dynamic output size)
  #   e.g. [1, 2, 3] -> [2, 3, 4]
  @tf.function(input_signature=[tf.TensorSpec([None], tf.int32)])
  def add_one(self, values):
    return tf.math.add(values, tf.constant(1, dtype=tf.int32))
%%capture
!python -m pip install iree-compiler iree-tools-tf -f https://github.com/google/iree/releases
#@title Import the TensorFlow program into IREE as MLIR
from IPython.display import clear_output
from iree.compiler import tf as tfc
compiler_module = tfc.compile_module(
    DynamicShapesModule(), import_only=True, 
    output_mlir_debuginfo=False)
clear_output()  # Skip over TensorFlow's output.
# Print the imported MLIR to see how the compiler views this program.
print("Dynamic Shapes MLIR:\n```\n%s```\n" % compiler_module.decode("utf-8"))
# Save the imported MLIR to disk.
imported_mlir_path = os.path.join(ARTIFACTS_DIR, "dynamic_shapes.mlir")
with open(imported_mlir_path, "wt") as output_file:
  output_file.write(compiler_module.decode("utf-8"))
print(f"Wrote MLIR to path '{imported_mlir_path}'")
Explanation: Create a program using TensorFlow and import it into IREE
NOTE: as in other domains, providing more information to a compiler allows it
to generate more efficient code. As a general rule, the slowest varying
dimensions of program data like batch index or timestep are safer to treat as
dynamic than faster varying dimensions like image x/y/channel. See
this paper for a discussion of the
challenges imposed by dynamic shapes and one project's approach to addressing
them.
End of explanation
%%capture
!python -m pip install iree-compiler -f https://github.com/google/iree/releases
#@title Compile the imported MLIR further into an IREE VM bytecode module
from iree.compiler import compile_str
# Note: we'll use the cpu (LLVM) backend since it has the best support
# for dynamic shapes among our compiler targets.
flatbuffer_blob = compile_str(compiler_module, target_backends=["cpu"], input_type="mhlo")
# Save the compiled program to disk.
flatbuffer_path = os.path.join(ARTIFACTS_DIR, "dynamic_shapes_cpu.vmfb")
with open(flatbuffer_path, "wb") as output_file:
  output_file.write(flatbuffer_blob)
print(f"Wrote compiled program to path '{flatbuffer_path}'")
%%capture
!python -m pip install iree-runtime -f https://github.com/google/iree/releases
#@title Test running the compiled VM module using IREE's runtime
from iree import runtime as ireert
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
config = ireert.Config("local-task")
ctx = ireert.SystemContext(config=config)
ctx.add_vm_module(vm_module)
import numpy as np
# Our @tf.functions are accessible by name on the module named 'module'
dynamic_shapes_program = ctx.modules.module
print(dynamic_shapes_program.reduce_sum_1d(np.array([1, 10, 100], dtype=np.int32)).to_host())
print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30]], dtype=np.int32)).to_host())
print(dynamic_shapes_program.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30], [100, 200, 300]], dtype=np.int32)).to_host())
print(dynamic_shapes_program.add_one(np.array([1, 10, 100], dtype=np.int32)).to_host())
Explanation: Test the imported program
Note: you can stop after each step and use intermediate outputs with other tools outside of Colab.
See the README for more details and example command line instructions.
The "imported MLIR" can be used by IREE's generic compiler tools
The "flatbuffer blob" can be saved and used by runtime applications
The specific point at which you switch from Python to native tools will depend on your project.
End of explanation
ARTIFACTS_ZIP = "/tmp/dynamic_shapes_colab_artifacts.zip"
print(f"Zipping '{ARTIFACTS_DIR}' to '{ARTIFACTS_ZIP}' for download...")
!cd {ARTIFACTS_DIR} && zip -r {ARTIFACTS_ZIP} .
# Note: you can also download files using Colab's file explorer
try:
  from google.colab import files
  print("Downloading the artifacts zip file...")
  files.download(ARTIFACTS_ZIP)  
except ImportError:
  print("Missing google_colab Python package, can't download files")
Explanation: Download compilation artifacts
End of explanation
<END_TASK>