Add basic AA pattern support for OpenCL example

5 files changed, 191 insertions, 57 deletions

diff --git a/lid_driven_cavity/opencl/AA.py b/lid_driven_cavity/opencl/AA.py
new file mode 100644
index 0000000..f317cd0
--- /dev/null
+++ b/lid_driven_cavity/opencl/AA.py
@@ -0,0 +1,105 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+import numpy
+
+class Memory:
+    def __init__(self, descriptor, geometry, context, float_type):
+        self.context    = context
+        self.float_type = float_type
+
+        self.size_x = geometry.size_x
+        self.size_y = geometry.size_y
+        self.size_z = geometry.size_z
+        self.volume = self.size_x * self.size_y * self.size_z
+
+        self.pop_size     = descriptor.q * self.volume * self.float_type(0).nbytes
+        self.moments_size = 3 * self.volume * self.float_type(0).nbytes
+
+        self.cl_pop = cl.Buffer(self.context, mf.READ_WRITE, size=self.pop_size)
+
+        self.cl_moments = cl.Buffer(self.context, mf.WRITE_ONLY, size=self.moments_size)
+
+    def gid(self, x, y, z = 0):
+        return z * (self.size_x*self.size_y) + y * self.size_x + x;
+
+class Lattice:
+    def __init__(self, geometry, kernel_src, descriptor, platform = 0, precision = 'single'):
+        self.geometry = geometry
+        self.descriptor = descriptor
+
+        self.float_type = {
+            'single': (numpy.float32, 'float'),
+            'double': (numpy.float64, 'double'),
+        }.get(precision, None)
+
+        self.platform = cl.get_platforms()[platform]
+        self.layout = None
+
+        self.context = cl.Context(
+            properties=[(cl.context_properties.PLATFORM, self.platform)])
+
+        self.queue = cl.CommandQueue(self.context)
+
+        self.memory = Memory(descriptor, self.geometry, self.context, self.float_type[0])
+        self.tick = False
+
+        self.compiler_args = {
+            'single': '-cl-single-precision-constant -cl-fast-relaxed-math',
+            'double': '-cl-fast-relaxed-math'
+        }.get(precision, None)
+
+        self.build_kernel(kernel_src)
+
+        self.program.equilibrilize_all(
+            self.queue, self.geometry.size(), self.layout, self.memory.cl_pop).wait()
+
+        self.tick_tasks = []
+        self.tock_tasks = []
+
+    def build_kernel(self, src):
+        self.program = cl.Program(self.context, src).build(self.compiler_args)
+
+    def schedule_tick(self, f, cells, *params):
+        self.tick_tasks += [ (eval("self.program.%s" % f), cells, params) ]
+
+    def schedule_tock(self, f, cells, *params):
+        self.tock_tasks += [ (eval("self.program.%s" % f), cells, params) ]
+
+    def evolve(self):
+        if self.tick:
+            self.tick = False
+            for f, cells, params in self.tick_tasks:
+                f(self.queue, cells.size(), self.layout, self.memory.cl_pop, cells.get(), *params)
+        else:
+            self.tick = True
+            for f, cells, params in self.tock_tasks:
+                f(self.queue, cells.size(), self.layout, self.memory.cl_pop, cells.get(), *params)
+
+    def sync(self):
+        self.queue.finish()
+
+    def get_moments(self):
+        moments = numpy.ndarray(shape=(self.memory.volume*(self.descriptor.d+1),1), dtype=self.float_type[0])
+
+        self.program.collect_moments_all(
+            self.queue, self.geometry.size(), self.layout, self.memory.cl_pop, self.memory.cl_moments)
+
+        cl.enqueue_copy(self.queue, moments, self.memory.cl_moments).wait();
+
+        return moments
+
+HelperTemplate = """
+__kernel void equilibrilize_all(__global ${float_type}* f_next)
+{
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
+    equilibrilize_tick(f_next, gid);
+}
+
+__kernel void collect_moments_all(__global ${float_type}* f,
+                                  __global ${float_type}* moments)
+{
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
+    collect_moments_tock(f, gid, moments);
+}
+"""
diff --git a/lid_driven_cavity/opencl/simulation.py b/lid_driven_cavity/opencl/AB.py
index 37aab5b..7cbf1c6 100644
--- a/lid_driven_cavity/opencl/simulation.py
+++ b/lid_driven_cavity/opencl/AB.py
@@ -24,20 +24,6 @@ class Memory:
     def gid(self, x, y, z = 0):
         return z * (self.size_x*self.size_y) + y * self.size_x + x;
 
-class CellList:
-    def __init__(self, context, queue, float_type, cells):
-        self.cl_cells = cl.Buffer(context, mf.READ_ONLY, size=len(cells) * numpy.uint32(0).nbytes)
-        self.np_cells = numpy.ndarray(shape=(len(cells), 1), dtype=numpy.uint32)
-        self.np_cells[:,0] = cells[:]
-
-        cl.enqueue_copy(queue, self.cl_cells, self.np_cells).wait();
-
-    def get(self):
-        return self.cl_cells
-
-    def size(self):
-        return (len(self.np_cells), 1, 1)
-
 class Lattice:
     def __init__(self, geometry, kernel_src, descriptor, platform = 0, precision = 'single'):
         self.geometry = geometry
@@ -103,3 +89,20 @@ class Lattice:
         cl.enqueue_copy(self.queue, moments, self.memory.cl_moments).wait();
 
         return moments
+
+HelperTemplate = """
+__kernel void equilibrilize_all(__global ${float_type}* f_next,
+                                __global ${float_type}* f_prev)
+{
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
+    equilibrilize(f_next, f_prev, gid);
+    equilibrilize(f_prev, f_next, gid);
+}
+
+__kernel void collect_moments_all(__global ${float_type}* f,
+                                  __global ${float_type}* moments)
+{
+    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
+    collect_moments(f, gid, moments);
+}
+"""
diff --git a/lid_driven_cavity/opencl/common.py b/lid_driven_cavity/opencl/common.py
new file mode 100644
index 0000000..717c68f
--- /dev/null
+++ b/lid_driven_cavity/opencl/common.py
@@ -0,0 +1,36 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+import numpy
+
+import matplotlib
+matplotlib.use('AGG')
+import matplotlib.pyplot as plt
+
+
+class CellList:
+    def __init__(self, context, queue, float_type, cells):
+        self.cl_cells = cl.Buffer(context, mf.READ_ONLY, size=len(cells) * numpy.uint32(0).nbytes)
+        self.np_cells = numpy.ndarray(shape=(len(cells), 1), dtype=numpy.uint32)
+        self.np_cells[:,0] = cells[:]
+
+        cl.enqueue_copy(queue, self.cl_cells, self.np_cells).wait();
+
+    def get(self):
+        return self.cl_cells
+
+    def size(self):
+        return (len(self.np_cells), 1, 1)
+
+def generate_moment_plots(lattice, moments):
+    for i, m in enumerate(moments):
+        print("Generating plot %d of %d." % (i+1, len(moments)))
+
+        gid = lattice.memory.gid
+        velocity = numpy.reshape(
+            [ numpy.sqrt(m[gid(x,y)*3+1]**2 + m[gid(x,y)*3+2]**2) for x, y in lattice.geometry.inner_cells() ],
+            lattice.geometry.inner_size())
+
+        plt.figure(figsize=(10, 10))
+        plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
+        plt.savefig("result/ldc_2d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
diff --git a/lid_driven_cavity/opencl/ldc_2d.py b/lid_driven_cavity/opencl/ldc_2d.py
index c080284..8d9e324 100644
--- a/lid_driven_cavity/opencl/ldc_2d.py
+++ b/lid_driven_cavity/opencl/ldc_2d.py
@@ -1,36 +1,16 @@
 import numpy
 import time
 
-import matplotlib
-matplotlib.use('AGG')
-import matplotlib.pyplot as plt
-
 from boltzgen import Generator, Geometry
 from boltzgen.lbm.lattice import D2Q9
 from boltzgen.lbm.model   import BGK
 
-from simulation import Lattice, CellList
-
-def MLUPS(cells, steps, time):
-    return cells * steps / time * 1e-6
-
-def generate_moment_plots(lattice, moments):
-    for i, m in enumerate(moments):
-        print("Generating plot %d of %d." % (i+1, len(moments)))
-
-        gid = lattice.memory.gid
-        velocity = numpy.reshape(
-            [ numpy.sqrt(m[gid(x,y)*3+1]**2 + m[gid(x,y)*3+2]**2) for x, y in lattice.geometry.inner_cells() ],
-            lattice.geometry.inner_size())
-
-        plt.figure(figsize=(10, 10))
-        plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
-        plt.savefig("result/ldc_2d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
+from common import CellList, generate_moment_plots
 
 nUpdates = 100000
 nStat    = 10000
 
-geometry = Geometry(256, 256)
+geometry = Geometry(512, 512)
 
 print("Generating kernel using boltzgen...\n")
 
@@ -38,37 +18,35 @@ functions = ['collide_and_stream', 'equilibrilize', 'collect_moments', 'momenta_
 extras    = ['cell_list_dispatch']
 
 precision = 'single'
+streaming = 'AA'
+
+import AA
+import AB
+
+Lattice = eval('%s.Lattice' % streaming)
+HelperTemplate = eval('%s.HelperTemplate' % streaming)
+
+def MLUPS(cells, steps, time):
+    return cells * steps / time * 1e-6
 
 generator = Generator(
-    model     = BGK(D2Q9, tau = 0.6),
+    model     = BGK(D2Q9, tau = 0.54),
     target    = 'cl',
     precision = precision,
+    streaming = streaming,
     index     = 'ZYX',
     layout    = 'SOA')
 
 kernel_src  = generator.kernel(geometry, functions, extras)
-kernel_src += generator.custom(geometry, """
-__kernel void equilibrilize_all(__global ${float_type}* f_next,
-                                __global ${float_type}* f_prev)
-{
-    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
-    equilibrilize(f_next, f_prev, gid);
-    equilibrilize(f_prev, f_next, gid);
-}
-
-__kernel void collect_moments_all(__global ${float_type}* f,
-                                  __global ${float_type}* moments)
-{
-    const unsigned int gid = ${index.gid('get_global_id(0)', 'get_global_id(1)')};
-    collect_moments(f, gid, moments);
-}
-""")
+kernel_src += generator.custom(geometry, HelperTemplate)
 
 print("Initializing simulation...\n")
 
 lattice = Lattice(geometry, kernel_src, D2Q9, precision = precision)
 gid = lattice.memory.gid
 
+ghost_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
+    [ gid(x,y) for x, y in geometry.cells() if x == 0 or y == 0 or x == geometry.size_x-1 or y == geometry.size_y-1 ])
 bulk_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
     [ gid(x,y) for x, y in geometry.inner_cells() if x > 1 and x < geometry.size_x-2 and y > 1 and y < geometry.size_y-2 ])
 wall_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
@@ -76,9 +54,21 @@ wall_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
 lid_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
     [ gid(x,y) for x, y in geometry.inner_cells() if y == geometry.size_y-2 ])
 
-lattice.schedule('collide_and_stream_cells', bulk_cells)
-lattice.schedule('velocity_momenta_boundary_cells', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
-lattice.schedule('velocity_momenta_boundary_cells', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+if streaming == 'AB':
+    lattice.schedule('collide_and_stream_cells', bulk_cells)
+    lattice.schedule('velocity_momenta_boundary_cells', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule('velocity_momenta_boundary_cells', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+
+elif streaming == 'AA':
+    lattice.schedule_tick('collide_and_stream_tick_cells', bulk_cells)
+    lattice.schedule_tick('velocity_momenta_boundary_tick_cells', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tick('velocity_momenta_boundary_tick_cells', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+
+    lattice.schedule_tock('equilibrilize_tick_cells', ghost_cells)
+    lattice.schedule_tock('collide_and_stream_tock_cells', bulk_cells)
+    lattice.schedule_tock('velocity_momenta_boundary_tock_cells', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tock('velocity_momenta_boundary_tock_cells', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+
 
 print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
 
diff --git a/shell.nix b/shell.nix
index c1806bf..4e36dce 100644
--- a/shell.nix
+++ b/shell.nix
@@ -11,7 +11,7 @@ pkgs.stdenvNoCC.mkDerivation rec {
 
       src = builtins.fetchGit {
         url = "https://code.kummerlaender.eu/boltzgen/";
-        rev = "02cb01c94fe26d425371ab74feeb50e8a9bf6bf6";
+        rev = "814e6253475c7955eb6a46d814e5a86974e58613";
       };
 
       propagatedBuildInputs = with pkgs.python37Packages; [