From de4f8b5d10d7045ceb0a1120b417f190d760b82e Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Thu, 30 Jan 2020 15:50:21 +0100
Subject: [PATCH] Append assignments to KernelFunction (for later analysis
etc.)
---
pystencils/astnodes.py | 3 ++-
pystencils/cpu/kernelcreation.py | 4 ++--
pystencils/gpucuda/kernelcreation.py | 11 +++++++++--
3 files changed, 13 insertions(+), 5 deletions(-)
diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py
index 8ec80917a..6d1501529 100644
--- a/pystencils/astnodes.py
+++ b/pystencils/astnodes.py
@@ -172,7 +172,7 @@ class KernelFunction(Node):
def field_name(self):
return self.fields[0].name
- def __init__(self, body, target, backend, compile_function, ghost_layers, function_name="kernel"):
+ def __init__(self, body, target, backend, compile_function, ghost_layers, function_name="kernel", assignments=None):
super(KernelFunction, self).__init__()
self._body = body
body.parent = self
@@ -186,6 +186,7 @@ class KernelFunction(Node):
self.instruction_set = None # used in `vectorize` function to tell the backend which i.s. (SSE,AVX) to use
# function that compiles the node to a Python callable, is set by the backends
self._compile_function = compile_function
+ self.assignments = assignments
@property
def target(self):
diff --git a/pystencils/cpu/kernelcreation.py b/pystencils/cpu/kernelcreation.py
index 9b119ea9a..38ce169af 100644
--- a/pystencils/cpu/kernelcreation.py
+++ b/pystencils/cpu/kernelcreation.py
@@ -66,7 +66,7 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
loop_node, ghost_layer_info = make_loop_over_domain(body, iteration_slice=iteration_slice,
ghost_layers=ghost_layers, loop_order=loop_order)
ast_node = KernelFunction(loop_node, 'cpu', 'c', compile_function=make_python_function,
- ghost_layers=ghost_layer_info, function_name=function_name)
+ ghost_layers=ghost_layer_info, function_name=function_name, assignments=assignments)
implement_interpolations(body)
if split_groups:
@@ -147,7 +147,7 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu
function_body = Block([loop_node])
ast_node = KernelFunction(function_body, "cpu", "c", make_python_function,
- ghost_layers=None, function_name=function_name)
+ ghost_layers=None, function_name=function_name, assignments=assignments)
fixed_coordinate_mapping = {f.name: coordinate_typed_symbols for f in non_index_fields}
diff --git a/pystencils/gpucuda/kernelcreation.py b/pystencils/gpucuda/kernelcreation.py
index d002cb0f0..33db3ad56 100644
--- a/pystencils/gpucuda/kernelcreation.py
+++ b/pystencils/gpucuda/kernelcreation.py
@@ -62,7 +62,13 @@ def create_cuda_kernel(assignments,
block = indexing.guard(block, common_shape)
unify_shape_symbols(block, common_shape=common_shape, fields=fields_without_buffers)
- ast = KernelFunction(block, 'gpu', 'gpucuda', make_python_function, ghost_layers, function_name)
+ ast = KernelFunction(block,
+ 'gpu',
+ 'gpucuda',
+ make_python_function,
+ ghost_layers,
+ function_name,
+ assignments=assignments)
ast.global_variables.update(indexing.index_variables)
implement_interpolations(ast, implement_by_texture_accesses=use_textures_for_interpolation)
@@ -137,7 +143,8 @@ def created_indexed_cuda_kernel(assignments,
function_body = Block(coordinate_symbol_assignments + assignments)
function_body = indexing.guard(function_body, get_common_shape(index_fields))
- ast = KernelFunction(function_body, 'gpu', 'gpucuda', make_python_function, None, function_name)
+ ast = KernelFunction(function_body, 'gpu', 'gpucuda', make_python_function,
+ None, function_name, assignments=assignments)
ast.global_variables.update(indexing.index_variables)
implement_interpolations(ast, implement_by_texture_accesses=use_textures_for_interpolation)
--
GitLab