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
Jonas Plewinski
pystencils
Commits
8d610054
Commit
8d610054
authored
Sep 23, 2019
by
Stephan Seitz
Browse files
Add pystencils.math_optimizations
parent
c3a4fb73
Changes
3
Hide whitespace changes
Inline
Side-by-side
pystencils/astnodes.py
View file @
8d610054
...
...
@@ -509,6 +509,13 @@ class SympyAssignment(Node):
self
.
lhs
=
fast_subs
(
self
.
lhs
,
subs_dict
)
self
.
rhs
=
fast_subs
(
self
.
rhs
,
subs_dict
)
def
optimize
(
self
,
optimizations
):
try
:
from
sympy.codegen.rewriting
import
optimize
self
.
rhs
=
optimize
(
self
.
rhs
,
optimizations
)
except
Exception
:
pass
@
property
def
args
(
self
):
return
[
self
.
_lhs_symbol
,
self
.
rhs
]
...
...
pystencils/math_optimizations.py
0 → 100644
View file @
8d610054
"""
Default Sympy optimizations applied in pystencils kernels using :func:`sympy.codegen.rewriting.optimize`.
See :func:`sympy.codegen.rewriting.optimize`.
"""
import
itertools
from
pystencils
import
Assignment
from
pystencils.astnodes
import
SympyAssignment
try
:
from
sympy.codegen.rewriting
import
optims_c99
,
optimize
from
sympy.codegen.rewriting
import
ReplaceOptim
HAS_REWRITING
=
True
# Evaluates all constant terms
evaluate_constant_terms
=
ReplaceOptim
(
lambda
e
:
hasattr
(
e
,
'is_constant'
)
and
e
.
is_constant
and
not
e
.
is_integer
,
lambda
p
:
p
.
evalf
()
)
optims_pystencils_cpu
=
[
evaluate_constant_terms
]
+
list
(
optims_c99
)
optims_pystencils_gpu
=
[
evaluate_constant_terms
]
+
list
(
optims_c99
)
except
ImportError
:
from
warnings
import
warn
warn
(
"Could not import ReplaceOptim, optims_c99, optimize from sympy.codegen.rewriting."
"Please update your sympy installation!"
)
optims_c99
=
[]
optims_pystencils_cpu
=
[]
optims_pystencils_gpu
=
[]
HAS_REWRITING
=
False
def
optimize_assignments
(
assignments
,
optimizations
):
if
HAS_REWRITING
:
assignments
=
[
Assignment
(
a
.
lhs
,
optimize
(
a
.
rhs
,
optimizations
))
if
hasattr
(
a
,
'lhs'
)
else
a
for
a
in
assignments
]
assignments_nodes
=
[
a
.
atoms
(
SympyAssignment
)
for
a
in
assignments
]
for
a
in
itertools
.
chain
.
from_iterable
(
assignments_nodes
):
a
.
optimize
(
optimizations
)
return
assignments
pystencils_tests/test_sympy_optimizations.py
0 → 100644
View file @
8d610054
import
pytest
import
sympy
as
sp
import
pystencils
from
pystencils.math_optimizations
import
HAS_REWRITING
,
optimize_assignments
,
optims_pystencils_cpu
@
pytest
.
mark
.
skipif
(
not
HAS_REWRITING
,
reason
=
"need sympy.codegen.rewriting"
)
def
test_sympy_optimizations
():
for
target
in
(
'cpu'
,
'gpu'
):
x
,
y
,
z
=
pystencils
.
fields
(
'x, y, z: float32[2d]'
)
# Triggers Sympy's expm1 optimization
assignments
=
pystencils
.
AssignmentCollection
({
x
[
0
,
0
]:
sp
.
exp
(
y
[
0
,
0
])
-
1
})
assignments
=
optimize_assignments
(
assignments
,
optims_pystencils_cpu
)
ast
=
pystencils
.
create_kernel
(
assignments
,
target
=
target
)
code
=
str
(
pystencils
.
show_code
(
ast
))
assert
'expm1('
in
code
@
pytest
.
mark
.
skipif
(
not
HAS_REWRITING
,
reason
=
"need sympy.codegen.rewriting"
)
def
test_evaluate_constant_terms
():
for
target
in
(
'cpu'
,
'gpu'
):
x
,
y
,
z
=
pystencils
.
fields
(
'x, y, z: float32[2d]'
)
# Triggers Sympy's expm1 optimization
assignments
=
pystencils
.
AssignmentCollection
({
x
[
0
,
0
]:
-
sp
.
cos
(
1
)
+
y
[
0
,
0
]
})
assignments
=
optimize_assignments
(
assignments
,
optims_pystencils_cpu
)
ast
=
pystencils
.
create_kernel
(
assignments
,
target
=
target
)
code
=
str
(
pystencils
.
show_code
(
ast
))
assert
'cos('
not
in
code
print
(
code
)
@
pytest
.
mark
.
skipif
(
not
HAS_REWRITING
,
reason
=
"need sympy.codegen.rewriting"
)
def
test_do_not_evaluate_constant_terms
():
optimizations
=
pystencils
.
math_optimizations
.
optims_pystencils_cpu
optimizations
.
remove
(
pystencils
.
math_optimizations
.
evaluate_constant_terms
)
for
target
in
(
'cpu'
,
'gpu'
):
x
,
y
,
z
=
pystencils
.
fields
(
'x, y, z: float32[2d]'
)
assignments
=
pystencils
.
AssignmentCollection
({
x
[
0
,
0
]:
-
sp
.
cos
(
1
)
+
y
[
0
,
0
]
})
optimize_assignments
(
assignments
,
optimizations
)
ast
=
pystencils
.
create_kernel
(
assignments
,
target
=
target
)
code
=
str
(
pystencils
.
show_code
(
ast
))
assert
'cos('
in
code
print
(
code
)
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