# coding: utf-8

# # Loopy: Controlling data layout
# 
# ## Setup code

# In[1]:

import numpy as np
import pyopencl as cl
import pyopencl.array
import pyopencl.clrandom
import loopy as lp


# In[2]:

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)


# ## A kernel on a structured array

# In[6]:

knl = lp.make_kernel(
    "{[el,dof, comp]: "
        "0<=el<nels "
        "and 0<=dof<14 "
        "and 0<=comp < 3}",
    "D[el, dof, comp] = eps[el] * E[el, dof, comp]")

knl = lp.set_options(knl, write_cl=True)


# In[7]:

eps = np.random.randn(500)
E = cl.clrandom.rand(queue, (500, 14, 3), dtype=np.float64)


# In[8]:

mknl = knl.copy()
evt, _ = mknl(queue, eps=eps, E=E)


# ## Changing the layout
# 
# `E` and `D` are currently laid out as AoS. What if I want SoA?

# In[9]:

mknl = knl

mknl  = lp.tag_data_axes(mknl, "E", "c,c,sep")
mknl = lp.tag_inames(mknl, {"comp": "unr"})
mknl = lp.set_loop_priority(mknl, "el,dof,comp")

# change data format of E
copy_knl = lp.make_copy_kernel("c,c,sep")
copy_knl = lp.fix_parameters(copy_knl, n2=3)
evt, E_new = copy_knl(queue, input=E)

evt, _ = mknl(queue, eps=eps, E=E_new)


# May want to add padding (demo).
# 
# ---
# 
# Grouped padding exists as well.