#@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: Variables and State
This notebook

Creates a TensorFlow program with basic tf.Variable use
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 simple "counter" TensorFlow module

import tensorflow as tf

class CounterModule(tf.Module):
  def __init__(self):
    super().__init__()
    self.counter = tf.Variable(0)

  @tf.function(input_signature=[])
  def get_value(self):
    return self.counter
    
  @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
  def set_value(self, new_value):
    self.counter.assign(new_value)
    
  @tf.function(input_signature=[tf.TensorSpec([], tf.int32)])
  def add_to_value(self, x):
    self.counter.assign(self.counter + x)

  @tf.function(input_signature=[])
  def reset_value(self):
    self.set_value(0)

%%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(
    CounterModule(), 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 TensorFlow program.
# Note IREE's `util.global` ops and the public (exported) functions.
print("Counter MLIR:\n```\n%s```\n" % compiler_module.decode("utf-8"))

# Save the imported MLIR to disk.
imported_mlir_path = os.path.join(ARTIFACTS_DIR, "counter.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
This program uses tf.Variable to track state internal to the program then exports functions which can be used to interact with that variable.
Note that each function we want to be callable from our compiled program needs
to use @tf.function with an input_signature specified.
References:

"Introduction to Variables" Guide
tf.Variable reference
tf.function reference
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

flatbuffer_blob = compile_str(compiler_module, target_backends=["vmvx"], input_type="mhlo")

# Save the compiled program to disk.
flatbuffer_path = os.path.join(ARTIFACTS_DIR, "counter_vmvx.vmfb")
with open(flatbuffer_path, "wb") as output_file:
  output_file.write(flatbuffer_blob)
print(f"Wrote .vmfb 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)

# Our @tf.functions are accessible by name on the module named 'module'
counter = ctx.modules.module

print(counter.get_value().to_host())
counter.set_value(101)
print(counter.get_value().to_host())

counter.add_to_value(20)
print(counter.get_value().to_host())
counter.add_to_value(-50)
print(counter.get_value().to_host())

counter.reset_value()
print(counter.get_value().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/variables_and_state_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
"""


