From 2d4418779f5793eee9d5ec5e317178e307515243 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 20:54:00 +0300 Subject: [PATCH 1/7] remove explicit force_device_scalars --- examples/curve-pot.py | 2 +- examples/expansion-toys.py | 2 +- examples/sym-exp-complexity.py | 2 +- sumpy/array_context.py | 2 +- test/test_fmm.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/curve-pot.py b/examples/curve-pot.py index c6402352f..c1c113662 100644 --- a/examples/curve-pot.py +++ b/examples/curve-pot.py @@ -63,7 +63,7 @@ def draw_pot_figure(aspect_ratio, from sumpy.array_context import PyOpenCLArrayContext ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) - actx = PyOpenCLArrayContext(queue, force_device_scalars=True) + actx = PyOpenCLArrayContext(queue) # {{{ make plot targets diff --git a/examples/expansion-toys.py b/examples/expansion-toys.py index c7052fe36..b61aafcf5 100644 --- a/examples/expansion-toys.py +++ b/examples/expansion-toys.py @@ -18,7 +18,7 @@ def main(): from sumpy.array_context import PyOpenCLArrayContext ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) - actx = PyOpenCLArrayContext(queue, force_device_scalars=True) + actx = PyOpenCLArrayContext(queue) tctx = t.ToyContext( actx.context, diff --git a/examples/sym-exp-complexity.py b/examples/sym-exp-complexity.py index 9dfbfb442..c0d0c8163 100644 --- a/examples/sym-exp-complexity.py +++ b/examples/sym-exp-complexity.py @@ -24,7 +24,7 @@ def find_flops(): from sumpy.array_context import PyOpenCLArrayContext ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) - actx = PyOpenCLArrayContext(queue, force_device_scalars=True) + actx = PyOpenCLArrayContext(queue) if 0: knl = LaplaceKernel(2) diff --git a/sumpy/array_context.py b/sumpy/array_context.py index eddfb2cc8..cc54048a8 100644 --- a/sumpy/array_context.py +++ b/sumpy/array_context.py @@ -64,7 +64,7 @@ def _acf(): ctx = cl.create_some_context() queue = cl.CommandQueue(ctx) - return PyOpenCLArrayContext(queue, force_device_scalars=True) + return PyOpenCLArrayContext(queue) class PytestPyOpenCLArrayContextFactory( diff --git a/test/test_fmm.py b/test/test_fmm.py index 525046e0c..d85eafe60 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -450,7 +450,7 @@ def test_sumpy_fmm_timing_data_collection(ctx_factory, use_fft, visualize=False) ctx = ctx_factory() queue = cl.CommandQueue(ctx, properties=cl.command_queue_properties.PROFILING_ENABLE) - actx = PyOpenCLArrayContext(queue, force_device_scalars=True) + actx = PyOpenCLArrayContext(queue) nsources = 500 dtype = np.float64 From 0c8576308d58674d00d7f42f112adabab51cb7c4 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 20:55:17 +0300 Subject: [PATCH 2/7] use parentheses instead of line breaks --- sumpy/codegen.py | 20 ++++++++++---------- sumpy/kernel.py | 4 ++-- sumpy/symbolic.py | 5 +++-- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 9917db731..4dd9c5a40 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -221,8 +221,8 @@ def _fp_contract_fast_preamble(preamble_info): def register_optimization_preambles(loopy_knl, device): if isinstance(loopy_knl.target, lp.PyOpenCLTarget): import pyopencl as cl - if device.platform.name == "Portable Computing Language" and \ - (device.type & cl.device_type.GPU): + if (device.platform.name == "Portable Computing Language" + and (device.type & cl.device_type.GPU)): loopy_knl = lp.register_preamble_generators(loopy_knl, [_fp_contract_fast_preamble]) return loopy_knl @@ -257,8 +257,8 @@ def __init__(self): self.bessel_j_arg_to_top_order = {} def map_call(self, expr, rec_self=None, *args): - if isinstance(expr.function, prim.Variable) \ - and expr.function.name == "bessel_j": + if (isinstance(expr.function, prim.Variable) + and expr.function.name == "bessel_j"): order, arg = expr.parameters self.rec(arg) assert isinstance(order, int) @@ -588,11 +588,12 @@ def map_sum(self, expr, *args): first_group.append(child) new_children = tuple(first_group + second_group) - if len(new_children) == len(expr.children) and \ - all(child is orig_child for child, orig_child in - zip(new_children, expr.children, strict=True)): + if (len(new_children) == len(expr.children) + and all(child is orig_child for child, orig_child + in zip(new_children, expr.children, strict=True))): return expr - return prim.Sum(tuple(first_group+second_group)) + + return prim.Sum(tuple(first_group + second_group)) map_common_subexpression_uncached = IdentityMapper.map_common_subexpression @@ -647,8 +648,7 @@ def combine_mappers(*mappers): class CombinedMapper(CSECachingIdentityMapper): def __init__(self, all_methods): self.all_methods = all_methods - map_common_subexpression_uncached = \ - IdentityMapper.map_common_subexpression + map_common_subexpression_uncached = IdentityMapper.map_common_subexpression def _map(method_name, self, expr, rec_self=None, *args): if method_name not in self.all_methods: diff --git a/sumpy/kernel.py b/sumpy/kernel.py index a7ca24866..4bf05fe46 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -1105,8 +1105,8 @@ def __init__(self, vec_name, additional_indices): def map_subscript(self, expr): from pymbolic.primitives import CommonSubexpression, cse_scope - if expr.aggregate.name == self.vec_name \ - and isinstance(expr.index, int): + if (expr.aggregate.name == self.vec_name + and isinstance(expr.index, int)): return CommonSubexpression( expr.aggregate.index((expr.index, *self.additional_indices)), prefix=None, scope=cse_scope.EVALUATION) diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index c3c58c68c..2870b50e0 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -316,8 +316,9 @@ def map_Mul(self, expr): # noqa: N802 num_args = [] den_args = [] for child in expr.args: - if isinstance(child, Pow) and isinstance(child.exp, Integer) \ - and child.exp < 0: + if (isinstance(child, Pow) + and isinstance(child.exp, Integer) + and child.exp < 0): den_args.append(self.rec(child.base)**(-self.rec(child.exp))) else: num_args.append(self.rec(child)) From ac354a771ccaab68225c947f8ca3003549148776 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 20:56:09 +0300 Subject: [PATCH 3/7] port to expr_dataclass Expressions --- sumpy/symbolic.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index 2870b50e0..d8bc011e5 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -275,17 +275,18 @@ def find_power_of(base, prod): return result[power] +@prim.expr_dataclass() class SpatialConstant(prim.Variable): """A symbolic constant to represent a symbolic variable that - is spatially constant, like for example the wave-number :math:`k` - in the setting of a constant-coefficient Helmholtz problem. - For use in :attr:`sumpy.kernel.ExpressionKernel.expression`. + is spatially constant. + + For example the wave-number :math:`k` in the setting of a constant-coefficient + Helmholtz problem. For use in :attr:`sumpy.kernel.ExpressionKernel.expression`. Any variable occurring there that is not a :class:`SpatialConstant` is assumed to have a spatial dependency. """ prefix = "_spatial_constant_" - mapper_method = "map_spatial_constant" def as_sympy(self): return sym.Symbol(f"{self.prefix}{self.name}") From 93440b2a87b5724dd49c88fe9c4d1909fe4d5caa Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 20:55:35 +0300 Subject: [PATCH 4/7] port away from Expression.index --- sumpy/codegen.py | 2 +- sumpy/kernel.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 4dd9c5a40..26e006128 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -548,7 +548,7 @@ def map_variable(self, expr, *args): name = match_obj.group(1) subscript = int(match_obj.group(2)) if name in self.name_whitelist: - return prim.Variable(name).index(subscript) + return prim.Variable(name)[subscript] else: return expr else: diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 4bf05fe46..6e889acd2 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -1108,7 +1108,7 @@ def map_subscript(self, expr): if (expr.aggregate.name == self.vec_name and isinstance(expr.index, int)): return CommonSubexpression( - expr.aggregate.index((expr.index, *self.additional_indices)), + expr.aggregate[(expr.index, *self.additional_indices)], prefix=None, scope=cse_scope.EVALUATION) else: return IdentityMapper.map_subscript(self, expr) From 67da3d293313e9dad5577fe3165a46254a793fba Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 20:55:50 +0300 Subject: [PATCH 5/7] port away from PersistentHashWalkMapper --- sumpy/kernel.py | 28 +++++++--------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 6e889acd2..76e24c620 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -394,15 +394,8 @@ def get_global_scaling_const(self): def update_persistent_hash(self, key_hash, key_builder): key_hash.update(type(self).__name__.encode("utf8")) - for name, value in zip(self.init_arg_names, self.__getinitargs__(), - strict=True): - if name in ["expression", "global_scaling_const"]: - from pymbolic.mapper.persistent_hash import ( - PersistentHashWalkMapper as PersistentHashWalkMapper, - ) - PersistentHashWalkMapper(key_hash)(value) - else: - key_builder.rec(key_hash, value) + for _, value in zip(self.init_arg_names, self.__getinitargs__(), strict=True): + key_builder.rec(key_hash, value) mapper_method = "map_expression_kernel" @@ -755,13 +748,11 @@ def __reduce__(self): return (ElasticityKernel, self.__getinitargs__()) def update_persistent_hash(self, key_hash, key_builder): - from pymbolic.mapper.persistent_hash import PersistentHashWalkMapper key_hash.update(type(self).__name__.encode()) key_builder.rec(key_hash, (self.dim, self.icomp, self.jcomp)) - mapper = PersistentHashWalkMapper(key_hash) - mapper(self.viscosity_mu) - mapper(self.poisson_ratio) + key_builder.rec(key_hash, self.viscosity_mu) + key_builder.rec(key_hash, self.poisson_ratio) def __repr__(self): return f"ElasticityKnl{self.dim}D_{self.icomp}{self.jcomp}" @@ -854,10 +845,7 @@ def __getinitargs__(self): def update_persistent_hash(self, key_hash, key_builder): key_hash.update(type(self).__name__.encode()) key_builder.rec(key_hash, (self.dim, self.icomp, self.jcomp, self.kcomp)) - - from pymbolic.mapper.persistent_hash import PersistentHashWalkMapper - mapper = PersistentHashWalkMapper(key_hash) - mapper(self.viscosity_mu) + key_builder.rec(key_hash, self.viscosity_mu) def __repr__(self): return f"StressletKnl{self.dim}D_{self.icomp}{self.jcomp}{self.kcomp}" @@ -930,12 +918,10 @@ def __getinitargs__(self): return (self.dim, self.axis, self.viscosity_mu, self.poisson_ratio) def update_persistent_hash(self, key_hash, key_builder): - from pymbolic.mapper.persistent_hash import PersistentHashWalkMapper key_hash.update(type(self).__name__.encode()) key_builder.rec(key_hash, (self.dim, self.axis)) - mapper = PersistentHashWalkMapper(key_hash) - mapper(self.viscosity_mu) - mapper(self.poisson_ratio) + key_builder.rec(key_hash, self.viscosity_mu) + key_builder.rec(key_hash, self.poisson_ratio) def __repr__(self): return f"LineOfCompressionKnl{self.dim}D_{self.axis}" From 9ef292ddf80745acb5c929bf9d919bbc15a28ac3 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Tue, 1 Oct 2024 21:30:14 +0300 Subject: [PATCH 6/7] ignore mypy errors --- sumpy/symbolic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index d8bc011e5..ccc07edbb 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -100,7 +100,7 @@ def _find_symbolic_backend(): else: import sympy as sym - from pymbolic.interop.sympy import ( + from pymbolic.interop.sympy import ( # type: ignore[assignment] PymbolicToSympyMapper as PymbolicToSympyMapperBase, SympyToPymbolicMapper as SympyToPymbolicMapperBase, ) From ef7b7a38e41e4976f701a57f61401c66512320d6 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Wed, 9 Oct 2024 10:01:18 +0300 Subject: [PATCH 7/7] docs: improve SpatialConstant docs --- sumpy/symbolic.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index ccc07edbb..d040971d3 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -39,9 +39,10 @@ import logging import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import numpy as np +import sympy as sp import pymbolic.primitives as prim from pymbolic.mapper import IdentityMapper as IdentityMapperBase @@ -277,22 +278,28 @@ def find_power_of(base, prod): @prim.expr_dataclass() class SpatialConstant(prim.Variable): - """A symbolic constant to represent a symbolic variable that - is spatially constant. + """A symbolic constant to represent a symbolic variable that is spatially constant. For example the wave-number :math:`k` in the setting of a constant-coefficient Helmholtz problem. For use in :attr:`sumpy.kernel.ExpressionKernel.expression`. - Any variable occurring there that is not a :class:`SpatialConstant` + Any variable occurring there that is not a :class:`~sumpy.symbolic.SpatialConstant` is assumed to have a spatial dependency. + + .. autoattribute:: prefix + .. automethod:: as_sympy + .. automethod:: from_sympy """ - prefix = "_spatial_constant_" + prefix: ClassVar[str] = "_spatial_constant_" + """Prefix used in code generation for variables of this type.""" - def as_sympy(self): + def as_sympy(self) -> sp.Symbol: + """Convert variable to a :mod:`sympy` expression.""" return sym.Symbol(f"{self.prefix}{self.name}") @classmethod - def from_sympy(cls, expr): + def from_sympy(cls, expr: sp.Symbol) -> SpatialConstant: + """Convert :mod:`sympy` expression to a constant.""" return cls(expr.name[len(cls.prefix):])