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
Files
6
 
from typing import List, Union
 
import sympy as sp
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
import pystencils.astnodes as ast
from pystencils.cpu.cpujit import make_python_function
from pystencils.assignment import Assignment
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]]
AssignmentOrAstNodeList = List[Union[Assignment, ast.Node]]
@@ -83,7 +91,6 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
@@ -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_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)
resolve_field_accesses(ast_node, read_only_fields, field_to_base_pointer_info=base_pointer_info)
move_constants_before_loop(ast_node)
move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node
return ast_node
@@ -148,7 +155,6 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu
@@ -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])
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)
resolve_field_accesses(ast_node, read_only_fields, field_to_fixed_coordinates=fixed_coordinate_mapping)
move_constants_before_loop(ast_node)
move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node
return ast_node
Loading