From f4629978b405ad6dc05a9d3de4019ea8ab92a2e0 Mon Sep 17 00:00:00 2001
From: Frederik Hennig <frederik.hennig@fau.de>
Date: Wed, 18 Dec 2024 14:47:12 +0100
Subject: [PATCH] Extend mdspan interface and fix mdspan memory layout mapping
---
docs/source/api/composer.rst | 6 +-
docs/source/api/lang.rst | 16 ++-
docs/source/conf.py | 7 +-
src/pystencilssfg/composer/basic_composer.py | 3 +-
src/pystencilssfg/emission/clang_format.py | 3 +-
src/pystencilssfg/lang/cpp/std_mdspan.py | 125 ++++++++++++++----
src/pystencilssfg/lang/expressions.py | 7 +-
tests/generator_scripts/index.yaml | 5 +
.../source/JacobiMdspan.harness.cpp | 12 +-
.../generator_scripts/source/JacobiMdspan.py | 6 +-
.../source/MdSpanFixedShape.harness.cpp | 35 +++++
.../source/MdSpanFixedShapeLayouts.py | 55 ++++++++
.../source/MdSpanLbStreaming.harness.cpp | 91 +++++++++++++
.../source/MdSpanLbStreaming.py | 62 +++++++++
14 files changed, 388 insertions(+), 45 deletions(-)
create mode 100644 tests/generator_scripts/source/MdSpanFixedShape.harness.cpp
create mode 100644 tests/generator_scripts/source/MdSpanFixedShapeLayouts.py
create mode 100644 tests/generator_scripts/source/MdSpanLbStreaming.harness.cpp
create mode 100644 tests/generator_scripts/source/MdSpanLbStreaming.py
diff --git a/docs/source/api/composer.rst b/docs/source/api/composer.rst
index a969d4f..667b0de 100644
--- a/docs/source/api/composer.rst
+++ b/docs/source/api/composer.rst
@@ -1,6 +1,6 @@
-***************************************
-Composer API (`pystencilssfg.composer`)
-***************************************
+*****************************************
+Composer API (``pystencilssfg.composer``)
+*****************************************
.. autoclass:: pystencilssfg.composer.SfgComposer
:members:
diff --git a/docs/source/api/lang.rst b/docs/source/api/lang.rst
index efc69d9..c83317e 100644
--- a/docs/source/api/lang.rst
+++ b/docs/source/api/lang.rst
@@ -9,8 +9,20 @@ Expressions
.. automodule:: pystencilssfg.lang.expressions
:members:
-C++ Standard Library (`pystencilssfg.lang.cpp`)
------------------------------------------------
+Header Files
+------------
+
+.. automodule:: pystencilssfg.lang.headers
+ :members:
+
+Data Types
+----------
+
+.. automodule:: pystencilssfg.lang.types
+ :members:
+
+C++ Standard Library (``pystencilssfg.lang.cpp``)
+-------------------------------------------------
Quick Access
^^^^^^^^^^^^
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 40d5a0f..4bdf700 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -65,12 +65,15 @@ html_theme_options = {
intersphinx_mapping = {
"python": ("https://docs.python.org/3.8", None),
- "numpy": ("https://docs.scipy.org/doc/numpy/", None),
- "matplotlib": ("https://matplotlib.org/", None),
+ "numpy": ("https://numpy.org/doc/stable/", None),
"sympy": ("https://docs.sympy.org/latest/", None),
"pystencils": ("https://da15siwa.pages.i10git.cs.fau.de/dev-docs/pystencils-nbackend/", None),
}
+# References
+
+# Treat `single-quoted` code blocks as references to any
+default_role = "any"
# Autodoc options
diff --git a/src/pystencilssfg/composer/basic_composer.py b/src/pystencilssfg/composer/basic_composer.py
index c48410c..90146e5 100644
--- a/src/pystencilssfg/composer/basic_composer.py
+++ b/src/pystencilssfg/composer/basic_composer.py
@@ -286,7 +286,7 @@ class SfgBasicComposer(SfgIComposer):
When using `call`, the given kernel will simply be called as a function.
To invoke a GPU kernel on a specified launch grid, use `cuda_invoke`
- or the interfaces of `pystencilssfg.extensions.sycl` instead.
+ or the interfaces of ``pystencilssfg.extensions.sycl`` instead.
Args:
kernel_handle: Handle to a kernel previously added to some kernel namespace.
@@ -300,6 +300,7 @@ class SfgBasicComposer(SfgIComposer):
threads_per_block: ExprLike,
stream: ExprLike | None,
):
+ """Dispatch a CUDA kernel to the device."""
num_blocks_str = str(num_blocks)
tpb_str = str(threads_per_block)
stream_str = str(stream) if stream is not None else None
diff --git a/src/pystencilssfg/emission/clang_format.py b/src/pystencilssfg/emission/clang_format.py
index 53540d3..1b15e8c 100644
--- a/src/pystencilssfg/emission/clang_format.py
+++ b/src/pystencilssfg/emission/clang_format.py
@@ -29,7 +29,8 @@ def invoke_clang_format(code: str, options: ClangFormatOptions) -> str:
binary = options.get_option("binary")
force = options.get_option("force")
- args = [binary, f"--style={options.code_style}"]
+ style = options.get_option("code_style")
+ args = [binary, f"--style={style}"]
if not shutil.which(binary):
if force:
diff --git a/src/pystencilssfg/lang/cpp/std_mdspan.py b/src/pystencilssfg/lang/cpp/std_mdspan.py
index 2ede108..5f80ad6 100644
--- a/src/pystencilssfg/lang/cpp/std_mdspan.py
+++ b/src/pystencilssfg/lang/cpp/std_mdspan.py
@@ -11,94 +11,164 @@ from pystencils.types import (
from pystencilssfg.lang.expressions import AugExpr
-from ...lang import SrcField, IFieldExtraction, cpptype, Ref, HeaderFile
+from ...lang import SrcField, IFieldExtraction, cpptype, Ref, HeaderFile, ExprLike
class StdMdspan(SrcField):
"""Represents an `std::mdspan` instance.
- **On Standard Library Adoption**
+ The `std::mdspan <https://en.cppreference.com/w/cpp/container/mdspan>`_
+ provides non-owning views into contiguous or strided n-dimensional arrays.
+ It has been added to the C++ STL with the C++23 standard.
+ As such, it is a natural data structure to target with pystencils kernels.
- Since `std::mdspan` is not yet widely adopted
+ **Concerning Headers and Namespaces**
+
+ Since ``std::mdspan`` is not yet widely adopted
(libc++ ships it as of LLVM 18, but GCC libstdc++ does not include it yet),
you might have to manually include an implementation in your project
- (you can get a reference implementation [here](https://github.com/kokkos/mdspan)).
+ (you can get a reference implementation at https://github.com/kokkos/mdspan).
However, when working with a non-standard mdspan implementation,
the path to its the header and the namespace it is defined in will likely be different.
- To tell pystencils-sfg which headers to include and which namespace to use for `mdspan`,
- use `StdMdspan.configure`.
+ To tell pystencils-sfg which headers to include and which namespace to use for ``mdspan``,
+ use `StdMdspan.configure`;
+ for instance, adding this call before creating any ``mdspan`` objects will
+ set their namespace to `std::experimental`, and require ``<experimental/mdspan>`` to be imported:
+
+ >>> from pystencilssfg.lang.cpp import std
+ >>> std.mdspan.configure("std::experimental", "<experimental/mdspan>")
+
+ **Creation from pystencils fields**
+
+ Using `from_field`, ``mdspan`` objects can be created directly from `Field <pystencils.Field>` instances.
+ The `extents`_ of the ``mdspan`` type will be inferred from the field;
+ each fixed entry in the field's shape will become a fixed entry of the ``mdspan``'s extents.
+
+ The ``mdspan``'s `layout_policy`_ defaults to `std::layout_stride`_,
+ which might not be the optimal choice depending on the memory layout of your fields.
+ You may therefore override this by specifying the name of the desired layout policy.
+ To map pystencils field layout identifiers to layout policies, consult the following table:
+
+ +------------------------+--------------------------+
+ | pystencils Layout Name | ``mdspan`` Layout Policy |
+ +========================+==========================+
+ | ``"fzyx"`` | `std::layout_left`_ |
+ | ``"soa"`` | |
+ | ``"f"`` | |
+ | ``"reverse_numpy"`` | |
+ +------------------------+--------------------------+
+ | ``"c"`` | `std::layout_right`_ |
+ | ``"numpy"`` | |
+ +------------------------+--------------------------+
+ | ``"zyxf"`` | `std::layout_stride`_ |
+ | ``"aos"`` | |
+ +------------------------+--------------------------+
+
+ The array-of-structures (``"aos"``, ``"zyxf"``) layout has no equivalent layout policy in the C++ standard,
+ so it can only be mapped onto ``layout_stride``.
+
+ .. _extents: https://en.cppreference.com/w/cpp/container/mdspan/extents
+ .. _layout_policy: https://en.cppreference.com/w/cpp/named_req/LayoutMappingPolicy
+ .. _std::layout_left: https://en.cppreference.com/w/cpp/container/mdspan/layout_left
+ .. _std::layout_right: https://en.cppreference.com/w/cpp/container/mdspan/layout_right
+ .. _std::layout_stride: https://en.cppreference.com/w/cpp/container/mdspan/layout_stride
+
+ Args:
+ T: Element type of the mdspan
"""
dynamic_extent = "std::dynamic_extent"
_namespace = "std"
- _template = cpptype("std::mdspan< {T}, {extents} >", "<mdspan>")
+ _template = cpptype("std::mdspan< {T}, {extents}, {layout_policy} >", "<mdspan>")
@classmethod
def configure(cls, namespace: str = "std", header: str | HeaderFile = "<mdspan>"):
- """Configure the namespace and header `mdspan` is defined in."""
+ """Configure the namespace and header ``std::mdspan`` is defined in."""
cls._namespace = namespace
- cls._template = cpptype(f"{namespace}::mdspan< {{T}}, {{extents}} >", header)
+ cls._template = cpptype(
+ f"{namespace}::mdspan< {{T}}, {{extents}}, {{layout_policy}} >", header
+ )
def __init__(
self,
T: UserTypeSpec,
extents: tuple[int | str, ...],
- extents_type: PsType = PsUnsignedIntegerType(64),
+ index_type: UserTypeSpec = PsUnsignedIntegerType(64),
+ layout_policy: str | None = None,
ref: bool = False,
const: bool = False,
):
T = create_type(T)
- extents_type_str = extents_type.c_string()
+ extents_type_str = create_type(index_type).c_string()
extents_str = f"{self._namespace}::extents< {extents_type_str}, {', '.join(str(e) for e in extents)} >"
- dtype = self._template(T=T, extents=extents_str, const=const)
+ if layout_policy is None:
+ layout_policy = f"{self._namespace}::layout_stride"
+ elif layout_policy in ("layout_left", "layout_right", "layout_stride"):
+ layout_policy = f"{self._namespace}::{layout_policy}"
+
+ dtype = self._template(
+ T=T, extents=extents_str, layout_policy=layout_policy, const=const
+ )
if ref:
dtype = Ref(dtype)
super().__init__(dtype)
- self._extents = extents
+ self._extents_type = extents_str
+ self._layout_type = layout_policy
self._dim = len(extents)
+ @property
+ def extents_type(self) -> str:
+ return self._extents_type
+
+ @property
+ def layout_type(self) -> str:
+ return self._layout_type
+
+ def extent(self, r: int | ExprLike) -> AugExpr:
+ return AugExpr.format("{}.extent({})", self, r)
+
+ def stride(self, r: int | ExprLike) -> AugExpr:
+ return AugExpr.format("{}.stride({})", self, r)
+
+ def data_handle(self) -> AugExpr:
+ return AugExpr.format("{}.data_handle()", self)
+
def get_extraction(self) -> IFieldExtraction:
mdspan = self
class Extraction(IFieldExtraction):
def ptr(self) -> AugExpr:
- return AugExpr.format("{}.data_handle()", mdspan)
+ return mdspan.data_handle()
def size(self, coordinate: int) -> AugExpr | None:
if coordinate > mdspan._dim:
return None
else:
- return AugExpr.format("{}.extents().extent({})", mdspan, coordinate)
+ return mdspan.extent(coordinate)
def stride(self, coordinate: int) -> AugExpr | None:
if coordinate > mdspan._dim:
return None
else:
- return AugExpr.format("{}.stride({})", mdspan, coordinate)
+ return mdspan.stride(coordinate)
return Extraction()
@staticmethod
def from_field(
field: Field,
- extents_type: PsType = PsUnsignedIntegerType(64),
+ extents_type: UserTypeSpec = PsUnsignedIntegerType(64),
+ layout_policy: str | None = None,
ref: bool = False,
const: bool = False,
):
"""Creates a `std::mdspan` instance for a given pystencils field."""
- from pystencils.field import layout_string_to_tuple
-
- if field.layout != layout_string_to_tuple("soa", field.spatial_dimensions):
- raise NotImplementedError(
- "mdspan mapping is currently only available for structure-of-arrays fields"
- )
-
extents: list[str | int] = []
for s in field.spatial_shape:
@@ -110,7 +180,12 @@ class StdMdspan(SrcField):
extents.append(StdMdspan.dynamic_extent if isinstance(s, Symbol) else s)
return StdMdspan(
- field.dtype, tuple(extents), extents_type=extents_type, ref=ref, const=const
+ field.dtype,
+ tuple(extents),
+ index_type=extents_type,
+ layout_policy=layout_policy,
+ ref=ref,
+ const=const,
).var(field.name)
diff --git a/src/pystencilssfg/lang/expressions.py b/src/pystencilssfg/lang/expressions.py
index ccaf2d1..53064bd 100644
--- a/src/pystencilssfg/lang/expressions.py
+++ b/src/pystencilssfg/lang/expressions.py
@@ -180,7 +180,7 @@ class AugExpr:
- Calling `var <AugExpr.var>` on an unbound `AugExpr` turns it into a *variable* with the given name.
This variable expression takes its set of required header files from the
- `required_headers <PsType.required_headers>` field of the data type of the `AugExpr`.
+ `required_headers <pystencils.types.PsType.required_headers>` field of the data type of the `AugExpr`.
- Using `bind <AugExpr.bind>`, an unbound `AugExpr` can be bound to an arbitrary string
of code. The `bind` method mirrors the interface of `str.format` to combine sub-expressions
and collect their dependencies.
@@ -206,6 +206,7 @@ class AugExpr:
self._is_variable = False
def var(self, name: str):
+ """Bind an unbound `AugExpr` instance as a new variable of given name."""
v = SfgVar(name, self.get_dtype())
expr = VarExpr(v)
return self._bind(expr)
@@ -220,6 +221,7 @@ class AugExpr:
return AugExpr().bind(fmt, *deps, **kwdeps)
def bind(self, fmt: str | AugExpr, *deps, **kwdeps):
+ """Bind an unbound `AugExpr` instance to an expression."""
if isinstance(fmt, AugExpr):
if bool(deps) or bool(kwdeps):
raise ValueError(
@@ -342,7 +344,8 @@ def asvar(var: VarLike) -> SfgVar:
Raises:
ValueError: If given a non-variable `AugExpr`,
- a `TypedSymbol` with a `DynamicType`,
+ a `TypedSymbol <pystencils.TypedSymbol>`
+ with a `DynamicType <pystencils.sympyextensions.typed_sympy.DynamicType>`,
or any non-variable-like object.
"""
match var:
diff --git a/tests/generator_scripts/index.yaml b/tests/generator_scripts/index.yaml
index 91bb77b..c78b335 100644
--- a/tests/generator_scripts/index.yaml
+++ b/tests/generator_scripts/index.yaml
@@ -54,6 +54,11 @@ ScaleKernel:
JacobiMdspan:
StlContainers1D:
+# std::mdspan
+
+MdSpanFixedShapeLayouts:
+MdSpanLbStreaming:
+
# SYCL
SyclKernels:
diff --git a/tests/generator_scripts/source/JacobiMdspan.harness.cpp b/tests/generator_scripts/source/JacobiMdspan.harness.cpp
index e2b8c71..529d12c 100644
--- a/tests/generator_scripts/source/JacobiMdspan.harness.cpp
+++ b/tests/generator_scripts/source/JacobiMdspan.harness.cpp
@@ -5,21 +5,21 @@
namespace stdex = std::experimental;
-using field_t = stdex::mdspan<double, stdex::extents< int64_t, std::dynamic_extent, std::dynamic_extent>>;
-using scalar_field_t = stdex::mdspan<double, stdex::extents< int64_t, std::dynamic_extent, std::dynamic_extent, 1>>;
+using field_t = stdex::mdspan<double, stdex::extents<int64_t, std::dynamic_extent, std::dynamic_extent>, stdex::layout_left>;
+using scalar_field_t = stdex::mdspan<double, stdex::extents<int64_t, std::dynamic_extent, std::dynamic_extent, 1>, stdex::layout_left>;
int main(void)
{
auto data_f = std::make_unique<double[]>(64);
- scalar_field_t f{ data_f.get(), 8, 8 };
+ scalar_field_t f{data_f.get(), 8, 8};
auto data_u = std::make_unique<double[]>(64);
- field_t u{ data_u.get(), 8, 8 };
+ field_t u{data_u.get(), 8, 8};
auto data_u_tmp = std::make_unique<double[]>(64);
- field_t u_tmp{ data_u_tmp.get(), 8, 8 };
+ field_t u_tmp{data_u_tmp.get(), 8, 8};
- double h { 1.0 / 7.0 };
+ double h{1.0 / 7.0};
gen::jacobi_smooth(f, h, u_tmp, u);
}
diff --git a/tests/generator_scripts/source/JacobiMdspan.py b/tests/generator_scripts/source/JacobiMdspan.py
index 41f149a..bbe95ac 100644
--- a/tests/generator_scripts/source/JacobiMdspan.py
+++ b/tests/generator_scripts/source/JacobiMdspan.py
@@ -20,8 +20,8 @@ with SourceFileGenerator() as sfg:
poisson_kernel = sfg.kernels.create(poisson_jacobi)
sfg.function("jacobi_smooth")(
- sfg.map_field(u_src, mdspan.from_field(u_src)),
- sfg.map_field(u_dst, mdspan.from_field(u_dst)),
- sfg.map_field(f, mdspan.from_field(f)),
+ sfg.map_field(u_src, mdspan.from_field(u_src, layout_policy="layout_left")),
+ sfg.map_field(u_dst, mdspan.from_field(u_dst, layout_policy="layout_left")),
+ sfg.map_field(f, mdspan.from_field(f, layout_policy="layout_left")),
sfg.call(poisson_kernel)
)
diff --git a/tests/generator_scripts/source/MdSpanFixedShape.harness.cpp b/tests/generator_scripts/source/MdSpanFixedShape.harness.cpp
new file mode 100644
index 0000000..cc979c8
--- /dev/null
+++ b/tests/generator_scripts/source/MdSpanFixedShape.harness.cpp
@@ -0,0 +1,35 @@
+#include "MdSpanLayouts.hpp"
+
+#include <concepts>
+#include <experimental/mdspan>
+
+namespace stdex = std::experimental;
+
+static_assert( std::is_same_v< gen::field_soa::layout_type, stdex::layout_left > );
+static_assert( std::is_same_v< gen::field_aos::layout_type, stdex::layout_stride > );
+static_assert( std::is_same_v< gen::field_c::layout_type, stdex::layout_right > );
+
+static_assert( gen::field_soa::static_extent(0) == 17 );
+static_assert( gen::field_soa::static_extent(1) == 19 );
+static_assert( gen::field_soa::static_extent(2) == 32 );
+static_assert( gen::field_soa::static_extent(3) == 9 );
+
+int main(void) {
+ gen::field_soa f_soa { nullptr };
+ gen::checkLayoutSoa(f_soa);
+
+ gen::field_aos::extents_type f_aos_extents { };
+ std::array< uint64_t, 4 > strides_aos {
+ /* stride(x) */ f_aos_extents.extent(3),
+ /* stride(y) */ f_aos_extents.extent(3) * f_aos_extents.extent(0),
+ /* stride(z) */ f_aos_extents.extent(3) * f_aos_extents.extent(0) * f_aos_extents.extent(1),
+ /* stride(f) */ 1
+ };
+
+ gen::field_aos::mapping_type f_aos_mapping { f_aos_extents, strides_aos };
+ gen::field_aos f_aos { nullptr, f_aos_mapping };
+ gen::checkLayoutAos(f_aos);
+
+ gen::field_c f_c { nullptr };
+ gen::checkLayoutC(f_c);
+}
diff --git a/tests/generator_scripts/source/MdSpanFixedShapeLayouts.py b/tests/generator_scripts/source/MdSpanFixedShapeLayouts.py
new file mode 100644
index 0000000..c89fe24
--- /dev/null
+++ b/tests/generator_scripts/source/MdSpanFixedShapeLayouts.py
@@ -0,0 +1,55 @@
+import pystencils as ps
+from pystencilssfg import SourceFileGenerator
+from pystencilssfg.lang.cpp import std
+from pystencilssfg.lang import strip_ptr_ref
+
+std.mdspan.configure(namespace="std::experimental", header="<experimental/mdspan>")
+
+with SourceFileGenerator() as sfg:
+ sfg.namespace("gen")
+ sfg.include("<cassert>")
+
+ def check_layout(field: ps.Field, mdspan: std.mdspan):
+ seq = []
+
+ for d in range(field.spatial_dimensions + field.index_dimensions):
+ seq += [
+ sfg.expr(
+ 'assert({} == {} && "Shape mismatch at coordinate {}");',
+ mdspan.extent(d),
+ field.shape[d],
+ d,
+ ),
+ sfg.expr(
+ 'assert({} == {} && "Stride mismatch at coordinate {}");',
+ mdspan.stride(d),
+ field.strides[d],
+ d,
+ ),
+ ]
+
+ return seq
+
+ f_soa = ps.fields("f_soa(9): double[17, 19, 32]", layout="soa")
+ f_soa_mdspan = std.mdspan.from_field(f_soa, layout_policy="layout_left", ref=True)
+
+ sfg.code(f"using field_soa = {strip_ptr_ref(f_soa_mdspan.dtype)};")
+ sfg.function("checkLayoutSoa")(
+ *check_layout(f_soa, f_soa_mdspan)
+ )
+
+ f_aos = ps.fields("f_aos(9): double[17, 19, 32]", layout="aos")
+ f_aos_mdspan = std.mdspan.from_field(f_aos, ref=True)
+ sfg.code(f"using field_aos = {strip_ptr_ref(f_aos_mdspan.dtype)};")
+
+ sfg.function("checkLayoutAos")(
+ *check_layout(f_aos, f_aos_mdspan)
+ )
+
+ f_c = ps.fields("f_c(9): double[17, 19, 32]", layout="c")
+ f_c_mdspan = std.mdspan.from_field(f_c, layout_policy="layout_right", ref=True)
+ sfg.code(f"using field_c = {strip_ptr_ref(f_c_mdspan.dtype)};")
+
+ sfg.function("checkLayoutC")(
+ *check_layout(f_c, f_c_mdspan)
+ )
diff --git a/tests/generator_scripts/source/MdSpanLbStreaming.harness.cpp b/tests/generator_scripts/source/MdSpanLbStreaming.harness.cpp
new file mode 100644
index 0000000..0b9b765
--- /dev/null
+++ b/tests/generator_scripts/source/MdSpanLbStreaming.harness.cpp
@@ -0,0 +1,91 @@
+#include "MdSpanLbStreaming.hpp"
+
+#include <concepts>
+#include <experimental/mdspan>
+#include <array>
+#include <cassert>
+#include <memory>
+#include <span>
+
+namespace stdex = std::experimental;
+
+static_assert(std::is_same_v<gen::field_fzyx::layout_type, stdex::layout_left>);
+static_assert(std::is_same_v<gen::field_zyxf::layout_type, stdex::layout_stride>);
+static_assert(std::is_same_v<gen::field_c::layout_type, stdex::layout_right>);
+
+using shape_type = stdex::extents< int64_t, std::dynamic_extent, std::dynamic_extent, std::dynamic_extent, 6 >;
+static_assert(std::is_same_v<gen::field_fzyx::extents_type, shape_type >);
+
+constexpr shape_type field_shape { 16l, 15l, 14l };
+
+constexpr std::array<std::array<int64_t, 3>, 2> slice{
+ {{3, 4, 5},
+ {7, 10, 12}}};
+
+template <typename Kernel, typename PdfField>
+void test_streaming(Kernel &kernel, PdfField &src_field, PdfField &dst_field)
+{
+ kernel.setZero(src_field);
+ kernel.setZero(dst_field);
+
+ for (int64_t z = slice[0][2]; z < slice[1][2]; ++z)
+ for (int64_t y = slice[0][1]; y < slice[1][1]; ++y)
+ for (int64_t x = slice[0][0]; x < slice[1][0]; ++x)
+ for (int64_t i = 0; i < int64_t(gen::STENCIL.size()); ++i)
+ {
+ src_field(x, y, z, i) = double(i);
+ }
+
+ kernel(dst_field, src_field);
+
+ for (int64_t z = slice[0][2]; z < slice[1][2]; ++z)
+ for (int64_t y = slice[0][1]; y < slice[1][1]; ++y)
+ for (int64_t x = slice[0][0]; x < slice[1][0]; ++x)
+ for (int64_t i = 0; i < int64_t(gen::STENCIL.size()); ++i)
+ {
+ const std::array<int64_t, 3> &offsets = gen::STENCIL[i];
+ assert((dst_field(x + offsets[0], y + offsets[1], z + offsets[2], i) == double(i)));
+ }
+}
+
+int main(void)
+{
+ constexpr size_t num_items { (size_t) field_shape.extent(0) * field_shape.extent(1) * field_shape.extent(2) * field_shape.extent(3) };
+
+ auto src_data = std::make_unique< double [] >( num_items );
+ auto dst_data = std::make_unique< double [] >( num_items );
+
+ // Structure-of-Arrays
+ {
+ gen::Kernel_fzyx kernel;
+ gen::field_fzyx src_arr { src_data.get(), field_shape };
+ gen::field_fzyx dst_arr { dst_data.get(), field_shape };
+ test_streaming(kernel, src_arr, dst_arr );
+ }
+
+ // Array-of-Structures
+ {
+ gen::Kernel_zyxf kernel;
+
+ std::array< uint64_t, 4 > strides_xyzf {
+ /* stride(x) */ field_shape.extent(3),
+ /* stride(y) */ field_shape.extent(3) * field_shape.extent(0),
+ /* stride(z) */ field_shape.extent(3) * field_shape.extent(0) * field_shape.extent(1),
+ /* stride(f) */ 1
+ };
+
+ gen::field_zyxf::mapping_type zyxf_mapping { field_shape, strides_xyzf };
+
+ gen::field_zyxf src_arr { src_data.get(), zyxf_mapping };
+ gen::field_zyxf dst_arr { dst_data.get(), zyxf_mapping };
+ test_streaming(kernel, src_arr, dst_arr );
+ }
+
+ // C Row-Major
+ {
+ gen::Kernel_c kernel;
+ gen::field_c src_arr { src_data.get(), field_shape };
+ gen::field_c dst_arr { dst_data.get(), field_shape };
+ test_streaming(kernel, src_arr, dst_arr );
+ }
+}
diff --git a/tests/generator_scripts/source/MdSpanLbStreaming.py b/tests/generator_scripts/source/MdSpanLbStreaming.py
new file mode 100644
index 0000000..60049a8
--- /dev/null
+++ b/tests/generator_scripts/source/MdSpanLbStreaming.py
@@ -0,0 +1,62 @@
+import numpy as np
+import pystencils as ps
+from pystencilssfg import SourceFileGenerator, SfgComposer
+from pystencilssfg.lang.cpp import std
+from pystencilssfg.lang import strip_ptr_ref
+
+std.mdspan.configure(namespace="std::experimental", header="<experimental/mdspan>")
+
+stencil = ((-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, 1), (0, 0, -1))
+
+
+def lbm_stream(sfg: SfgComposer, field_layout: str, layout_policy: str):
+ src, dst = ps.fields("src(6), dst(6): double[3D]", layout=field_layout)
+
+ src_mdspan = std.mdspan.from_field(src, layout_policy=layout_policy, extents_type="int64", ref=True)
+ dst_mdspan = std.mdspan.from_field(dst, layout_policy=layout_policy, extents_type="int64", ref=True)
+
+ asms = []
+ asms_zero = []
+
+ for i, dir in enumerate(stencil):
+ asms.append(ps.Assignment(dst.center(i), src[-np.array(dir)](i)))
+ asms_zero.append(ps.Assignment(dst.center(i), 0))
+
+ khandle = sfg.kernels.create(asms, f"stream_{field_layout}")
+ khandle_zero = sfg.kernels.create(asms_zero, f"zero_{field_layout}")
+
+ sfg.code(f"using field_{field_layout} = {strip_ptr_ref(src_mdspan.get_dtype())};")
+
+ sfg.klass(f"Kernel_{field_layout}")(
+ sfg.public(
+ sfg.method("operator()")(
+ sfg.map_field(src, src_mdspan),
+ sfg.map_field(dst, dst_mdspan),
+ sfg.call(khandle),
+ ),
+
+ sfg.method("setZero")(
+ sfg.map_field(dst, dst_mdspan),
+ sfg.call(khandle_zero),
+ )
+ )
+ )
+
+
+with SourceFileGenerator() as sfg:
+ sfg.namespace("gen")
+ sfg.include("<cassert>")
+ sfg.include("<array>")
+
+ stencil_code = (
+ "{{"
+ + ", ".join("{" + ", ".join(str(ci) for ci in c) + "}" for c in stencil)
+ + "}}"
+ )
+ sfg.code(
+ f"constexpr std::array< std::array< int64_t, 3 >, 6 > STENCIL = {stencil_code};"
+ )
+
+ lbm_stream(sfg, "fzyx", "layout_left")
+ lbm_stream(sfg, "c", "layout_right")
+ lbm_stream(sfg, "zyxf", "layout_stride")
--
GitLab