BifrostDemo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tayN_DCBZgmn"
      },
      "source": [
        "\n",
        "## Instructions\n",
        "1. Work on a copy of this notebook: _File_ > _Save a copy in Drive_ (you will need a Google account). Alternatively, you can download the notebook using _File_ > _Download .ipynb_, then upload it to [Colab](https://colab.research.google.com/).\n",
        "2. Turn on the GPU with Edit->Notebook settings->Hardware accelerator->GPU\n",
        "3. Execute the following cell (click on it and press Ctrl+Enter) to install dependencies.\n",
        "4. Continue to the next section.\n",
        "\n",
        "_Notes_:\n",
        "* If your Colab Runtime gets reset (e.g., due to inactivity), repeat steps 3, 4.\n",
        "* After installation, if you want to activate/deactivate the GPU, you will need to reset the Runtime: _Runtime_ > _Factory reset runtime_ and repeat steps 2-4."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "i_Ec54bF6mVj",
        "outputId": "79bec2be-3986-4f80-a450-6e5edcf9a8c4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": []
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# @title Install C++ deps\n",
        "%%shell\n",
        "\n",
        "sudo apt-get -qq install exuberant-ctags libopenblas-dev software-properties-common build-essential"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "45aDEOKc6JGW",
        "outputId": "cd058f85-f376-4915-e9f5-825a9b51ad45"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[?25l\r\u001b[K     |█▋                              | 10 kB 35.4 MB/s eta 0:00:01\r\u001b[K     |███▏                            | 20 kB 32.4 MB/s eta 0:00:01\r\u001b[K     |████▊                           | 30 kB 19.6 MB/s eta 0:00:01\r\u001b[K     |██████▎                         | 40 kB 16.2 MB/s eta 0:00:01\r\u001b[K     |███████▉                        | 51 kB 8.9 MB/s eta 0:00:01\r\u001b[K     |█████████▍                      | 61 kB 9.2 MB/s eta 0:00:01\r\u001b[K     |███████████                     | 71 kB 8.8 MB/s eta 0:00:01\r\u001b[K     |████████████▌                   | 81 kB 9.6 MB/s eta 0:00:01\r\u001b[K     |██████████████                  | 92 kB 10.1 MB/s eta 0:00:01\r\u001b[K     |███████████████▋                | 102 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |█████████████████▏              | 112 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |██████████████████▊             | 122 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |████████████████████▎           | 133 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |█████████████████████▉          | 143 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |███████████████████████▍        | 153 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████       | 163 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▋     | 174 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▏   | 184 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▊  | 194 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |███████████████████████████████▎| 204 kB 8.4 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 209 kB 8.4 MB/s \n",
            "\u001b[?25h  Building wheel for ctypesgen (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        },
        {
          "data": {
            "text/plain": []
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# @title Install python deps\n",
        "%%shell\n",
        "\n",
        "pip install -q contextlib2 pint simplejson ctypesgen==1.0.2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WQL4Psci6S92"
      },
      "outputs": [],
      "source": [
        "# @title Build and Install Bifrost\n",
        "%%shell\n",
        "cd \"${HOME}\"\n",
        "if [ -d \"${HOME}/bifrost_repo\" ]; then\n",
        "    echo \"Already cloned.\"\n",
        "else\n",
        "    git clone https://github.com/ledatelescope/bifrost bifrost_repo\n",
        "fi\n",
        "cd \"${HOME}/bifrost_repo\"\n",
        "git pull --all\n",
        "\n",
        "./configure\n",
        "\n",
        "# Build and install:\n",
        "make -j all\n",
        "sudo make install\n",
        "export LD_LIBRARY_PATH=/usr/local/lib:${LD_LIBRARY_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q8q8ZMix8kF6"
      },
      "source": [
        "Now, let's create and test a pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "YZOGJ-HhBAxg"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Environment path doesn't propagate, so add it manually:\n",
        "if \"/usr/local/lib\" not in os.environ['LD_LIBRARY_PATH']:\n",
        "    os.environ['LD_LIBRARY_PATH'] += \":/usr/local/lib\"\n",
        "\n",
        "import bifrost as bf\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i2oRp6B1drwd"
      },
      "source": [
        "Let's first create a simple CUDA kernel within Bifrost.\n",
        "\n",
        "We will generate 1000 integers, feed them into Bifrost as a CUDA array, perform a kernel operation `x * 3`, then copy them back."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "yxm87l5acwIt"
      },
      "outputs": [],
      "source": [
        "x = np.random.randint(256, size=1000)\n",
        "\n",
        "x_orig = x\n",
        "x = bf.asarray(x, 'cuda')\n",
        "y = bf.empty_like(x)\n",
        "x.flags['WRITEABLE'] = False\n",
        "x.bf.immutable = True\n",
        "for _ in range(3):\n",
        "    bf.map(\"y = x * 3\", {'x': x, 'y': y})\n",
        "x = x.copy('system')\n",
        "y = y.copy('system')\n",
        "if isinstance(x_orig, bf.ndarray):\n",
        "    x_orig = x\n",
        "np.testing.assert_equal(y, x_orig * 3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_jHnukaHdzUD"
      },
      "source": [
        "Now, let's generate a full pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "E36NHJDzdyi9"
      },
      "outputs": [],
      "source": [
        "from bifrost.block import Pipeline, NumpyBlock, NumpySourceBlock\n",
        "\n",
        "def generate_different_arrays():\n",
        "    \"\"\"Yield four different groups of two arrays\"\"\"\n",
        "    dtypes = ['float32', 'float64', 'complex64', 'int8']\n",
        "    shapes = [(4,), (4, 5), (4, 5, 6), (2,) * 8]\n",
        "    for array_index in range(4):\n",
        "        yield np.ones(\n",
        "            shape=shapes[array_index],\n",
        "            dtype=dtypes[array_index])\n",
        "        yield 2 * np.ones(\n",
        "            shape=shapes[array_index],\n",
        "            dtype=dtypes[array_index])\n",
        "\n",
        "def switch_types(array):\n",
        "    \"\"\"Return two copies of the array, one with a different type\"\"\"\n",
        "    return np.copy(array), np.copy(array).astype(np.complex128)\n",
        "\n",
        "occurences = 0\n",
        "def compare_arrays(array1, array2):\n",
        "    \"\"\"Make sure that all arrays coming in are equal\"\"\"\n",
        "    global occurences\n",
        "    occurences += 1\n",
        "    np.testing.assert_almost_equal(array1, array2)\n",
        "\n",
        "blocks = [\n",
        "    (NumpySourceBlock(generate_different_arrays), {'out_1': 0}),\n",
        "    (NumpyBlock(switch_types, outputs=2), {'in_1': 0, 'out_1': 1, 'out_2': 2}),\n",
        "    (NumpyBlock(np.fft.fft), {'in_1': 2, 'out_1': 3}),\n",
        "    (NumpyBlock(np.fft.ifft), {'in_1': 3, 'out_1': 4}),\n",
        "    (NumpyBlock(compare_arrays, inputs=2, outputs=0), {'in_1': 1, 'in_2': 4})]\n",
        "\n",
        "Pipeline(blocks).main()\n",
        "\n",
        "# The function `compare_arrays` should be hit 8 times:\n",
        "assert occurences == 8"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "BifrostDemo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}