Match OpenCL and CUDA cell list dispatch templates

4 files changed, 101 insertions, 99 deletions

diff --git a/boltzgen/kernel/template/collect_moments.cl.mako b/boltzgen/kernel/template/collect_moments.cl.mako
index 63387b7..e3ce166 100644
--- a/boltzgen/kernel/template/collect_moments.cl.mako
+++ b/boltzgen/kernel/template/collect_moments.cl.mako
@@ -16,6 +16,18 @@ moments_subexpr, moments_assignment = model.moments()
 % endfor
 </%call>
 
+<%call expr="pattern.functor_with_domain_dispatch('collect_moments_domain', [('__global %s*' % float_type, 'm')])">
+% for i, expr in enumerate(moments_subexpr):
+    const ${float_type} ${expr[0]} = ${sympy.ccode(expr[1])};
+% endfor
+
+    __global ${float_type}* preshifted_m = m + gid*${descriptor.d+1};
+
+% for i, expr in enumerate(moments_assignment):
+    preshifted_m[${i}] = ${sympy.ccode(expr.rhs)};
+% endfor
+</%call>
+
 % if 'opencl_gl_interop' in extras:
 <%call expr="pattern.functor_with_domain_dispatch('collect_moments_to_texture', [('__write_only %s' % {2: 'image2d_t', 3: 'image3d_t'}.get(descriptor.d), 'm')])">
 % for i, expr in enumerate(moments_subexpr):
diff --git a/boltzgen/kernel/template/equilibrilize.cl.mako b/boltzgen/kernel/template/equilibrilize.cl.mako
index 8a423ae..36ec48b 100644
--- a/boltzgen/kernel/template/equilibrilize.cl.mako
+++ b/boltzgen/kernel/template/equilibrilize.cl.mako
@@ -5,3 +5,9 @@
     const ${float_type} f_next_${i} = ${w_i.evalf()};
 % endfor
 </%call>
+
+<%call expr="pattern.operator_with_domain_dispatch('equilibrilize_domain')">
+% for i, w_i in enumerate(descriptor.w):
+    const ${float_type} f_next_${i} = ${w_i.evalf()};
+% endfor
+</%call>
diff --git a/boltzgen/kernel/template/pattern/AA.cl.mako b/boltzgen/kernel/template/pattern/AA.cl.mako
index 2a88755..74ad2fa 100644
--- a/boltzgen/kernel/template/pattern/AA.cl.mako
+++ b/boltzgen/kernel/template/pattern/AA.cl.mako
@@ -1,13 +1,20 @@
 <%def name="operator(name, params = None)">
 __kernel void ${name}_tick(
       __global ${float_type}* f
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i, c_i in enumerate(descriptor.c):
@@ -23,13 +30,20 @@ __kernel void ${name}_tick(
 
 __kernel void ${name}_tock(
       __global ${float_type}* f
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i, c_i in enumerate(descriptor.c):
@@ -42,60 +56,71 @@ __kernel void ${name}_tock(
     preshifted_f[${layout.pop_offset(i) + layout.neighbor_offset(c_i)}] = f_next_${i};
 % endfor
 }
+</%def>
 
-% if 'cell_list_dispatch' in extras:
-__kernel void ${name}_cells_tick(
+<%def name="operator_with_domain_dispatch(name, params = None)">
+__kernel void ${name}_tick(
       __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}
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)', 'get_global_id(2)')};
+    __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
-% endif
-    );
 }
 
-__kernel void ${name}_cells_tock(
+__kernel void ${name}_tock(
       __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}
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)', 'get_global_id(2)')};
+    __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
-% endif
-    );
 }
-% endif
 </%def>
 
 <%def name="functor(name, params = None)">
 __kernel void ${name}_tick(
       __global ${float_type}* f
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i, c_i in enumerate(descriptor.c):
@@ -107,13 +132,20 @@ __kernel void ${name}_tick(
 
 __kernel void ${name}_tock(
       __global ${float_type}* f
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i, c_i in enumerate(descriptor.c):
@@ -122,48 +154,6 @@ __kernel void ${name}_tock(
 
     ${caller.body()}
 }
-
-% if 'cell_list_dispatch' in extras:
-__kernel void ${name}_cells_tick(
-      __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}_cells_tock(
-      __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_with_domain_dispatch(name, params = None)">
diff --git a/boltzgen/kernel/template/pattern/AB.cl.mako b/boltzgen/kernel/template/pattern/AB.cl.mako
index 8e6421e..8fa941b 100644
--- a/boltzgen/kernel/template/pattern/AB.cl.mako
+++ b/boltzgen/kernel/template/pattern/AB.cl.mako
@@ -2,13 +2,20 @@
 __kernel void ${name}(
       __global ${float_type}* f_next
     , __global ${float_type}* f_prev
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
     __global ${float_type}* preshifted_f_prev = f_prev + ${layout.cell_preshift('gid')};
 
@@ -22,42 +29,51 @@ __kernel void ${name}(
     preshifted_f_next[${layout.pop_offset(i)}] = f_next_${i};
 % endfor
 }
+</%def>
 
-% if 'cell_list_dispatch' in extras:
-__kernel void ${name}_cells(
+<%def name="operator_with_domain_dispatch(name, params = None)">
+__kernel void ${name}(
       __global ${float_type}* f_next
     , __global ${float_type}* f_prev
-    , __global unsigned int* cells
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
-    ${name}(
-          f_next
-        , f_prev
-        , cells[get_global_id(0)]
-% if params is not None:
-% for param_type, param_name in params:
-        , ${param_name}
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)', 'get_global_id(2)')};
+    __global ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
+    __global ${float_type}* preshifted_f_prev = f_prev + ${layout.cell_preshift('gid')};
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${i} = preshifted_f_prev[${layout.pop_offset(i) + layout.neighbor_offset(-c_i)}];
+% endfor
+
+    ${caller.body()}
+
+% for i, _ in enumerate(descriptor.c):
+    preshifted_f_next[${layout.pop_offset(i)}] = f_next_${i};
 % endfor
-% endif
-    );
 }
-% endif
 </%def>
 
 <%def name="functor(name, params = None)">
 __kernel void ${name}(
       __global ${float_type}* f
+% if 'cell_list_dispatch' in extras:
+    , __global unsigned int* cells
+% else:
     , unsigned int gid
+% endif
 % if params is not None:
 % for param_type, param_name in params:
     , ${param_type} ${param_name}
 % endfor
 % endif
 ) {
+% if 'cell_list_dispatch' in extras:
+    const unsigned int gid = cells[get_global_id(0)];
+% endif
     __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i in range(0,descriptor.q):
@@ -66,28 +82,6 @@ __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>
 
 <%def name="functor_with_domain_dispatch(name, params = None)">