Extract cell indexing function

8 files changed, 62 insertions, 73 deletions

diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py
index 5a1a37d..da32a76 100644
--- a/boltzgen/kernel/generator.py
+++ b/boltzgen/kernel/generator.py
@@ -32,16 +32,24 @@ class Generator:
             extras = extras
         )
 
-    def kernel(self, target, precision, layout, geometry, functions, extras = []):
-        layout_impl = eval("boltzgen.kernel.target.layout.%s.%s" % (target, layout))
+    def kernel(self, target, precision, layout, cell_index, geometry, functions, extras = []):
+        cell_index_impl = eval("boltzgen.kernel.target.cell_index.%s" % cell_index)
+        if cell_index_impl is None:
+            raise Exception("There is no cell indexing scheme '%s'" % (target, layout))
+        else:
+            cell_index_impl = cell_index_impl(geometry)
+
+        layout_impl = eval("boltzgen.kernel.target.layout.%s" % layout)
         if layout_impl is None:
-            raise Exception("Target '%s' doesn't support layout '%s'" % (target, layout))
+            raise Exception("There is no layout '%s'" % (target, layout))
         else:
-            layout_impl = layout_impl(self.descriptor, geometry)
+            layout_impl = layout_impl(self.descriptor, cell_index_impl, geometry)
 
         if geometry.dimension() != self.descriptor.d:
             raise Exception('Geometry dimension must match descriptor dimension')
 
         float_type = eval("boltzgen.kernel.target.precision.%s" % target).get_float_type(precision)
 
-        return "\n".join(map(lambda f: self.instantiate(target, f, float_type, layout_impl, geometry, extras), functions))
+        return "\n".join(map(
+            lambda f: self.instantiate(target, f, float_type, layout_impl, geometry, extras),
+            functions))
diff --git a/boltzgen/kernel/target/__init__.py b/boltzgen/kernel/target/__init__.py
index e3f8dad..d39f48b 100644
--- a/boltzgen/kernel/target/__init__.py
+++ b/boltzgen/kernel/target/__init__.py
@@ -1,2 +1,3 @@
 from . import layout
 from . import precision
+from . import cell_index
diff --git a/boltzgen/kernel/target/cell_index.py b/boltzgen/kernel/target/cell_index.py
new file mode 100644
index 0000000..7e74dc5
--- /dev/null
+++ b/boltzgen/kernel/target/cell_index.py
@@ -0,0 +1,18 @@
+class CellIndexer:
+    def __init__(self, geometry):
+        self.geometry   = geometry
+
+    def neighbor(self, c_i):
+        return self.gid(*[ c for _, c in enumerate(c_i) ])
+
+class XYZ(CellIndexer):
+    pass
+
+    def gid(self, x, y, z=0):
+        return x*self.geometry.size_y*self.geometry.size_z + y*self.geometry.size_z + z
+
+class ZYX(CellIndexer):
+    pass
+
+    def gid(self, x, y, z=0):
+        return z*self.geometry.size_x*self.geometry.size_y + y*self.geometry.size_x + x
diff --git a/boltzgen/kernel/target/layout/AOS.py b/boltzgen/kernel/target/layout/AOS.py
new file mode 100644
index 0000000..c603869
--- /dev/null
+++ b/boltzgen/kernel/target/layout/AOS.py
@@ -0,0 +1,14 @@
+class AOS:
+    def __init__(self, descriptor, cell_index, geometry):
+        self.descriptor = descriptor
+        self.cell_index = cell_index
+        self.geometry = geometry
+
+    def cell_preshift(self, gid):
+        return "(%s)*%d" % (gid, self.descriptor.q)
+
+    def pop_offset(self, i):
+        return i
+
+    def neighbor_offset(self, c_i):
+        return self.descriptor.q * self.cell_index.neighbor(c_i)
diff --git a/boltzgen/kernel/target/layout/SOA.py b/boltzgen/kernel/target/layout/SOA.py
new file mode 100644
index 0000000..8f126ff
--- /dev/null
+++ b/boltzgen/kernel/target/layout/SOA.py
@@ -0,0 +1,14 @@
+class SOA:
+    def __init__(self, descriptor, cell_index, geometry):
+        self.descriptor = descriptor
+        self.cell_index = cell_index
+        self.geometry = geometry
+
+    def cell_preshift(self, gid):
+        return gid
+
+    def pop_offset(self, i):
+        return i * self.geometry.volume
+
+    def neighbor_offset(self, c_i):
+        return self.cell_index.neighbor(c_i)
diff --git a/boltzgen/kernel/target/layout/__init__.py b/boltzgen/kernel/target/layout/__init__.py
index 44e92e6..731fbe8 100644
--- a/boltzgen/kernel/target/layout/__init__.py
+++ b/boltzgen/kernel/target/layout/__init__.py
@@ -1,2 +1,2 @@
-from . import cl
-from . import cpp
+from .AOS import AOS
+from .SOA import SOA
diff --git a/boltzgen/kernel/target/layout/cl.py b/boltzgen/kernel/target/layout/cl.py
deleted file mode 100644
index e144faf..0000000
--- a/boltzgen/kernel/target/layout/cl.py
+++ /dev/null
@@ -1,33 +0,0 @@
-class SOA:
-    def __init__(self, descriptor, geometry):
-        self.descriptor = descriptor
-        self.geometry = geometry
-
-    def cell_preshift(self, gid):
-        return gid
-
-    def pop_offset(self, i):
-        return i * self.geometry.volume
-
-    def neighbor_offset(self, c_i):
-        return {
-            2: lambda:                                                    c_i[1]*self.geometry.size_x + c_i[0],
-            3: lambda: c_i[2]*self.geometry.size_x*self.geometry.size_y + c_i[1]*self.geometry.size_x + c_i[0]
-        }.get(self.descriptor.d)()
-
-class AOS:
-    def __init__(self, descriptor, geometry):
-        self.descriptor = descriptor
-        self.geometry = geometry
-
-    def cell_preshift(self, gid):
-        return "(%s)*%d" % (gid, self.descriptor.q)
-
-    def pop_offset(self, i):
-        return i
-
-    def neighbor_offset(self, c_i):
-        return self.descriptor.q * {
-            2: lambda:                                                    c_i[1]*self.geometry.size_x + c_i[0],
-            3: lambda: c_i[2]*self.geometry.size_x*self.geometry.size_y + c_i[1]*self.geometry.size_x + c_i[0]
-        }.get(self.descriptor.d)()
diff --git a/boltzgen/kernel/target/layout/cpp.py b/boltzgen/kernel/target/layout/cpp.py
deleted file mode 100644
index 2a8c89e..0000000
--- a/boltzgen/kernel/target/layout/cpp.py
+++ /dev/null
@@ -1,33 +0,0 @@
-class SOA:
-    def __init__(self, descriptor, geometry):
-        self.descriptor = descriptor
-        self.geometry = geometry
-
-    def cell_preshift(self, gid):
-        return gid
-
-    def pop_offset(self, i):
-        return i * self.geometry.volume
-
-    def neighbor_offset(self, c_i):
-        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 AOS:
-    def __init__(self, descriptor, geometry):
-        self.descriptor = descriptor
-        self.geometry = geometry
-
-    def cell_preshift(self, gid):
-        return "(%s)*%d" % (gid, self.descriptor.q)
-
-    def pop_offset(self, i):
-        return i
-
-    def neighbor_offset(self, c_i):
-        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)()