diff --git a/pystencils/field.py b/pystencils/field.py
index e92cac4046400061b19d851755af63db76dea110..e56cc3bf1322af053ba6f516b07c60dee4690ead 100644
--- a/pystencils/field.py
+++ b/pystencils/field.py
@@ -5,7 +5,7 @@ import pickle
import re
from enum import Enum
from itertools import chain
-from typing import List, Optional, Sequence, Set, Tuple
+from typing import List, Optional, Sequence, Set, Tuple, Union
import numpy as np
import sympy as sp
@@ -69,74 +69,6 @@ class FieldType(Enum):
return field.field_type == FieldType.STAGGERED_FLUX
-def fields(description=None, index_dimensions=0, layout=None, field_type=FieldType.GENERIC, **kwargs):
- """Creates pystencils fields from a string description.
-
- Examples:
- Create a 2D scalar and vector field:
- >>> s, v = fields("s, v(2): double[2D]")
- >>> assert s.spatial_dimensions == 2 and s.index_dimensions == 0
- >>> assert (v.spatial_dimensions, v.index_dimensions, v.index_shape) == (2, 1, (2,))
-
- Create an integer field of shape (10, 20):
- >>> f = fields("f : int32[10, 20]")
- >>> f.has_fixed_shape, f.shape
- (True, (10, 20))
-
- Numpy arrays can be used as template for shape and data type of field:
- >>> arr_s, arr_v = np.zeros([20, 20]), np.zeros([20, 20, 2])
- >>> s, v = fields("s, v(2)", s=arr_s, v=arr_v)
- >>> assert s.index_dimensions == 0 and s.dtype.numpy_dtype == arr_s.dtype
- >>> assert v.index_shape == (2,)
-
- Format string can be left out, field names are taken from keyword arguments.
- >>> fields(f1=arr_s, f2=arr_s)
- [f1: double[20,20], f2: double[20,20]]
-
- The keyword names ``index_dimension`` and ``layout`` have special meaning, don't use them for field names
- >>> f = fields(f=arr_v, index_dimensions=1)
- >>> assert f.index_dimensions == 1
- >>> f = fields("pdfs(19) : float32[3D]", layout='fzyx')
- >>> f.layout
- (2, 1, 0)
- """
- result = []
- if description:
- field_descriptions, dtype, shape = _parse_description(description)
- layout = 'numpy' if layout is None else layout
- for field_name, idx_shape in field_descriptions:
- if field_name in kwargs:
- arr = kwargs[field_name]
- idx_shape_of_arr = () if not len(idx_shape) else arr.shape[-len(idx_shape):]
- assert idx_shape_of_arr == idx_shape
- f = Field.create_from_numpy_array(field_name, kwargs[field_name], index_dimensions=len(idx_shape),
- field_type=field_type)
- elif isinstance(shape, tuple):
- f = Field.create_fixed_size(field_name, shape + idx_shape, dtype=dtype,
- index_dimensions=len(idx_shape), layout=layout, field_type=field_type)
- elif isinstance(shape, int):
- f = Field.create_generic(field_name, spatial_dimensions=shape, dtype=dtype,
- index_shape=idx_shape, layout=layout, field_type=field_type)
- elif shape is None:
- f = Field.create_generic(field_name, spatial_dimensions=2, dtype=dtype,
- index_shape=idx_shape, layout=layout, field_type=field_type)
- else:
- assert False
- result.append(f)
- else:
- assert layout is None, "Layout can not be specified when creating Field from numpy array"
- for field_name, arr in kwargs.items():
- result.append(Field.create_from_numpy_array(field_name, arr, index_dimensions=index_dimensions,
- field_type=field_type))
-
- if len(result) == 0:
- return None
- elif len(result) == 1:
- return result[0]
- else:
- return result
-
-
class Field:
"""
With fields one can formulate stencil-like update rules on structured grids.
@@ -875,6 +807,75 @@ class Field:
return f"{n}[{offset_str}]"
+def fields(description=None, index_dimensions=0, layout=None,
+ field_type=FieldType.GENERIC, **kwargs) -> Union[Field, List[Field]]:
+ """Creates pystencils fields from a string description.
+
+ Examples:
+ Create a 2D scalar and vector field:
+ >>> s, v = fields("s, v(2): double[2D]")
+ >>> assert s.spatial_dimensions == 2 and s.index_dimensions == 0
+ >>> assert (v.spatial_dimensions, v.index_dimensions, v.index_shape) == (2, 1, (2,))
+
+ Create an integer field of shape (10, 20):
+ >>> f = fields("f : int32[10, 20]")
+ >>> f.has_fixed_shape, f.shape
+ (True, (10, 20))
+
+ Numpy arrays can be used as template for shape and data type of field:
+ >>> arr_s, arr_v = np.zeros([20, 20]), np.zeros([20, 20, 2])
+ >>> s, v = fields("s, v(2)", s=arr_s, v=arr_v)
+ >>> assert s.index_dimensions == 0 and s.dtype.numpy_dtype == arr_s.dtype
+ >>> assert v.index_shape == (2,)
+
+ Format string can be left out, field names are taken from keyword arguments.
+ >>> fields(f1=arr_s, f2=arr_s)
+ [f1: double[20,20], f2: double[20,20]]
+
+ The keyword names ``index_dimension`` and ``layout`` have special meaning, don't use them for field names
+ >>> f = fields(f=arr_v, index_dimensions=1)
+ >>> assert f.index_dimensions == 1
+ >>> f = fields("pdfs(19) : float32[3D]", layout='fzyx')
+ >>> f.layout
+ (2, 1, 0)
+ """
+ result = []
+ if description:
+ field_descriptions, dtype, shape = _parse_description(description)
+ layout = 'numpy' if layout is None else layout
+ for field_name, idx_shape in field_descriptions:
+ if field_name in kwargs:
+ arr = kwargs[field_name]
+ idx_shape_of_arr = () if not len(idx_shape) else arr.shape[-len(idx_shape):]
+ assert idx_shape_of_arr == idx_shape
+ f = Field.create_from_numpy_array(field_name, kwargs[field_name], index_dimensions=len(idx_shape),
+ field_type=field_type)
+ elif isinstance(shape, tuple):
+ f = Field.create_fixed_size(field_name, shape + idx_shape, dtype=dtype,
+ index_dimensions=len(idx_shape), layout=layout, field_type=field_type)
+ elif isinstance(shape, int):
+ f = Field.create_generic(field_name, spatial_dimensions=shape, dtype=dtype,
+ index_shape=idx_shape, layout=layout, field_type=field_type)
+ elif shape is None:
+ f = Field.create_generic(field_name, spatial_dimensions=2, dtype=dtype,
+ index_shape=idx_shape, layout=layout, field_type=field_type)
+ else:
+ assert False
+ result.append(f)
+ else:
+ assert layout is None, "Layout can not be specified when creating Field from numpy array"
+ for field_name, arr in kwargs.items():
+ result.append(Field.create_from_numpy_array(field_name, arr, index_dimensions=index_dimensions,
+ field_type=field_type))
+
+ if len(result) == 0:
+ raise ValueError("Could not parse field description")
+ elif len(result) == 1:
+ return result[0]
+ else:
+ return result
+
+
def get_layout_from_strides(strides: Sequence[int], index_dimension_ids: Optional[List[int]] = None):
index_dimension_ids = [] if index_dimension_ids is None else index_dimension_ids
coordinates = list(range(len(strides)))
diff --git a/pystencils/gpucuda/kernelcreation.py b/pystencils/gpucuda/kernelcreation.py
index a50953b64e720a50f41b83fea0f6b376834fd257..595cf8cb53a7bcd30c3cb09b80ac0f393bd53fd0 100644
--- a/pystencils/gpucuda/kernelcreation.py
+++ b/pystencils/gpucuda/kernelcreation.py
@@ -91,10 +91,8 @@ def create_cuda_kernel(assignments: Union[AssignmentCollection, NodeCollection],
coord_mapping = {f.name: cell_idx_symbols for f in all_fields}
- loop_strides = list(fields_without_buffers)[0].shape
-
if any(FieldType.is_buffer(f) for f in all_fields):
- resolve_buffer_accesses(ast, get_base_buffer_index(ast, indexing.coordinates, loop_strides), read_only_fields)
+ resolve_buffer_accesses(ast, get_base_buffer_index(ast, indexing.coordinates, common_shape), read_only_fields)
resolve_field_accesses(ast, read_only_fields, field_to_base_pointer_info=base_pointer_info,
field_to_fixed_coordinates=coord_mapping)
diff --git a/pystencils_tests/test_buffer_gpu.py b/pystencils_tests/test_buffer_gpu.py
index 39750301a43288ada788f94cc469e055cb55f749..dd20acd5d6e4bb080980d7ddab83e38e820f7099 100644
--- a/pystencils_tests/test_buffer_gpu.py
+++ b/pystencils_tests/test_buffer_gpu.py
@@ -1,10 +1,12 @@
"""Tests for the (un)packing (from)to buffers on a CUDA GPU."""
+from dataclasses import replace
import numpy as np
import pytest
import pystencils
-from pystencils import Assignment, Field, FieldType, CreateKernelConfig, create_kernel
+from pystencils import Assignment, Field, FieldType, Target, CreateKernelConfig, create_kernel, fields
+from pystencils.bit_masks import flag_cond
from pystencils.field import create_numpy_array_with_layout, layout_string_to_tuple
from pystencils.slicing import (
add_ghost_layers, get_ghost_region_slice, get_slice_before_ghost_layer)
@@ -240,3 +242,35 @@ def test_field_layouts():
unpack_kernel = unpack_ast.compile()
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
+
+
+def test_buffer_indexing():
+ src_field, dst_field = fields(f'pdfs_src(19), pdfs_dst(19) :double[3D]')
+ mask_field = fields(f'mask : uint32 [3D]')
+ buffer = Field.create_generic('buffer', spatial_dimensions=1, field_type=FieldType.BUFFER,
+ dtype="float64",
+ index_shape=(19,))
+
+ src_field_size = src_field.spatial_shape
+ mask_field_size = mask_field.spatial_shape
+
+ up = Assignment(buffer(0), flag_cond(1, mask_field.center, src_field[0, 1, 0](1)))
+ iteration_slice = tuple(slice(None, None, 2) for _ in range(3))
+ config = CreateKernelConfig(target=Target.GPU)
+ config = replace(config, iteration_slice=iteration_slice, ghost_layers=0)
+
+ ast = create_kernel(up, config=config)
+ parameters = ast.get_parameters()
+
+ spatial_shape_symbols = [p.symbol for p in parameters if p.is_field_shape]
+
+ # The loop counters as well as the resolved field access should depend on one common spatial shape
+ if spatial_shape_symbols[0] in mask_field_size:
+ for s in spatial_shape_symbols:
+ assert s in mask_field_size
+
+ if spatial_shape_symbols[0] in src_field_size:
+ for s in spatial_shape_symbols:
+ assert s in src_field_size
+
+ assert len(spatial_shape_symbols) <= 3