Move OpenCL buffers into Memory class

5 files changed, 69 insertions, 64 deletions

diff --git a/implosion.py b/implosion.py
index 370ef92..245f2bb 100644
--- a/implosion.py
+++ b/implosion.py
@@ -19,7 +19,7 @@ def generate_moment_plots(lattice, moments):
 
         velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))
         for x, y in lattice.geometry.inner_cells():
-            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.gid(x,y)]**2 + m[2,lattice.gid(x,y)]**2)
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.memory.gid(x,y)]**2 + m[2,lattice.memory.gid(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
diff --git a/ldc_2d.py b/ldc_2d.py
index 52f0912..5fae3fe 100644
--- a/ldc_2d.py
+++ b/ldc_2d.py
@@ -23,7 +23,7 @@ def generate_moment_plots(lattice, moments):
 
         velocity = numpy.ndarray(shape=tuple(reversed(lattice.geometry.inner_size())))
         for x, y in lattice.geometry.inner_cells():
-            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.gid(x,y)]**2 + m[2,lattice.gid(x,y)]**2)
+            velocity[y-1,x-1] = numpy.sqrt(m[1,lattice.memory.gid(x,y)]**2 + m[2,lattice.memory.gid(x,y)]**2)
 
         plt.figure(figsize=(10, 10))
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
@@ -61,10 +61,10 @@ lbm = LBM(D2Q9)
 
 lattice = Lattice(
     descriptor = D2Q9,
-    geometry   = Geometry(300, 300),
+    geometry   = Geometry(600, 600),
 
     layout  = (30,1),
-    padding = (30,1,1),
+    padding = (30,1),
     align   = True,
 
     moments = lbm.moments(optimize = False),
diff --git a/ldc_2d_gl_interop.py b/ldc_2d_gl_interop.py
index 1c3439f..e5ea866 100644
--- a/ldc_2d_gl_interop.py
+++ b/ldc_2d_gl_interop.py
@@ -138,13 +138,13 @@ def on_display():
 
     glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
-    lattice.gl_moments.bind()
+    lattice.memory.gl_moments.bind()
     glEnableClientState(GL_VERTEX_ARRAY)
 
     shaders.glUseProgram(shader_program)
     glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
 
-    glVertexPointer(4, GL_FLOAT, 0, lattice.gl_moments)
+    glVertexPointer(4, GL_FLOAT, 0, lattice.memory.gl_moments)
 
     glPointSize(cells_per_pixel)
     glDrawArrays(GL_POINTS, 0, lattice.geometry.volume)
diff --git a/ldc_3d.py b/ldc_3d.py
index 9149cbc..a41f0ae 100644
--- a/ldc_3d.py
+++ b/ldc_3d.py
@@ -34,13 +34,13 @@ def generate_moment_plots(lattice, moments):
         # extract x-z-plane
         y_slice = lattice.geometry.size_y//2
         for z, x in numpy.ndindex(lattice.geometry.size_z-2, lattice.geometry.size_x-2):
-            gid = lattice.gid(x+1,y_slice,z+1)
+            gid = lattice.memory.gid(x+1,y_slice,z+1)
             velocity[z,y_slice,x] = numpy.sqrt(m[1,gid]**2 + m[2,gid]**2 + m[3,gid]**2)
 
         # extract y-z-plane
         x_slice = lattice.geometry.size_x//2
         for z, y in numpy.ndindex(lattice.geometry.size_z-2, lattice.geometry.size_y-2):
-            gid = lattice.gid(x_slice,y+1,z+1)
+            gid = lattice.memory.gid(x_slice,y+1,z+1)
             velocity[z,y,x_slice] = numpy.sqrt(m[1,gid]**2 + m[2,gid]**2 + m[3,gid]**2)
 
         plt.figure(figsize=(20, 10))
@@ -87,7 +87,7 @@ lbm = LBM(D3Q19)
 
 lattice = Lattice(
     descriptor = D3Q19,
-    geometry   = Geometry(128, 128, 128),
+    geometry   = Geometry(64, 64, 64),
 
     moments = lbm.moments(optimize = False),
     collide = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.56),
diff --git a/simulation.py b/simulation.py
index d3b8641..039486c 100644
--- a/simulation.py
+++ b/simulation.py
@@ -11,7 +11,6 @@ from pyopencl.tools import get_gl_sharing_context_properties
 import OpenGL.GL as gl
 from OpenGL.arrays import vbo
 
-
 class Geometry:
     def __init__(self, size_x, size_y, size_z = 1):
         self.size_x = size_x
@@ -36,6 +35,9 @@ class Geometry:
             return (self.size_x-2, self.size_y-2, self.size_z-2)
 
 
+def pad(n, m):
+    return (n // m + min(1,n % m)) * m
+
 class Grid:
     def __init__(self, geometry, padding = None):
         if padding == None:
@@ -43,9 +45,12 @@ class Grid:
             self.size_y = geometry.size_y
             self.size_z = geometry.size_z
         else:
-            self.size_x = (geometry.size_x // padding[0] + min(1,geometry.size_x % padding[0])) * padding[0]
-            self.size_y = (geometry.size_y // padding[1] + min(1,geometry.size_y % padding[1])) * padding[1]
-            self.size_z = (geometry.size_z // padding[2] + min(1,geometry.size_z % padding[2])) * padding[2]
+            self.size_x = pad(geometry.size_x, padding[0])
+            self.size_y = pad(geometry.size_y, padding[1])
+            if geometry.size_z == 1:
+                self.size_z = geometry.size_z
+            else:
+                self.size_z = pad(geometry.size_z, padding[2])
 
         self.volume = self.size_x * self.size_y * self.size_z
 
@@ -55,11 +60,14 @@ class Grid:
         else:
             return (self.size_x, self.size_y, self.size_z)
 
-
 class Memory:
-    def __init__(self, grid, align = True):
+    def __init__(self, descriptor, grid, context, float_type, align, opengl):
+        self.descriptor = descriptor
+        self.context    = context
+        self.float_type = float_type
+
         if align:
-            self.size_x = (grid.size_x // 32 + min(1,grid.size_x % 32)) * 32
+            self.size_x = pad(grid.size_x, 32)
         else:
             self.size_x = grid.size_x
 
@@ -68,6 +76,25 @@ class Memory:
 
         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 = (descriptor.d+1) * 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)
+
+        if opengl:
+            self.np_moments = numpy.ndarray(shape=(self.volume, 4), dtype=self.float_type)
+            self.gl_moments = vbo.VBO(data=self.np_moments, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
+            self.gl_moments.bind()
+            self.cl_gl_moments  = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_moments))
+        else:
+            self.cl_moments  = cl.Buffer(self.context, mf.WRITE_ONLY, size=self.moments_size)
+
+        self.cl_material = cl.Buffer(self.context, mf.READ_ONLY, size=self.volume * numpy.int32(0).nbytes)
+
+    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,
@@ -77,28 +104,13 @@ class Lattice:
     ):
         self.descriptor = descriptor
         self.geometry   = geometry
-
-        self.grid   = Grid(self.geometry, padding)
-        self.memory = Memory(self.grid, align)
-
-        self.moments = moments
-        self.collide = collide
-
-        self.pop_eq_src = pop_eq_src
-        self.boundary_src = boundary_src
+        self.grid       = Grid(self.geometry, padding)
 
         self.float_type = {
             'single': (numpy.float32, 'float'),
             'double': (numpy.float64, 'double'),
         }.get(precision, None)
 
-        self.layout = layout
-
-        self.compiler_args = {
-            'single': '-cl-single-precision-constant -cl-fast-relaxed-math',
-            'double': '-cl-fast-relaxed-math'
-        }.get(precision, None)
-
         self.platform = cl.get_platforms()[platform]
 
         if opengl:
@@ -110,45 +122,40 @@ class Lattice:
 
         self.queue = cl.CommandQueue(self.context)
 
-        self.pop_size     = descriptor.q     * self.memory.volume * self.float_type[0](0).nbytes
-        self.moments_size = (descriptor.d+1) * self.memory.volume * self.float_type[0](0).nbytes
+        self.memory = Memory(self.descriptor, self.grid, self.context, self.float_type[0], align, opengl)
+        self.tick = False
 
-        self.tick = True
-        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.moments = moments
+        self.collide = collide
 
-        if opengl:
-            self.np_moments = numpy.ndarray(shape=(self.memory.volume, 4), dtype=self.float_type[0])
-            self.gl_moments = vbo.VBO(data=self.np_moments, usage=gl.GL_DYNAMIC_DRAW, target=gl.GL_ARRAY_BUFFER)
-            self.gl_moments.bind()
-            self.cl_gl_moments  = cl.GLBuffer(self.context, mf.READ_WRITE, int(self.gl_moments))
-        else:
-            self.cl_moments  = cl.Buffer(self.context, mf.WRITE_ONLY, size=self.moments_size)
+        self.pop_eq_src = pop_eq_src
+        self.boundary_src = boundary_src
 
-        self.cl_material = cl.Buffer(self.context, mf.READ_ONLY, size=self.memory.volume * numpy.int32(0).nbytes)
+        self.layout = layout
+
+        self.compiler_args = {
+            'single': '-cl-single-precision-constant -cl-fast-relaxed-math',
+            'double': '-cl-fast-relaxed-math'
+        }.get(precision, None)
 
         self.build_kernel()
 
         self.program.equilibrilize(
-            self.queue, self.grid.size(), self.layout, self.cl_pop_a, self.cl_pop_b).wait()
-
-    def gid(self, x, y, z = 0):
-        return z * (self.memory.size_x*self.memory.size_y) + y * self.memory.size_x + x;
+            self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_pop_b).wait()
 
     def setup_geometry(self, material_at):
         material = numpy.ndarray(shape=(self.memory.volume, 1), dtype=numpy.int32)
         material[:,:] = 0
 
         for idx in self.geometry.inner_cells():
-            material[self.gid(*idx)] = material_at(self.geometry, *idx)
+            material[self.memory.gid(*idx)] = material_at(self.geometry, *idx)
 
-        cl.enqueue_copy(self.queue, self.cl_material, material).wait();
+        cl.enqueue_copy(self.queue, self.memory.cl_material, material).wait();
 
     def build_kernel(self):
         program_src = Template(filename = str(Path(__file__).parent/'template/kernel.mako')).render(
             descriptor = self.descriptor,
             geometry   = self.geometry,
-            grid       = self.grid,
             memory     = self.memory,
 
             moments_subexpr    = self.moments[0],
@@ -161,16 +168,14 @@ class Lattice:
             pop_eq_src = Template(self.pop_eq_src).render(
                 descriptor = self.descriptor,
                 geometry   = self.geometry,
-                grid       = self.grid,
                 memory     = self.memory,
-                float_type = self.float_type[1]
+                float_type = self.float_type[1],
             ),
             boundary_src = Template(self.boundary_src).render(
                 descriptor = self.descriptor,
                 geometry   = self.geometry,
-                grid       = self.grid,
                 memory     = self.memory,
-                float_type = self.float_type[1]
+                float_type = self.float_type[1],
             ),
 
             ccode = sympy.ccode
@@ -181,11 +186,11 @@ class Lattice:
         if self.tick:
             self.tick = False
             self.program.collide_and_stream(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_a, self.cl_pop_b, self.cl_material)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_pop_b, self.memory.cl_material)
         else:
             self.tick = True
             self.program.collide_and_stream(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_b, self.cl_pop_a, self.cl_material)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_pop_a, self.memory.cl_material)
 
     def sync(self):
         self.queue.finish()
@@ -195,24 +200,24 @@ class Lattice:
 
         if self.tick:
             self.program.collect_moments(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_b, self.cl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_moments)
         else:
             self.program.collect_moments(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_a, self.cl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_moments)
 
-        cl.enqueue_copy(self.queue, moments, self.cl_moments).wait();
+        cl.enqueue_copy(self.queue, moments, self.memory.cl_moments).wait();
 
         return moments
 
     def sync_gl_moments(self):
-        cl.enqueue_acquire_gl_objects(self.queue, [self.cl_gl_moments])
+        cl.enqueue_acquire_gl_objects(self.queue, [self.memory.cl_gl_moments])
 
         if self.tick:
             self.program.collect_gl_moments(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_b, self.cl_gl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_gl_moments)
         else:
             self.program.collect_gl_moments(
-                self.queue, self.grid.size(), self.layout, self.cl_pop_a, self.cl_gl_moments)
+                self.queue, self.grid.size(), self.layout, self.memory.cl_pop_a, self.memory.cl_gl_moments)
 
         #cl.enqueue_release_gl_objects(self.queue, [self.cl_gl_moments])
         self.sync()