Compiled backends integration

.github/workflows/tox.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   tox-tests:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         python-version: ["3.10"]
@@ -34,6 +34,12 @@ jobs:
     - name: Install python dependencies
       run: |
         python -m pip install --upgrade pip setuptools wheel
+    - name: Set up C++ devtools
+      uses: seanmiddleditch/gha-setup-ninja@v3
+    - name: Install C++ dependencies
+      run: |
+        sudo apt install pybind11-dev
+        sudo apt install libboost-all-dev
     - name: Test with tox
       run: |
         pyversion_no_dot="${{ matrix.python-version }}"

src/functional/fencil_processors/callables/cache.py

@@ -0,0 +1,38 @@
+import enum
+import hashlib
+import pathlib
+import tempfile
+
+
+class Strategy(enum.Enum):
+    SESSION = 1
+    PERSISTENT = 2
+
+
+_session_cache_dir = tempfile.TemporaryDirectory(prefix="gt4py_session_")
+
+_session_cache_dir_path = pathlib.Path(_session_cache_dir.name)
+_persistent_cache_dir_path = (
+    pathlib.Path(tempfile.tempdir) / "gt4py_cache" if tempfile.tempdir else _session_cache_dir_path
+)
+
+
+def _cache_folder_name(module_name: str, module_src: str) -> str:
+    hashed = hashlib.sha256(module_src.encode(encoding="utf-8"))
+    hashed_str = hashed.hexdigest()
+    return module_name + "_" + hashed_str
+
+
+def get_cache_folder(module_name: str, module_src: str, strategy: Strategy) -> pathlib.Path:
+    folder_name = _cache_folder_name(module_name, module_src)
+
+    base_path = {
+        Strategy.SESSION: _session_cache_dir_path,
+        Strategy.PERSISTENT: _persistent_cache_dir_path,
+    }[strategy]
+    base_path.mkdir(exist_ok=True)
+
+    complete_path = base_path / folder_name
+    complete_path.mkdir(exist_ok=True)
+
+    return complete_path

src/functional/fencil_processors/callables/cpp/__init__.py

@@ -0,0 +1,4 @@
+from .callable import create_callable
+
+
+__all__ = ["create_callable"]

src/functional/fencil_processors/callables/cpp/bindings.py

@@ -0,0 +1,220 @@
+import os
+import typing
+from typing import Any, Sequence, Type, TypeVar
+
+import jinja2
+
+import eve.codegen
+import functional.fencil_processors.cpp as cpp
+import functional.fencil_processors.defs as defs
+from eve import Node
+from eve.codegen import JinjaTemplate, TemplatedGenerator
+
+
+T = TypeVar("T")
+
+
+class CommaSeparatedList(Node):
+    items: list
+
+
+class NewlineSeparatedList(Node):
+    items: list
+
+
+class BindingFunction(Node):
+    exported_name: str
+    wrapper_name: str
+    doc: str
+
+
+class BindingModule(Node):
+    name: str
+    doc: str
+    functions: NewlineSeparatedList
+
+
+class FunctionParameter(Node):
+    name: str
+    ndim: int
+    dtype: Type
+
+
+class FunctionCall(Node):
+    target: defs.Function
+
+
+class ReturnStmt(Node):
+    expr: Any
+
+
+class GridToolsSidExpr(Node):
+    name: str
+    num_dims: int
+    dtype: typing.Type
+    dim_config: int
+
+
+class SidExpr(Node):
+    identifier: str
+
+
+class WrapperFunction(Node):
+    name: str
+    parameters: CommaSeparatedList
+    body: Any
+
+
+class BindingFile(Node):
+    callee_header_file: str
+    header_files: list[str]
+    wrapper: WrapperFunction
+    binding_module: BindingModule
+
+
+class BindingCodeGenerator(TemplatedGenerator):
+    BindingFile = JinjaTemplate(
+        """\
+        #include "{{callee_header_file}}"
+
+        {% for header_file in header_files: %}\
+        #include <{{header_file}}>
+        {% endfor %}
+
+        {{wrapper}}
+
+        {{binding_module}}\
+        """
+    )
+
+    WrapperFunction = JinjaTemplate(
+        """\
+        decltype(auto) {{name}}(
+            {{parameters}}
+        )
+        {
+            {{body}}
+        }\
+        """
+    )
+
+    def visit_FunctionParameter(self, param: FunctionParameter):
+        if param.ndim > 0:
+            type_str = "pybind11::buffer"
+        else:
+            type_str = cpp.render_python_type(param.dtype)
+        return type_str + " " + param.name
+
+    FunctionArgument = JinjaTemplate("""{{expr}}""")
+
+    ReturnStmt = JinjaTemplate("""return {{expr}};""")
+
+    def visit_GridToolsSidExpr(self, expr: GridToolsSidExpr):
+        template = jinja2.Template(
+            """\
+            gridtools::as_sid<{{dtype}},\
+                              {{num_dims}},\
+                              gridtools::integral_constant<int, {{dim_config}}>,\
+                              999'999'999>({{name}})\
+            """
+        )
+        return template.render(
+            name=expr.name,
+            num_dims=expr.num_dims,
+            dtype=cpp.render_python_type(expr.dtype),
+            dim_config=expr.dim_config,
+        )
+
+    BindingModule = JinjaTemplate(
+        """\
+        PYBIND11_MODULE({{name}}, module) {
+            module.doc() = "{{doc}}";
+            {{functions}}
+        }\
+        """
+    )
+
+    BindingFunction = JinjaTemplate(
+        """module.def("{{exported_name}}", &{{wrapper_name}}, "{{doc}}");"""
+    )
+
+    def visit_CommaSeparatedList(self, items: NewlineSeparatedList):
+        str_items = [self.visit(item) for item in items.items]
+        return ", ".join(str_items)
+
+    def visit_NewlineSeparatedList(self, items: NewlineSeparatedList):
+        str_items = [self.visit(item) for item in items.items]
+        return str(os.linesep).join(str_items)
+
+    def visit_FunctionCall(self, call: FunctionCall):
+        args = [render_argument(index, param) for index, param in enumerate(call.target.parameters)]
+        return cpp.render_function_call(call.target, args)
+
+
+def make_parameter_list(
+    parameters: Sequence[defs.ScalarParameter | defs.BufferParameter],
+) -> Sequence[FunctionParameter]:
+    def make_parameter(parameter: defs.ScalarParameter | defs.BufferParameter):
+        if isinstance(parameter, defs.ScalarParameter):
+            return FunctionParameter(name=parameter.name, ndim=0, dtype=parameter.type_)
+        else:
+            return FunctionParameter(
+                name=parameter.name, ndim=parameter.num_dimensions, dtype=parameter.scalar_type
+            )
+
+    regulars = [make_parameter(param) for param in parameters]
+    return regulars
+
+
+def render_argument(index: int, param: defs.ScalarParameter | defs.BufferParameter) -> str:
+    if isinstance(param, defs.ScalarParameter):
+        return """{name}""".format(name=param.name)
+    else:
+        template = jinja2.Template(
+            """\
+            gridtools::as_sid<{{dtype}},\
+                              {{ndims}},\
+                              gridtools::integral_constant<int, {{dim_config}}>,\
+                              999'999'999>({{name}})\
+            """
+        )
+        return template.render(
+            name=param.name,
+            ndims=param.num_dimensions,
+            dtype=cpp.render_python_type(param.scalar_type),
+            dim_config=index,
+        )
+
+
+def create_bindings(target: defs.Function, target_header: str) -> defs.BindingCodeModule:
+    wrapper_name = target.name + "_wrapper"
+
+    file_binding = BindingFile(
+        callee_header_file=target_header,
+        header_files=[
+            "pybind11/pybind11.h",
+            "pybind11/stl.h",
+            "gridtools/storage/adapter/python_sid_adapter.hpp",
+            "gridtools/common/defs.hpp",
+            "gridtools/fn/unstructured.hpp",
+            "gridtools/fn/cartesian.hpp",
+            "gridtools/fn/backend/naive.hpp",
+        ],
+        wrapper=WrapperFunction(
+            name=wrapper_name,
+            parameters=CommaSeparatedList(items=make_parameter_list(target.parameters)),
+            body=ReturnStmt(expr=FunctionCall(target=target)),
+        ),
+        binding_module=BindingModule(
+            name=target.name,
+            doc="",
+            functions=NewlineSeparatedList(
+                items=[
+                    BindingFunction(exported_name=target.name, wrapper_name=wrapper_name, doc="")
+                ]
+            ),
+        ),
+    )
+
+    src = eve.codegen.format_source("cpp", BindingCodeGenerator.apply(file_binding), style="LLVM")
+    return defs.BindingCodeModule(src, [defs.LibraryDependency("pybind11", "2.9.2")])