diff --git a/grudge/symbolic/compiler.py b/grudge/symbolic/compiler.py index 270de648b3fa439545e1999785d58cbd5229b3c4..bd19a6c3d1b8246d5145dbdc6db454c391b85199 100644 --- a/grudge/symbolic/compiler.py +++ b/grudge/symbolic/compiler.py @@ -24,6 +24,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import numpy as np import six # noqa from six.moves import zip, reduce @@ -880,6 +881,59 @@ class ToLoopyExpressionMapper(mappers.IdentityMapper): "compilation process") +# {{{ bessel handling + +BESSEL_PREAMBLE = """//CL// +#include +#include +""" + + +def bessel_preamble_generator(preamble_info): + from loopy.target.pyopencl import PyOpenCLTarget + if not isinstance(preamble_info.kernel.target, PyOpenCLTarget): + raise NotImplementedError("Only the PyOpenCLTarget is supported as of now") + + if any(func.name in ["bessel_j", "bessel_y"] + for func in preamble_info.seen_functions): + yield ("50-grudge-bessel", BESSEL_PREAMBLE) + + +def bessel_function_mangler(kernel, name, arg_dtypes): + from loopy.types import NumpyType + if name == "bessel_j" and len(arg_dtypes) == 2: + n_dtype, x_dtype, = arg_dtypes + + # *technically* takes a float, but let's not worry about that. + if n_dtype.numpy_dtype.kind != "i": + raise TypeError("%s expects an integer first argument") + + from loopy.kernel.data import CallMangleInfo + return CallMangleInfo( + "bessel_jv", + (NumpyType(np.float64),), + (NumpyType(np.int32), NumpyType(np.float64)), + ) + + elif name == "bessel_y" and len(arg_dtypes) == 2: + n_dtype, x_dtype, = arg_dtypes + + # *technically* takes a float, but let's not worry about that. + if n_dtype.numpy_dtype.kind != "i": + raise TypeError("%s expects an integer first argument") + + from loopy.kernel.data import CallMangleInfo + return CallMangleInfo( + "bessel_yn", + (NumpyType(np.float64),), + (NumpyType(np.int32), NumpyType(np.float64)), + ) + + return None + +# }}} + + class ToLoopyInstructionMapper(object): def __init__(self, dd_inference_mapper): self.dd_inference_mapper = dd_inference_mapper @@ -908,7 +962,11 @@ class ToLoopyInstructionMapper(object): lp.Assignment( expr_mapper(var(name)), expr_mapper(expr), - temp_var_type=lp.auto if dnr else None)) + temp_var_type=lp.auto if dnr else None, + no_sync_with=frozenset([ + ("*", "any"), + ]), + )) if not expr_mapper.non_scalar_vars: return insn @@ -916,7 +974,11 @@ class ToLoopyInstructionMapper(object): knl = lp.make_kernel( "{[%s]: 0 <= %s < %s}" % (iname, iname, size_name), insns, - default_offset=lp.auto) + default_offset=lp.auto, + + # Single-insn kernels may have their no_sync_with resolve to an + # empty set, that's OK. + options=lp.Options(check_dep_resolution=False)) knl = lp.set_options(knl, return_dict=True) knl = lp.split_iname(knl, iname, 128, outer_tag="g.0", inner_tag="l.0") @@ -926,6 +988,11 @@ class ToLoopyInstructionMapper(object): self.dd_inference_mapper(expr) for expr in insn.exprs) + knl = lp.register_preamble_generators(knl, + [bessel_preamble_generator]) + knl = lp.register_function_manglers(knl, + [bessel_function_mangler]) + return LoopyKernelInstruction( LoopyKernelDescriptor( loopy_kernel=knl, diff --git a/grudge/symbolic/primitives.py b/grudge/symbolic/primitives.py index 00f7b8d386af6130810cf900bbaaebb0027f5adc..471645d62317f9521e2ce09c4d8d600a09cfdc78 100644 --- a/grudge/symbolic/primitives.py +++ b/grudge/symbolic/primitives.py @@ -78,6 +78,13 @@ Symbols .. autoclass:: make_sym_mv .. autoclass:: CFunction +.. function :: sqrt(arg) +.. function :: exp(arg) +.. function :: sin(arg) +.. function :: cos(arg) +.. function :: besesl_j(n, arg) +.. function :: besesl_y(n, arg) + Helpers ^^^^^^^ @@ -345,12 +352,12 @@ class CFunction(pymbolic.primitives.Variable): from grudge.symbolic.mappers import StringifyMapper return StringifyMapper - def __call__(self, expr): - from pytools.obj_array import with_object_array_or_scalar + def __call__(self, *exprs): + from pytools.obj_array import with_object_array_or_scalar_n_args from functools import partial - return with_object_array_or_scalar( + return with_object_array_or_scalar_n_args( partial(pymbolic.primitives.Expression.__call__, self), - expr) + *exprs) mapper_method = "map_c_function" @@ -359,6 +366,8 @@ sqrt = CFunction("sqrt") exp = CFunction("exp") sin = CFunction("sin") cos = CFunction("cos") +bessel_j = CFunction("bessel_j") +bessel_y = CFunction("bessel_y") # }}} diff --git a/test/test_grudge.py b/test/test_grudge.py index 7cebbd4b8ab7eca46834f6b068b99519606ccc21..9537a31d6f0eb658d6b131407c57fc5cd6711f2d 100644 --- a/test/test_grudge.py +++ b/test/test_grudge.py @@ -485,6 +485,35 @@ def test_foreign_points(ctx_factory): bind(pdiscr, sym.nodes(dim)**2)(queue) +def test_bessel(ctx_factory): + cl_ctx = cl.create_some_context() + queue = cl.CommandQueue(cl_ctx) + + dims = 2 + + from meshmode.mesh.generation import generate_regular_rect_mesh + mesh = generate_regular_rect_mesh( + a=(0.1,)*dims, + b=(1.0,)*dims, + n=(8,)*dims) + + discr = DGDiscretizationWithBoundaries(cl_ctx, mesh, order=3) + + nodes = sym.nodes(dims) + r = sym.cse(sym.sqrt(nodes[0]**2 + nodes[1]**2)) + + # https://dlmf.nist.gov/10.6.1 + n = 3 + bessel_zero = ( + sym.bessel_j(n+1, r) + + sym.bessel_j(n-1, r) + - 2*n/r * sym.bessel_j(n, r)) + + z = bind(discr, sym.norm(2, bessel_zero))(queue) + + assert z < 1e-15 + + # You can test individual routines by typing # $ python test_grudge.py 'test_routine()'