{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Loopy: Controlling data layout\n",
        "\n",
        "## Setup code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import pyopencl as cl\n",
        "import pyopencl.array\n",
        "import pyopencl.clrandom\n",
        "import loopy as lp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ctx = cl.create_some_context()\n",
        "queue = cl.CommandQueue(ctx)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## A kernel on a structured array"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "knl = lp.make_kernel(\n",
        "    \"{[el,dof, comp]: \"\n",
        "        \"0<=el<nels \"\n",
        "        \"and 0<=dof<14 \"\n",
        "        \"and 0<=comp < 3}\",\n",
        "    \"D[el, dof, comp] = eps[el] * E[el, dof, comp]\")\n",
        "\n",
        "knl = lp.set_options(knl, write_cl=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "eps = np.random.randn(500)\n",
        "E = cl.clrandom.rand(queue, (500, 14, 3), dtype=np.float64)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "#if __OPENCL_C_VERSION__ < 120\n",
            "#pragma OPENCL EXTENSION cl_khr_fp64: enable\n",
            "#endif\n",
            "#define lid(N) ((int) get_local_id(N))\n",
            "#define gid(N) ((int) get_group_id(N))\n",
            "\n",
            "__kernel void __attribute__ ((reqd_work_group_size(1, 1, 1))) loopy_kernel(__global double *restrict D, __global double const *restrict E, __global double const *restrict eps, int const nels)\n",
            "{\n",
            "\n",
            "  for (int comp = 0; comp <= 2; ++comp)\n",
            "    if (-1 + nels >= 0)\n",
            "      for (int dof = 0; dof <= 13; ++dof)\n",
            "        for (int el = 0; el <= -1 + nels; ++el)\n",
            "          D[42 * el + 3 * dof + comp] = eps[el] * E[42 * el + 3 * dof + comp];\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "mknl = knl.copy()\n",
        "evt, _ = mknl(queue, eps=eps, E=E)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Changing the layout\n",
        "\n",
        "`E` and `D` are currently laid out as AoS. What if I want SoA?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "#if __OPENCL_C_VERSION__ < 120\n",
            "#pragma OPENCL EXTENSION cl_khr_fp64: enable\n",
            "#endif\n",
            "#define lid(N) ((int) get_local_id(N))\n",
            "#define gid(N) ((int) get_group_id(N))\n",
            "\n",
            "__kernel void __attribute__ ((reqd_work_group_size(1, 1, 1))) loopy_kernel(__global double *restrict D, __global double const *restrict E_s0, __global double const *restrict E_s1, __global double const *restrict E_s2, __global double const *restrict eps, int const nels)\n",
            "{\n",
            "\n",
            "  for (int el = 0; el <= -1 + nels; ++el)\n",
            "    for (int dof = 0; dof <= 13; ++dof)\n",
            "    {\n",
            "      D[42 * el + 3 * dof + 0] = eps[el] * E_s0[14 * el + dof];\n",
            "      D[42 * el + 3 * dof + 1] = eps[el] * E_s1[14 * el + dof];\n",
            "      D[42 * el + 3 * dof + 2] = eps[el] * E_s2[14 * el + dof];\n",
            "    }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "mknl = knl\n",
        "\n",
        "mknl  = lp.tag_data_axes(mknl, \"E\", \"c,c,sep\")\n",
        "mknl = lp.tag_inames(mknl, {\"comp\": \"unr\"})\n",
        "mknl = lp.set_loop_priority(mknl, \"el,dof,comp\")\n",
        "\n",
        "# change data format of E\n",
        "copy_knl = lp.make_copy_kernel(\"c,c,sep\")\n",
        "copy_knl = lp.fix_parameters(copy_knl, n2=3)\n",
        "evt, E_new = copy_knl(queue, input=E)\n",
        "\n",
        "evt, _ = mknl(queue, eps=eps, E=E_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "May want to add padding (demo).\n",
        "\n",
        "---\n",
        "\n",
        "Grouped padding exists as well."
      ]
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 0
}