Fix iree-samples buildkite pipeline. (#8142)
Might drop or rewrite the notebooks using TensorFlow soon anyways (https://github.com/google/iree/issues/8086)
diff --git a/colab/edge_detection.ipynb b/colab/edge_detection.ipynb
index 3d59dbc..b82c3a4 100644
--- a/colab/edge_detection.ipynb
+++ b/colab/edge_detection.ipynb
@@ -63,7 +63,7 @@
"colab": {
"base_uri": "https://localhost:8080/"
},
- "outputId": "91312c08-0bee-4835-b1c6-7a13a72d7fcc"
+ "outputId": "bb1e5d0b-ef67-4f25-d4a1-61bc2b7d8356"
},
"source": [
"!python -m pip install iree-compiler iree-runtime iree-tools-tf -f https://github.com/google/iree/releases"
@@ -72,21 +72,23 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Looking in links: https://github.com/google/iree/releases\n",
"Collecting iree-compiler\n",
- "\u001b[?25l Downloading https://github.com/google/iree/releases/download/snapshot-20210608.328/iree_compiler_snapshot-20210608.328-py3-none-manylinux2014_x86_64.whl (33.6MB)\n",
- "\u001b[K |████████████████████████████████| 33.6MB 109kB/s \n",
+ " Downloading https://github.com/google/iree/releases/download/candidate-20220119.22/iree_compiler-20220119.22-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (42.6 MB)\n",
+ "\u001b[K |████████████████████████████████| 42.6 MB 1.3 MB/s \n",
"\u001b[?25hCollecting iree-runtime\n",
- "\u001b[?25l Downloading https://github.com/google/iree/releases/download/snapshot-20210608.328/iree_runtime_snapshot-20210608.328-cp37-cp37m-manylinux2014_x86_64.whl (521kB)\n",
- "\u001b[K |████████████████████████████████| 522kB 44.2MB/s \n",
+ " Downloading https://github.com/google/iree/releases/download/candidate-20220119.22/iree_runtime-20220119.22-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.7 MB)\n",
+ "\u001b[K |████████████████████████████████| 1.7 MB 45.2 MB/s \n",
"\u001b[?25hCollecting iree-tools-tf\n",
- "\u001b[?25l Downloading https://github.com/google/iree/releases/download/snapshot-20210608.328/iree_tools_tf_snapshot-20210608.328-py3-none-manylinux2014_x86_64.whl (53.6MB)\n",
- "\u001b[K |████████████████████████████████| 53.6MB 133kB/s \n",
- "\u001b[?25hInstalling collected packages: iree-compiler, iree-runtime, iree-tools-tf\n",
- "Successfully installed iree-compiler-20210608.328 iree-runtime-20210608.328 iree-tools-tf-20210608.328\n"
- ],
- "name": "stdout"
+ " Downloading https://github.com/google/iree/releases/download/candidate-20220119.22/iree_tools_tf-20220119.22-py3-none-linux_x86_64.whl (49.1 MB)\n",
+ "\u001b[K |████████████████████████████████| 49.1 MB 94 kB/s \n",
+ "\u001b[?25hRequirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from iree-compiler) (3.13)\n",
+ "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from iree-compiler) (1.19.5)\n",
+ "Installing collected packages: iree-tools-tf, iree-runtime, iree-compiler\n",
+ "Successfully installed iree-compiler-20220119.22 iree-runtime-20220119.22 iree-tools-tf-20220119.22\n"
+ ]
}
]
},
@@ -163,10 +165,10 @@
"cellView": "form",
"colab": {
"base_uri": "https://localhost:8080/",
- "height": 281
+ "height": 299
},
"id": "Q2no8DO_s125",
- "outputId": "31fffb8c-89e6-48ae-a972-9305cbfadb0c"
+ "outputId": "96b9aaa4-7d77-4525-d487-347b29aa3c11"
},
"source": [
"#@title Load a test image of a [labrador](https://commons.wikimedia.org/wiki/File:YellowLabradorLooking_new.jpg) and run the module with TF\n",
@@ -206,11 +208,12 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg\n",
- "90112/83281 [================================] - 0s 0us/step\n"
- ],
- "name": "stdout"
+ "90112/83281 [================================] - 0s 0us/step\n",
+ "98304/83281 [===================================] - 0s 0us/step\n"
+ ]
},
{
"output_type": "display_data",
@@ -221,7 +224,6 @@
]
},
"metadata": {
- "tags": [],
"needs_background": "light"
}
}
@@ -256,10 +258,10 @@
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
- "height": 366
+ "height": 264
},
"id": "2OPBr0eGroE0",
- "outputId": "589ea782-a09d-4c3e-919d-1dcf09f8eb77"
+ "outputId": "a0cc5e49-2ef1-483d-eecc-81d5b29b3e1d"
},
"source": [
"#@title Compile and Run the EdgeDetectionModule with IREE.\n",
@@ -273,22 +275,10 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
- "INFO:tensorflow:Assets written to: /tmp/tmp9wpoq1ug.sm/assets\n"
- ],
- "name": "stdout"
- },
- {
- "output_type": "stream",
- "text": [
- "2021-06-09 17:02:22.734737: I external/org_tensorflow/tensorflow/cc/saved_model/bundle_v2.cc:37] Reading SavedModel from: /tmp/tmp9wpoq1ug.sm\n",
- "2021-06-09 17:02:22.735196: I external/org_tensorflow/tensorflow/cc/saved_model/bundle_v2.cc:60] Reading SavedModel debug info (if present) from: /tmp/tmp9wpoq1ug.sm\n",
- "2021-06-09 17:02:22.774953: I external/org_tensorflow/tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA\n",
- "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
- "Created IREE driver vmvx: <iree.runtime.binding.HalDriver object at 0x7f685cfc18f0>\n",
- "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7f685cfc18f0>\n"
- ],
- "name": "stderr"
+ "INFO:tensorflow:Assets written to: /tmp/tmpitk5g1ug.sm/assets\n"
+ ]
},
{
"output_type": "display_data",
@@ -299,7 +289,6 @@
]
},
"metadata": {
- "tags": [],
"needs_background": "light"
}
}
@@ -328,7 +317,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "6YGqN2uqP_7P",
- "outputId": "8514efd3-3863-46e6-f16f-dd10439eb51d"
+ "outputId": "ce5ec981-446b-4e22-fe7c-eb5afc0ea42d"
},
"source": [
"#@title Construct a module containing the edge detection function\n",
@@ -353,91 +342,68 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
- "INFO:tensorflow:Assets written to: /tmp/tmplefnnpwo.sm/assets\n",
- "Edge Detection MLIR: #loc0 = loc(unknown)\n",
- "#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>\n",
- "module {\n",
- " func @edge_detect_sobel_operator(%arg0: tensor<1x128x128x1xf32> {tf._user_specified_name = \"image\"} loc(unknown)) -> tensor<1x128x128x1xf32> attributes {iree.module.export, iree.reflection = {abi = \"sip\", abiv = 1 : i32, sip = \"I8!S5!k0_0R3!_0\"}, tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf.shape<1x128x128x1>]} {\n",
- " %cst = arith.constant 0.000000e+00 : f32 loc(#loc1)\n",
- " %c0 = arith.constant 0 : index loc(#loc1)\n",
- " %c1 = arith.constant 1 : index loc(#loc1)\n",
- " %cst_0 = arith.constant dense<[[[-1.000000e+00], [0.000000e+00], [1.000000e+00]], [[-2.000000e+00], [0.000000e+00], [2.000000e+00]], [[-1.000000e+00], [0.000000e+00], [1.000000e+00]]]> : tensor<3x3x1xf32> loc(#loc1)\n",
- " %cst_1 = arith.constant dense<[[[1.000000e+00], [2.000000e+00], [1.000000e+00]], [[0.000000e+00], [0.000000e+00], [0.000000e+00]], [[-1.000000e+00], [-2.000000e+00], [-1.000000e+00]]]> : tensor<3x3x1xf32> loc(#loc2)\n",
- " %0 = linalg.pad_tensor %arg0 low[%c0, %c1, %c1, %c0] high[%c0, %c1, %c1, %c0] {\n",
- " ^bb0(%arg1: index loc(unknown), %arg2: index loc(unknown), %arg3: index loc(unknown), %arg4: index loc(unknown)): // no predecessors\n",
- " linalg.yield %cst : f32 loc(#loc1)\n",
- " } : tensor<1x128x128x1xf32> to tensor<1x130x130x1xf32> loc(#loc1)\n",
- " %1 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc1)\n",
- " %2 = linalg.fill(%1, %cst) : tensor<1x128x128x1xf32>, f32 -> tensor<1x128x128x1xf32> loc(#loc1)\n",
- " %3 = linalg.depthwise_conv_2d_input_nhwc_filter_hwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%0, %cst_0 : tensor<1x130x130x1xf32>, tensor<3x3x1xf32>) outs(%2 : tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc1)\n",
- " %4 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc3)\n",
- " %5 = linalg.generic {indexing_maps = [#map, #map], iterator_types = [\"parallel\", \"parallel\", \"parallel\", \"parallel\"]} ins(%3 : tensor<1x128x128x1xf32>) outs(%4 : tensor<1x128x128x1xf32>) {\n",
- " ^bb0(%arg1: f32 loc(unknown), %arg2: f32 loc(unknown)): // no predecessors\n",
- " %15 = arith.mulf %arg1, %arg1 : f32 loc(#loc3)\n",
- " linalg.yield %15 : f32 loc(#loc3)\n",
- " } -> tensor<1x128x128x1xf32> loc(#loc3)\n",
- " %6 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc2)\n",
- " %7 = linalg.fill(%6, %cst) : tensor<1x128x128x1xf32>, f32 -> tensor<1x128x128x1xf32> loc(#loc2)\n",
- " %8 = linalg.depthwise_conv_2d_input_nhwc_filter_hwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%0, %cst_1 : tensor<1x130x130x1xf32>, tensor<3x3x1xf32>) outs(%7 : tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc2)\n",
- " %9 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc4)\n",
- " %10 = linalg.generic {indexing_maps = [#map, #map], iterator_types = [\"parallel\", \"parallel\", \"parallel\", \"parallel\"]} ins(%8 : tensor<1x128x128x1xf32>) outs(%9 : tensor<1x128x128x1xf32>) {\n",
- " ^bb0(%arg1: f32 loc(unknown), %arg2: f32 loc(unknown)): // no predecessors\n",
- " %15 = arith.mulf %arg1, %arg1 : f32 loc(#loc4)\n",
- " linalg.yield %15 : f32 loc(#loc4)\n",
- " } -> tensor<1x128x128x1xf32> loc(#loc4)\n",
- " %11 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc5)\n",
- " %12 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = [\"parallel\", \"parallel\", \"parallel\", \"parallel\"]} ins(%5, %10 : tensor<1x128x128x1xf32>, tensor<1x128x128x1xf32>) outs(%11 : tensor<1x128x128x1xf32>) {\n",
- " ^bb0(%arg1: f32 loc(unknown), %arg2: f32 loc(unknown), %arg3: f32 loc(unknown)): // no predecessors\n",
- " %15 = arith.addf %arg1, %arg2 : f32 loc(#loc5)\n",
- " linalg.yield %15 : f32 loc(#loc5)\n",
- " } -> tensor<1x128x128x1xf32> loc(#loc5)\n",
- " %13 = linalg.init_tensor [1, 128, 128, 1] : tensor<1x128x128x1xf32> loc(#loc6)\n",
- " %14 = linalg.generic {indexing_maps = [#map, #map], iterator_types = [\"parallel\", \"parallel\", \"parallel\", \"parallel\"]} ins(%12 : tensor<1x128x128x1xf32>) outs(%13 : tensor<1x128x128x1xf32>) {\n",
- " ^bb0(%arg1: f32 loc(unknown), %arg2: f32 loc(unknown)): // no predecessors\n",
- " %15 = math.sqrt %arg1 : f32 loc(#loc6)\n",
- " linalg.yield %15 : f32 loc(#loc6)\n",
- " } -> tensor<1x128x128x1xf32> loc(#loc6)\n",
- " return %14 : tensor<1x128x128x1xf32> loc(#loc0)\n",
- " } loc(#loc0)\n",
- "} loc(#loc0)\n",
- "#loc1 = loc(callsite(\"Conv2D@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/nn_ops.py\":2169:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":206:0 at callsite(\"<ipython-input-5-814f817fc65c>\":15:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "#loc2 = loc(callsite(\"Conv2D_1@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/nn_ops.py\":2169:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":206:0 at callsite(\"<ipython-input-5-814f817fc65c>\":16:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "#loc3 = loc(callsite(\"mul@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1575:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1234:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "#loc4 = loc(callsite(\"mul_1@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1575:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1234:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "#loc5 = loc(callsite(\"add@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":206:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1234:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "#loc6 = loc(callsite(\"Sqrt@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":5086:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":206:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":982:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3971:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":999:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3289:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3444:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3050:0))))))))))\n",
- "\n",
- "\n",
- "Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlir'\n"
- ],
- "name": "stdout"
+ "INFO:tensorflow:Assets written to: /tmp/tmp9t5tvru4.sm/assets\n"
+ ]
},
{
"output_type": "stream",
+ "name": "stderr",
"text": [
- "2021-06-09 17:02:23.470129: I external/org_tensorflow/tensorflow/cc/saved_model/bundle_v2.cc:37] Reading SavedModel from: /tmp/tmplefnnpwo.sm\n",
- "2021-06-09 17:02:23.470639: I external/org_tensorflow/tensorflow/cc/saved_model/bundle_v2.cc:60] Reading SavedModel debug info (if present) from: /tmp/tmplefnnpwo.sm\n",
- "2021-06-09 17:02:23.506427: I external/org_tensorflow/tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA\n",
- "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
- ],
- "name": "stderr"
+ "INFO:tensorflow:Assets written to: /tmp/tmp9t5tvru4.sm/assets\n"
+ ]
+ },
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Edge Detection MLIR: #loc0 = loc(unknown)\n",
+ "\"builtin.module\"() ( {\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view loc(unknown)): // no predecessors\n",
+ " %0 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<1x128x128x1xf32> loc(#loc0)\n",
+ " %1 = \"std.call\"(%0) {callee = @__inference_edge_detect_sobel_operator_1110} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc0)\n",
+ " %2 = \"iree_input.cast.tensor_to_buffer_view\"(%1) : (tensor<1x128x128x1xf32>) -> !iree_input.buffer_view loc(#loc0)\n",
+ " \"std.return\"(%2) : (!iree_input.buffer_view) -> () loc(#loc0)\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22f32\\22,4,1,128,128,1]],\\22r\\22:[[\\22ndarray\\22,\\22f32\\22,4,1,128,128,1]],\\22v\\22:1}\", sym_name = \"edge_detect_sobel_operator\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> () loc(#loc0)\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: tensor<1x128x128x1xf32> loc(unknown)): // no predecessors\n",
+ " %0 = \"mhlo.constant\"() {value = dense<[[[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]], [[[-2.000000e+00]], [[0.000000e+00]], [[2.000000e+00]]], [[[-1.000000e+00]], [[0.000000e+00]], [[1.000000e+00]]]]> : tensor<3x3x1x1xf32>} : () -> tensor<3x3x1x1xf32> loc(#loc1)\n",
+ " %1 = \"mhlo.constant\"() {value = dense<[[[[1.000000e+00]], [[2.000000e+00]], [[1.000000e+00]]], [[[0.000000e+00]], [[0.000000e+00]], [[0.000000e+00]]], [[[-1.000000e+00]], [[-2.000000e+00]], [[-1.000000e+00]]]]> : tensor<3x3x1x1xf32>} : () -> tensor<3x3x1x1xf32> loc(#loc2)\n",
+ " %2 = \"mhlo.convolution\"(%arg0, %0) {batch_group_count = 1 : i64, dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, feature_group_count = 1 : i64, padding = dense<1> : tensor<2x2xi64>, rhs_dilation = dense<1> : tensor<2xi64>, window_strides = dense<1> : tensor<2xi64>} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc3)\n",
+ " %3 = \"chlo.broadcast_multiply\"(%2, %2) : (tensor<1x128x128x1xf32>, tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc4)\n",
+ " %4 = \"mhlo.convolution\"(%arg0, %1) {batch_group_count = 1 : i64, dimension_numbers = #mhlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>, feature_group_count = 1 : i64, padding = dense<1> : tensor<2x2xi64>, rhs_dilation = dense<1> : tensor<2xi64>, window_strides = dense<1> : tensor<2xi64>} : (tensor<1x128x128x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc5)\n",
+ " %5 = \"chlo.broadcast_multiply\"(%4, %4) : (tensor<1x128x128x1xf32>, tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc6)\n",
+ " %6 = \"chlo.broadcast_add\"(%3, %5) : (tensor<1x128x128x1xf32>, tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc7)\n",
+ " %7 = \"mhlo.sqrt\"(%6) : (tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32> loc(#loc8)\n",
+ " \"std.return\"(%7) : (tensor<1x128x128x1xf32>) -> () loc(#loc0)\n",
+ " }) {arg_attrs = [{tf._user_specified_name = \"image\"}], sym_name = \"__inference_edge_detect_sobel_operator_1110\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<1x128x128x1>], type = (tensor<1x128x128x1xf32>) -> tensor<1x128x128x1xf32>} : () -> () loc(#loc0)\n",
+ "}) : () -> () loc(#loc0)\n",
+ "#loc1 = loc(fused[\"Const:\", callsite(\"Const@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\":294:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\":272:0 at callsite(\"<ipython-input-5-814f817fc65c>\":7:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc2 = loc(fused[\"Const:\", callsite(\"Const_1@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\":294:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\":272:0 at callsite(\"<ipython-input-5-814f817fc65c>\":11:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc3 = loc(fused[\"Conv2D:\", callsite(\"Conv2D@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":1096:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":15:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc4 = loc(fused[\"Mul:\", callsite(\"mul@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1383:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc5 = loc(fused[\"Conv2D:\", callsite(\"Conv2D_1@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":1096:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":16:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc6 = loc(fused[\"Mul:\", callsite(\"mul_1@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1383:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc7 = loc(fused[\"AddV2:\", callsite(\"add@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py\":1383:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "#loc8 = loc(fused[\"Sqrt:\", callsite(\"Sqrt@__inference_edge_detect_sobel_operator_111\"(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py\":1096:0) at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\":150:0 at callsite(\"<ipython-input-5-814f817fc65c>\":17:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1125:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":4095:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\":672:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py\":1143:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3402:0 at callsite(\"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3557:0 at \"/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\":3157:0)))))))))])\n",
+ "\n",
+ "\n",
+ "Wrote MLIR to path '/tmp/iree/colab_artifacts/edge_detection.mlir'\n"
+ ]
}
]
},
{
"cell_type": "code",
"metadata": {
- "colab": {
- "base_uri": "https://localhost:8080/"
- },
- "id": "Ytvb5Gx_EFJl",
- "outputId": "b4fe8f83-db13-4380-d86b-a3595bbf3db1"
+ "id": "Ytvb5Gx_EFJl"
},
"source": [
"#@title Compile and prepare to test the edge detection module\n",
"\n",
- "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"])\n",
+ "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"], input_type=\"mhlo\")\n",
"vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)\n",
"\n",
"# Register the module with a runtime context.\n",
@@ -445,17 +411,8 @@
"ctx = ireert.SystemContext(config=config)\n",
"ctx.add_vm_module(vm_module)"
],
- "execution_count": 10,
- "outputs": [
- {
- "output_type": "stream",
- "text": [
- "Created IREE driver vmvx: <iree.runtime.binding.HalDriver object at 0x7f68a40586b0>\n",
- "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7f68a40586b0>\n"
- ],
- "name": "stderr"
- }
- ]
+ "execution_count": 11,
+ "outputs": []
},
{
"cell_type": "code",
@@ -465,7 +422,7 @@
"height": 247
},
"id": "OUUXxol7wl-f",
- "outputId": "99f32904-33e6-41ed-ac6e-0f8e15f12430"
+ "outputId": "463c4b3b-40f5-4569-ce4a-53accbd88040"
},
"source": [
"edge_detect_sobel_operator_f = ctx.modules.module[\"edge_detect_sobel_operator\"]\n",
@@ -474,7 +431,7 @@
"\n",
"show_images(image, low_level_iree_edges)"
],
- "execution_count": 11,
+ "execution_count": 12,
"outputs": [
{
"output_type": "display_data",
@@ -485,7 +442,6 @@
]
},
"metadata": {
- "tags": [],
"needs_background": "light"
}
}
diff --git a/colab/test_notebooks.py b/colab/test_notebooks.py
index 2de3de6..af545b8 100644
--- a/colab/test_notebooks.py
+++ b/colab/test_notebooks.py
@@ -14,7 +14,13 @@
NOTEBOOKS_TO_SKIP = []
-NOTEBOOKS_EXPECTED_TO_FAIL = []
+NOTEBOOKS_EXPECTED_TO_FAIL = [
+ # Logs: https://buildkite.com/iree/iree-samples/builds/212#2d108f5a-b7d1-46de-93d3-85ce60250514
+ # error: 'util.global.store' op global type mismatch;
+ # global __iree_flow___sm_node14__optimizer.iter is 'tensor<i64>' but
+ # store is 'tensor<i32>'
+ "mnist_training.ipynb",
+]
class ColabNotebookTests(absltest.TestCase):
diff --git a/iree/samples/dynamic_shapes/README.md b/iree/samples/dynamic_shapes/README.md
index 6678425..aafebcf 100644
--- a/iree/samples/dynamic_shapes/README.md
+++ b/iree/samples/dynamic_shapes/README.md
@@ -84,6 +84,7 @@
../iree-build/iree/tools/iree-translate \
-iree-mlir-to-vm-bytecode-module \
-iree-hal-target-backends=dylib-llvm-aot \
+ -iree-input-type=mhlo \
dynamic_shapes.mlir -o dynamic_shapes_dylib.vmfb
```
diff --git a/iree/samples/dynamic_shapes/dynamic_shapes.ipynb b/iree/samples/dynamic_shapes/dynamic_shapes.ipynb
index 8b1b6da..ab67a75 100644
--- a/iree/samples/dynamic_shapes/dynamic_shapes.ipynb
+++ b/iree/samples/dynamic_shapes/dynamic_shapes.ipynb
@@ -62,7 +62,7 @@
"colab": {
"base_uri": "https://localhost:8080/"
},
- "outputId": "1fb6eb27-ae8b-4259-a38e-0b92ea307f73"
+ "outputId": "bfab242f-9bb2-44bd-de62-24fd6df58be4"
},
"source": [
"#@title General setup\n",
@@ -78,10 +78,10 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Using artifacts directory '/tmp/iree/colab_artifacts'\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -141,7 +141,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-compiler-snapshot iree-tools-tf-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-compiler iree-tools-tf -f https://github.com/google/iree/releases"
],
"execution_count": 4,
"outputs": []
@@ -153,7 +153,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "3nSXZiZ_X8-P",
- "outputId": "321b77e3-5890-4d95-a149-83d86ea6f039"
+ "outputId": "f193b051-1935-4fcc-ba4a-8a5d7a2f3c51"
},
"source": [
"#@title Import the TensorFlow program into IREE as MLIR\n",
@@ -180,76 +180,54 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Dynamic Shapes MLIR:\n",
"```\n",
- "#map0 = affine_map<(d0) -> ()>\n",
- "#map1 = affine_map<(d0) -> (d0)>\n",
- "#map2 = affine_map<(d0, d1) -> (d1, d0)>\n",
- "#map3 = affine_map<(d0, d1) -> (d0)>\n",
- "module {\n",
- " func @add_one(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22v\\22:1}\"} {\n",
- " %c0 = arith.constant 0 : index\n",
- " %0 = hal.buffer_view.dim %arg0, 0 : index\n",
- " %1 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<?xi32>{%0}\n",
- " %2 = call @__inference_add_one_70(%1) : (tensor<?xi32>) -> tensor<?xi32>\n",
- " %3 = memref.dim %2, %c0 : tensor<?xi32>\n",
- " %4 = hal.tensor.export %2 : tensor<?xi32>{%3} -> !hal.buffer_view\n",
- " return %4 : !hal.buffer_view\n",
- " }\n",
- " func private @__inference_add_one_70(%arg0: tensor<?xi32> {tf._user_specified_name = \"values\"}) -> tensor<?xi32> attributes {tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf.shape<?>]} {\n",
- " %cst = arith.constant dense<1> : tensor<i32>\n",
- " %c0 = arith.constant 0 : index\n",
- " %0 = memref.dim %arg0, %c0 : tensor<?xi32>\n",
- " %1 = linalg.init_tensor [%0] : tensor<?xi32>\n",
- " %2 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = [\"parallel\"]} ins(%cst : tensor<i32>) outs(%1 : tensor<?xi32>) {\n",
- " ^bb0(%arg1: i32, %arg2: i32): // no predecessors\n",
- " linalg.yield %arg1 : i32\n",
- " } -> tensor<?xi32>\n",
- " %3 = memref.dim %arg0, %c0 : tensor<?xi32>\n",
- " %4 = linalg.init_tensor [%3] : tensor<?xi32>\n",
- " %5 = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = [\"parallel\"]} ins(%arg0, %2 : tensor<?xi32>, tensor<?xi32>) outs(%4 : tensor<?xi32>) {\n",
- " ^bb0(%arg1: i32, %arg2: i32, %arg3: i32): // no predecessors\n",
- " %6 = arith.addi %arg1, %arg2 : i32\n",
- " linalg.yield %6 : i32\n",
- " } -> tensor<?xi32>\n",
- " return %5 : tensor<?xi32>\n",
- " }\n",
- " func @reduce_sum_1d(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22v\\22:1}\"} {\n",
- " %c0_i32 = arith.constant 0 : i32\n",
- " %0 = hal.buffer_view.dim %arg0, 0 : index\n",
- " %1 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<?xi32>{%0}\n",
- " %2 = linalg.init_tensor [] : tensor<i32>\n",
- " %3 = linalg.fill(%2, %c0_i32) : tensor<i32>, i32 -> tensor<i32> \n",
- " %4 = linalg.generic {indexing_maps = [#map1, #map0], iterator_types = [\"reduction\"]} ins(%1 : tensor<?xi32>) outs(%3 : tensor<i32>) {\n",
- " ^bb0(%arg1: i32, %arg2: i32): // no predecessors\n",
- " %6 = arith.addi %arg1, %arg2 : i32\n",
- " linalg.yield %6 : i32\n",
- " } -> tensor<i32>\n",
- " %5 = hal.tensor.export %4 : tensor<i32> -> !hal.buffer_view\n",
- " return %5 : !hal.buffer_view\n",
- " }\n",
- " func @reduce_sum_2d(%arg0: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,2,null,3]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,3]],\\22v\\22:1}\"} {\n",
- " %c0_i32 = arith.constant 0 : i32\n",
- " %0 = hal.buffer_view.dim %arg0, 0 : index\n",
- " %1 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<?x3xi32>{%0}\n",
- " %2 = linalg.init_tensor [3] : tensor<3xi32>\n",
- " %3 = linalg.fill(%2, %c0_i32) : tensor<3xi32>, i32 -> tensor<3xi32> \n",
- " %4 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = [\"parallel\", \"reduction\"]} ins(%1 : tensor<?x3xi32>) outs(%3 : tensor<3xi32>) {\n",
- " ^bb0(%arg1: i32, %arg2: i32): // no predecessors\n",
- " %6 = arith.addi %arg1, %arg2 : i32\n",
- " linalg.yield %6 : i32\n",
- " } -> tensor<3xi32>\n",
- " %5 = hal.tensor.export %4 : tensor<3xi32> -> !hal.buffer_view\n",
- " return %5 : !hal.buffer_view\n",
- " }\n",
- "}\n",
+ "\"builtin.module\"() ( {\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n",
+ " %0 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?xi32>\n",
+ " %1 = \"std.call\"(%0) {callee = @__inference_add_one_70} : (tensor<?xi32>) -> tensor<?xi32>\n",
+ " %2 = \"iree_input.cast.tensor_to_buffer_view\"(%1) : (tensor<?xi32>) -> !iree_input.buffer_view\n",
+ " \"std.return\"(%2) : (!iree_input.buffer_view) -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22v\\22:1}\", sym_name = \"add_one\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: tensor<?xi32>): // no predecessors\n",
+ " %0 = \"mhlo.constant\"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>\n",
+ " %1 = \"chlo.broadcast_add\"(%arg0, %0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>\n",
+ " \"std.return\"(%1) : (tensor<?xi32>) -> ()\n",
+ " }) {arg_attrs = [{tf._user_specified_name = \"values\"}], sym_name = \"__inference_add_one_70\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<?>], type = (tensor<?xi32>) -> tensor<?xi32>} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n",
+ " %0 = \"mhlo.constant\"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>\n",
+ " %1 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?xi32>\n",
+ " %2 = \"mhlo.reduce\"(%1, %0) ( {\n",
+ " ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>): // no predecessors\n",
+ " %4 = \"mhlo.add\"(%arg1, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i32>\n",
+ " \"mhlo.return\"(%4) : (tensor<i32>) -> ()\n",
+ " }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<?xi32>, tensor<i32>) -> tensor<i32>\n",
+ " %3 = \"iree_input.cast.tensor_to_buffer_view\"(%2) : (tensor<i32>) -> !iree_input.buffer_view\n",
+ " \"std.return\"(%3) : (!iree_input.buffer_view) -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,1,null]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22v\\22:1}\", sym_name = \"reduce_sum_1d\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n",
+ " %0 = \"mhlo.constant\"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>\n",
+ " %1 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<?x3xi32>\n",
+ " %2 = \"mhlo.reduce\"(%1, %0) ( {\n",
+ " ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>): // no predecessors\n",
+ " %4 = \"mhlo.add\"(%arg1, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i32>\n",
+ " \"mhlo.return\"(%4) : (tensor<i32>) -> ()\n",
+ " }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<?x3xi32>, tensor<i32>) -> tensor<3xi32>\n",
+ " %3 = \"iree_input.cast.tensor_to_buffer_view\"(%2) : (tensor<3xi32>) -> !iree_input.buffer_view\n",
+ " \"std.return\"(%3) : (!iree_input.buffer_view) -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,2,null,3]],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,1,3]],\\22v\\22:1}\", sym_name = \"reduce_sum_2d\", type = (!iree_input.buffer_view) -> !iree_input.buffer_view} : () -> ()\n",
+ "}) : () -> ()\n",
"\n",
"```\n",
"\n",
"Wrote MLIR to path '/tmp/iree/colab_artifacts/dynamic_shapes.mlir'\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -278,7 +256,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-compiler-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-compiler -f https://github.com/google/iree/releases"
],
"execution_count": 6,
"outputs": []
@@ -296,7 +274,7 @@
"# Note: we'll use the dylib-llvm-aot backend since it has the best support\n",
"# for dynamic shapes among our compiler targets.\n",
"\n",
- "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"dylib-llvm-aot\"])\n",
+ "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"dylib-llvm-aot\"], input_type=\"mhlo\")\n",
"\n",
"# Note: the dylib-llvm-aot target produces platform-specific code. Since you\n",
"# may need to recompile it yourself using the appropriate\n",
@@ -312,7 +290,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-runtime-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-runtime -f https://github.com/google/iree/releases"
],
"execution_count": 8,
"outputs": []
@@ -320,11 +298,7 @@
{
"cell_type": "code",
"metadata": {
- "colab": {
- "base_uri": "https://localhost:8080/"
- },
- "id": "h8cmF6nAfza0",
- "outputId": "0e89eecd-ad88-45bb-9a81-c8cc9519eccb"
+ "id": "h8cmF6nAfza0"
},
"source": [
"#@title Test running the compiled VM module using IREE's runtime\n",
@@ -337,16 +311,7 @@
"ctx.add_vm_module(vm_module)"
],
"execution_count": 9,
- "outputs": [
- {
- "output_type": "stream",
- "text": [
- "Created IREE driver dylib: <iree.runtime.binding.HalDriver object at 0x7fc97f1611f0>\n",
- "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7fc97f1611f0>\n"
- ],
- "name": "stderr"
- }
- ]
+ "outputs": []
},
{
"cell_type": "code",
@@ -355,7 +320,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "CQffg1iQatkb",
- "outputId": "f1729cac-4bbc-420e-fe22-bd47b36d8ad4"
+ "outputId": "dab42bd6-0291-41fa-d664-3739fbfa1908"
},
"source": [
"import numpy as np\n",
@@ -372,13 +337,13 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"111\n",
"[11 22 33]\n",
"[111 222 333]\n",
"[ 2 11 101]\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -394,7 +359,12 @@
{
"cell_type": "code",
"metadata": {
- "id": "bUaNUkS2ohRj"
+ "id": "bUaNUkS2ohRj",
+ "outputId": "0382c67c-72a9-4310-a3a4-550708c515aa",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 69
+ }
},
"source": [
"ARTIFACTS_ZIP = \"/tmp/dynamic_shapes_colab_artifacts.zip\"\n",
@@ -410,8 +380,82 @@
"except ImportError:\n",
" print(\"Missing google_colab Python package, can't download files\")"
],
- "execution_count": null,
- "outputs": []
+ "execution_count": 11,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Zipping '/tmp/iree/colab_artifacts' to '/tmp/dynamic_shapes_colab_artifacts.zip' for download...\n",
+ " adding: dynamic_shapes.mlir (deflated 80%)\n",
+ "Downloading the artifacts zip file...\n"
+ ]
+ },
+ {
+ "output_type": "display_data",
+ "data": {
+ "application/javascript": [
+ "\n",
+ " async function download(id, filename, size) {\n",
+ " if (!google.colab.kernel.accessAllowed) {\n",
+ " return;\n",
+ " }\n",
+ " const div = document.createElement('div');\n",
+ " const label = document.createElement('label');\n",
+ " label.textContent = `Downloading \"${filename}\": `;\n",
+ " div.appendChild(label);\n",
+ " const progress = document.createElement('progress');\n",
+ " progress.max = size;\n",
+ " div.appendChild(progress);\n",
+ " document.body.appendChild(div);\n",
+ "\n",
+ " const buffers = [];\n",
+ " let downloaded = 0;\n",
+ "\n",
+ " const channel = await google.colab.kernel.comms.open(id);\n",
+ " // Send a message to notify the kernel that we're ready.\n",
+ " channel.send({})\n",
+ "\n",
+ " for await (const message of channel.messages) {\n",
+ " // Send a message to notify the kernel that we're ready.\n",
+ " channel.send({})\n",
+ " if (message.buffers) {\n",
+ " for (const buffer of message.buffers) {\n",
+ " buffers.push(buffer);\n",
+ " downloaded += buffer.byteLength;\n",
+ " progress.value = downloaded;\n",
+ " }\n",
+ " }\n",
+ " }\n",
+ " const blob = new Blob(buffers, {type: 'application/binary'});\n",
+ " const a = document.createElement('a');\n",
+ " a.href = window.URL.createObjectURL(blob);\n",
+ " a.download = filename;\n",
+ " div.appendChild(a);\n",
+ " a.click();\n",
+ " div.remove();\n",
+ " }\n",
+ " "
+ ],
+ "text/plain": [
+ "<IPython.core.display.Javascript object>"
+ ]
+ },
+ "metadata": {}
+ },
+ {
+ "output_type": "display_data",
+ "data": {
+ "application/javascript": [
+ "download(\"download_f673a666-f490-4fc5-865b-7d6da53b46a9\", \"dynamic_shapes_colab_artifacts.zip\", 831)"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Javascript object>"
+ ]
+ },
+ "metadata": {}
+ }
+ ]
}
]
-}
+}
\ No newline at end of file
diff --git a/iree/samples/dynamic_shapes/test.sh b/iree/samples/dynamic_shapes/test.sh
index d78a0ec..d7c19c6 100755
--- a/iree/samples/dynamic_shapes/test.sh
+++ b/iree/samples/dynamic_shapes/test.sh
@@ -29,6 +29,7 @@
${BUILD_DIR}/iree/tools/iree-translate \
-iree-mlir-to-vm-bytecode-module \
-iree-hal-target-backends=dylib-llvm-aot \
+ -iree-input-type=mhlo \
${ARTIFACTS_DIR}/dynamic_shapes.mlir -o ${ARTIFACTS_DIR}/dynamic_shapes_dylib.vmfb
# 4. Build the `iree_samples_dynamic_shapes` CMake target.
diff --git a/iree/samples/variables_and_state/README.md b/iree/samples/variables_and_state/README.md
index de84608..dc0e086 100644
--- a/iree/samples/variables_and_state/README.md
+++ b/iree/samples/variables_and_state/README.md
@@ -94,6 +94,7 @@
../iree-build/iree/tools/iree-translate \
-iree-mlir-to-vm-bytecode-module \
-iree-hal-target-backends=dylib-llvm-aot \
+ -iree-input-type=mhlo \
counter.mlir -o counter_dylib.vmfb
```
diff --git a/iree/samples/variables_and_state/variables_and_state.ipynb b/iree/samples/variables_and_state/variables_and_state.ipynb
index 733f9f6..6f537f5 100644
--- a/iree/samples/variables_and_state/variables_and_state.ipynb
+++ b/iree/samples/variables_and_state/variables_and_state.ipynb
@@ -62,7 +62,7 @@
"colab": {
"base_uri": "https://localhost:8080/"
},
- "outputId": "a0d86566-9ae9-4b2a-bbf0-893ff72980d1"
+ "outputId": "e2d13e53-c3ba-4258-d247-4b64a376182e"
},
"source": [
"#@title General setup\n",
@@ -78,10 +78,10 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Using artifacts directory '/tmp/iree/colab_artifacts'\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -146,7 +146,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-compiler-snapshot iree-tools-tf-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-compiler iree-tools-tf -f https://github.com/google/iree/releases"
],
"execution_count": 4,
"outputs": []
@@ -158,7 +158,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "3nSXZiZ_X8-P",
- "outputId": "fbd021ef-3da9-475d-87b9-00e0e8139bf3"
+ "outputId": "d266ec8b-cc3c-498b-f212-b0efe1f50284"
},
"source": [
"#@title Import the TensorFlow program into IREE as MLIR\n",
@@ -185,50 +185,64 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Counter MLIR:\n",
"```\n",
- "#map = affine_map<() -> ()>\n",
- "module {\n",
- " util.global private mutable @counter = dense<0> : tensor<i32>\n",
- " func @add_to_value(%arg0: !hal.buffer_view) attributes {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22r\\22:[],\\22v\\22:1}\"} {\n",
- " %0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<i32>\n",
- " call @__inference_add_to_value_100(%0) : (tensor<i32>) -> ()\n",
- " return\n",
- " }\n",
- " func private @__inference_add_to_value_100(%arg0: tensor<i32> {tf._user_specified_name = \"x\"}) attributes {tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf.shape<>, #tf.shape<>], tf.signature.is_stateful} {\n",
- " %0 = util.global.load @counter : tensor<i32>\n",
- " %1 = linalg.init_tensor [] : tensor<i32>\n",
- " %2 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} ins(%0, %arg0 : tensor<i32>, tensor<i32>) outs(%1 : tensor<i32>) {\n",
- " ^bb0(%arg1: i32, %arg2: i32, %arg3: i32): // no predecessors\n",
- " %3 = arith.addi %arg1, %arg2 : i32\n",
- " linalg.yield %3 : i32\n",
- " } -> tensor<i32>\n",
- " util.global.store %2, @counter : tensor<i32>\n",
- " return\n",
- " }\n",
- " func @get_value() -> !hal.buffer_view attributes {iree.abi = \"{\\22a\\22:[],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22v\\22:1}\"} {\n",
- " %0 = util.global.load @counter : tensor<i32>\n",
- " %1 = hal.tensor.export %0 : tensor<i32> -> !hal.buffer_view\n",
- " return %1 : !hal.buffer_view\n",
- " }\n",
- " func @reset_value() attributes {iree.abi = \"{\\22a\\22:[],\\22r\\22:[],\\22v\\22:1}\"} {\n",
- " %cst = arith.constant dense<0> : tensor<i32>\n",
- " util.global.store %cst, @counter : tensor<i32>\n",
- " return\n",
- " }\n",
- " func @set_value(%arg0: !hal.buffer_view) attributes {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22r\\22:[],\\22v\\22:1}\"} {\n",
- " %0 = hal.tensor.import %arg0 : !hal.buffer_view -> tensor<i32>\n",
- " util.global.store %0, @counter : tensor<i32>\n",
- " return\n",
- " }\n",
- "}\n",
+ "\"builtin.module\"() ( {\n",
+ " \"iree_input.global\"() {initial_value = dense<0> : tensor<i32>, is_mutable, sym_name = \"counter\", sym_visibility = \"private\", type = tensor<i32>} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n",
+ " %0 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<i32>\n",
+ " \"std.call\"(%0) {callee = @__inference_add_to_value_100} : (tensor<i32>) -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22r\\22:[],\\22v\\22:1}\", sym_name = \"add_to_value\", type = (!iree_input.buffer_view) -> ()} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: tensor<i32>): // no predecessors\n",
+ " %0 = \"iree_input.global.address\"() {global = @counter} : () -> !iree_input.ptr<tensor<i32>>\n",
+ " %1 = \"iree_input.global.load.indirect\"(%0) : (!iree_input.ptr<tensor<i32>>) -> tensor<i32>\n",
+ " %2 = \"chlo.broadcast_add\"(%1, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>\n",
+ " \"iree_input.global.store.indirect\"(%2, %0) : (tensor<i32>, !iree_input.ptr<tensor<i32>>) -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {arg_attrs = [{tf._user_specified_name = \"x\"}], sym_name = \"__inference_add_to_value_100\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<>, #tf_type.shape<>], tf.signature.is_stateful, type = (tensor<i32>) -> ()} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " %0 = \"std.call\"() {callee = @__inference_get_value_160} : () -> tensor<i32>\n",
+ " %1 = \"iree_input.cast.tensor_to_buffer_view\"(%0) : (tensor<i32>) -> !iree_input.buffer_view\n",
+ " \"std.return\"(%1) : (!iree_input.buffer_view) -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[],\\22r\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22v\\22:1}\", sym_name = \"get_value\", type = () -> !iree_input.buffer_view} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " %0 = \"iree_input.global.address\"() {global = @counter} : () -> !iree_input.ptr<tensor<i32>>\n",
+ " %1 = \"iree_input.global.load.indirect\"(%0) : (!iree_input.ptr<tensor<i32>>) -> tensor<i32>\n",
+ " \"std.return\"(%1) : (tensor<i32>) -> ()\n",
+ " }) {sym_name = \"__inference_get_value_160\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, type = () -> tensor<i32>} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " \"std.call\"() {callee = @__inference_reset_value_270} : () -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[],\\22r\\22:[],\\22v\\22:1}\", sym_name = \"reset_value\", type = () -> ()} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " %0 = \"mhlo.constant\"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>\n",
+ " %1 = \"iree_input.global.address\"() {global = @counter} : () -> !iree_input.ptr<tensor<i32>>\n",
+ " \"iree_input.global.store.indirect\"(%0, %1) : (tensor<i32>, !iree_input.ptr<tensor<i32>>) -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {sym_name = \"__inference_reset_value_270\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, type = () -> ()} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: !iree_input.buffer_view): // no predecessors\n",
+ " %0 = \"iree_input.cast.buffer_view_to_tensor\"(%arg0) : (!iree_input.buffer_view) -> tensor<i32>\n",
+ " \"std.call\"(%0) {callee = @__sm_exported___inference_set_value_230} : (tensor<i32>) -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {iree.abi = \"{\\22a\\22:[[\\22ndarray\\22,\\22i32\\22,0]],\\22r\\22:[],\\22v\\22:1}\", sym_name = \"set_value\", type = (!iree_input.buffer_view) -> ()} : () -> ()\n",
+ " \"builtin.func\"() ( {\n",
+ " ^bb0(%arg0: tensor<i32>): // no predecessors\n",
+ " %0 = \"iree_input.global.address\"() {global = @counter} : () -> !iree_input.ptr<tensor<i32>>\n",
+ " \"iree_input.global.store.indirect\"(%arg0, %0) : (tensor<i32>, !iree_input.ptr<tensor<i32>>) -> ()\n",
+ " \"std.return\"() : () -> ()\n",
+ " }) {arg_attrs = [{tf._user_specified_name = \"new_value\"}], sym_name = \"__sm_exported___inference_set_value_230\", sym_visibility = \"private\", tf._construction_context = \"kEagerRuntime\", tf._input_shapes = [#tf_type.shape<>, #tf_type.shape<>], tf.signature.is_stateful, type = (tensor<i32>) -> ()} : () -> ()\n",
+ "}) : () -> ()\n",
"\n",
"```\n",
"\n",
"Wrote MLIR to path '/tmp/iree/colab_artifacts/counter.mlir'\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -257,7 +271,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-compiler-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-compiler -f https://github.com/google/iree/releases"
],
"execution_count": 6,
"outputs": []
@@ -269,14 +283,14 @@
"colab": {
"base_uri": "https://localhost:8080/"
},
- "outputId": "850d68d0-36e5-498c-fa83-b93e90d1adbc"
+ "outputId": "a8cd9ce6-4b7f-40c0-9628-b18c5c4748d6"
},
"source": [
"#@title Compile the imported MLIR further into an IREE VM bytecode module\n",
"\n",
"from iree.compiler import compile_str\n",
"\n",
- "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"])\n",
+ "flatbuffer_blob = compile_str(compiler_module, target_backends=[\"vmvx\"], input_type=\"mhlo\")\n",
"\n",
"# Save the imported MLIR to disk.\n",
"flatbuffer_path = os.path.join(ARTIFACTS_DIR, \"counter_vmvx.vmfb\")\n",
@@ -288,10 +302,10 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"Wrote .vmfb to path '/tmp/iree/colab_artifacts/counter_vmvx.vmfb'\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -302,7 +316,7 @@
},
"source": [
"%%capture\n",
- "!python -m pip install iree-runtime-snapshot -f https://github.com/google/iree/releases"
+ "!python -m pip install iree-runtime -f https://github.com/google/iree/releases"
],
"execution_count": 8,
"outputs": []
@@ -310,11 +324,7 @@
{
"cell_type": "code",
"metadata": {
- "colab": {
- "base_uri": "https://localhost:8080/"
- },
- "id": "h8cmF6nAfza0",
- "outputId": "62636478-965a-4b5c-c170-28a5c84f4af1"
+ "id": "h8cmF6nAfza0"
},
"source": [
"#@title Test running the compiled VM module using IREE's runtime\n",
@@ -327,16 +337,7 @@
"ctx.add_vm_module(vm_module)"
],
"execution_count": 9,
- "outputs": [
- {
- "output_type": "stream",
- "text": [
- "Created IREE driver vmvx: <iree.runtime.binding.HalDriver object at 0x7fb8d992d070>\n",
- "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7fb8d992d070>\n"
- ],
- "name": "stderr"
- }
- ]
+ "outputs": []
},
{
"cell_type": "code",
@@ -345,7 +346,7 @@
"base_uri": "https://localhost:8080/"
},
"id": "CQffg1iQatkb",
- "outputId": "689e2314-6523-470d-9faf-93acc4362e56"
+ "outputId": "f5daa64a-50d2-4ecf-94cc-b609ff09ddad"
},
"source": [
"# Our @tf.functions are accessible by name on the module named 'module'\n",
@@ -367,14 +368,14 @@
"outputs": [
{
"output_type": "stream",
+ "name": "stdout",
"text": [
"0\n",
"101\n",
"121\n",
"71\n",
"0\n"
- ],
- "name": "stdout"
+ ]
}
]
},
@@ -390,7 +391,12 @@
{
"cell_type": "code",
"metadata": {
- "id": "bUaNUkS2ohRj"
+ "id": "bUaNUkS2ohRj",
+ "outputId": "afe4caff-18b2-4f62-b48e-2633813ffa98",
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 86
+ }
},
"source": [
"ARTIFACTS_ZIP = \"/tmp/variables_and_state_colab_artifacts.zip\"\n",
@@ -406,8 +412,83 @@
"except ImportError:\n",
" print(\"Missing google_colab Python package, can't download files\")"
],
- "execution_count": null,
- "outputs": []
+ "execution_count": 11,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Zipping '/tmp/iree/colab_artifacts' to '/tmp/variables_and_state_colab_artifacts.zip' for download...\n",
+ " adding: counter_vmvx.vmfb (deflated 61%)\n",
+ " adding: counter.mlir (deflated 82%)\n",
+ "Downloading the artifacts zip file...\n"
+ ]
+ },
+ {
+ "output_type": "display_data",
+ "data": {
+ "application/javascript": [
+ "\n",
+ " async function download(id, filename, size) {\n",
+ " if (!google.colab.kernel.accessAllowed) {\n",
+ " return;\n",
+ " }\n",
+ " const div = document.createElement('div');\n",
+ " const label = document.createElement('label');\n",
+ " label.textContent = `Downloading \"${filename}\": `;\n",
+ " div.appendChild(label);\n",
+ " const progress = document.createElement('progress');\n",
+ " progress.max = size;\n",
+ " div.appendChild(progress);\n",
+ " document.body.appendChild(div);\n",
+ "\n",
+ " const buffers = [];\n",
+ " let downloaded = 0;\n",
+ "\n",
+ " const channel = await google.colab.kernel.comms.open(id);\n",
+ " // Send a message to notify the kernel that we're ready.\n",
+ " channel.send({})\n",
+ "\n",
+ " for await (const message of channel.messages) {\n",
+ " // Send a message to notify the kernel that we're ready.\n",
+ " channel.send({})\n",
+ " if (message.buffers) {\n",
+ " for (const buffer of message.buffers) {\n",
+ " buffers.push(buffer);\n",
+ " downloaded += buffer.byteLength;\n",
+ " progress.value = downloaded;\n",
+ " }\n",
+ " }\n",
+ " }\n",
+ " const blob = new Blob(buffers, {type: 'application/binary'});\n",
+ " const a = document.createElement('a');\n",
+ " a.href = window.URL.createObjectURL(blob);\n",
+ " a.download = filename;\n",
+ " div.appendChild(a);\n",
+ " a.click();\n",
+ " div.remove();\n",
+ " }\n",
+ " "
+ ],
+ "text/plain": [
+ "<IPython.core.display.Javascript object>"
+ ]
+ },
+ "metadata": {}
+ },
+ {
+ "output_type": "display_data",
+ "data": {
+ "application/javascript": [
+ "download(\"download_52327d87-b758-4770-b3cf-8a4af0381082\", \"variables_and_state_colab_artifacts.zip\", 4007)"
+ ],
+ "text/plain": [
+ "<IPython.core.display.Javascript object>"
+ ]
+ },
+ "metadata": {}
+ }
+ ]
}
]
-}
+}
\ No newline at end of file
