| { |
| "cells": [ |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "colab_type": "text", |
| "id": "vQTT2EYu4q_W" |
| }, |
| "source": [ |
| "##### Copyright 2020 Google LLC.\n", |
| "\n", |
| "Licensed under the Apache License, Version 2.0 (the \"License\");" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": null, |
| "metadata": { |
| "cellView": "form", |
| "colab": {}, |
| "colab_type": "code", |
| "id": "BgQ7yyp84qDj" |
| }, |
| "outputs": [], |
| "source": [ |
| "#@title License header\n", |
| "# Copyright 2020 Google LLC\n", |
| "#\n", |
| "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", |
| "# you may not use this file except in compliance with the License.\n", |
| "# You may obtain a copy of the License at\n", |
| "#\n", |
| "# https://www.apache.org/licenses/LICENSE-2.0\n", |
| "#\n", |
| "# Unless required by applicable law or agreed to in writing, software\n", |
| "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", |
| "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", |
| "# See the License for the specific language governing permissions and\n", |
| "# limitations under the License." |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "metadata": { |
| "colab_type": "text", |
| "id": "h5s6ncerSpc5" |
| }, |
| "source": [ |
| "# Defines a simple TF module, saves it and loads it in IREE.\n", |
| "\n", |
| "## Start kernel:\n", |
| "* [Install a TensorFlow2 nightly pip](https://www.tensorflow.org/install) (or bring your own)\n", |
| "* Enable IREE/TF integration by adding to your user.bazelrc: `build --define=iree_tensorflow=true`\n", |
| "* *Optional:* Prime the build: `bazel build bindings/python/pyiree`\n", |
| "* Start colab by running `python colab/start_colab_kernel.py` (see that file for initial setup instructions)" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 1, |
| "metadata": { |
| "colab": {}, |
| "colab_type": "code", |
| "executionInfo": { |
| "elapsed": 6652, |
| "status": "ok", |
| "timestamp": 1598480165652, |
| "user": { |
| "displayName": "", |
| "photoUrl": "", |
| "userId": "" |
| }, |
| "user_tz": 420 |
| }, |
| "id": "s2bScbYkP6VZ" |
| }, |
| "outputs": [], |
| "source": [ |
| "import tensorflow as tf\n", |
| "from pyiree.tf import compiler as ireec" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 2, |
| "metadata": { |
| "colab": {}, |
| "colab_type": "code", |
| "executionInfo": { |
| "elapsed": 623, |
| "status": "ok", |
| "timestamp": 1598480319071, |
| "user": { |
| "displayName": "", |
| "photoUrl": "", |
| "userId": "" |
| }, |
| "user_tz": 420 |
| }, |
| "id": "6YGqN2uqP_7P" |
| }, |
| "outputs": [], |
| "source": [ |
| "class MyModule(tf.Module):\n", |
| "\n", |
| " def __init__(self):\n", |
| " self.v = tf.Variable([4], dtype=tf.float32)\n", |
| "\n", |
| " @tf.function(input_signature=[\n", |
| " tf.TensorSpec([4], tf.float32),\n", |
| " tf.TensorSpec([4], tf.float32)\n", |
| " ])\n", |
| " def add(self, a, b):\n", |
| " return tf.tanh(self.v * a + b)" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 10, |
| "metadata": { |
| "colab": { |
| "height": 326 |
| }, |
| "colab_type": "code", |
| "executionInfo": { |
| "elapsed": 954, |
| "status": "ok", |
| "timestamp": 1598480413200, |
| "user": { |
| "displayName": "", |
| "photoUrl": "", |
| "userId": "" |
| }, |
| "user_tz": 420 |
| }, |
| "id": "r2H2BOpn2SpG", |
| "outputId": "9d6ad89a-d491-4dc4-f6b5-a26a8ef24fdb" |
| }, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "LOADED ASM: \n", |
| "\n", |
| "module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 504 : i32}, tf_saved_model.semantics} {\n", |
| " \"tf_saved_model.global_tensor\"() {is_mutable, sym_name = \"__sm_node1__v\", tf_saved_model.exported_names = [\"v\"], type = tensor\u003c1xf32\u003e, value = dense\u003c4.000000e+00\u003e : tensor\u003c1xf32\u003e} : () -\u003e ()\n", |
| " func @__inference_add_160(%arg0: tensor\u003c4xf32\u003e {tf._user_specified_name = \"a\", tf_saved_model.index_path = [0]}, %arg1: tensor\u003c4xf32\u003e {tf._user_specified_name = \"b\", tf_saved_model.index_path = [1]}, %arg2: tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e {tf_saved_model.bound_input = @__sm_node1__v}) -\u003e (tensor\u003c4xf32\u003e {tf_saved_model.index_path = []}) attributes {tf._input_shapes = [#tf.shape\u003c4\u003e, #tf.shape\u003c4\u003e, #tf.shape\u003c*\u003e], tf.signature.is_stateful, tf_saved_model.exported_names = [\"add\"]} {\n", |
| " %0 = \"tf.Cast\"(%arg2) {Truncate = false} : (tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e) -\u003e tensor\u003c*x!tf.resource\u003e\n", |
| " %1 = tf_executor.graph {\n", |
| " %outputs, %control = tf_executor.island wraps \"tf.ReadVariableOp\"(%0) {device = \"\"} : (tensor\u003c*x!tf.resource\u003e) -\u003e tensor\u003c1xf32\u003e\n", |
| " %outputs_0, %control_1 = tf_executor.island wraps \"tf.Mul\"(%outputs, %arg0) {device = \"\"} : (tensor\u003c1xf32\u003e, tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %outputs_2, %control_3 = tf_executor.island wraps \"tf.AddV2\"(%outputs_0, %arg1) {device = \"\"} : (tensor\u003c4xf32\u003e, tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %outputs_4, %control_5 = tf_executor.island wraps \"tf.Tanh\"(%outputs_2) {device = \"\"} : (tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %outputs_6, %control_7 = tf_executor.island wraps \"tf.Identity\"(%outputs_4) {device = \"\"} : (tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " tf_executor.fetch %outputs_6 : tensor\u003c4xf32\u003e\n", |
| " }\n", |
| " return %1 : tensor\u003c4xf32\u003e\n", |
| " }\n", |
| "}\n" |
| ] |
| } |
| ], |
| "source": [ |
| "#@title Compile to MLIR (mhlo).\n", |
| "compiler_module = ireec.tf_module_to_compiler_module(MyModule(),\n", |
| " pass_pipeline=())\n", |
| "print('LOADED ASM:', compiler_module.to_asm())" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 12, |
| "metadata": { |
| "colab": { |
| "height": 258 |
| }, |
| "colab_type": "code", |
| "executionInfo": { |
| "elapsed": 82, |
| "status": "ok", |
| "timestamp": 1598480427027, |
| "user": { |
| "displayName": "", |
| "photoUrl": "", |
| "userId": "" |
| }, |
| "user_tz": 420 |
| }, |
| "id": "S7IrzODx2RIF", |
| "outputId": "a96b729d-e64d-497e-bd0f-7c0e2fa31d4c" |
| }, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "LOWERED TF ASM: \n", |
| "\n", |
| "module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 504 : i32}, tf_saved_model.semantics} {\n", |
| " \"tf_saved_model.global_tensor\"() {is_mutable, sym_name = \"__sm_node1__v\", tf_saved_model.exported_names = [\"v\"], type = tensor\u003c1xf32\u003e, value = dense\u003c4.000000e+00\u003e : tensor\u003c1xf32\u003e} : () -\u003e ()\n", |
| " func @__inference_add_160(%arg0: tensor\u003c4xf32\u003e {tf._user_specified_name = \"a\", tf_saved_model.index_path = [0]}, %arg1: tensor\u003c4xf32\u003e {tf._user_specified_name = \"b\", tf_saved_model.index_path = [1]}, %arg2: tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e {tf_saved_model.bound_input = @__sm_node1__v}) -\u003e (tensor\u003c4xf32\u003e {tf_saved_model.index_path = []}) attributes {tf._input_shapes = [#tf.shape\u003c4\u003e, #tf.shape\u003c4\u003e, #tf.shape\u003c*\u003e], tf.signature.is_stateful, tf_saved_model.exported_names = [\"add\"]} {\n", |
| " %0 = \"tf.ReadVariableOp\"(%arg2) : (tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e) -\u003e tensor\u003c1xf32\u003e\n", |
| " %1 = \"tf.Mul\"(%0, %arg0) {device = \"\"} : (tensor\u003c1xf32\u003e, tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %2 = \"tf.AddV2\"(%1, %arg1) {device = \"\"} : (tensor\u003c4xf32\u003e, tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %3 = \"tf.Tanh\"(%2) {device = \"\"} : (tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %4 = \"tf.Identity\"(%3) {device = \"\"} : (tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " return %4 : tensor\u003c4xf32\u003e\n", |
| " }\n", |
| "}\n" |
| ] |
| } |
| ], |
| "source": [ |
| "#@title Canonicalize the TF import.\n", |
| "compiler_module.run_pass_pipeline([\n", |
| " \"tf-executor-graph-pruning\",\n", |
| " \"tf-standard-pipeline\",\n", |
| " \"canonicalize\",\n", |
| "])\n", |
| "print(\"LOWERED TF ASM:\", compiler_module.to_asm())" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 13, |
| "metadata": { |
| "colab": { |
| "height": 445 |
| }, |
| "colab_type": "code", |
| "executionInfo": { |
| "elapsed": 47, |
| "status": "ok", |
| "timestamp": 1598480427942, |
| "user": { |
| "displayName": "", |
| "photoUrl": "", |
| "userId": "" |
| }, |
| "user_tz": 420 |
| }, |
| "id": "BVmeZrRx2Uh9", |
| "outputId": "cfaa4c01-8f47-40ed-edbe-19f7cbf2c43c" |
| }, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "XLA ASM: \n", |
| "\n", |
| "module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 504 : i32}, tf_saved_model.semantics} {\n", |
| " \"tf_saved_model.global_tensor\"() {is_mutable, sym_name = \"__sm_node1__v\", tf_saved_model.exported_names = [\"v\"], type = tensor\u003c1xf32\u003e, value = dense\u003c4.000000e+00\u003e : tensor\u003c1xf32\u003e} : () -\u003e ()\n", |
| " func @__inference_add_160(%arg0: tensor\u003c4xf32\u003e {tf._user_specified_name = \"a\", tf_saved_model.index_path = [0]}, %arg1: tensor\u003c4xf32\u003e {tf._user_specified_name = \"b\", tf_saved_model.index_path = [1]}, %arg2: tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e {tf_saved_model.bound_input = @__sm_node1__v}) -\u003e (tensor\u003c4xf32\u003e {tf_saved_model.index_path = []}) attributes {tf._input_shapes = [#tf.shape\u003c4\u003e, #tf.shape\u003c4\u003e, #tf.shape\u003c*\u003e], tf.signature.is_stateful, tf_saved_model.exported_names = [\"add\"]} {\n", |
| " %0 = \"tf.ReadVariableOp\"(%arg2) : (tensor\u003c!tf.resource\u003ctensor\u003c1xf32\u003e\u003e\u003e) -\u003e tensor\u003c1xf32\u003e\n", |
| " %1 = shape.shape_of %0 : tensor\u003c1xf32\u003e -\u003e tensor\u003c?xindex\u003e\n", |
| " %2 = shape.shape_of %arg0 : tensor\u003c4xf32\u003e -\u003e tensor\u003c?xindex\u003e\n", |
| " %3 = shape.cstr_broadcastable %1, %2 : tensor\u003c?xindex\u003e, tensor\u003c?xindex\u003e\n", |
| " %4 = shape.assuming %3 -\u003e (tensor\u003c4xf32\u003e) {\n", |
| " %7 = shape.const_shape [1] : tensor\u003c?xindex\u003e\n", |
| " %8 = shape.const_shape [4] : tensor\u003c?xindex\u003e\n", |
| " %9 = shape.const_shape [4] : tensor\u003c?xindex\u003e\n", |
| " %10 = tensor_cast %9 : tensor\u003c?xindex\u003e to tensor\u003c1xindex\u003e\n", |
| " %11 = \"mhlo.dynamic_broadcast_in_dim\"(%0, %10) {broadcast_dimensions = dense\u003c0\u003e : tensor\u003c1xi64\u003e} : (tensor\u003c1xf32\u003e, tensor\u003c1xindex\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %12 = \"mhlo.dynamic_broadcast_in_dim\"(%arg0, %10) {broadcast_dimensions = dense\u003c0\u003e : tensor\u003c1xi64\u003e} : (tensor\u003c4xf32\u003e, tensor\u003c1xindex\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " %13 = mhlo.multiply %11, %12 : tensor\u003c4xf32\u003e\n", |
| " shape.assuming_yield %13 : tensor\u003c4xf32\u003e\n", |
| " }\n", |
| " %5 = mhlo.add %4, %arg1 : tensor\u003c4xf32\u003e\n", |
| " %6 = \"mhlo.tanh\"(%5) : (tensor\u003c4xf32\u003e) -\u003e tensor\u003c4xf32\u003e\n", |
| " return %6 : tensor\u003c4xf32\u003e\n", |
| " }\n", |
| "}\n" |
| ] |
| } |
| ], |
| "source": [ |
| "#@title Legalize to XLA (high-level).\n", |
| "compiler_module.run_pass_pipeline([\n", |
| " \"xla-legalize-tf{allow-partial-conversion=true}\",\n", |
| "])\n", |
| "print(\"XLA ASM:\", compiler_module.to_asm())" |
| ] |
| } |
| ], |
| "metadata": { |
| "colab": { |
| "collapsed_sections": [], |
| "last_runtime": { |
| "build_target": "", |
| "kind": "local" |
| }, |
| "name": "simple_tensorflow_module_import.ipynb", |
| "provenance": [] |
| }, |
| "kernelspec": { |
| "display_name": "Python 3", |
| "name": "python3" |
| } |
| }, |
| "nbformat": 4, |
| "nbformat_minor": 0 |
| } |
