Implement AA pattern for OpenCL target

Works well but function naming is getting kind of clunky, e.g. "velocity_momenta_boundary_tick_cells"

This could be hidden to a degree by proving branching wrappers for
the odd and even time step implementations. However this would not
vectorize when targeting Intel via OpenCL.

3 files changed, 189 insertions, 9 deletions

diff --git a/boltzgen/kernel/template/collect_moments.cl.mako b/boltzgen/kernel/template/collect_moments.cl.mako
index 2b5bf69..67b2d68 100644
--- a/boltzgen/kernel/template/collect_moments.cl.mako
+++ b/boltzgen/kernel/template/collect_moments.cl.mako
@@ -15,12 +15,3 @@ moments_subexpr, moments_assignment = model.moments()
     preshifted_m[${i}] = ${sympy.ccode(expr.rhs)};
 % endfor
 </%call>
-
-% if 'cell_list_dispatch' in extras:
-__kernel void collect_moments_cells(__global ${float_type}* f,
-                                    __global ${float_type}* m,
-                                    __global unsigned int*  cells)
-{
-    collect_moments(f, cells[get_global_id(0)], m);
-}
-% endif
diff --git a/boltzgen/kernel/template/pattern/AA.cl.mako b/boltzgen/kernel/template/pattern/AA.cl.mako
new file mode 100644
index 0000000..786b019
--- /dev/null
+++ b/boltzgen/kernel/template/pattern/AA.cl.mako
@@ -0,0 +1,167 @@
+<%def name="operator(name, params = None)">
+__kernel void ${name}_tick(
+      __global ${float_type}* f
+    , unsigned int gid
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i)}];
+% endfor
+
+    ${caller.body()}
+
+% for i, c_i in enumerate(descriptor.c):
+    preshifted_f[${layout.pop_offset(i)}] = f_next_${descriptor.c.index(-c_i)};
+% endfor
+}
+
+__kernel void ${name}_tock(
+      __global ${float_type}* f
+    , unsigned int gid
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${descriptor.c.index(-c_i)} = preshifted_f[${layout.pop_offset(i) + layout.neighbor_offset(c_i)}];
+% endfor
+
+    ${caller.body()}
+
+% for i, c_i in enumerate(descriptor.c):
+    preshifted_f[${layout.pop_offset(i) + layout.neighbor_offset(c_i)}] = f_next_${i};
+% endfor
+}
+
+% if 'cell_list_dispatch' in extras:
+__kernel void ${name}_tick_cells(
+      __global ${float_type}* f
+    , __global unsigned int* cells
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    ${name}_tick(
+          f
+        , cells[get_global_id(0)]
+% if params is not None:
+% for param_type, param_name in params:
+        , ${param_name}
+% endfor
+% endif
+    );
+}
+
+__kernel void ${name}_tock_cells(
+      __global ${float_type}* f
+    , __global unsigned int* cells
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    ${name}_tock(
+          f
+        , cells[get_global_id(0)]
+% if params is not None:
+% for param_type, param_name in params:
+        , ${param_name}
+% endfor
+% endif
+    );
+}
+% endif
+</%def>
+
+<%def name="functor(name, params = None)">
+__kernel void ${name}_tick(
+      __global ${float_type}* f
+    , unsigned int gid
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(descriptor.c.index(-c_i))}];
+% endfor
+
+    ${caller.body()}
+}
+
+__kernel void ${name}_tock(
+      __global ${float_type}* f
+    , unsigned int gid
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i) + layout.neighbor_offset(c_i)}];
+% endfor
+
+    ${caller.body()}
+}
+
+% if 'cell_list_dispatch' in extras:
+__kernel void ${name}_tick_cells(
+      __global ${float_type}* f
+    , __global unsigned int* cells
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    ${name}_tick(
+          f
+        , cells[get_global_id(0)]
+% if params is not None:
+% for param_type, param_name in params:
+        , ${param_name}
+% endfor
+% endif
+    );
+}
+
+__kernel void ${name}_tock_cells(
+      __global ${float_type}* f
+    , __global unsigned int* cells
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    ${name}_tock(
+          f
+        , cells[get_global_id(0)]
+% if params is not None:
+% for param_type, param_name in params:
+        , ${param_name}
+% endfor
+% endif
+    );
+}
+% endif
+</%def>
diff --git a/boltzgen/kernel/template/pattern/AB.cl.mako b/boltzgen/kernel/template/pattern/AB.cl.mako
index d2d8b8a..33836ea 100644
--- a/boltzgen/kernel/template/pattern/AB.cl.mako
+++ b/boltzgen/kernel/template/pattern/AB.cl.mako
@@ -66,4 +66,26 @@ __kernel void ${name}(
 
     ${caller.body()}
 }
+
+% if 'cell_list_dispatch' in extras:
+__kernel void ${name}_cells(
+      __global ${float_type}* f
+    , __global unsigned int* cells
+% if params is not None:
+% for param_type, param_name in params:
+    , ${param_type} ${param_name}
+% endfor
+% endif
+) {
+    ${name}(
+          f
+        , cells[get_global_id(0)]
+% if params is not None:
+% for param_type, param_name in params:
+        , ${param_name}
+% endfor
+% endif
+    );
+}
+% endif
 </%def>