pack info codegen WIP

src/walberla/codegen/communication/__init__.py

+from .pack_infos import GpuFieldPackInfo
+
+__all__ = ["GpuFieldPackInfo"]

src/walberla/codegen/communication/pack_infos.py

+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pystencils import (
+    Field,
+    FieldType,
+    Assignment,
+    CreateKernelConfig,
+    DynamicType,
+    Target,
+)
+from pystencils.stencil import offset_to_direction_string
+from lbmpy import LBStencil
+
+from pystencilssfg import SfgComposer
+from pystencilssfg.composer.basic_composer import KernelsAdder
+from pystencilssfg.composer.custom import CustomGenerator
+from pystencilssfg.ir import SfgKernelHandle
+from pystencilssfg.lang.cpp import std
+from pystencilssfg.lang.gpu import CudaAPI, HipAPI
+
+from ..api import GpuFieldPtr, Direction, CellInterval
+from ..build_config import get_build_config
+
+
+@dataclass
+class PackingKernelsContext:
+    sfg: SfgComposer
+    kns: KernelsAdder
+    cfg: CreateKernelConfig
+
+
+class GpuFieldPackInfo(CustomGenerator):
+    def __init__(self, name: str, stencil: LBStencil, field: Field):
+        if field.index_dimensions > 1:
+            raise ValueError(
+                "GpuFieldPackInfo currently does not support higher-order tensor fields"
+            )
+
+        if isinstance(field.dtype, DynamicType):
+            raise ValueError(
+                "Cannot generate GpuFieldPackInfo for a dynamically-typed field"
+            )
+
+        self._name = name
+        self._stencil = stencil
+        self._field = field
+        self._dtype = field.dtype
+
+    def generate(self, sfg: SfgComposer) -> None:
+        base_class = f"walberla::experimental::communication::UniformGpuFieldPackInfoImpl< {self._name } >"
+
+        build_config = get_build_config(sfg)
+
+        pkc = PackingKernelsContext(
+            sfg,
+            kns=sfg.kernel_namespace(f"{self._name}_kernels"),
+            cfg=build_config.get_pystencils_config(),
+        )
+
+        # GpuAPI: type[ProvidesGpuRuntimeAPI]
+        match pkc.cfg.get_target():
+            case Target.CUDA:
+                GpuAPI = CudaAPI
+            case Target.HIP:
+                GpuAPI = HipAPI
+            case other:
+                raise ValueError(
+                    f"Invalid target for generating GpuFieldPackInfo: {other}"
+                )
+
+        pack_kernels: dict[tuple[int, int, int], SfgKernelHandle] = dict()
+        unpack_kernels: dict[tuple[int, int, int], SfgKernelHandle] = dict()
+
+        for comm_dir in self._stencil:
+            pack_kernels[comm_dir] = self._do_pack(pkc, comm_dir)
+            unpack_kernels[comm_dir] = self._do_unpack(pkc, comm_dir)
+
+        gpu_field = GpuFieldPtr.create(self._field)
+        buffer_span = std.span(self._dtype).var("buffer")
+        dir = Direction().var("dir")
+        ci = CellInterval().var("ci")
+        stream = GpuAPI.stream_t().var("stream")
+
+        common_buffer = self._buffer(1)
+
+        sfg.klass(self._name, bases=[f"public {base_class}"])(
+            sfg.public(
+                f"using Field_T = {gpu_field.get_dtype().c_string()};",
+                sfg.method("doPack").params(gpu_field, buffer_span, dir, ci, stream)(
+                    sfg.map_field(self._field, gpu_field),
+                    sfg.map_field(common_buffer, buffer_span),
+                    sfg.switch(dir).cases(
+                        {
+                            f"walberla::stencil::Direction::{offset_to_direction_string(comm_dir)}": sfg.gpu_invoke(
+                                pack_kernels[comm_dir], stream=stream
+                            )
+                            for comm_dir in self._stencil
+                        }
+                    ),
+                ),
+                sfg.method("doUnpack").params(gpu_field, buffer_span, dir, ci, stream)(
+                    sfg.map_field(self._field, gpu_field),
+                    sfg.map_field(common_buffer, buffer_span),
+                    sfg.switch(dir).cases(
+                        {
+                            f"walberla::stencil::Direction::{offset_to_direction_string(comm_dir)}": sfg.gpu_invoke(
+                                unpack_kernels[comm_dir], stream=stream
+                            )
+                            for comm_dir in self._stencil
+                        }
+                    ),
+                ),
+            )
+        )
+
+    def _pack_accesses(self, comm_dir: tuple[int, int, int]):
+        return [self._field.center(i) for i in range(self._field.index_shape[0])]
+
+    def _do_pack(
+        self, pkc: PackingKernelsContext, comm_dir: tuple[int, int, int]
+    ) -> SfgKernelHandle:
+        pack_accs = self._pack_accesses(comm_dir)
+        buffer = self._buffer(len(pack_accs))
+        asms = [Assignment(buffer(i), acc) for i, acc in enumerate(pack_accs)]
+        dir_str = offset_to_direction_string(comm_dir)
+        return pkc.kns.create(asms, f"pack{dir_str}", pkc.cfg)
+
+    def _do_unpack(
+        self, pkc: PackingKernelsContext, comm_dir: tuple[int, int, int]
+    ) -> SfgKernelHandle:
+        pack_accs = self._pack_accesses(comm_dir)
+        buffer = self._buffer(len(pack_accs))
+        asms = [Assignment(acc, buffer(i)) for i, acc in enumerate(pack_accs)]
+        dir_str = offset_to_direction_string(comm_dir)
+        return pkc.kns.create(asms, f"unpack{dir_str}", pkc.cfg)
+
+    def _buffer(self, num_elems: int):
+        return Field.create_generic(
+            "buffer",
+            1,
+            field_type=FieldType.BUFFER,
+            dtype=self._field.dtype,
+            index_shape=(num_elems,),
+        )

tests/BasicLbmScenarios/PackInfo.py

+import pystencils as ps
+from lbmpy import Stencil, LBStencil
+from pystencilssfg import SourceFileGenerator
+from walberla.codegen.communication import GpuFieldPackInfo
+from walberla.codegen.build_config import DEBUG_MOCK_CMAKE
+
+DEBUG_MOCK_CMAKE.use_hip_default()
+
+with SourceFileGenerator() as sfg:
+    field = ps.fields("f(3): double[3D]")
+    stencil = LBStencil(Stencil.D3Q19)
+    sfg.generate(GpuFieldPackInfo("PackInfo", stencil, field))