colab/low_level_invoke_function.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "low_level_invoke_function.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uMVh8_lZDRa7"
      },
      "source": [
        "This notebook shows off some concepts of the low level IREE python bindings."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Go2Nw7BgIHYU",
        "outputId": "dbafcd78-3cc4-4fa0-f408-1b8e6270084c",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "!python -m pip install iree-compiler-snapshot iree-runtime-snapshot -f https://github.com/google/iree/releases"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Looking in links: https://github.com/google/iree/releases\n",
            "Collecting iree-compiler-snapshot\n",
            "\u001b[?25l  Downloading https://github.com/google/iree/releases/download/snapshot-20210105.14/iree_compiler_snapshot-20210105.14-py3-none-manylinux2014_x86_64.whl (32.1MB)\n",
            "\u001b[K     |████████████████████████████████| 32.1MB 155kB/s \n",
            "\u001b[?25hCollecting iree-runtime-snapshot\n",
            "\u001b[?25l  Downloading https://github.com/google/iree/releases/download/snapshot-20210105.14/iree_runtime_snapshot-20210105.14-cp36-cp36m-manylinux2014_x86_64.whl (1.0MB)\n",
            "\u001b[K     |████████████████████████████████| 1.0MB 46.9MB/s \n",
            "\u001b[?25hInstalling collected packages: iree-compiler-snapshot, iree-runtime-snapshot\n",
            "Successfully installed iree-compiler-snapshot-20210105.14 iree-runtime-snapshot-20210105.14\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1F144M4wAFPz"
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "from iree import runtime as ireert\n",
        "from iree.compiler import compile_str"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2Rq-JdzMAFPU"
      },
      "source": [
        "# Compile a module.\n",
        "SIMPLE_MUL_ASM = \"\"\"\n",
        "  module @arithmetic {\n",
        "    func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>\n",
        "          attributes { iree.module.export } {\n",
        "        %0 = \"mhlo.multiply\"(%arg0, %arg1) {name = \"mul.1\"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>\n",
        "        return %0 : tensor<4xf32>\n",
        "    } \n",
        "  }\n",
        "\"\"\"\n",
        "\n",
        "compiled_flatbuffer = compile_str(SIMPLE_MUL_ASM, target_backends=[\"vmla\"])\n",
        "vm_module = ireert.VmModule.from_flatbuffer(compiled_flatbuffer)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TNQiNeOU_cpK",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "42df9878-4af7-458d-b23f-92791873f9ab"
      },
      "source": [
        "# Register the module with a runtime context.\n",
        "# Use the CPU interpreter (which has the most implementation done):\n",
        "config = ireert.Config(\"vmla\")\n",
        "ctx = ireert.SystemContext(config=config)\n",
        "ctx.add_module(vm_module)\n",
        "\n",
        "# Invoke the function and print the result.\n",
        "print(\"INVOKE simple_mul\")\n",
        "arg0 = np.array([1., 2., 3., 4.], dtype=np.float32)\n",
        "arg1 = np.array([4., 5., 6., 7.], dtype=np.float32)\n",
        "f = ctx.modules.arithmetic[\"simple_mul\"]\n",
        "results = f(arg0, arg1)\n",
        "print(\"Results:\", results)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INVOKE simple_mul\n",
            "Results: [ 4. 10. 18. 28.]\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Created IREE driver vmla: <iree.runtime.binding.HalDriver object at 0x7f6201491458>\n",
            "SystemContext driver=<iree.runtime.binding.HalDriver object at 0x7f6201491458>\n"
          ],
          "name": "stderr"
        }
      ]
    }
  ]
}