Interface: add support for mathematical special functions (real inputs only)

pykokkos/core/visitors/pykokkos_visitor.py

@@ -427,6 +427,23 @@ def visit_Call(self, node: ast.Call) -> cppast.CallExpr:
         if visitors_util.is_math_function(name) or name in ["printf", "abs", "Kokkos::PerTeam", "Kokkos::PerThread", "Kokkos::fence"]:
             return cppast.CallExpr(function, args)
 
+        if visitors_util.is_math_special_function(name):
+            if name in ["expint1", "erfcx"]:
+                if len(args) != 1:
+                    self.error(node, f"pk.{name}() accepts only one argument")
+                s = cppast.Serializer()
+                math_call = cppast.CallExpr(cppast.DeclRefExpr(f"Kokkos::Experimental::{name}<double>"), args)
+
+                return math_call
+            else:
+                if len(args) != 1:
+                    self.error(node, f"pk.{name}() accepts only one argument")
+                s = cppast.Serializer()
+                arg_str = s.serialize(args[0])
+                math_call = cppast.CallExpr(cppast.DeclRefExpr(f"Kokkos::Experimental::{name}<Kokkos::complex<decltype({arg_str})>, double, int>"), args)
+                real_number_call = cppast.MemberCallExpr(math_call, cppast.DeclRefExpr("real"), [])
+                return real_number_call
+
         if function in self.kokkos_functions:
             if "PK_RESTRICT" in os.environ:
                 return adjust_kokkos_function_call(function, args, self.restrict_views, self.views)

pykokkos/core/visitors/visitors_util.py

@@ -116,6 +116,23 @@ def pretty_print(node):
     "nan",
 }
 
+math_special_functions: Set = {
+    "expint1",
+    "erfcx",
+    "cyl_bessel_j0", 
+    "cyl_bessel_y0", 
+    "cyl_bessel_i0", 
+    "cyl_bessel_k0",
+    "cyl_bessel_j1", 
+    "cyl_bessel_y1", 
+    "cyl_bessel_i1", 
+    "cyl_bessel_k1", 
+    "cyl_bessel_h10", 
+    "cyl_bessel_h11", 
+    "cyl_bessel_h20", 
+    "cyl_bessel_h21",
+}
+
 math_constants: Dict[str, str] = {
     "e": "M_E",
     "pi": "M_PI",
@@ -171,6 +188,8 @@ def get_allowed_type_str(python_type: str) -> str:
 def is_math_function(function: str) -> bool:
     return function in math_functions
 
+def is_math_special_function(function: str) -> bool:
+    return function in math_special_functions
 
 def get_node_name(node: Union[ast.Attribute, ast.Name]) -> str:
     name: str = ""

pykokkos/core/visitors/workunit_visitor.py

@@ -307,9 +307,22 @@ def visit_Call(self, node: ast.Call) -> cppast.CallExpr:
 
             return rand_call
 
-        if name in {"cyl_bessel_j0", "cyl_bessel_j1"}:
+        if name in {"expint1", "erfcx"}:
             if len(args) != 1:
-                self.error(node, "pk.cyl_bessel_j0/j1 accepts only one argument")
+                self.error(node, f"pk.{name}() accepts only one argument")
+
+            s = cppast.Serializer()
+            math_call = cppast.CallExpr(cppast.DeclRefExpr(f"Kokkos::Experimental::{name}<double>"), args)
+
+            return math_call
+
+        if name in {
+                "cyl_bessel_j0", "cyl_bessel_y0", "cyl_bessel_i0", "cyl_bessel_k0",  
+                "cyl_bessel_j1", "cyl_bessel_y1", "cyl_bessel_i1", "cyl_bessel_k1", 
+                "cyl_bessel_h10", "cyl_bessel_h11", "cyl_bessel_h20", "cyl_bessel_h21"
+                }:
+            if len(args) != 1:
+                self.error(node, f"pk.{name}() accepts only one argument")
 
             s = cppast.Serializer()
             arg_str = s.serialize(args[0])

pykokkos/interface/__init__.py

@@ -33,7 +33,10 @@
     AUTO, TeamMember, PerTeam, PerThread, single
 )
 from .mathematical_special_functions import (
-    cyl_bessel_j0, cyl_bessel_j1
+    expint1, erfcx, 
+    cyl_bessel_j0, cyl_bessel_y0, cyl_bessel_k0, cyl_bessel_i0,  
+    cyl_bessel_j1, cyl_bessel_y1, cyl_bessel_k1, cyl_bessel_i1, 
+    cyl_bessel_h10, cyl_bessel_h11, cyl_bessel_h20, cyl_bessel_h21
 )
 from .memory_space import MemorySpace, get_default_memory_space
 from .parallel_dispatch import (

pykokkos/interface/mathematical_special_functions.py

@@ -1,6 +1,41 @@
+def expint1(input: float) -> float:
+    pass
+
+def erfcx(input: float) -> float:
+    pass
 
 def cyl_bessel_j0(input: float) -> float:
     pass
 
+def cyl_bessel_y0(input: float) -> float:
+    pass
+
+def cyl_bessel_i0(input: float) -> float:
+    pass
+
+def cyl_bessel_k0(input: float) -> float:
+    pass
+
 def cyl_bessel_j1(input: float) -> float:
-    pass
+    pass
+
+def cyl_bessel_y1(input: float) -> float:
+    pass
+
+def cyl_bessel_i1(input: float) -> float:
+    pass
+
+def cyl_bessel_k1(input: float) -> float:
+    pass
+
+def cyl_bessel_h10(input: float) -> float:
+    pass
+
+def cyl_bessel_h11(input: float) -> float:
+    pass
+
+def cyl_bessel_h20(input: float) -> float:
+    pass
+
+def cyl_bessel_h21(input: float) -> float:
+    pass