update gpu pdf field pack info: restrict to streaming PDFs

src/walberla/codegen/communication/__init__.py

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

src/walberla/codegen/communication/pack_infos.py

@@ -12,7 +12,7 @@ from pystencils import (
 )
 from pystencils.stencil import offset_to_direction_string
 from pystencils.codegen.properties import FieldBasePtr
-from lbmpy import LBStencil
+from lbmpy import Stencil, LBStencil
 
 from pystencilssfg import SfgComposer
 from pystencilssfg.composer.basic_composer import KernelsAdder
@@ -55,7 +55,26 @@ class CaptureBufferPointer(SfgDeferredNode):
         return SfgEmptyNode()
 
 
-class GpuFieldPackInfo(CustomGenerator):
+class GpuPdfFieldPackInfo(CustomGenerator):
+    """Pack Info for lattice Boltzmann Gpu PDF fields.
+
+    Generate a ghost layer exchange pack info for communicating lattice Boltzmann populations
+    streaming across a block boundary,
+    for use with `gpu::GpuField` and `gpu::UniformGpuScheme`.
+
+    For a given velocity set, this pack info will only communicate those populations f_i
+    from block A to a neighbor block B which are being advected, by the streaming step,
+    from a cell in A to an adjacent cell in B.
+
+    .. note::
+        For the time being, this pack info is restricted to the *pull* streaming pattern.
+
+    Args:
+        name: Name of the generated pack info class
+        stencil: Velocity set of the lattice Boltzmann method
+        field: Symbolic representation of the PDF field
+    """
+
     def __init__(self, name: str, stencil: LBStencil, field: Field):
         if field.index_dimensions > 1:
             raise ValueError(
@@ -69,6 +88,18 @@ class GpuFieldPackInfo(CustomGenerator):
 
         self._name = name
         self._stencil = stencil
+        self._full_stencil = (
+            LBStencil(Stencil.D3Q27)
+            if self._stencil.D == 3
+            else LBStencil(Stencil.D2Q9)
+        )
+
+        #   Map storing the set of communicated populations for each communication direction
+        self._communication_sets: dict[tuple[int, int, int], list[int]] = dict()
+        for comm_dir in self._full_stencil:
+            if indices := self._get_streaming_indices(comm_dir):
+                self._communication_sets[comm_dir] = indices
+
         self._field = field
         self._dtype = field.dtype
         self._src_field = Field.new_field_with_different_name(
@@ -107,7 +138,7 @@ class GpuFieldPackInfo(CustomGenerator):
         unpack_kernels: dict[tuple[int, int, int], SfgKernelHandle] = dict()
         local_copy_kernels: dict[tuple[int, int, int], SfgKernelHandle] = dict()
 
-        comm_dirs = [c for c in self._stencil if not all(x == 0 for x in c)]
+        comm_dirs = self._communication_sets.keys()
 
         for comm_dir in comm_dirs:
             if not all(c == 0 for c in comm_dir):
@@ -200,10 +231,20 @@ class GpuFieldPackInfo(CustomGenerator):
         )
 
     def _pack_accesses(self, comm_dir: tuple[int, int, int]) -> list[Field.Access]:
-        return [self._src_field.center(i) for i in range(self._field.index_shape[0])]
+        return [self._src_field.center(i) for i in self._communication_sets[comm_dir]]
 
     def _unpack_accesses(self, comm_dir: tuple[int, int, int]) -> list[Field.Access]:
-        return [self._dst_field.center(i) for i in range(self._field.index_shape[0])]
+        return [self._dst_field.center(i) for i in self._communication_sets[comm_dir]]
+
+    def _get_streaming_indices(self, comm_dir) -> list[int]:
+        if all(d == 0 for d in comm_dir):
+            return []
+        else:
+            from lbmpy.advanced_streaming.communication import _extend_dir
+
+            directions = set(_extend_dir(comm_dir)) & set(self._stencil)
+            indices = sorted(self._stencil.index(d) for d in directions)
+            return indices
 
     def _do_pack(
         self, pkc: PackingKernelsContext, comm_dir: tuple[int, int, int]

tests/BasicLbmScenarios/LbmAlgorithms.py

@@ -16,7 +16,7 @@ from lbmpy.macroscopic_value_kernels import macroscopic_values_setter
 
 from walberla.codegen import Sweep, get_build_config
 from walberla.codegen.boundaries import NoSlip, FreeSlip
-from walberla.codegen.communication import GpuFieldPackInfo
+from walberla.codegen.communication import GpuPdfFieldPackInfo
 
 from walberla.codegen.build_config import DEBUG_MOCK_CMAKE
 
@@ -123,5 +123,5 @@ with SourceFileGenerator(keep_unknown_argv=True) as sfg:
 
     if build_config.override.target.is_gpu():
         with sfg.namespace("comm"):
-            pack_info = GpuFieldPackInfo("GpuPdfsPackInfo", stencil, f)
+            pack_info = GpuPdfFieldPackInfo("GpuPdfsPackInfo", stencil, f)
             sfg.generate(pack_info)

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.communication import GpuPdfFieldPackInfo
 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))
+    field = ps.fields(f"f({stencil.Q}): double[{stencil.D}D]")
+    sfg.generate(GpuPdfFieldPackInfo("PackInfo", stencil, field))