samples/colab/low_level_invoke_function.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2019 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
        "# See https://llvm.org/LICENSE.txt for license information.\n",
        "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Low Level Invoke Function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This notebook shows off some concepts of the low level IREE python bindings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python -m pip install --pre iree-base-compiler iree-base-runtime -f https://iree.dev/pip-release-links.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "from iree import runtime as ireert\n",
        "from iree.compiler import compile_str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Compile a module.\n",
        "SIMPLE_MUL_ASM = \"\"\"\n",
        "  module @arithmetic {\n",
        "    func.func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {\n",
        "      %0 = arith.mulf %arg0, %arg1 : tensor<4xf32>\n",
        "      return %0 : tensor<4xf32>\n",
        "    } \n",
        "  }\n",
        "\"\"\"\n",
        "\n",
        "# Compile using the vmvx (reference) target:\n",
        "compiled_flatbuffer = compile_str(SIMPLE_MUL_ASM, target_backends=[\"vmvx\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Register the module with a runtime context.\n",
        "# Use the \"local-task\" CPU driver, which can load the vmvx executable:\n",
        "config = ireert.Config(\"local-task\")\n",
        "ctx = ireert.SystemContext(config=config)\n",
        "vm_module = ireert.VmModule.from_flatbuffer(ctx.instance, compiled_flatbuffer)\n",
        "ctx.add_vm_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).to_host()\n",
        "print(\"Results:\", results)"
      ]
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 0
}