<jupyter_start><jupyter_text>Code generation with `pystella` and `loopy`<jupyter_code>import numpy as np
import pyopencl as cl
import pyopencl.array as cla
import pyopencl.clrandom as clr
import loopy as lp
from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2<jupyter_output><empty_output><jupyter_text>Four ways to an OpenCL kernelWe're going to create (and run!) an OpenCL kernel that computes$$a(\mathbf{x}) = b(\mathbf{x})^2 \cdot c(\mathbf{x}) + z$$in four different ways.First, we'll generate data and expected results with `numpy`.<jupyter_code>n = 64  # the grid size in each dimension

b_h = np.random.rand(n, n, n).astype(np.float64)
c_h = np.random.rand(n, n, n).astype(np.float64)
z = np.array(3.2)

a_true_h = b_h**2 * c_h + z<jupyter_output><empty_output><jupyter_text>1. `pyopencl` arrays methods First, we need an OpenCL "context" (the umbrella construct for running programs with OpenCL) and a "queue" (to which kernels will be submitted to execute on a device).Check out `pyopencl`'s [docs](https://documen.tician.de/pyopencl/) for examples and details.<jupyter_code>ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)<jupyter_output><empty_output><jupyter_text>`pyopencl` has a very convenient `Array` construct, which emulates `numpy` arrays---but with memory residing on the device.We'll copy the data to the device and try it out.<jupyter_code>b = cla.to_device(queue, b_h)
c = cla.to_device(queue, c_h)
a = cla.zeros_like(b)
a_true = cla.to_device(queue, a_true_h)
a[:] = b**2 * c + z<jupyter_output><empty_output><jupyter_text>To compare results, compute the maximum of `a - a_true`:<jupyter_code>difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>Note that we had to call `difference.get()`, which returns a `numpy.ndarray` on the "host" (the CPU) with data copied from `difference` (on the GPU).We can also use `pyopencl`'s `max` method:<jupyter_code>cla.max(difference)<jupyter_output><empty_output><jupyter_text>2. OpenCL kernel generation with `loopy` Refer to `loopy`'s [tutorial](https://documen.tician.de/loopy/tutorial.html) to get started.Let's create a kernel which computes the above for $i \in [0, N_x)$, $j \in [0, N_y)$, and $k \in [0, N_z)$:<jupyter_code>knl = lp.make_kernel(
    "{[i, j, k]: 0 <= i < Nx and 0 <= j < Ny and 0 <= k < Nz}",
    """
    a[i, j, k] = b[i, j, k]**2 * c[i, j, k] + z
    """
)<jupyter_output><empty_output><jupyter_text>Inspect your kernel to see if it appears correct by printing it.How did `make_kernel` interpret the un-indexed scalar variable `z`?<jupyter_code>print(knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: <auto/runtime>
Ny: ValueArg, type: <auto/runtime>
Nz: ValueArg, type: <auto/runtime>
a: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
b: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
c: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
z: ValueArg, type: <auto/runtime>
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [i, j, k] : 0 <= i < Nx and 0 <= j < Ny and 0 <= k < Nz }
---------------------------------------------------------------------------
INAME IMPLEMENTATION TAGS:
i: None
j: None
k: None
-----------------------------------------[...]<jupyter_text>We can now test our kernel by directly calling the `knl` we created above:<jupyter_code>evt, _ = knl(queue, a=a, b=b, c=c, z=z)<jupyter_output><empty_output><jupyter_text>Note that `z` needs to be a `numpy.array` so that `loopy` can infer its datatype. To compare results, compute the maximum of `a - a_true`:<jupyter_code>difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>Parallelization Now, GPUs are parallel, and the kernel we just wrote isn't making use of any parallelism.First, let's check what OpenCL code was produced by setting the kernel option `write_cl` to `True`:<jupyter_code>knl = lp.set_options(knl, write_cl=True)<jupyter_output><empty_output><jupyter_text>If we run the kernel now, it will print OpenCL code:<jupyter_code>evt, _ = knl(queue, a=a, b=b, c=c, z=z)<jupyter_output>[36m#[39;49;00m[36mdefine lid(N) ((int) get_local_id(N))[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mdefine gid(N) ((int) get_group_id(N))[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mif __OPENCL_C_VERSION__ < 120[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mpragma OPENCL EXTENSION cl_khr_fp64: enable[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mendif[39;49;00m[36m[39;49;00m

__kernel [36mvoid[39;49;00m [32m__attribute__[39;49;00m ((reqd_work_group_size([34m1[39;49;00m, [34m1[39;49;00m, [34m1[39;49;00m))) loopy_kernel([36mint[39;49;00m [34mconst[39;49;00m Nx, [36mint[39;49;00m [34mconst[39;49;00m Ny, [36mint[39;49;00m [34mconst[39;49;00m Nz, __global [36mdouble[39;49;00m *__restrict__ a, __global [36mdouble[39;49;00m [34mconst[39;49;00m *__restrict__ b, __global [36mdouble[39;49;00m [34mconst[39;49;00m *__restrict__ c, [36mdouble[39;49;00m [34mconst[39;49;00m z)
{
  [34mfor[39;49;00m ([36mint[39;49;00m k = [34m0[39;49;00m[...]<jupyter_text>It looks correct, array indexing and all.But that's a lot of sequential loops!`loopy` enables *code transformations* that (aim to) optimize the performance of a given kernel.For instance, mapping the `k` index to the "0" index of the local and global OpenCL thread dimensions is accomplished via `loopy.split_iname`:<jupyter_code>knl = lp.split_iname(knl, "k", 32, outer_tag="g.0", inner_tag="l.0")<jupyter_output><empty_output><jupyter_text>Let's see what this did to the kernel:<jupyter_code>print(knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: <auto/runtime>
Ny: ValueArg, type: <auto/runtime>
Nz: ValueArg, type: <auto/runtime>
a: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
b: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
c: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
z: ValueArg, type: <auto/runtime>
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [i, j, k_outer, k_inner] : 0 <= i < Nx and 0 <= j < Ny and k_inner >= 0 and -32k_outer <= k_inner <= 31 and k_inner < Nz - 32k_outer }
---------------------------------------------------------------------------
INAME IMPLEMENT[...]<jupyter_text>The above splits the loop over the "iname" `k` into (a yet-undetermined number of) blocks of 32 threads each. The "iname" (index name) `k` is gone, repalced by the combination `k_inner + k_outer * 32`.Observe also that the "implementation" of these new inames has been tagged to map to axes of global and local parallelization (as we specified).If we run the kernel now (enabling `write_cl` again), we see that the sequential loop over `k` is gone, and the indexing of `k` has been replaced by `32 * gid(0) + lid(0)`.<jupyter_code>knl = lp.set_options(knl, write_cl=True)
evt, _ = knl(queue, a=a, b=b, c=c, z=z)<jupyter_output>[36m#[39;49;00m[36mdefine lid(N) ((int) get_local_id(N))[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mdefine gid(N) ((int) get_group_id(N))[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mif __OPENCL_C_VERSION__ < 120[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mpragma OPENCL EXTENSION cl_khr_fp64: enable[39;49;00m[36m[39;49;00m
[36m#[39;49;00m[36mendif[39;49;00m[36m[39;49;00m

__kernel [36mvoid[39;49;00m [32m__attribute__[39;49;00m ((reqd_work_group_size([34m32[39;49;00m, [34m1[39;49;00m, [34m1[39;49;00m))) loopy_kernel([36mint[39;49;00m [34mconst[39;49;00m Nx, [36mint[39;49;00m [34mconst[39;49;00m Ny, [36mint[39;49;00m [34mconst[39;49;00m Nz, __global [36mdouble[39;49;00m *__restrict__ a, __global [36mdouble[39;49;00m [34mconst[39;49;00m *__restrict__ b, __global [36mdouble[39;49;00m [34mconst[39;49;00m *__restrict__ c, [36mdouble[39;49;00m [34mconst[39;49;00m z)
{
  [34mif[39;49;00m (-[34m1[39;49;00m + Nx >= [34m0[39;49;[...]<jupyter_text>Note that no loops over `k_inner` nor `k_outer` appear. They have been mapped to the "hardware" axes of parallelization: the kernel implicitly runs over a bunch of work groups (one for each value of `gid(0)`), each with 32 work items (each with their own index `lid(0)`). And the result is still correct!<jupyter_code>difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>We can achieve more parallelism by "tagging" `j` and `i` as, say, global indices 1 and 2.<jupyter_code>knl = lp.tag_inames(knl, {'j': 'g.1', 'i': 'g.2'})
knl = lp.set_options(knl, write_cl=True)
evt, _ = knl(queue, a=a, b=b, c=c, z=z)
difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>Observe the pesky `if` statement, which is ensuring that no out-of-bounds array elements are accessed by the replacement of `k` with `32 * gid(0) + lid(0)`.If we are *sure* that this won't happen (namely, that `Nz` is divisble by 32), we can add this as an assumption:<jupyter_code>knl = lp.assume(knl, 'Nz mod 32 = 0')
knl = lp.set_options(knl, write_cl=True)
evt, (x,) = knl(queue, a=a, b=b, c=c, z=z)
difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>Symbolic representation of code First, let's inspect the actual instruction `knl` is executing as represented by `loopy` kernel objects:<jupyter_code>knl.instructions<jupyter_output><empty_output><jupyter_text>This has quite a lot of details - but we're interested in the "assignee" and the "expression" of the first (and only) instruction:<jupyter_code>knl.instructions[0].assignee
knl.instructions[0].expression<jupyter_output><empty_output><jupyter_text>This is an *expression tree*. We can actually see what's going on if we `print(assignee, '=', expression)`:<jupyter_code>print(knl.instructions[0].assignee, '=', knl.instructions[0].expression)<jupyter_output>a[i, j, k_inner + k_outer*32] = b[i, j, k_inner + k_outer*32]**2*c[i, j, k_inner + k_outer*32] + z<jupyter_text>This is exactly the statement that appears if we print `knl` itself:<jupyter_code>print(knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: <auto/runtime>
Ny: ValueArg, type: <auto/runtime>
Nz: ValueArg, type: <auto/runtime>
a: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
b: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
c: type: <auto/runtime>, shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1) aspace: global
z: ValueArg, type: <auto/runtime>
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [i, j, k_outer, k_inner] : 0 <= i < Nx and 0 <= j < Ny and k_inner >= 0 and -32k_outer <= k_inner <= 31 and k_inner < Nz - 32k_outer }
---------------------------------------------------------------------------
INAME IMPLEMENT[...]<jupyter_text>Let's reproduce this instruction using `pymbolic`. First we need to import all the "primitive" objects:<jupyter_code>import pymbolic.primitives as p<jupyter_output><empty_output><jupyter_text>Now create some named "Variables" (prepending their name with an underscore so they don't overwrite our `pyopencl` arrays):<jupyter_code>_a = p.Variable('a')
_b = p.Variable('b')
_c = p.Variable('c')
_z = p.Variable('z')<jupyter_output><empty_output><jupyter_text>We also need some index variables:<jupyter_code>i = p.Variable('i')
j = p.Variable('j')
k = p.Variable('k')<jupyter_output><empty_output><jupyter_text>If we index (or "subscript") a `Variable`, we get a `Subscript` object:<jupyter_code>_a[i, j, k]<jupyter_output><empty_output><jupyter_text>This matches the `assignee` of the instruction above. Let's try the `expression`:<jupyter_code>_b[i, j, k]**2 * _c[i, j, k] + _z
print(_b[i, j, k]**2 * _c[i, j, k] + _z)<jupyter_output>b[i, j, k]**2*c[i, j, k] + z<jupyter_text>Looks good. The lesson here is that `pymbolic` provides a symbolic way to generate code.Rather than inputting strings of instructions to `loopy.make_kernel` (which, as we saw above, are parsed to `pymbolic` expressions by `loopy` behind the scenes!), we can work with the symbolic code directly.This unlocks a lot of potential to actually use python as a scripting language to generate the code (which `loopy` uses to subsequently generate OpenCL code).`pymbolic` can be thought of as a very simple computer algebra system (it can take derivatives, for instance), but geared toward manipulating and generating code.  3. `pystella.ElementWiseMap` Let's see how `pystella` provides a simpler interface to `loopy` to turn `pymbolic` expressions into kernels.Instructions are represented by lists of pairs of assignees (the left-hand side) and expressions (the right-hand side).These can be input as dictionaries (with keys and values corresponding to assignees and expressions, respectively) or lists of tuples. First, we'll recreate our same kernel again.<jupyter_code>import pystella as ps
map_instructions = {
    _a[i, j, k]: _b[i, j, k]**2 * _c[i, j, k] + _z
}
ewmap = ps.ElementWiseMap(map_instructions, dtype='float64', halo_shape=0)
print(ewmap.knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: np:dtype('int64')
Ny: ValueArg, type: np:dtype('int64')
Nz: ValueArg, type: np:dtype('int64')
a: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
b: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
c: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
z: ValueArg, type: np:dtype('float64')
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [k_outer, k_inner, j_outer, j_inner, i_outer, i_inner] : i_inner = 0 an[...]<jupyter_text>It's the same kernel! Already parallelized---`ElementWiseMap` implements a default parallelization that works well for these types of operations. Let's check the results:<jupyter_code>evt, _ = ewmap(queue, a=a, b=b, c=c, z=z)
difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>4. Using `pystella.Field`'s as input to `pystella.ElementWiseMap` `pystella.Field`'s can make our life even easier.Constantly indexing with `[i, j, k]` can get pretty annoying, and can be automate with `pymbolic`'s mapping methods.<jupyter_code>_a = ps.Field('a')
_b = ps.Field('b')
_c = ps.Field('c')
_z = p.Variable('z')

map_instructions = [
    (_a, _b**2 * _c + _z)
]

ewmap = ps.ElementWiseMap(map_instructions, dtype='float64', halo_shape=0)
print(ewmap.knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: np:dtype('int64')
Ny: ValueArg, type: np:dtype('int64')
Nz: ValueArg, type: np:dtype('int64')
a: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
b: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
c: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
z: ValueArg, type: np:dtype('float64')
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [k_outer, k_inner, j_outer, j_inner, i_outer, i_inner] : i_inner = 0 an[...]<jupyter_text>We made no mention of indices or subscripts, yet the kernels are identical.<jupyter_code>evt, _ = ewmap(queue, a=a, b=b, c=c, z=z)
difference = a - a_true
np.max(difference.get())<jupyter_output><empty_output><jupyter_text>To further illustrate why `pystella.Field`'s are useful, consider the (extremely common) case where arrays are padded in each direction.This is implemented by passing a value for `offset`.<jupyter_code>_a = ps.Field('a', offset='h')
_b = ps.Field('b', offset='h')
_c = ps.Field('c', offset='h')
_z = p.Variable('z')

map_instructions = {
    _a: _b**2 * _c + _z
}

ewmap = ps.ElementWiseMap(map_instructions, dtype='float64', halo_shape=1)
print(ewmap.knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: np:dtype('int64')
Ny: ValueArg, type: np:dtype('int64')
Nz: ValueArg, type: np:dtype('int64')
a: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
b: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
c: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
z: ValueArg, type: np:dtype('float64')
---------------------------------------------------------------------------
DOMAINS:
[Nx, [...]<jupyter_text>This would get quite cumbersome to type out manually, and it's easy to forget which arrays should be padded.From experience, it can be difficult to see errors in array indexing.  Indexer Behind the scenes, `ElementWiseMap` is calling `index_fields`:<jupyter_code>_f = ps.Field('f', offset=('h', 3, 2))
print(_f)
print(ps.index_fields(_f))<jupyter_output>f
f[i + h, j + 3, k + 2]<jupyter_text>Exercise: computing spatial gradients with `pystella.Field` and the `pystella.Stencil` kernel generator `Field`'s indices may be shifted by `pystella.field.shift_fields`, which does what it sounds like:<jupyter_code>from pystella.field import shift_fields
print(ps.index_fields(shift_fields(_a, (1, 0, 0))))<jupyter_output>a[i + h + 1, j + h, k + h]<jupyter_text>For this type of kernel, `pystella.Stencil` provides good parallelization (by allowing arrays to be *prefetched* into so-called "shared" memory).<jupyter_code>f = ps.Field('f', offset=1)

dfdx = ps.Field('dfdx', offset=0)
dfdy = ps.Field('dfdy', offset=0)
dfdz = ps.Field('dfdz', offset=0)

dx = p.Variable('dx')<jupyter_output><empty_output><jupyter_text>Fill in `map_instructions` below to compute the second-order centered-difference approximation to the gradient of `f`:<jupyter_code>map_instructions = {
    dfdx: (shift_fields(f, (1, 0, 0)) - shift_fields(f, (-1, 0, 0))) / 2 / dx,
    dfdy: (shift_fields(f, (0, 1, 0)) - shift_fields(f, (0, -1, 0))) / 2 / dx,
    dfdz: (shift_fields(f, (0, 0, 1)) - shift_fields(f, (0, 0, -1))) / 2 / dx,
}
stencil = ps.Stencil(map_instructions, prefetch_args=['f'], halo_shape=1, dtype='float64')
print(stencil.knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: np:dtype('int64')
Ny: ValueArg, type: np:dtype('int64')
Nz: ValueArg, type: np:dtype('int64')
dfdx: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
dfdy: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
dfdz: type: np:dtype('float64'), shape: (Nx, Ny, Nz), dim_tags: (N2:stride:Nz*Ny, N1:stride:Nz, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
dx: ValueArg, type: np:dtype('float64')
f: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.[...]<jupyter_text>We can also do something more complicated by inputting `tmp_instructions`, which computes temporary values (that don't get stored in global arrays) before executing the assignments specified by `map_instructions`:<jupyter_code>f = ps.Field('f', offset=1)
g = ps.Field('g', offset=1)
tmp = p.Variable('tmp')
tmp_instructions = {}
for i in range(3):
    shift = [0, 0, 0]
    shift[i] = 1
    expr = shift_fields(f, tuple(shift))
    shift[i] = - 1
    expr += shift_fields(f, tuple(shift))
    tmp_instructions[tmp[i]] = expr<jupyter_output><empty_output><jupyter_text>Let's check what we just did:<jupyter_code>for key, value in tmp_instructions.items():
    print(key, '=', ps.index_fields(value))
map_instructions = {
    g: tmp[0] * tmp[1] * tmp[2]
}
stencil = ps.Stencil(map_instructions, tmp_instructions=tmp_instructions, prefetch_args=['f'], halo_shape=1, dtype='float64')

print(stencil.knl)<jupyter_output>---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
Nx: ValueArg, type: np:dtype('int64')
Ny: ValueArg, type: np:dtype('int64')
Nz: ValueArg, type: np:dtype('int64')
f: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
g: type: np:dtype('float64'), shape: (Nx + 2, Ny + 2, Nz + 2), dim_tags: (N2:stride:(Nz + 2)*(Ny + 2), N1:stride:Nz + 2, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
---------------------------------------------------------------------------
DOMAINS:
[Nx, Ny, Nz] -> { [k_outer, k_inner, j_outer, j_inner, i_outer, i_inner, f_dim_0, f_dim_1, f_dim_2] : k_outer >= 0 and 0 <= k_inner <= 7 and k_inner < Nz - 8k_outer and j_outer >= 0 and 0 <= j_inner <= 7 and j_inner < Ny - 8j_outer and[...]
{'fraction_non_alphanumeric': 0.236776253343688, 'fraction_numerical': 0.04146173095176461, 'mean_word_length': 3.7459049959049957, 'pattern_counts': {'":': 0, '<': 209, '<?xml version=': 0, '>': 185, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '38', 'dir': 'jupyter-structured-clean-dedup', 'id': '477456', 'license': 'MIT', 'n_tokens_mistral': 9224, 'n_tokens_neox': 8511, 'n_words': 2198, 'path': 'examples/codegen-tutorial.ipynb', 'repo_name': 'zachjweiner/pystella', 'size': '23178'}