Skip to content
Snippets Groups Projects

Improve Vectorisation

Merged Markus Holzer requested to merge holzer/pystencils:ImproveVec into master
3 unresolved threads
2 files
+ 8
5
Compare changes
  • Side-by-side
  • Inline
Files
2
@@ -43,14 +43,15 @@ class NodeCollection:
@@ -43,14 +43,15 @@ class NodeCollection:
def evaluate_terms(self):
def evaluate_terms(self):
evaluate_constant_terms = ReplaceOptim(
evaluate_constant_terms = ReplaceOptim(
lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
lambda p: p.evalf())
lambda p: p.evalf()
 
)
evaluate_pow = ReplaceOptim(
evaluate_pow = ReplaceOptim(
lambda e: e.is_Pow and e.exp.is_Integer and abs(e.exp) <= 8,
lambda e: e.is_Pow and e.exp.is_Integer and abs(e.exp) <= 8,
lambda p: (
lambda p: sp.UnevaluatedExpr(sp.Mul(*([p.base] * +p.exp), evaluate=False)) if p.exp > 0 else
sp.UnevaluatedExpr(sp.Mul(*([p.base] * +p.exp), evaluate=False)) if p.exp > 0 else
(DivFunc(sp.Integer(1), p.base) if p.exp == -1 else
DivFunc(sp.Integer(1), sp.UnevaluatedExpr(sp.Mul(*([p.base] * -p.exp), evaluate=False)))
DivFunc(sp.Integer(1), sp.UnevaluatedExpr(sp.Mul(*([p.base] * -p.exp), evaluate=False))))
))
)
sympy_optimisations = [evaluate_constant_terms, evaluate_pow]
sympy_optimisations = [evaluate_constant_terms, evaluate_pow]
if self.is_Nodes:
if self.is_Nodes:
@@ -65,6 +66,7 @@ class NodeCollection:
@@ -65,6 +66,7 @@ class NodeCollection:
return optimize(node, sympy_optimisations)
return optimize(node, sympy_optimisations)
else:
else:
raise NotImplementedError(f'{node} {type(node)} has no valid visitor')
raise NotImplementedError(f'{node} {type(node)} has no valid visitor')
 
self.all_assignments = [visitor(assignment) for assignment in self.all_assignments]
self.all_assignments = [visitor(assignment) for assignment in self.all_assignments]
else:
else:
self.all_assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
self.all_assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
Loading