Extract offset helper into target and layout specific classes

1 files changed, 10 insertions, 2 deletions

diff --git a/boltzgen/kernel/generator.py b/boltzgen/kernel/generator.py
index c9fdb27..e897325 100644
--- a/boltzgen/kernel/generator.py
+++ b/boltzgen/kernel/generator.py
@@ -1,9 +1,11 @@
 import sympy
 
 from mako.template import Template
-
 from pathlib import Path
 
+import kernel.target.cl
+import kernel.target.cpp
+
 class Generator:
     def __init__(self, descriptor, moments, collision, boundary = ''):
         self.descriptor = descriptor
@@ -16,10 +18,16 @@ class Generator:
         if not template_path.exists():
             raise Exception("Target '%s' not supported" % target)
 
+        layout_impl = eval("kernel.target.%s.%s" % (target, layout))
+        if layout_impl is None:
+            raise Exception("Target '%s' doesn't support layout '%s'" % (target, layout))
+        else:
+            layout_impl = layout_impl(self.descriptor, geometry)
+
         return Template(filename = str(template_path)).render(
             descriptor = self.descriptor,
             geometry   = geometry,
-            layout     = layout,
+            layout     = layout_impl,
 
             moments_subexpr    = self.moments[0],
             moments_assignment = self.moments[1],