Attach assignments to KernelFunction

pystencils/kernelcreation.py

@@ -130,6 +130,7 @@ def create_kernel(assignments,
     if use_auto_for_assignments:
         for a in ast.atoms(SympyAssignment):
             a.use_auto = True
+    ast.assignments = assignments
 
     return ast
 
@@ -184,6 +185,7 @@ def create_indexed_kernel(assignments,
         assignments = [assignments]
     elif isinstance(assignments, AssignmentCollection):
         assignments = assignments.all_assignments
+
     if target == 'cpu':
         from pystencils.cpu import create_indexed_kernel
         from pystencils.cpu import add_openmp
@@ -191,7 +193,6 @@ def create_indexed_kernel(assignments,
                                     coordinate_names=coordinate_names)
         if cpu_openmp:
             add_openmp(ast, num_threads=cpu_openmp)
-        return ast
     elif target == 'llvm':
         raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend")
     elif target == 'gpu' or target == 'opencl':
@@ -209,10 +210,12 @@ def create_indexed_kernel(assignments,
             ast.compile = functools.partial(make_python_function, ast, opencl_queue, opencl_ctx)
             ast._target = 'opencl'
             ast._backend = 'opencl'
-        return ast
     else:
         raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,))
 
+    ast.assignments = assignments
+    return ast
+
 
 def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=False, **kwargs):
     """Kernel that updates a staggered field.