<jupyter_start><jupyter_text>Copyright 2021 The IREE Authors<jupyter_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<jupyter_output><empty_output><jupyter_text>Variables and StateThis notebook1. Creates a TensorFlow program with basic tf.Variable use2. Imports that program into IREE's compiler3. Compiles the imported program to an IREE VM bytecode module4. Tests running the compiled VM module using IREE's runtime5. Downloads compilation artifacts for use with the native (C API) sample application<jupyter_code>#@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}'")<jupyter_output>Using artifacts directory '/tmp/iree/colab_artifacts'<jupyter_text>Create a program using TensorFlow and import it into IREEThis 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 needsto use `@tf.function` with an `input_signature` specified.References:* ["Introduction to Variables" Guide](https://www.tensorflow.org/guide/variable)* [`tf.Variable` reference](https://www.tensorflow.org/api_docs/python/tf/Variable)* [`tf.function` reference](https://www.tensorflow.org/api_docs/python/tf/function)<jupyter_code>#@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-snapshot iree-tools-tf-snapshot -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)
clear_output()  # Skip over TensorFlow's output.

# Print the imported MLIR to see how the compiler views this TensorFlow program.
# Note IREE's `flow.variable` ops and the public (exported) functions.
print("Counter MLIR:\n```\n", compiler_module.decode("utf-8"), "```\n")

# 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}'")<jupyter_output>Counter MLIR:
```
 #loc0 = loc(unknown)
#map = affine_map<() -> ()>
module  {
  flow.variable @counter mutable dense<0> : tensor<i32> attributes {sym_visibility = "private"} loc(#loc0)
  func @add_to_value(%arg0: tensor<i32> {tf._user_specified_name = "x"} loc(unknown)) attributes {iree.module.export, iree.reflection = {abi = "sip", abiv = 1 : i32, sip = "I8!S5!k0_0R1!"}, tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf.shape<>, #tf.shape<>], tf.signature.is_stateful} {
    %0 = flow.variable.load @counter : tensor<i32> loc(#loc1)
    %1 = linalg.init_tensor [] : tensor<i32> loc(#loc2)
    %2 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} ins(%0, %arg0 : tensor<i32>, tensor<i32>) outs(%1 : tensor<i32>) {
    ^bb0(%arg1: i32 loc(unknown), %arg2: i32 loc(unknown), %arg3: i32 loc(unknown)):  // no predecessors
      %3 = addi %arg1, %arg2 : i32 loc(#loc2)
      linalg.yield %3 : i32 loc(#loc2)
    } -> tensor<i32> loc(#loc2)
    flow.variable[...]<jupyter_text>Test the imported program_Note: you can stop after each step and use intermediate outputs with other tools outside of Colab.__See the [README](https://github.com/google/iree/tree/main/iree/samples/variables_and_statechanging-compilation-options) 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._<jupyter_code>%%capture
!python -m pip install iree-compiler-snapshot -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"])

# Save the imported MLIR 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-snapshot -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("vmvx")
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())
counter.set_value(101)
print(counter.get_value())

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

counter.reset_value()
print(counter.get_value())<jupyter_output>0
101
121
71
0<jupyter_text>Download compilation artifacts<jupyter_code>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
from google.colab import files
print("Downloading the artifacts zip file...")
files.download(ARTIFACTS_ZIP)<jupyter_output>Zipping '/tmp/iree/colab_artifacts' to '/tmp/variables_and_state_colab_artifacts.zip' for download...
  adding: counter_vmvx.vmfb (deflated 67%)
  adding: counter.mlir (deflated 89%)
Downloading the artifacts zip file...
{'fraction_non_alphanumeric': 0.11982903373530758, 'fraction_numerical': 0.015112196611204396, 'mean_word_length': 5.733812949640288, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 31, 'https://': 8, 'lorem ipsum': 0, 'www.': 3, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '15019', 'license': 'Apache-2.0', 'n_tokens_mistral': 2235, 'n_tokens_neox': 2145, 'n_words': 574, 'path': 'iree/samples/variables_and_state/variables_and_state.ipynb', 'repo_name': 'pzeballos/iree', 'size': '6551'}