Add OpenCL GL interop example

5 files changed, 501 insertions, 2 deletions

diff --git a/lid_driven_cavity/opencl_gl_interop/AA.py b/lid_driven_cavity/opencl_gl_interop/AA.py
new file mode 100644
index 0000000..ca9a132
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/AA.py
@@ -0,0 +1,123 @@
+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);
+}
+"""
+
+class MomentsTexture(MomentsTextureBase):
+    pass
+
+    def collect(self):
+        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
+
+        if self.lattice.tick:
+            self.lattice.program.collect_moments_to_texture_tick(
+                self.lattice.queue,
+                self.lattice.geometry.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop,
+                self.cl_gl_moments)
+        else:
+            self.lattice.program.collect_moments_to_texture_tock(
+                self.lattice.queue,
+                self.lattice.geometry.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop,
+                self.cl_gl_moments)
diff --git a/lid_driven_cavity/opencl_gl_interop/AB.py b/lid_driven_cavity/opencl_gl_interop/AB.py
new file mode 100644
index 0000000..f6b437c
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/AB.py
@@ -0,0 +1,126 @@
+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);
+}
+"""
+
+class MomentsTexture(MomentsTextureBase):
+    pass
+
+    def collect(self):
+        cl.enqueue_acquire_gl_objects(self.lattice.queue, [self.cl_gl_moments])
+
+        if self.lattice.tick:
+            self.lattice.program.collect_moments_to_texture(
+                self.lattice.queue,
+                self.lattice.geometry.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_a,
+                self.cl_gl_moments)
+        else:
+            self.lattice.program.collect_moments_to_texture(
+                self.lattice.queue,
+                self.lattice.geometry.size(),
+                self.lattice.layout,
+                self.lattice.memory.cl_pop_b,
+                self.cl_gl_moments)
diff --git a/lid_driven_cavity/opencl_gl_interop/common.py b/lid_driven_cavity/opencl_gl_interop/common.py
new file mode 100644
index 0000000..808ec0c
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/common.py
@@ -0,0 +1,53 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+from OpenGL.GL import *
+from OpenGL.arrays import vbo
+
+import numpy
+
+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 MomentsTextureBase:
+    def __init__(self, lattice):
+        self.lattice = lattice
+        self.gl_texture_buffer = numpy.ndarray(shape=(self.lattice.memory.volume, 4), dtype=self.lattice.memory.float_type)
+        self.gl_texture_buffer[:,:] = 0.0
+
+        self.gl_moments = glGenTextures(1)
+        self.gl_texture_type = {2: GL_TEXTURE_2D, 3: GL_TEXTURE_3D}.get(self.lattice.descriptor.d)
+        glBindTexture(self.gl_texture_type, self.gl_moments)
+
+        if self.gl_texture_type == GL_TEXTURE_3D:
+            glTexImage3D(self.gl_texture_type, 0, GL_RGBA32F, self.lattice.memory.size_x, self.lattice.memory.size_y, self.lattice.memory.size_z, 0, GL_RGBA, GL_FLOAT, self.gl_texture_buffer)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_T,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_S,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_R,     GL_CLAMP_TO_EDGE)
+            self.cl_gl_moments  = cl.GLTexture(self.lattice.context, mf.READ_WRITE, self.gl_texture_type, 0, self.gl_moments, 3)
+        elif self.gl_texture_type == GL_TEXTURE_2D:
+            glTexImage2D(self.gl_texture_type, 0, GL_RGBA32F, self.lattice.memory.size_x, self.lattice.memory.size_y, 0, GL_RGBA, GL_FLOAT, self.gl_texture_buffer)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_T,     GL_CLAMP_TO_EDGE)
+            glTexParameteri(self.gl_texture_type, GL_TEXTURE_WRAP_S,     GL_CLAMP_TO_EDGE)
+            self.cl_gl_moments  = cl.GLTexture(self.lattice.context, mf.READ_WRITE, self.gl_texture_type, 0, self.gl_moments, 2)
+
+    def bind(self, location = GL_TEXTURE0):
+        glEnable(self.gl_texture_type)
+        glActiveTexture(location);
+        glBindTexture(self.gl_texture_type, self.gl_moments)
diff --git a/lid_driven_cavity/opencl_gl_interop/ldc_2d.py b/lid_driven_cavity/opencl_gl_interop/ldc_2d.py
new file mode 100644
index 0000000..daac486
--- /dev/null
+++ b/lid_driven_cavity/opencl_gl_interop/ldc_2d.py
@@ -0,0 +1,180 @@
+import numpy
+import time
+from string import Template
+
+from boltzgen import Generator, Geometry
+from boltzgen.lbm.lattice import D2Q9
+from boltzgen.lbm.model   import BGK
+
+from common import CellList
+
+from OpenGL.GL   import *
+from OpenGL.GLUT import *
+from OpenGL.GL import shaders
+from pyrr import matrix44
+
+geometry = Geometry(512, 512)
+
+functions = ['collide_and_stream', 'equilibrilize', 'collect_moments', 'momenta_boundary']
+extras    = ['cell_list_dispatch', 'opencl_gl_interop']
+
+precision = 'single'
+streaming = 'AA'
+
+import AA
+import AB
+
+def glut_window(fullscreen = False):
+    glutInit(sys.argv)
+    glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
+
+    if fullscreen:
+        window = glutEnterGameMode()
+    else:
+        glutInitWindowSize(800, 500)
+        glutInitWindowPosition(0, 0)
+        window = glutCreateWindow("LDC 2D")
+
+    return window
+
+window = glut_window(fullscreen = False)
+
+Lattice = eval('%s.Lattice' % streaming)
+HelperTemplate = eval('%s.HelperTemplate' % streaming)
+MomentsTexture = eval('%s.MomentsTexture' % streaming)
+
+generator = Generator(
+    model     = BGK(D2Q9, tau = 0.53),
+    target    = 'cl',
+    precision = precision,
+    streaming = streaming,
+    index     = 'ZYX',
+    layout    = 'SOA')
+
+kernel_src  = generator.kernel(geometry, functions, extras)
+kernel_src += generator.custom(geometry, HelperTemplate)
+
+lattice = Lattice(geometry, kernel_src, D2Q9, precision = precision)
+moments = MomentsTexture(lattice)
+
+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,
+    [ gid(x,y) for x, y in geometry.inner_cells() if x == 1 or y == 1 or x == geometry.size_x-2 ])
+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]))
+
+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]))
+
+
+def get_projection(width, height):
+    world_height = geometry.size_y
+    world_width  = world_height / height * width
+
+    projection  = matrix44.create_orthogonal_projection(-world_width/2, world_width/2, -world_height/2, world_height/2, -1, 1)
+    translation = matrix44.create_from_translation([-geometry.size_x/2, -geometry.size_y/2, 0])
+
+    point_size = width / world_width
+
+    return numpy.matmul(translation, projection), point_size
+
+vertex_shader = shaders.compileShader("""
+#version 430
+
+layout (location=0) in vec4 vertex;
+                   out vec2 frag_pos;
+
+uniform mat4 projection;
+
+void main() {
+    gl_Position = projection * vertex;
+    frag_pos    = vertex.xy;
+}""", GL_VERTEX_SHADER)
+
+fragment_shader = shaders.compileShader(Template("""
+#version 430
+
+in vec2 frag_pos;
+
+uniform sampler2D moments;
+
+out vec4 result;
+
+vec2 unit(vec2 v) {
+    return vec2(v[0] / $size_x, v[1] / $size_y);
+}
+
+vec3 blueRedPalette(float x) {
+    return mix(
+        vec3(0.0, 0.0, 1.0),
+        vec3(1.0, 0.0, 0.0),
+        x
+    );
+}
+
+void main(){
+    const vec2 sample_pos = unit(frag_pos);
+    const vec4 data = texture(moments, sample_pos);
+    result.a = 1.0;
+    result.rgb = blueRedPalette(data[3] / 0.1);
+}
+""").substitute({
+    "size_x": geometry.size_x,
+    "size_y": geometry.size_y,
+}), GL_FRAGMENT_SHADER)
+
+shader_program = shaders.compileProgram(vertex_shader, fragment_shader)
+projection_id = shaders.glGetUniformLocation(shader_program, 'projection')
+
+def on_display():
+    for i in range(0,100):
+        lattice.evolve()
+
+    moments.collect()
+
+    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
+
+    shaders.glUseProgram(shader_program)
+    glUniformMatrix4fv(projection_id, 1, False, numpy.asfortranarray(projection))
+    moments.bind()
+
+    glBegin(GL_POLYGON)
+    glVertex(0,0,0)
+    glVertex(lattice.geometry.size_x,0,0)
+    glVertex(lattice.geometry.size_x,lattice.geometry.size_y,0)
+    glVertex(0,lattice.geometry.size_y,0)
+    glEnd()
+
+    glutSwapBuffers()
+
+def on_reshape(width, height):
+    global projection, point_size
+    glViewport(0,0,width,height)
+    projection, point_size = get_projection(width, height)
+
+def on_timer(t):
+    glutTimerFunc(t, on_timer, t)
+    glutPostRedisplay()
+
+glutDisplayFunc(on_display)
+glutReshapeFunc(on_reshape)
+glutTimerFunc(10, on_timer, 10)
+
+glutMainLoop()
diff --git a/shell.nix b/shell.nix
index 4e36dce..e3e52d4 100644
--- a/shell.nix
+++ b/shell.nix
@@ -5,13 +5,29 @@ pkgs.stdenvNoCC.mkDerivation rec {
   env = pkgs.buildEnv { name = name; paths = buildInputs; };
 
   buildInputs = let
+    custom-python = let
+      packageOverrides = self: super: {
+        pyopencl = super.pyopencl.overridePythonAttrs(old: rec {
+          buildInputs = with pkgs; [
+            opencl-headers ocl-icd python37Packages.pybind11
+            libGLU_combined
+          ];
+        # Enable OpenGL integration and fix build
+          preBuild = ''
+            python configure.py --cl-enable-gl
+            export HOME=/tmp/pyopencl
+          '';
+        });
+      };
+    in pkgs.python3.override { inherit packageOverrides; };
+
     boltzgen = pkgs.python3.pkgs.buildPythonPackage rec {
       pname = "boltzgen";
       version = "0.1";
 
       src = builtins.fetchGit {
         url = "https://code.kummerlaender.eu/boltzgen/";
-        rev = "814e6253475c7955eb6a46d814e5a86974e58613";
+        rev = "286e243a171c8bcdfc91b5b6dcdd937ac95b0b7b";
       };
 
       propagatedBuildInputs = with pkgs.python37Packages; [
@@ -23,10 +39,11 @@ pkgs.stdenvNoCC.mkDerivation rec {
       doCheck = false;
     };
 
-    local-python = pkgs.python3.withPackages (python-packages: with python-packages; [
+    local-python = custom-python.withPackages (python-packages: with python-packages; [
       boltzgen
       numpy
       pyopencl setuptools
+      pyopengl pyrr
       matplotlib
     ]);