From e2b00f4ec963060be98939c7b0d12d6c00e50a02 Mon Sep 17 00:00:00 2001 From: Adrian Kummerlaender Date: Thu, 31 Oct 2019 13:13:00 +0100 Subject: Call symbolic generator inside code templates This paves the way for dropping in other LBM collision models. As a side benefit the default momenta calulcation is now fully inlined where possible. --- boltzgen.py | 5 +--- boltzgen/kernel/generator.py | 17 ++++--------- .../kernel/template/bounce_back_boundary.cl.mako | 14 ++++------- .../kernel/template/bounce_back_boundary.cpp.mako | 14 ++++------- boltzgen/kernel/template/collect_moments.cl.mako | 4 +++ boltzgen/kernel/template/collect_moments.cpp.mako | 4 +++ .../kernel/template/collide_and_stream.cl.mako | 14 ++++------- .../kernel/template/collide_and_stream.cpp.mako | 16 +++++------- boltzgen/kernel/template/momenta_boundary.cl.mako | 5 ++++ boltzgen/kernel/template/momenta_boundary.cpp.mako | 5 ++++ boltzgen/lbm/__init__.py | 29 ++++++++++++++++------ boltzgen/utility/optimizations.py | 12 +++++++-- 12 files changed, 77 insertions(+), 62 deletions(-) diff --git a/boltzgen.py b/boltzgen.py index b4a9fe8..82adafe 100755 --- a/boltzgen.py +++ b/boltzgen.py @@ -26,11 +26,8 @@ lattice = eval("lbm.model.%s" % args.lattice) if args.index is None: args.index = 'XYZ' -lbm = LBM(lattice) generator = Generator( - descriptor = lattice, - moments = lbm.moments(optimize = not args.disable_cse), - collision = lbm.bgk(f_eq = lbm.equilibrium(), tau = float(args.tau), optimize = not args.disable_cse), + model = LBM(lattice, tau = float(args.tau), optimize = not args.disable_cse), target = args.target, precision = args.precision, index = args.index, diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py index 92ceb72..4fd3c25 100644 --- a/boltzgen/kernel/generator.py +++ b/boltzgen/kernel/generator.py @@ -6,10 +6,9 @@ from pathlib import Path from . import memory class Generator: - def __init__(self, descriptor, moments, collision, target, precision, index, layout): - self.descriptor = descriptor - self.moments = moments - self.collision = collision + def __init__(self, model, target, precision, index, layout): + self.model = model + self.descriptor = self.model.descriptor self.target = target self.float_type = eval("memory.precision.%s" % target).get_float_type(precision) @@ -30,14 +29,11 @@ class Generator: return Template(filename = str(template_path)).render( descriptor = self.descriptor, + model = self.model, geometry = geometry, index = self.index_impl(geometry), layout = self.layout_impl(self.descriptor, self.index_impl, geometry), - moments_subexpr = self.moments[0], - moments_assignment = self.moments[1], - collision_subexpr = self.collision[0], - collision_assignment = self.collision[1], ccode = sympy.ccode, float_type = self.float_type, @@ -54,14 +50,11 @@ class Generator: def custom(self, geometry, source): return Template(text = source).render( descriptor = self.descriptor, + model = self.model, geometry = geometry, index = self.index_impl(geometry), layout = self.layout_impl(self.descriptor, self.index_impl, geometry), - moments_subexpr = self.moments[0], - moments_assignment = self.moments[1], - collision_subexpr = self.collision[0], - collision_assignment = self.collision[1], ccode = sympy.ccode, float_type = self.float_type, diff --git a/boltzgen/kernel/template/bounce_back_boundary.cl.mako b/boltzgen/kernel/template/bounce_back_boundary.cl.mako index cb59cbd..8bb4f6c 100644 --- a/boltzgen/kernel/template/bounce_back_boundary.cl.mako +++ b/boltzgen/kernel/template/bounce_back_boundary.cl.mako @@ -9,19 +9,15 @@ __kernel void bounce_back_boundary_gid(__global ${float_type}* f_next, const ${float_type} f_curr_${i} = preshifted_f_prev[${layout.pop_offset(i) + layout.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 +<% + subexpr, assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = True)) +%> -% for i, expr in enumerate(collision_subexpr): +% for i, expr in enumerate(subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor -% for i, expr in enumerate(collision_assignment): +% for i, expr in enumerate(assignment): const ${float_type} ${ccode(expr)} % endfor diff --git a/boltzgen/kernel/template/bounce_back_boundary.cpp.mako b/boltzgen/kernel/template/bounce_back_boundary.cpp.mako index c7abd2a..91dcfa0 100644 --- a/boltzgen/kernel/template/bounce_back_boundary.cpp.mako +++ b/boltzgen/kernel/template/bounce_back_boundary.cpp.mako @@ -9,19 +9,15 @@ void bounce_back_boundary( ${float_type}* f_next, const ${float_type} f_curr_${i} = preshifted_f_prev[${layout.pop_offset(i) + layout.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 +<% + subexpr, assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = True)) +%> -% for i, expr in enumerate(collision_subexpr): +% for i, expr in enumerate(subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor -% for i, expr in enumerate(collision_assignment): +% for i, expr in enumerate(assignment): const ${float_type} ${ccode(expr)} % endfor diff --git a/boltzgen/kernel/template/collect_moments.cl.mako b/boltzgen/kernel/template/collect_moments.cl.mako index ddf5102..587f58a 100644 --- a/boltzgen/kernel/template/collect_moments.cl.mako +++ b/boltzgen/kernel/template/collect_moments.cl.mako @@ -9,6 +9,10 @@ __kernel void collect_moments_gid(__global ${float_type}* f, const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i)}]; % endfor +<% + moments_subexpr, moments_assignment = model.moments() +%> + % for i, expr in enumerate(moments_subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor diff --git a/boltzgen/kernel/template/collect_moments.cpp.mako b/boltzgen/kernel/template/collect_moments.cpp.mako index d72ff5d..1605084 100644 --- a/boltzgen/kernel/template/collect_moments.cpp.mako +++ b/boltzgen/kernel/template/collect_moments.cpp.mako @@ -9,6 +9,10 @@ void collect_moments(const ${float_type}* f, const ${float_type} f_curr_${i} = preshifted_f[${layout.pop_offset(i)}]; % endfor +<% + moments_subexpr, moments_assignment = model.moments() +%> + % for i, expr in enumerate(moments_subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor diff --git a/boltzgen/kernel/template/collide_and_stream.cl.mako b/boltzgen/kernel/template/collide_and_stream.cl.mako index c586884..303b21b 100644 --- a/boltzgen/kernel/template/collide_and_stream.cl.mako +++ b/boltzgen/kernel/template/collide_and_stream.cl.mako @@ -9,19 +9,15 @@ __kernel void collide_and_stream_gid(__global ${float_type}* f_next, const ${float_type} f_curr_${i} = preshifted_f_prev[${layout.pop_offset(i) + layout.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 +<% + subexpr, assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = True)) +%> -% for i, expr in enumerate(collision_subexpr): +% for i, expr in enumerate(subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor -% for i, expr in enumerate(collision_assignment): +% for i, expr in enumerate(assignment): const ${float_type} ${ccode(expr)} % endfor diff --git a/boltzgen/kernel/template/collide_and_stream.cpp.mako b/boltzgen/kernel/template/collide_and_stream.cpp.mako index 20fc529..47dd6dd 100644 --- a/boltzgen/kernel/template/collide_and_stream.cpp.mako +++ b/boltzgen/kernel/template/collide_and_stream.cpp.mako @@ -9,23 +9,19 @@ void collide_and_stream( ${float_type}* f_next, const ${float_type} f_curr_${i} = preshifted_f_prev[${layout.pop_offset(i) + layout.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 +<% + subexpr, assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = True)) +%> -% for i, expr in enumerate(collision_subexpr): +% for i, expr in enumerate(subexpr): const ${float_type} ${expr[0]} = ${ccode(expr[1])}; % endfor -% for i, expr in enumerate(collision_assignment): +% for i, expr in enumerate(assignment): const ${float_type} ${ccode(expr)} % endfor -% for i, expr in enumerate(collision_assignment): +% for i, expr in enumerate(assignment): preshifted_f_next[${layout.pop_offset(i)}] = f_next_${i}; % endfor } diff --git a/boltzgen/kernel/template/momenta_boundary.cl.mako b/boltzgen/kernel/template/momenta_boundary.cl.mako index f5edb61..4e29601 100644 --- a/boltzgen/kernel/template/momenta_boundary.cl.mako +++ b/boltzgen/kernel/template/momenta_boundary.cl.mako @@ -1,3 +1,8 @@ +<% +moments_subexpr, moments_assignment = model.moments() +collision_subexpr, collision_assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = False)) +%> + <%def name="momenta_boundary(name, param)"> __kernel void ${name}_momenta_boundary_gid( __global ${float_type}* f_next, diff --git a/boltzgen/kernel/template/momenta_boundary.cpp.mako b/boltzgen/kernel/template/momenta_boundary.cpp.mako index 3f51087..6d05288 100644 --- a/boltzgen/kernel/template/momenta_boundary.cpp.mako +++ b/boltzgen/kernel/template/momenta_boundary.cpp.mako @@ -1,3 +1,8 @@ +<% +moments_subexpr, moments_assignment = model.moments() +collision_subexpr, collision_assignment = model.bgk(f_eq = model.equilibrium(resolve_moments = False)) +%> + <%def name="momenta_boundary(name, param)"> void ${name}_momenta_boundary( ${float_type}* f_next, diff --git a/boltzgen/lbm/__init__.py b/boltzgen/lbm/__init__.py index 8206baa..3f608cb 100644 --- a/boltzgen/lbm/__init__.py +++ b/boltzgen/lbm/__init__.py @@ -9,8 +9,14 @@ def assign(names, definitions): return list(map(lambda x: Assignment(*x), zip(names, definitions))) class LBM: - def __init__(self, descriptor): + def __init__(self, descriptor, tau, optimize = True): self.descriptor = descriptor + self.tau = tau + self.optimize = optimize + + if self.tau <= 0.5: + raise Exception('Relaxation time must be larger than 0.5') + self.f_next = symarray('f_next', descriptor.q) self.f_curr = symarray('f_curr', descriptor.q) @@ -20,7 +26,10 @@ class LBM: if not hasattr(descriptor, 'c_s'): self.descriptor.c_s = c_s(descriptor.d, descriptor.c, self.descriptor.w) - def moments(self, optimize = True): + def moments(self, optimize = None): + if optimize is None: + optimize = self.optimize + rho = symbols('rho') u = Matrix(symarray('u', self.descriptor.d)) @@ -35,10 +44,16 @@ class LBM: else: return ([], exprs) - def equilibrium(self): + def equilibrium(self, resolve_moments = False): rho = symbols('rho') u = Matrix(symarray('u', self.descriptor.d)) + if resolve_moments: + moments = self.moments(optimize = False)[1] + rho = moments[0].rhs + for i, m in enumerate(moments[1:]): + u[i] = m.rhs + f_eq = [] for i, c_i in enumerate(self.descriptor.c): @@ -50,11 +65,11 @@ class LBM: return f_eq - def bgk(self, tau, f_eq, optimize = True): - if tau <= 0.5: - raise Exception('Relaxation time must be larger than 0.5') + def bgk(self, f_eq, optimize = None): + if optimize is None: + optimize = self.optimize - exprs = [ self.f_curr[i] + 1/tau * (f_eq_i - self.f_curr[i]) for i, f_eq_i in enumerate(f_eq) ] + exprs = [ self.f_curr[i] + 1/self.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) diff --git a/boltzgen/utility/optimizations.py b/boltzgen/utility/optimizations.py index 93dad09..6dc23e9 100644 --- a/boltzgen/utility/optimizations.py +++ b/boltzgen/utility/optimizations.py @@ -2,9 +2,17 @@ from sympy import * from sympy.codegen.rewriting import ReplaceOptim -expand_square = ReplaceOptim( +expand_pos_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 +expand_neg_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_pos_square, expand_pos_square), + (expand_neg_square, expand_neg_square) +] + cse_main.basic_optimizations -- cgit v1.2.3