Unify AOS, SOA specific cell preshift between targets

SOA and AOS should not be target specific, neighbor offset calculation /
bijection between gid and cell coordinates should be customizable.

12 files changed, 32 insertions, 38 deletions

diff --git a/boltzgen/kernel/target/layout/cl.py b/boltzgen/kernel/target/layout/cl.py
index 9707cb3..e144faf 100644
--- a/boltzgen/kernel/target/layout/cl.py
+++ b/boltzgen/kernel/target/layout/cl.py
@@ -3,11 +3,8 @@ class SOA:
         self.descriptor = descriptor
         self.geometry = geometry
 
-    def gid(self):
-        return {
-            2: 'get_global_id(1)*%d + get_global_id(0)' % self.geometry.size_x,
-            3: 'get_global_id(2)*%d + get_global_id(1)*%d + get_global_id(0)' % (self.geometry.size_x*self.geometry.size_y, self.geometry.size_x)
-        }.get(self.descriptor.d)
+    def cell_preshift(self, gid):
+        return gid
 
     def pop_offset(self, i):
         return i * self.geometry.volume
@@ -23,11 +20,8 @@ class AOS:
         self.descriptor = descriptor
         self.geometry = geometry
 
-    def gid(self):
-        return {
-            2: '%d*(get_global_id(1)*%d + get_global_id(0))' % (self.descriptor.q, self.geometry.size_x),
-            3: '%d*(get_global_id(2)*%d + get_global_id(1)*%d + get_global_id(0))' % (self.descriptor.q, self.geometry.size_x*self.geometry.size_y, self.geometry.size_x)
-        }.get(self.descriptor.d)
+    def cell_preshift(self, gid):
+        return "(%s)*%d" % (gid, self.descriptor.q)
 
     def pop_offset(self, i):
         return i
diff --git a/boltzgen/kernel/target/layout/cpp.py b/boltzgen/kernel/target/layout/cpp.py
index bb5dba0..2a8c89e 100644
--- a/boltzgen/kernel/target/layout/cpp.py
+++ b/boltzgen/kernel/target/layout/cpp.py
@@ -1,33 +1,33 @@
-class AOS:
+class SOA:
     def __init__(self, descriptor, geometry):
         self.descriptor = descriptor
         self.geometry = geometry
 
-    def gid_offset(self):
-        return self.descriptor.q
+    def cell_preshift(self, gid):
+        return gid
 
     def pop_offset(self, i):
-        return i
+        return i * self.geometry.volume
 
     def neighbor_offset(self, c_i):
-        return self.descriptor.q * {
+        return {
             2: lambda:                                                    c_i[0]*self.geometry.size_y + c_i[1],
             3: lambda: c_i[0]*self.geometry.size_y*self.geometry.size_z + c_i[1]*self.geometry.size_z + c_i[2]
         }.get(self.descriptor.d)()
 
-class SOA:
+class AOS:
     def __init__(self, descriptor, geometry):
         self.descriptor = descriptor
         self.geometry = geometry
 
-    def gid_offset(self):
-        return 1
+    def cell_preshift(self, gid):
+        return "(%s)*%d" % (gid, self.descriptor.q)
 
     def pop_offset(self, i):
-        return i * self.geometry.volume
+        return i
 
     def neighbor_offset(self, c_i):
-        return {
+        return self.descriptor.q * {
             2: lambda:                                                    c_i[0]*self.geometry.size_y + c_i[1],
             3: lambda: c_i[0]*self.geometry.size_y*self.geometry.size_z + c_i[1]*self.geometry.size_z + c_i[2]
         }.get(self.descriptor.d)()
diff --git a/boltzgen/kernel/template/bounce_back_boundary.cl.mako b/boltzgen/kernel/template/bounce_back_boundary.cl.mako
index 0762a09..cb59cbd 100644
--- a/boltzgen/kernel/template/bounce_back_boundary.cl.mako
+++ b/boltzgen/kernel/template/bounce_back_boundary.cl.mako
@@ -2,8 +2,8 @@ __kernel void bounce_back_boundary_gid(__global ${float_type}* f_next,
                                        __global ${float_type}* f_prev,
                                        unsigned int gid)
 {
-    __global ${float_type}* preshifted_f_next = f_next + gid;
-    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+    __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)}];
diff --git a/boltzgen/kernel/template/bounce_back_boundary.cpp.mako b/boltzgen/kernel/template/bounce_back_boundary.cpp.mako
index 06dd718..c7abd2a 100644
--- a/boltzgen/kernel/template/bounce_back_boundary.cpp.mako
+++ b/boltzgen/kernel/template/bounce_back_boundary.cpp.mako
@@ -2,8 +2,8 @@ void bounce_back_boundary(      ${float_type}* f_next,
                           const ${float_type}* f_prev,
                           std::size_t gid)
 {
-          ${float_type}* preshifted_f_next = f_next + gid*${layout.gid_offset()};
-    const ${float_type}* preshifted_f_prev = f_prev + gid*${layout.gid_offset()};
+          ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
+    const ${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)}];
diff --git a/boltzgen/kernel/template/collect_moments.cl.mako b/boltzgen/kernel/template/collect_moments.cl.mako
index 0ab42d1..06b92c9 100644
--- a/boltzgen/kernel/template/collect_moments.cl.mako
+++ b/boltzgen/kernel/template/collect_moments.cl.mako
@@ -2,7 +2,7 @@ __kernel void collect_moments_gid(__global ${float_type}* f,
                                   __global ${float_type}* moments,
                                   unsigned int gid)
 {
-    __global ${float_type}* preshifted_f = f + gid;
+    __global ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i in range(0,descriptor.q):
     const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i)}];
diff --git a/boltzgen/kernel/template/collect_moments.cpp.mako b/boltzgen/kernel/template/collect_moments.cpp.mako
index 8c37db2..d72ff5d 100644
--- a/boltzgen/kernel/template/collect_moments.cpp.mako
+++ b/boltzgen/kernel/template/collect_moments.cpp.mako
@@ -3,7 +3,7 @@ void collect_moments(const ${float_type}* f,
                      ${float_type}& rho,
                      ${float_type} u[${descriptor.d}])
 {
-    const ${float_type}* preshifted_f = f + gid*${layout.gid_offset()};
+    const ${float_type}* preshifted_f = f + ${layout.cell_preshift('gid')};
 
 % for i in range(0,descriptor.q):
     const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i)}];
diff --git a/boltzgen/kernel/template/collide_and_stream.cl.mako b/boltzgen/kernel/template/collide_and_stream.cl.mako
index a8fe532..c586884 100644
--- a/boltzgen/kernel/template/collide_and_stream.cl.mako
+++ b/boltzgen/kernel/template/collide_and_stream.cl.mako
@@ -2,8 +2,8 @@ __kernel void collide_and_stream_gid(__global ${float_type}* f_next,
                                      __global ${float_type}* f_prev,
                                      unsigned int gid)
 {
-    __global ${float_type}* preshifted_f_next = f_next + gid;
-    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+    __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)}];
diff --git a/boltzgen/kernel/template/collide_and_stream.cpp.mako b/boltzgen/kernel/template/collide_and_stream.cpp.mako
index 71c62b7..20fc529 100644
--- a/boltzgen/kernel/template/collide_and_stream.cpp.mako
+++ b/boltzgen/kernel/template/collide_and_stream.cpp.mako
@@ -2,8 +2,8 @@ void collide_and_stream(      ${float_type}* f_next,
                         const ${float_type}* f_prev,
                         std::size_t gid)
 {
-          ${float_type}* preshifted_f_next = f_next + gid*${layout.gid_offset()};
-    const ${float_type}* preshifted_f_prev = f_prev + gid*${layout.gid_offset()};
+          ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
+    const ${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)}];
diff --git a/boltzgen/kernel/template/equilibrilize.cl.mako b/boltzgen/kernel/template/equilibrilize.cl.mako
index 4b9b984..4ee8d41 100644
--- a/boltzgen/kernel/template/equilibrilize.cl.mako
+++ b/boltzgen/kernel/template/equilibrilize.cl.mako
@@ -2,8 +2,8 @@ __kernel void equilibrilize_gid(__global ${float_type}* f_next,
                                 __global ${float_type}* f_prev,
                                 unsigned int gid)
 {
-    __global ${float_type}* preshifted_f_next = f_next + gid;
-    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+    __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, w_i in enumerate(descriptor.w):
     preshifted_f_next[${layout.pop_offset(i)}] = ${w_i}.f;
diff --git a/boltzgen/kernel/template/equilibrilize.cpp.mako b/boltzgen/kernel/template/equilibrilize.cpp.mako
index 9f082a1..3b95a31 100644
--- a/boltzgen/kernel/template/equilibrilize.cpp.mako
+++ b/boltzgen/kernel/template/equilibrilize.cpp.mako
@@ -2,8 +2,8 @@ void equilibrilize(${float_type}* f_next,
                    ${float_type}* f_prev,
                    std::size_t gid)
 {
-    ${float_type}* preshifted_f_next = f_next + gid*${layout.gid_offset()};
-    ${float_type}* preshifted_f_prev = f_prev + gid*${layout.gid_offset()};
+    ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
+    ${float_type}* preshifted_f_prev = f_prev + ${layout.cell_preshift('gid')};
 
 % for i, w_i in enumerate(descriptor.w):
     preshifted_f_next[${layout.pop_offset(i)}] = ${w_i.evalf()};
diff --git a/boltzgen/kernel/template/momenta_boundary.cl.mako b/boltzgen/kernel/template/momenta_boundary.cl.mako
index e4a8ff3..f5edb61 100644
--- a/boltzgen/kernel/template/momenta_boundary.cl.mako
+++ b/boltzgen/kernel/template/momenta_boundary.cl.mako
@@ -4,8 +4,8 @@ __kernel void ${name}_momenta_boundary_gid(
    __global ${float_type}* f_prev,
    unsigned int gid, ${param})
 {
-    __global ${float_type}* preshifted_f_next = f_next + gid;
-    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+    __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)}];
diff --git a/boltzgen/kernel/template/momenta_boundary.cpp.mako b/boltzgen/kernel/template/momenta_boundary.cpp.mako
index 5475241..3f51087 100644
--- a/boltzgen/kernel/template/momenta_boundary.cpp.mako
+++ b/boltzgen/kernel/template/momenta_boundary.cpp.mako
@@ -4,8 +4,8 @@ void ${name}_momenta_boundary(
     const ${float_type}* f_prev,
     std::size_t gid, ${param})
 {
-          ${float_type}* preshifted_f_next = f_next + gid*${layout.gid_offset()};
-    const ${float_type}* preshifted_f_prev = f_prev + gid*${layout.gid_offset()};
+          ${float_type}* preshifted_f_next = f_next + ${layout.cell_preshift('gid')};
+    const ${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)}];