Skip to content
Snippets Groups Projects
Commit f4629978 authored by Frederik Hennig's avatar Frederik Hennig Committed by Christoph Alt
Browse files

Extend mdspan interface and fix mdspan memory layout mapping

parent af3aaea8
Branches
No related tags found
1 merge request!5Extend mdspan interface and fix mdspan memory layout mapping
Showing
with 388 additions and 45 deletions
***************************************
Composer API (`pystencilssfg.composer`)
***************************************
*****************************************
Composer API (``pystencilssfg.composer``)
*****************************************
.. autoclass:: pystencilssfg.composer.SfgComposer
:members:
......
......@@ -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
^^^^^^^^^^^^
......
......@@ -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
......
......@@ -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
......
......@@ -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:
......
......@@ -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)
......
......@@ -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:
......
......@@ -54,6 +54,11 @@ ScaleKernel:
JacobiMdspan:
StlContainers1D:
# std::mdspan
MdSpanFixedShapeLayouts:
MdSpanLbStreaming:
# SYCL
SyclKernels:
......
......@@ -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);
}
......@@ -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)
)
#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);
}
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)
)
#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 );
}
}
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")
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment