Extract cell indexing function

9 files changed, 72 insertions, 79 deletions

diff --git a/boltzgen.py b/boltzgen.py
index 5224942..34d641d 100755
--- a/boltzgen.py
+++ b/boltzgen.py
@@ -8,11 +8,12 @@ argparser = argparse.ArgumentParser(
 
 argparser.add_argument('language', help = 'Target language (currently either "cl" or "cpp")')
 
-argparser.add_argument('--lattice',   required = True, help = 'Lattice type (D2Q9, D3Q7, D3Q19, D3Q27)')
-argparser.add_argument('--layout',    required = True, help = 'Memory layout ("AOS" or "SOA")')
-argparser.add_argument('--precision', required = True, help = 'Floating precision ("single" or "double")')
-argparser.add_argument('--geometry',  required = True, help = 'Size of the block geometry ("x:y(:z)")')
-argparser.add_argument('--tau',       required = True, help = 'BGK relaxation time')
+argparser.add_argument('--lattice',   required = True,  help = 'Lattice type (D2Q9, D3Q7, D3Q19, D3Q27)')
+argparser.add_argument('--layout',    required = True,  help = 'Memory layout ("AOS" or "SOA")')
+argparser.add_argument('--indexing',  required = False, help = 'Cell indexing ("XYZ" or "ZYX")')
+argparser.add_argument('--precision', required = True,  help = 'Floating precision ("single" or "double")')
+argparser.add_argument('--geometry',  required = True,  help = 'Size of the block geometry ("x:y(:z)")')
+argparser.add_argument('--tau',       required = True,  help = 'BGK relaxation time')
 
 argparser.add_argument('--disable-cse', action = 'store_const', const = True, help = 'Disable common subexpression elimination')
 argparser.add_argument('--functions',   action = 'append', nargs = '+', default = [], help = 'Function templates to be generated')
@@ -28,6 +29,9 @@ generator = Generator(
     moments    = lbm.moments(optimize = not args.disable_cse),
     collision  = lbm.bgk(f_eq = lbm.equilibrium(), tau = float(args.tau), optimize = not args.disable_cse))
 
+if args.indexing is None:
+    args.indexing = 'XYZ'
+
 geometry = Geometry.parse(args.geometry)
 
 functions = sum(args.functions, [])
@@ -40,5 +44,5 @@ if 'default' in functions:
 
 extras = sum(args.extras, [])
 
-src = generator.kernel(args.language, args.precision, args.layout, geometry, functions, extras)
+src = generator.kernel(args.language, args.precision, args.layout, args.indexing, geometry, functions, extras)
 print(src)
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)()