Share Lattice implementation between plain and interop OpenCL example

8 files changed, 61 insertions, 281 deletions

diff --git a/lid_driven_cavity/opencl/AA.py b/lid_driven_cavity/opencl/AA.py
index 74979c5..7ea3417 100644
--- a/lid_driven_cavity/opencl/AA.py
+++ b/lid_driven_cavity/opencl/AA.py
@@ -24,7 +24,7 @@ class Memory:
         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'):
+    def __init__(self, geometry, kernel_src, descriptor, context, queue, precision = 'single'):
         self.geometry = geometry
         self.descriptor = descriptor
 
@@ -33,13 +33,10 @@ class Lattice:
             '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.context = context
+        self.queue   = queue
 
         self.memory = Memory(descriptor, self.geometry, self.context, self.float_type[0])
         self.tick = False
@@ -51,7 +48,7 @@ class Lattice:
 
         self.build_kernel(kernel_src)
 
-        self.program.equilibrilize_all(
+        self.program.equilibrilize_domain_tick(
             self.queue, self.geometry.size(), self.layout, self.memory.cl_pop).wait()
 
         self.tick_tasks = []
@@ -82,24 +79,9 @@ class Lattice:
     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.program.collect_moments_domain_tock(
             self.queue, self.geometry.inner_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)+1', 'get_global_id(1)+1')};
-    collect_moments_tock(f, gid, moments);
-}
-"""
diff --git a/lid_driven_cavity/opencl/AB.py b/lid_driven_cavity/opencl/AB.py
index 7cbf1c6..797bc83 100644
--- a/lid_driven_cavity/opencl/AB.py
+++ b/lid_driven_cavity/opencl/AB.py
@@ -25,7 +25,7 @@ class Memory:
         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'):
+    def __init__(self, geometry, kernel_src, descriptor, context, queue, precision = 'single'):
         self.geometry = geometry
         self.descriptor = descriptor
 
@@ -34,13 +34,10 @@ class Lattice:
             '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.context = context
+        self.queue   = queue
 
         self.memory = Memory(descriptor, self.geometry, self.context, self.float_type[0])
         self.tick = False
@@ -52,8 +49,10 @@ class Lattice:
 
         self.build_kernel(kernel_src)
 
-        self.program.equilibrilize_all(
+        self.program.equilibrilize_domain(
             self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_pop_b).wait()
+        self.program.equilibrilize_domain(
+            self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_pop_a).wait()
 
         self.tasks = []
 
@@ -80,29 +79,12 @@ class Lattice:
         moments = numpy.ndarray(shape=(self.memory.volume*(self.descriptor.d+1),1), dtype=self.float_type[0])
 
         if self.tick:
-            self.program.collect_moments_all(
+            self.program.collect_moments_domain(
                 self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_moments)
         else:
-            self.program.collect_moments_all(
+            self.program.collect_moments_domain(
                 self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_a, 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,
-                                __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/ldc_2d.py b/lid_driven_cavity/opencl/ldc_2d.py
index 7380aa9..1ea80aa 100644
--- a/lid_driven_cavity/opencl/ldc_2d.py
+++ b/lid_driven_cavity/opencl/ldc_2d.py
@@ -1,16 +1,18 @@
 import numpy
 import time
 
+import pyopencl as cl
+
 from boltzgen import Generator, Geometry
 from boltzgen.lbm.lattice import D2Q9
 from boltzgen.lbm.model   import BGK
 
 from common import CellList, generate_moment_plots
 
-nUpdates = 100000
-nStat    = 10000
+nUpdates = 20000
+nStat    = 1000
 
-geometry = Geometry(256, 256)
+geometry = Geometry(512, 512)
 
 print("Generating kernel using boltzgen...\n")
 
@@ -24,7 +26,6 @@ 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
@@ -37,12 +38,17 @@ generator = Generator(
     index     = 'ZYX',
     layout    = 'SOA')
 
-kernel_src  = generator.kernel(geometry, functions, extras)
-kernel_src += generator.custom(geometry, HelperTemplate)
+kernel_src = generator.kernel(geometry, functions, extras)
+
+print("Constructing OpenCL context...\n")
+
+cl_platform = cl.get_platforms()[0]
+cl_context  = cl.Context(properties=[(cl.context_properties.PLATFORM, cl_platform)])
+cl_queue    = cl.CommandQueue(cl_context)
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(geometry, kernel_src, D2Q9, precision = precision)
+lattice = Lattice(geometry, kernel_src, D2Q9, cl_context, cl_queue, precision = precision)
 gid = lattice.memory.gid
 
 ghost_cells = CellList(lattice.context, lattice.queue, lattice.float_type,
@@ -55,19 +61,19 @@ 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 ])
 
 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]))
+    lattice.schedule('collide_and_stream', bulk_cells)
+    lattice.schedule('velocity_momenta_boundary', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule('velocity_momenta_boundary', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
 
 elif streaming == 'AA':
-    lattice.schedule_tick('collide_and_stream_cells_tick', bulk_cells)
-    lattice.schedule_tick('velocity_momenta_boundary_cells_tick', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
-    lattice.schedule_tick('velocity_momenta_boundary_cells_tick', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
-
-    lattice.schedule_tock('equilibrilize_cells_tick', ghost_cells)
-    lattice.schedule_tock('collide_and_stream_cells_tock', bulk_cells)
-    lattice.schedule_tock('velocity_momenta_boundary_cells_tock', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
-    lattice.schedule_tock('velocity_momenta_boundary_cells_tock', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tick('collide_and_stream_tick', bulk_cells)
+    lattice.schedule_tick('velocity_momenta_boundary_tick', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tick('velocity_momenta_boundary_tick', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+
+    lattice.schedule_tock('equilibrilize_tick', ghost_cells)
+    lattice.schedule_tock('collide_and_stream_tock', bulk_cells)
+    lattice.schedule_tock('velocity_momenta_boundary_tock', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tock('velocity_momenta_boundary_tock', 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/lid_driven_cavity/opencl_gl_interop/AA.py b/lid_driven_cavity/opencl_gl_interop/AA.py
index ca9a132..1322856 100644
--- a/lid_driven_cavity/opencl_gl_interop/AA.py
+++ b/lid_driven_cavity/opencl_gl_interop/AA.py
@@ -1,105 +1,8 @@
 import pyopencl as cl
-mf = cl.mem_flags
-
-from pyopencl.tools import get_gl_sharing_context_properties
-
-import numpy
 
 from common import MomentsTextureBase
 
-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)] + get_gl_sharing_context_properties())
-
-        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);
-}
-"""
+from lattice.AA import Lattice
 
 class MomentsTexture(MomentsTextureBase):
     pass
diff --git a/lid_driven_cavity/opencl_gl_interop/AB.py b/lid_driven_cavity/opencl_gl_interop/AB.py
index f6b437c..fce10d9 100644
--- a/lid_driven_cavity/opencl_gl_interop/AB.py
+++ b/lid_driven_cavity/opencl_gl_interop/AB.py
@@ -1,108 +1,8 @@
 import pyopencl as cl
-mf = cl.mem_flags
-
-from pyopencl.tools import get_gl_sharing_context_properties
-
-import numpy
 
 from common import MomentsTextureBase
 
-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_a = cl.Buffer(self.context, mf.READ_WRITE, size=self.pop_size)
-        self.cl_pop_b = 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)] + get_gl_sharing_context_properties())
-
-        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_a, self.memory.cl_pop_b).wait()
-
-        self.tasks = []
-
-    def build_kernel(self, src):
-        self.program = cl.Program(self.context, src).build(self.compiler_args)
-
-    def schedule(self, f, cells, *params):
-        self.tasks += [ (eval("self.program.%s" % f), cells, params) ]
-
-    def evolve(self):
-        if self.tick:
-            self.tick = False
-            for f, cells, params in self.tasks:
-                f(self.queue, cells.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_pop_b, cells.get(), *params)
-        else:
-            self.tick = True
-            for f, cells, params in self.tasks:
-                f(self.queue, cells.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_pop_a, 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])
-
-        if self.tick:
-            self.program.collect_moments_all(
-                self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_moments)
-        else:
-            self.program.collect_moments_all(
-                self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_a, 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,
-                                __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);
-}
-"""
+from lattice.AB import Lattice
 
 class MomentsTexture(MomentsTextureBase):
     pass
diff --git a/lid_driven_cavity/opencl_gl_interop/lattice/AA.py b/lid_driven_cavity/opencl_gl_interop/lattice/AA.py
new file mode 120000
index 0000000..7ca466f
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/lattice/AA.py
@@ -0,0 +1 @@
+../../opencl/AA.py
\ No newline at end of file
diff --git a/lid_driven_cavity/opencl_gl_interop/lattice/AB.py b/lid_driven_cavity/opencl_gl_interop/lattice/AB.py
new file mode 120000
index 0000000..769b0aa
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/lattice/AB.py
@@ -0,0 +1 @@
+../../opencl/AB.py
\ No newline at end of file
diff --git a/lid_driven_cavity/opencl_gl_interop/ldc_2d.py b/lid_driven_cavity/opencl_gl_interop/ldc_2d.py
index daac486..ac8119b 100644
--- a/lid_driven_cavity/opencl_gl_interop/ldc_2d.py
+++ b/lid_driven_cavity/opencl_gl_interop/ldc_2d.py
@@ -2,6 +2,9 @@ import numpy
 import time
 from string import Template
 
+import pyopencl as cl
+from pyopencl.tools import get_gl_sharing_context_properties
+
 from boltzgen import Generator, Geometry
 from boltzgen.lbm.lattice import D2Q9
 from boltzgen.lbm.model   import BGK
@@ -40,7 +43,6 @@ def glut_window(fullscreen = False):
 window = glut_window(fullscreen = False)
 
 Lattice = eval('%s.Lattice' % streaming)
-HelperTemplate = eval('%s.HelperTemplate' % streaming)
 MomentsTexture = eval('%s.MomentsTexture' % streaming)
 
 generator = Generator(
@@ -51,10 +53,13 @@ generator = Generator(
     index     = 'ZYX',
     layout    = 'SOA')
 
-kernel_src  = generator.kernel(geometry, functions, extras)
-kernel_src += generator.custom(geometry, HelperTemplate)
+kernel_src = generator.kernel(geometry, functions, extras)
+
+cl_platform = cl.get_platforms()[0]
+cl_context  = cl.Context(properties=[(cl.context_properties.PLATFORM, cl_platform)] + get_gl_sharing_context_properties())
+cl_queue    = cl.CommandQueue(cl_context)
 
-lattice = Lattice(geometry, kernel_src, D2Q9, precision = precision)
+lattice = Lattice(geometry, kernel_src, D2Q9, cl_context, cl_queue, precision = precision)
 moments = MomentsTexture(lattice)
 
 gid = lattice.memory.gid
@@ -69,19 +74,19 @@ 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 ])
 
 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]))
+    lattice.schedule('collide_and_stream', bulk_cells)
+    lattice.schedule('velocity_momenta_boundary', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule('velocity_momenta_boundary', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
 
 elif streaming == 'AA':
-    lattice.schedule_tick('collide_and_stream_cells_tick', bulk_cells)
-    lattice.schedule_tick('velocity_momenta_boundary_cells_tick', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
-    lattice.schedule_tick('velocity_momenta_boundary_cells_tick', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
-
-    lattice.schedule_tock('equilibrilize_cells_tick', ghost_cells)
-    lattice.schedule_tock('collide_and_stream_cells_tock', bulk_cells)
-    lattice.schedule_tock('velocity_momenta_boundary_cells_tock', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
-    lattice.schedule_tock('velocity_momenta_boundary_cells_tock', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tick('collide_and_stream_tick', bulk_cells)
+    lattice.schedule_tick('velocity_momenta_boundary_tick', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tick('velocity_momenta_boundary_tick', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
+
+    lattice.schedule_tock('equilibrilize_tick', ghost_cells)
+    lattice.schedule_tock('collide_and_stream_tock', bulk_cells)
+    lattice.schedule_tock('velocity_momenta_boundary_tock', wall_cells, numpy.array([0.0, 0.0], dtype=lattice.float_type[0]))
+    lattice.schedule_tock('velocity_momenta_boundary_tock', lid_cells,  numpy.array([0.1, 0.0], dtype=lattice.float_type[0]))
 
 
 def get_projection(width, height):