Tidy up symbolic kernel generation

5 files changed, 60 insertions, 33 deletions

diff --git a/D2Q9.py b/D2Q9.py
index 88bd88c..27ae375 100644
--- a/D2Q9.py
+++ b/D2Q9.py
@@ -12,12 +12,16 @@ class Lattice:
     def idx(self, x, y):
         return y * self.nX + x;
 
-    def __init__(self, nX, nY, tau, geometry, pop_eq_src = '', boundary_src = ''):
+    def __init__(self, nX, nY, geometry, moments, collide, pop_eq_src = '', boundary_src = ''):
         self.nX = nX
         self.nY = nY
         self.nCells = nX * nY
-        self.tau = tau
-        self.tick = True
+
+        self.moments = moments
+        self.collide = collide
+
+        self.pop_eq_src = pop_eq_src
+        self.boundary_src = boundary_src
 
         self.platform = cl.get_platforms()[0]
         self.context  = cl.Context(properties=[(cl.context_properties.PLATFORM, self.platform)])
@@ -26,9 +30,7 @@ class Lattice:
         self.np_material = numpy.ndarray(shape=(self.nCells, 1), dtype=numpy.int32)
         self.setup_geometry(geometry)
 
-        self.pop_eq_src = pop_eq_src
-        self.boundary_src = boundary_src
-
+        self.tick = True
         self.cl_pop_a = cl.Buffer(self.context, mf.READ_WRITE, size=9*self.nCells*numpy.float32(0).nbytes)
         self.cl_pop_b = cl.Buffer(self.context, mf.READ_WRITE, size=9*self.nCells*numpy.float32(0).nbytes)
 
@@ -49,11 +51,10 @@ class Lattice:
             nX     = self.nX,
             nY     = self.nY,
             nCells = self.nCells,
-            tau    = self.tau,
-            moments_helper     = D2Q9.moments_opt[0],
-            moments_assignment = D2Q9.moments_opt[1],
-            collide_helper     = D2Q9.collide_opt[0],
-            collide_assignment = D2Q9.collide_opt[1],
+            moments_helper     = self.moments[0],
+            moments_assignment = self.moments[1],
+            collide_helper     = self.collide[0],
+            collide_assignment = self.collide[1],
             c     = D2Q9.c,
             w     = D2Q9.w,
             ccode = sympy.ccode,
@@ -69,7 +70,7 @@ class Lattice:
                 d = D2Q9.d
             )
         )
-        self.program = cl.Program(self.context, program_src).build()
+        self.program = cl.Program(self.context, program_src).build('-cl-single-precision-constant')
 
     def evolve(self):
         if self.tick:
diff --git a/implosion.py b/implosion.py
index e5e87bc..fb101b4 100644
--- a/implosion.py
+++ b/implosion.py
@@ -7,6 +7,8 @@ matplotlib.use('AGG')
 
 from D2Q9 import Lattice
 
+import symbolic.D2Q9 as D2Q9
+
 def MLUPS(cells, steps, time):
     return cells * steps / time * 1e-6
 
@@ -56,7 +58,13 @@ moments = []
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(nX = 1024, nY = 1024, tau = 0.8, geometry = box, pop_eq_src = pop_eq, boundary_src = boundary)
+lattice = Lattice(
+    nX = 1024, nY = 1024,
+    geometry = box,
+    moments = D2Q9.moments(optimize = False),
+    collide = D2Q9.bgk(tau = 0.8),
+    pop_eq_src   = pop_eq,
+    boundary_src = boundary)
 
 print("Starting simulation using %d cells...\n" % lattice.nCells)
 
diff --git a/lid_driven_cavity.py b/lid_driven_cavity.py
index db51e73..a4de67d 100644
--- a/lid_driven_cavity.py
+++ b/lid_driven_cavity.py
@@ -7,6 +7,8 @@ matplotlib.use('AGG')
 
 from D2Q9 import Lattice
 
+import symbolic.D2Q9 as D2Q9
+
 def MLUPS(cells, steps, time):
     return cells * steps / time * 1e-6
 
@@ -23,7 +25,7 @@ def generate_moment_plots(lattice, moments):
         plt.imshow(velocity, origin='lower', cmap=plt.get_cmap('seismic'))
         plt.savefig("result/velocity_" + str(i) + ".png", bbox_inches='tight', pad_inches=0)
 
-def box(nX, nY, x, y):
+def cavity(nX, nY, x, y):
     if x == 1 or y == 1 or x == nX-2:
         return 2
     elif y == nY-2:
@@ -49,7 +51,12 @@ moments = []
 
 print("Initializing simulation...\n")
 
-lattice = Lattice(nX = 256, nY = 256, tau = 0.56, geometry = box, boundary_src = boundary)
+lattice = Lattice(
+    nX = 256, nY = 256,
+    geometry = cavity,
+    moments = D2Q9.moments(optimize = False),
+    collide = D2Q9.bgk(tau = 0.56),
+    boundary_src = boundary)
 
 print("Starting simulation using %d cells...\n" % lattice.nCells)
 
diff --git a/symbolic/D2Q9.py b/symbolic/D2Q9.py
index a3c2503..8d42245 100644
--- a/symbolic/D2Q9.py
+++ b/symbolic/D2Q9.py
@@ -9,29 +9,42 @@ w = [ Rational(*x) for x in [(1,36), (1,9), (1,36), (1,9), (4,9), (1,9), (1,36),
 
 c_s = sqrt(Rational(1,3))
 
-rho, tau = symbols('rho tau')
-
 f_next = symarray('f_next', q)
 f_curr = symarray('f_curr', q)
 
-u = Matrix(symarray('u', d))
+def moments(f = f_curr, optimize = True):
+    rho = symbols('rho')
+    u   = Matrix(symarray('u', d))
+
+    exprs = [ Assignment(rho, sum(f)) ]
+
+    for i, u_i in enumerate(u):
+        exprs.append(Assignment(u_i, sum([ (c_j*f[j])[i] for j, c_j in enumerate(c) ]) / sum(f)))
 
-moments = [ Assignment(rho, sum(f_curr)) ]
+    if optimize:
+        return cse(exprs, optimizations='basic', symbols=numbered_symbols(prefix='m'))
+    else:
+        return ([], exprs)
 
-for i, u_i in enumerate(u):
-    moments.append(Assignment(u_i, sum([ (c_j*f_curr[j])[i] for j, c_j in enumerate(c) ]) / sum(f_curr)))
+def equilibrium():
+    rho = symbols('rho')
+    u   = Matrix(symarray('u', d))
 
-moments_opt = cse(moments, optimizations='basic', symbols=numbered_symbols(prefix='m'))
+    f_eq = []
 
-f_eq = []
+    for i, c_i in enumerate(c):
+        f_eq_i = w[i] * rho * (  1
+                               + c_i.dot(u)    /    c_s**2
+                               + c_i.dot(u)**2 / (2*c_s**4)
+                               - u.dot(u)      / (2*c_s**2) )
+        f_eq.append(f_eq_i)
 
-for i, c_i in enumerate(c):
-    f_eq_i = w[i] * rho * (  1
-                           + c_i.dot(u)    /    c_s**2
-                           + c_i.dot(u)**2 / (2*c_s**4)
-                           - u.dot(u)      / (2*c_s**2) )
-    f_eq.append(f_eq_i)
+    return f_eq
 
-collide = [ Assignment(f_next[i], f_curr[i] + 1/tau * ( f_eq_i - f_curr[i] )) for i, f_eq_i in enumerate(f_eq) ]
+def bgk(tau, f_eq = equilibrium(), optimize = True):
+    exprs = [ Assignment(f_next[i], f_curr[i] + 1/tau * ( f_eq_i - f_curr[i] )) for i, f_eq_i in enumerate(f_eq) ]
 
-collide_opt = cse(collide, optimizations='basic')
+    if optimize:
+        return cse(exprs, optimizations='basic')
+    else:
+        return ([], exprs)
diff --git a/template/kernel.mako b/template/kernel.mako
index b81d432..6f82b20 100644
--- a/template/kernel.mako
+++ b/template/kernel.mako
@@ -1,5 +1,3 @@
-__constant float tau = ${tau};
-
 __kernel void equilibrilize(__global __write_only float* f_a,
                             __global __write_only float* f_b)
 {