Skip to content
Snippets Groups Projects

Add support for sycl accessors

Merged Christoph Alt requested to merge ob28imeq/pystencils-sfg:sycl_accessor into master
All threads resolved!
Viewing commit 764f56e4
Show latest version
2 files
+ 53
25
Preferences
Compare changes
Files
2
import math
from ...lang import SrcField, IFieldExtraction
from ...ir.source_components import SfgHeaderInclude
from typing import Sequence
from pystencils import Field
from pystencils.types import (
@@ -17,22 +15,17 @@ class SyclAccessor(SrcField):
self,
T: PsType,
dimensions: int,
index_shape: Sequence[int],
reference: bool = False,
):
cpp_typestr = T.c_string()
if dimensions not in [1, 2, 3]:
if 3 < dimensions:
raise ValueError("sycl accessors can only have dims 1, 2 or 3")
typestring = (
f"sycl::accessor< {cpp_typestr}, {dimensions} > {'&' if reference else ''}"
)
super().__init__(PsCustomType(typestring))
self._spatial_dimensions = dimensions
self._index_dimensions = len(index_shape)
self._index_shape = index_shape
self._index_size = math.prod(index_shape)
self._total_dimensions_ = self._spatial_dimensions + self._index_dimensions
self._dim = dimensions
self._inner_stride = 1
@property
def required_includes(self) -> set[SfgHeaderInclude]:
@@ -49,7 +42,7 @@ class SyclAccessor(SrcField):
)
def size(self, coordinate: int) -> AugExpr | None:
if coordinate > accessor._spatial_dimensions:
if coordinate > accessor._dim:
return None
else:
return AugExpr.format(
@@ -57,36 +50,30 @@ class SyclAccessor(SrcField):
)
def stride(self, coordinate: int) -> AugExpr | None:
if coordinate > accessor._total_dimensions_:
if coordinate > accessor._dim:
return None
elif coordinate >= accessor._spatial_dimensions - 1:
start = (coordinate - accessor._spatial_dimensions) + 1
return AugExpr.format(
"{}", math.prod(accessor._index_shape[start:])
)
elif coordinate == accessor._dim - 1:
return AugExpr.format("{}", accessor._inner_stride)
else:
exprs = []
args = []
for d in range(coordinate + 1, accessor._spatial_dimensions):
for d in range(coordinate + 1, accessor._dim):
args.extend([accessor, d])
exprs.append("{}.get_range().get({})")
expr = " * ".join(exprs)
expr += " * {}"
return AugExpr.format(expr, *args, accessor._index_size)
return AugExpr.format(expr, *args, accessor._inner_stride)
return Extraction()
def sycl_accessor_ref(field: Field):
"""Creates a `sycl::accessor &` for a given pystencils field."""
# Sycl accesors allow only at max 3 dimensions:
# So also mapping the index dimens to the sycl accesor we only can have 2D LBM stuff
# In principle it would be possible to map it to something like sycl::buffer<std::array<double, 19>, 3>
# but then would need to generate kernels that have sycl accessors as arguments
# Sycl Accessor do not expose information about strides, so the linearization is like here
# https://registry.khronos.org/SYCL/specs/sycl-2020/html/sycl-2020.html#_multi_dimensional_objects_and_linearization
return SyclAccessor(
field.dtype,
field.spatial_dimensions,
field.index_shape,
field.spatial_dimensions + field.index_dimensions,
reference=True,
).var(field.name)