Pull in basics from symlbm_playground

It's time to extract the generator-part of my GPU LBM playground and turn it
into a nice reusable library. The goal is to produce a framework that can be
used to generate collision and streaming programs from symbolic descriptions.

i.e. it should be possible to select a LB model, the desired boundary
conditions as well as a data structure / streaming model and use this
information to automatically generate matching OpenCL / CUDA / C++
programs.

14 files changed, 400 insertions, 0 deletions

diff --git a/boltzgen/__init__.py b/boltzgen/__init__.py
new file mode 100644
index 0000000..8896c65
--- /dev/null
+++ b/boltzgen/__init__.py
@@ -0,0 +1,9 @@
+from lbm import LBM
+
+import lbm.model.D2Q9  as D2Q9
+import lbm.model.D3Q19 as D3Q19
+import lbm.model.D3Q27 as D3Q27
+
+from geometry import Geometry
+
+from kernel.generator import source
diff --git a/boltzgen/geometry.py b/boltzgen/geometry.py
new file mode 100644
index 0000000..4b61982
--- /dev/null
+++ b/boltzgen/geometry.py
@@ -0,0 +1,24 @@
+from utility.ndindex import ndindex
+
+class Geometry:
+    def __init__(self, size_x, size_y, size_z = 1):
+        self.size_x = size_x
+        self.size_y = size_y
+        self.size_z = size_z
+        self.volume = size_x * size_y * size_z
+
+    def inner_cells(self):
+        for idx in numpy.ndindex(self.inner_size()):
+            yield tuple(map(lambda i: i + 1, idx))
+
+    def size(self):
+        if self.size_z == 1:
+            return (self.size_x, self.size_y)
+        else:
+            return (self.size_x, self.size_y, self.size_z)
+
+    def inner_size(self):
+        if self.size_z == 1:
+            return (self.size_x-2, self.size_y-2)
+        else:
+            return (self.size_x-2, self.size_y-2, self.size_z-2)
diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py
new file mode 100644
index 0000000..bd1bb86
--- /dev/null
+++ b/boltzgen/kernel/generator.py
@@ -0,0 +1,25 @@
+import sympy
+
+from mako.template import Template
+from pathlib import Path
+
+def source(descriptor, moments, collide, boundary_src, float_type, geometry):
+    return Template(filename = str(Path(__file__).parent/'template/kernel.mako')).render(
+        descriptor = descriptor,
+        geometry   = geometry,
+
+        moments_subexpr    = moments[0],
+        moments_assignment = moments[1],
+        collide_subexpr    = collide[0],
+        collide_assignment = collide[1],
+
+        float_type = float_type,
+
+        boundary_src = Template(boundary_src).render(
+            descriptor = descriptor,
+            geometry   = geometry,
+            float_type = float_type
+        ),
+
+        ccode = sympy.ccode
+    )
diff --git a/boltzgen/kernel/template/kernel.mako b/boltzgen/kernel/template/kernel.mako
new file mode 100644
index 0000000..5ddf64c
--- /dev/null
+++ b/boltzgen/kernel/template/kernel.mako
@@ -0,0 +1,104 @@
+% if float_type == 'double':
+#if defined(cl_khr_fp64)
+#pragma OPENCL EXTENSION cl_khr_fp64 : enable
+#elif defined(cl_amd_fp64)
+#pragma OPENCL EXTENSION cl_amd_fp64 : enable
+#endif
+% endif
+
+<%
+def gid():
+    return {
+        2: 'get_global_id(1)*%d + get_global_id(0)' % geometry.size_x,
+        3: 'get_global_id(2)*%d + get_global_id(1)*%d + get_global_id(0)' % (geometry.size_x*geometry.size_y, geometry.size_x)
+    }.get(descriptor.d)
+
+def pop_offset(i):
+    return i * geometry.volume
+%>
+
+__kernel void equilibrilize(__global ${float_type}* f_next,
+                            __global ${float_type}* f_prev)
+{
+    const unsigned int gid = ${gid()};
+
+    __global ${float_type}* preshifted_f_next = f_next + gid;
+    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+
+% for i, w_i in enumerate(descriptor.w):
+    preshifted_f_next[${pop_offset(i)}] = ${w_i}.f;
+    preshifted_f_prev[${pop_offset(i)}] = ${w_i}.f;
+% endfor
+}
+
+<%
+def neighbor_offset(c_i):
+    return {
+        2: lambda:                                          c_i[1]*geometry.size_x + c_i[0],
+        3: lambda: c_i[2]*geometry.size_x*geometry.size_y + c_i[1]*geometry.size_x + c_i[0]
+    }.get(descriptor.d)()
+
+%>
+
+__kernel void collide_and_stream(__global ${float_type}* f_next,
+                                 __global ${float_type}* f_prev,
+                                 __global int* material,
+                                 unsigned int time)
+{
+    const unsigned int gid = ${gid()};
+
+    const int m = material[gid];
+
+    if ( m == 0 ) {
+        return;
+    }
+
+    __global ${float_type}* preshifted_f_next = f_next + gid;
+    __global ${float_type}* preshifted_f_prev = f_prev + gid;
+
+% for i, c_i in enumerate(descriptor.c):
+    const ${float_type} f_curr_${i} = preshifted_f_prev[${pop_offset(i) + neighbor_offset(-c_i)}];
+% endfor
+
+% for i, expr in enumerate(moments_subexpr):
+    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+% endfor
+
+% for i, expr in enumerate(moments_assignment):
+    ${float_type} ${ccode(expr)}
+% endfor
+
+  ${boundary_src}
+
+% for i, expr in enumerate(collide_subexpr):
+    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+% endfor
+
+% for i, expr in enumerate(collide_assignment):
+    const ${float_type} ${ccode(expr)}
+% endfor
+
+% for i in range(0,descriptor.q):
+    preshifted_f_next[${pop_offset(i)}] = f_next_${i};
+% endfor
+}
+
+__kernel void collect_moments(__global ${float_type}* f,
+                              __global ${float_type}* moments)
+{
+    const unsigned int gid = ${gid()};
+
+    __global ${float_type}* preshifted_f = f + gid;
+
+% for i in range(0,descriptor.q):
+    const ${float_type} f_curr_${i} = preshifted_f[${pop_offset(i)}];
+% endfor
+
+% for i, expr in enumerate(moments_subexpr):
+    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+% endfor
+
+% for i, expr in enumerate(moments_assignment):
+    moments[${pop_offset(i)} + gid] = ${ccode(expr.rhs)};
+% endfor
+}
diff --git a/boltzgen/lbm/__init__.py b/boltzgen/lbm/__init__.py
new file mode 100644
index 0000000..f80feaa
--- /dev/null
+++ b/boltzgen/lbm/__init__.py
@@ -0,0 +1,60 @@
+from sympy import *
+from sympy.codegen.ast import Assignment
+
+import utility.optimizations as optimizations
+from lbm.model.characteristics import weights, c_s
+
+
+def assign(names, definitions):
+    return list(map(lambda x: Assignment(*x), zip(names, definitions)))
+
+class LBM:
+    def __init__(self, descriptor):
+        self.descriptor = descriptor
+        self.f_next = symarray('f_next', descriptor.q)
+        self.f_curr = symarray('f_curr', descriptor.q)
+
+        if not hasattr(descriptor, 'w'):
+            self.descriptor.w = weights(descriptor.d, descriptor.c)
+
+        if not hasattr(descriptor, 'c_s'):
+            self.descriptor.c_s = c_s(descriptor.d, descriptor.c, self.descriptor.w)
+
+    def moments(self, optimize = True):
+        rho = symbols('rho')
+        u   = Matrix(symarray('u', self.descriptor.d))
+
+        exprs = [ Assignment(rho, sum(self.f_curr)) ]
+
+        for i, u_i in enumerate(u):
+            exprs.append(
+                Assignment(u_i, sum([ (c_j*self.f_curr[j])[i] for j, c_j in enumerate(self.descriptor.c) ]) / sum(self.f_curr)))
+
+        if optimize:
+            return cse(exprs, optimizations=optimizations.custom, symbols=numbered_symbols(prefix='m'))
+        else:
+            return ([], exprs)
+
+    def equilibrium(self):
+        rho = symbols('rho')
+        u   = Matrix(symarray('u', self.descriptor.d))
+
+        f_eq = []
+
+        for i, c_i in enumerate(self.descriptor.c):
+            f_eq_i = self.descriptor.w[i] * rho * ( 1
+                                                  + c_i.dot(u)    /    self.descriptor.c_s**2
+                                                  + c_i.dot(u)**2 / (2*self.descriptor.c_s**4)
+                                                  - u.dot(u)      / (2*self.descriptor.c_s**2) )
+            f_eq.append(f_eq_i)
+
+        return f_eq
+
+    def bgk(self, tau, f_eq, optimize = True):
+        exprs = [ self.f_curr[i] + 1/tau * (f_eq_i - self.f_curr[i]) for i, f_eq_i in enumerate(f_eq) ]
+
+        if optimize:
+            subexprs, f = cse(exprs, optimizations=optimizations.custom)
+            return (subexprs, assign(self.f_next, f))
+        else:
+            return ([], assign(self.f_next, exprs))
diff --git a/boltzgen/lbm/model/D2Q9.py b/boltzgen/lbm/model/D2Q9.py
new file mode 100644
index 0000000..e3ac9de
--- /dev/null
+++ b/boltzgen/lbm/model/D2Q9.py
@@ -0,0 +1,7 @@
+from sympy import Matrix
+from itertools import product
+
+d = 2
+q = 9
+
+c = [ Matrix(x) for x in product([-1,0,1], repeat=d) ]
diff --git a/boltzgen/lbm/model/D3Q19.py b/boltzgen/lbm/model/D3Q19.py
new file mode 100644
index 0000000..e9e6eec
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q19.py
@@ -0,0 +1,18 @@
+from sympy import Matrix, Rational, sqrt
+
+d = 3
+q = 19
+
+c = [ Matrix(x) for x in [
+    ( 0, 1, 1), (-1, 0, 1), ( 0, 0, 1), ( 1, 0, 1), ( 0, -1, 1),
+    (-1, 1, 0), ( 0, 1, 0), ( 1, 1, 0), (-1, 0, 0), ( 0, 0, 0), ( 1, 0, 0), (-1,-1, 0), ( 0, -1, 0), ( 1, -1, 0),
+    ( 0, 1,-1), (-1, 0,-1), ( 0, 0,-1), ( 1, 0,-1), ( 0, -1,-1)
+]]
+
+w = [Rational(*x) for x in [
+    (1,36), (1,36), (1,18), (1,36), (1,36),
+    (1,36), (1,18), (1,36), (1,18), (1,3), (1,18), (1,36), (1,18), (1,36),
+    (1,36), (1,36), (1,18), (1,36), (1,36)
+]]
+
+c_s = sqrt(Rational(1,3))
diff --git a/boltzgen/lbm/model/D3Q27.py b/boltzgen/lbm/model/D3Q27.py
new file mode 100644
index 0000000..6fb1f80
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q27.py
@@ -0,0 +1,7 @@
+from sympy import Matrix
+from itertools import product
+
+d = 3
+q = 27
+
+c = [ Matrix(x) for x in product([-1,0,1], repeat=d) ]
diff --git a/boltzgen/lbm/model/D3Q7.py b/boltzgen/lbm/model/D3Q7.py
new file mode 100644
index 0000000..04e16a3
--- /dev/null
+++ b/boltzgen/lbm/model/D3Q7.py
@@ -0,0 +1,18 @@
+from sympy import *
+
+q = 7
+d = 3
+
+c = [ Matrix(x) for x in [
+    ( 0, 0, 1),
+    ( 0, 1, 0), (-1, 0, 0), ( 0, 0, 0), ( 1, 0, 0), ( 0, -1, 0),
+    ( 0, 0,-1)
+]]
+
+w = [Rational(*x) for x in [
+    (1,8),
+    (1,8), (1,8), (1,4), (1,8), (1,8),
+    (1,8)
+]]
+
+c_s = sqrt(Rational(1,4))
diff --git a/boltzgen/lbm/model/characteristics.py b/boltzgen/lbm/model/characteristics.py
new file mode 100644
index 0000000..b68afeb
--- /dev/null
+++ b/boltzgen/lbm/model/characteristics.py
@@ -0,0 +1,27 @@
+from sympy import *
+
+# copy of `sympy.integrals.quadrature.gauss_hermite` sans evaluation
+def gauss_hermite(n):
+    x = Dummy("x")
+    p  = hermite_poly(n, x, polys=True)
+    p1 = hermite_poly(n-1, x, polys=True)
+    xi = []
+    w  = []
+    for r in p.real_roots():
+        xi.append(r)
+        w.append(((2**(n-1) * factorial(n) * sqrt(pi))/(n**2 * p1.subs(x, r)**2)))
+    return xi, w
+
+# determine weights of a d-dimensional LBM model on velocity set c
+# (only works for velocity sets that result into NSE-recovering LB models when
+#  plugged into Gauss-Hermite quadrature without any additional arguments
+#  i.e. D2Q9 and D3Q27 but not D3Q19)
+def weights(d, c):
+    _, omegas = gauss_hermite(3)
+    return list(map(lambda c_i: Mul(*[ omegas[1+c_i[iDim]] for iDim in range(0,d) ]) / pi**(d/2), c))
+
+# determine lattice speed of sound using directions and their weights
+def c_s(d, c, w):
+    speeds = set([ sqrt(sum([ w[i] * c_i[j]**2 for i, c_i in enumerate(c) ])) for j in range(0,d) ])
+    assert len(speeds) == 1 # verify isotropy
+    return speeds.pop()
diff --git a/boltzgen/utility/ndindex.py b/boltzgen/utility/ndindex.py
new file mode 100644
index 0000000..0c5ed9b
--- /dev/null
+++ b/boltzgen/utility/ndindex.py
@@ -0,0 +1,14 @@
+import numpy
+from numpy.lib.stride_tricks import as_strided
+import numpy.core.numeric as _nx
+
+class ndindex(numpy.ndindex):
+    pass
+
+    def __init__(self, *shape, order):
+        if len(shape) == 1 and isinstance(shape[0], tuple):
+            shape = shape[0]
+        x = as_strided(_nx.zeros(1), shape=shape,
+                       strides=_nx.zeros_like(shape))
+        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
+                              order=order)
diff --git a/boltzgen/utility/optimizations.py b/boltzgen/utility/optimizations.py
new file mode 100644
index 0000000..93dad09
--- /dev/null
+++ b/boltzgen/utility/optimizations.py
@@ -0,0 +1,10 @@
+from sympy import *
+
+from sympy.codegen.rewriting import ReplaceOptim
+
+expand_square = ReplaceOptim(
+    lambda e: e.is_Pow and e.exp.is_integer and e.exp == 2,
+    lambda p: UnevaluatedExpr(Mul(p.base, p.base, evaluate = False))
+)
+
+custom = [ (expand_square, expand_square) ] + cse_main.basic_optimizations
diff --git a/shell.nix b/shell.nix
new file mode 100644
index 0000000..b91c772
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,62 @@
+{ pkgs ? import <nixpkgs> { }, ... }:
+
+pkgs.stdenvNoCC.mkDerivation rec {
+  name = "boltzgen-env";
+  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; };
+
+    pyevtk = pkgs.python3.pkgs.buildPythonPackage rec {
+      pname = "PyEVTK";
+      version = "1.2.1";
+
+      src = pkgs.fetchFromGitHub {
+        owner  = "paulo-herrera";
+        repo   = "PyEVTK";
+        rev    = "v1.2.1";
+        sha256 = "1p2459dqvgakywvy5d31818hix4kic6ks9j4m582ypxyk5wj1ksz";
+      };
+
+      buildInputs = with pkgs.python37Packages; [
+        numpy
+      ];
+
+      doCheck = false;
+    };
+
+    local-python = custom-python.withPackages (python-packages: with python-packages; [
+      numpy
+      sympy
+      pyopencl
+      pyopengl
+      pyrr
+      Mako
+      pyevtk
+    ]);
+
+  in [
+    local-python
+    pkgs.opencl-info
+  ];
+
+  shellHook = ''
+    export NIX_SHELL_NAME="${name}"
+    export PYOPENCL_COMPILER_OUTPUT=1
+    export PYTHONPATH="$PWD/boltzgen:$PYTHONPATH"
+  '';
+}
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..b5831aa
--- /dev/null
+++ b/test.py
@@ -0,0 +1,15 @@
+from boltzgen import *
+
+lbm = LBM(D2Q9)
+geometry = Geometry(32,32)
+
+src = source(
+    D2Q9,
+    lbm.moments(),
+    lbm.bgk(f_eq = lbm.equilibrium(), tau = 0.6),
+    "",
+    'float',
+    geometry
+)
+
+print(src)