Basic 2D LDC using boltzgen for kernel generation

Using cell lists as parameters for multiple non-branching kernels seems
to reduce performance by ~50 MLUPS (for single precision D2Q9).
This might be alleviated by padding the cell lists to enable thread
layout control or by improved kernel dispatching.

On the upside this OpenCL program runs not only on GPUs but is also vectorized on Intel
CPUs yielding about 180 MLUPS (single precision) and - anticlimactically - 85 MLUPS for
double precision on a i7-4790K.
However both these values compare well to the performance of established CPU LBM codes.

3 files changed, 242 insertions, 0 deletions

diff --git a/ldc_2d.py b/ldc_2d.py
new file mode 100644
index 0000000..14e8a5c
--- /dev/null
+++ b/ldc_2d.py
@@ -0,0 +1,94 @@
+import numpy
+import time
+from string import Template
+
+import matplotlib
+matplotlib.use('AGG')
+import matplotlib.pyplot as plt
+
+from boltzgen import LBM, Generator, Geometry
+from boltzgen.lbm.model import D2Q9
+
+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)))
+
+        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.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'))
+        plt.savefig("result/ldc_2d_%02d.png" % i, bbox_inches='tight', pad_inches=0)
+
+nUpdates = 100000
+nStat    = 5000
+
+geometry = Geometry(512, 512)
+
+print("Generating kernel using boltzgen...\n")
+
+lbm = LBM(D2Q9)
+generator = Generator(
+    descriptor = D2Q9,
+    moments    = lbm.moments(),
+    collision  = lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.6))
+
+functions = ['collide_and_stream', 'equilibrilize', 'collect_moments', 'momenta_boundary']
+extras    = ['cell_list_dispatch']
+
+kernel_src = generator.kernel('cl', 'single', 'SOA', geometry, functions, extras) + Template("""
+__kernel void equilibrilize(__global $float_type* f_next,
+                            __global $float_type* f_prev)
+{
+    const unsigned int gid = get_global_id(1)*$size_x + get_global_id(0);
+    equilibrilize_gid(f_next, f_prev, gid);
+}
+
+__kernel void collect_moments(__global $float_type* f,
+                              __global $float_type* moments)
+{
+    const unsigned int gid = get_global_id(1)*$size_x + get_global_id(0);
+    collect_moments_gid(f, moments, gid);
+}
+""").substitute(float_type = 'float', size_x = geometry.size_x)
+
+print("Initializing simulation...\n")
+
+lattice = Lattice(geometry, kernel_src)
+gid = lattice.memory.gid
+
+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 ])
+
+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]))
+
+print("Starting simulation using %d cells...\n" % lattice.geometry.volume)
+
+moments = []
+
+lastStat = time.time()
+
+for i in range(1,nUpdates+1):
+    lattice.evolve()
+
+    if i % nStat == 0:
+        lattice.sync()
+        print("i = %4d; %3.0f MLUPS" % (i, MLUPS(lattice.geometry.volume, nStat, time.time() - lastStat)))
+        moments.append(lattice.get_moments())
+        lastStat = time.time()
+
+print("\nConcluded simulation.\n")
+
+generate_moment_plots(lattice, moments)
diff --git a/shell.nix b/shell.nix
new file mode 100644
index 0000000..99d794b
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,43 @@
+{ pkgs ? import <nixpkgs> { }, ... }:
+
+pkgs.stdenvNoCC.mkDerivation rec {
+  name = "pycl-env";
+  env = pkgs.buildEnv { name = name; paths = buildInputs; };
+
+  buildInputs = let
+    boltzgen = pkgs.python3.pkgs.buildPythonPackage rec {
+      pname = "boltzgen";
+      version = "0.1";
+
+      src = pkgs.fetchFromGitHub {
+        owner  = "KnairdA";
+        repo   = "boltzgen";
+        rev    = "v0.1";
+        sha256 = "072kx4jrzd0g9rn63hjb0yic7qhbga47lp2vbz7rq3gvkqv1hz4d";
+      };
+
+      propagatedBuildInputs = with pkgs.python37Packages; [
+        sympy
+        numpy
+        Mako
+      ];
+    };
+
+    local-python = pkgs.python3.withPackages (python-packages: with python-packages; [
+      boltzgen
+      numpy
+      pyopencl setuptools
+      matplotlib
+    ]);
+
+  in [
+    local-python
+    pkgs.opencl-info
+  ];
+
+  shellHook = ''
+    export NIX_SHELL_NAME="${name}"
+    export PYOPENCL_COMPILER_OUTPUT=1
+    export PYTHONPATH="$PWD:$PYTHONPATH"
+  '';
+}
diff --git a/simulation.py b/simulation.py
new file mode 100644
index 0000000..d2f24c9
--- /dev/null
+++ b/simulation.py
@@ -0,0 +1,105 @@
+import pyopencl as cl
+mf = cl.mem_flags
+
+import numpy
+
+class Memory:
+    def __init__(self, grid, context, float_type):
+        self.context    = context
+        self.float_type = float_type
+
+        self.size_x = grid.size_x
+        self.size_y = grid.size_y
+        self.size_z = grid.size_z
+
+        self.volume = self.size_x * self.size_y * self.size_z
+
+        self.pop_size     = 9 * 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 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, platform = 0, precision = 'single'):
+        self.geometry = geometry
+
+        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(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(
+            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=(3, self.memory.volume), dtype=self.float_type[0])
+
+        if self.tick:
+            self.program.collect_moments(
+                self.queue, self.geometry.size(), self.layout, self.memory.cl_pop_b, self.memory.cl_moments)
+        else:
+            self.program.collect_moments(
+                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