From c7a4918aa4dc1317e41e843b7640e5774a4939c8 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Tue, 10 Dec 2019 14:27:40 +0100
Subject: [PATCH] Attach assignments to KernelFunction
---
pystencils/kernelcreation.py | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py
index 73b45b844..43b239ef5 100644
--- a/pystencils/kernelcreation.py
+++ b/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.
--
GitLab