diff --git a/src/pystencilssfg/lang/cpp/sycl_accessor.py b/src/pystencilssfg/lang/cpp/sycl_accessor.py index e8618a5febf760126d63f999eca7a5234116c92a..f704477b3a92925acf7e318095c8fb67af32cfb2 100644 --- a/src/pystencilssfg/lang/cpp/sycl_accessor.py +++ b/src/pystencilssfg/lang/cpp/sycl_accessor.py @@ -1,7 +1,5 @@ -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) diff --git a/tests/lang/test_sycl.py b/tests/lang/test_sycl.py new file mode 100644 index 0000000000000000000000000000000000000000..88f958fae19839fe5a6e8794339aaef895d928f0 --- /dev/null +++ b/tests/lang/test_sycl.py @@ -0,0 +1,41 @@ +import pytest + + +from pystencilssfg.lang.cpp.sycl_accessor import sycl_accessor_ref +import pystencils as ps + + +@pytest.mark.parametrize("data_type", ["double", "float"]) +@pytest.mark.parametrize("dim", [1, 2, 3]) +def test_spatial_field(data_type, dim): + f = ps.fields(f"f:{data_type}[{dim}D]") + ref = sycl_accessor_ref(f) + assert ( + f"sycl::accessor< {data_type}, {dim} > &" in ref.as_variable().name_and_type() + ) + + +@pytest.mark.parametrize("data_type", ["double", "float"]) +def test_to_large_dim(data_type): + dim = 4 + f = ps.fields(f"f:{data_type}[{dim}D]") + with pytest.raises(ValueError): + sycl_accessor_ref(f) + + +@pytest.mark.parametrize("data_type", ["double", "float"]) +@pytest.mark.parametrize("spatial_dim", [1, 2, 3]) +@pytest.mark.parametrize("index_dims", [1, 2, 3]) +def test_index_field(data_type, spatial_dim, index_dims): + index_shape = ("19",) * index_dims + total_dims = spatial_dim + index_dims + f = ps.fields(f"f({', '.join(index_shape)}):{data_type}[{spatial_dim}D]") + if total_dims <= 3: + ref = sycl_accessor_ref(f) + assert ( + f"sycl::accessor< {data_type}, {total_dims} > &" + in ref.as_variable().name_and_type() + ) + else: + with pytest.raises(ValueError): + sycl_accessor_ref(f)