diff --git a/src/walberla/codegen/communication/__init__.py b/src/walberla/codegen/communication/__init__.py
index 5ed6fdb7f3b34b538169f058cd4d9c38544cf3a5..9a514927bd1a1247b9d782080cc5414ab7d0b5a1 100644
--- a/src/walberla/codegen/communication/__init__.py
+++ b/src/walberla/codegen/communication/__init__.py
@@ -1,3 +1,3 @@
-from .pack_infos import GpuFieldPackInfo
+from .pack_infos import GpuPdfFieldPackInfo
-__all__ = ["GpuFieldPackInfo"]
+__all__ = ["GpuPdfFieldPackInfo"]
diff --git a/src/walberla/codegen/communication/pack_infos.py b/src/walberla/codegen/communication/pack_infos.py
index 229bc9dc8cbf5d08a74bc20718140e1bbe25d3b6..84a746608e2e1c969c64c5f89f3b1d455174a2ae 100644
--- a/src/walberla/codegen/communication/pack_infos.py
+++ b/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]
diff --git a/tests/BasicLbmScenarios/LbmAlgorithms.py b/tests/BasicLbmScenarios/LbmAlgorithms.py
index 9658792ce11f159b680eb62e5514d03f53a1a655..c3faa284b767f9e1de7102791faca527e3581528 100644
--- a/tests/BasicLbmScenarios/LbmAlgorithms.py
+++ b/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)
diff --git a/tests/BasicLbmScenarios/PackInfo.py b/tests/BasicLbmScenarios/PackInfo.py
index aa88327f7c538c06611260a8cc539cec553604a6..df3ccdd2729d51922d83bb65ddb55c6362149f9d 100644
--- a/tests/BasicLbmScenarios/PackInfo.py
+++ b/tests/BasicLbmScenarios/PackInfo.py
@@ -1,12 +1,12 @@
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))