Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Jan Hönig
pystencils
Commits
c1ca74ad
Commit
c1ca74ad
authored
Nov 05, 2021
by
Jan Hönig
Committed by
Markus Holzer
Nov 05, 2021
Browse files
Kernel decorator fix
parent
157c4f11
Changes
4
Hide whitespace changes
Inline
Side-by-side
.gitignore
View file @
c1ca74ad
...
...
@@ -20,7 +20,9 @@ pystencils/boundaries/createindexlistcython.*.so
pystencils_tests/tmp
pystencils_tests/kerncraft_inputs/.2d-5pt.c_kerncraft/
pystencils_tests/kerncraft_inputs/.3d-7pt.c_kerncraft/
report.xml
coverage_report/
# macOS
**/.DS_Store
\ No newline at end of file
**/.DS_Store
pystencils/__init__.py
View file @
c1ca74ad
...
...
@@ -7,7 +7,7 @@ from .data_types import TypedSymbol
from
.datahandling
import
create_data_handling
from
.display_utils
import
get_code_obj
,
get_code_str
,
show_code
,
to_dot
from
.field
import
Field
,
FieldType
,
fields
from
.kernel_decorator
import
kernel
from
.kernel_decorator
import
kernel
,
kernel_config
from
.kernelcreation
import
(
CreateKernelConfig
,
create_domain_kernel
,
create_indexed_kernel
,
create_kernel
,
create_staggered_kernel
)
from
.simp
import
AssignmentCollection
...
...
@@ -34,7 +34,7 @@ __all__ = ['Field', 'FieldType', 'fields',
'assignment_from_stencil'
,
'SymbolCreator'
,
'create_data_handling'
,
'kernel'
,
'kernel'
,
'kernel_config'
,
'x_'
,
'y_'
,
'z_'
,
'x_staggered'
,
'y_staggered'
,
'z_staggered'
,
'x_vector'
,
'x_staggered_vector'
,
...
...
pystencils/kernel_decorator.py
View file @
c1ca74ad
import
ast
import
inspect
import
textwrap
from
typing
import
Callable
,
Union
,
List
,
Dict
from
typing
import
Callable
,
Union
,
List
,
Dict
,
Tuple
import
sympy
as
sp
from
pystencils.assignment
import
Assignment
from
pystencils.sympyextensions
import
SymbolCreator
from
pystencils.kernelcreation
import
CreateKernelConfig
__all__
=
[
'kernel'
]
__all__
=
[
'kernel'
,
'kernel_config'
]
def
kernel
(
func
:
Callable
[...,
None
],
return_config
:
bool
=
False
,
**
kwargs
)
->
Union
[
List
[
Assignment
],
Dict
]:
"""Decorator to simplify generation of pystencils Assignments.
Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
sympy Piecewise.
The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
a SymbolCreator()
func: the decorated function
return_config: Specify whether to return the list with assignments, or a dictionary containing additional settings
like func_name
Examples:
>>> import pystencils as ps
>>> @kernel
... def my_kernel(s):
... f, g = ps.fields('f, g: [2D]')
... s.neighbors @= f[0,1] + f[1,0]
... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
>>> f, g = ps.fields('f, g: [2D]')
>>> assert my_kernel[0].rhs == f[0,1] + f[1,0]
def
_kernel
(
func
:
Callable
[...,
None
],
**
kwargs
)
->
Tuple
[
List
[
Assignment
],
str
]:
"""
Convenient function for kernel decorator to prevent code duplication
Args:
func: decorated function
**kwargs: kwargs for the function
Returns:
assignments, function_name
"""
source
=
inspect
.
getsource
(
func
)
source
=
textwrap
.
dedent
(
source
)
...
...
@@ -55,10 +42,74 @@ def kernel(func: Callable[..., None], return_config: bool = False, **kwargs) ->
if
's'
in
args
and
's'
not
in
kwargs
:
kwargs
[
's'
]
=
SymbolCreator
()
func
(
**
kwargs
)
if
return_config
:
return
{
'assignments'
:
assignments
,
'function_name'
:
func
.
__name__
}
else
:
return
assignments
return
assignments
,
func
.
__name__
def
kernel
(
func
:
Callable
[...,
None
],
**
kwargs
)
->
List
[
Assignment
]:
"""Decorator to simplify generation of pystencils Assignments.
Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
sympy Piecewise.
The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
a SymbolCreator()
Args:
func: decorated function
**kwargs: kwargs for the function
Examples:
>>> import pystencils as ps
>>> @kernel
... def my_kernel(s):
... f, g = ps.fields('f, g: [2D]')
... s.neighbors @= f[0,1] + f[1,0]
... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
>>> f, g = ps.fields('f, g: [2D]')
>>> assert my_kernel[0].rhs == f[0,1] + f[1,0]
"""
assignments
,
_
=
_kernel
(
func
,
**
kwargs
)
return
assignments
def
kernel_config
(
config
:
CreateKernelConfig
,
**
kwargs
)
->
Callable
[...,
Dict
]:
"""Decorator to simplify generation of pystencils Assignments, which takes a configuration
and updates the function name accordingly.
Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
sympy Piecewise.
The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
a SymbolCreator()
Args:
config: Specify whether to return the list with assignments, or a dictionary containing additional settings
like func_name
Returns:
decorator with config
Examples:
>>> import pystencils as ps
>>> config = ps.CreateKernelConfig()
>>> @kernel_config(config)
... def my_kernel(s):
... f, g = ps.fields('f, g: [2D]')
... s.neighbors @= f[0,1] + f[1,0]
... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
>>> f, g = ps.fields('f, g: [2D]')
>>> assert my_kernel['assignments'][0].rhs == f[0,1] + f[1,0]
"""
def
decorator
(
func
:
Callable
[...,
None
])
->
Union
[
List
[
Assignment
],
Dict
]:
"""
Args:
func: decorated function
Returns:
Dict for unpacking into create_kernel
"""
assignments
,
func_name
=
_kernel
(
func
,
**
kwargs
)
config
.
function_name
=
func_name
return
{
'assignments'
:
assignments
,
'config'
:
config
}
return
decorator
# noinspection PyMethodMayBeStatic
...
...
pystencils_tests/test_create_kernel_config.py
View file @
c1ca74ad
import
numpy
as
np
import
pystencils
as
ps
...
...
@@ -15,3 +16,14 @@ def test_create_kernel_config():
c
=
ps
.
CreateKernelConfig
(
backend
=
ps
.
Backend
.
CUDA
)
assert
c
.
target
==
ps
.
Target
.
CPU
assert
c
.
backend
==
ps
.
Backend
.
CUDA
def
test_kernel_decorator_config
():
config
=
ps
.
CreateKernelConfig
()
a
,
b
,
c
=
ps
.
fields
(
a
=
np
.
ones
(
100
),
b
=
np
.
ones
(
100
),
c
=
np
.
ones
(
100
))
@
ps
.
kernel_config
(
config
)
def
test
():
a
[
0
]
@=
b
[
0
]
+
c
[
0
]
ps
.
create_kernel
(
**
test
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment