Skip to content
Snippets Groups Projects

Add `pystencils.make_python_function` used for KernelFunction.compile

Closed Stephan Seitz requested to merge seitz/pystencils:make_python_function into master
6 files
+ 64
32
Compare changes
  • Side-by-side
  • Inline
Files
6
  • `KernelFunction.compile = None` is currently set by the
    `create_kernel` function of each respective backend as partial function
    of `<backend>.make_python_function`.
    
    The code would be clearer with a unified `make_python_function`.
    `KernelFunction.compile` can then be implemented  as a call to this
    function with the respective backend.
from typing import List, Union
import sympy as sp
from functools import partial
from pystencils.astnodes import SympyAssignment, Block, LoopOverCoordinate, KernelFunction
from pystencils.transformations import resolve_buffer_accesses, resolve_field_accesses, make_loop_over_domain, \
add_types, get_optimal_loop_ordering, parse_base_pointer_info, move_constants_before_loop, \
split_inner_loop, get_base_buffer_index, filtered_tree_iteration
from pystencils.data_types import TypedSymbol, BasicType, StructType, create_type
from pystencils.field import Field, FieldType
import pystencils.astnodes as ast
from pystencils.cpu.cpujit import make_python_function
from pystencils.assignment import Assignment
from typing import List, Union
from pystencils.astnodes import (Block, KernelFunction, LoopOverCoordinate,
SympyAssignment)
from pystencils.data_types import (BasicType, StructType, TypedSymbol,
create_type)
from pystencils.field import Field, FieldType
from pystencils.transformations import (add_types, filtered_tree_iteration,
get_base_buffer_index,
get_optimal_loop_ordering,
make_loop_over_domain,
move_constants_before_loop,
parse_base_pointer_info,
resolve_buffer_accesses,
resolve_field_accesses,
split_inner_loop)
AssignmentOrAstNodeList = List[Union[Assignment, ast.Node]]
@@ -83,7 +91,6 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
resolve_buffer_accesses(ast_node, get_base_buffer_index(ast_node), read_only_fields)
resolve_field_accesses(ast_node, read_only_fields, field_to_base_pointer_info=base_pointer_info)
move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node
@@ -148,7 +155,6 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu
read_only_fields = set([f.name for f in fields_read - fields_written])
resolve_field_accesses(ast_node, read_only_fields, field_to_fixed_coordinates=fixed_coordinate_mapping)
move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node
Loading