tests/frontend/test_typed_sympy.py

+import numpy as np
+import pickle
+import sympy as sp
+from sympy.logic import boolalg
+
+from pystencils.sympyextensions.typed_sympy import (
+    TypedSymbol,
+    tcast,
+    TypeAtom,
+    DynamicType,
+)
+from pystencils.types import create_type
+from pystencils.types.quick import UInt, Ptr
+
+
+def test_type_atoms():
+    atom1 = TypeAtom(create_type("int32"))
+    atom2 = TypeAtom(create_type(np.int32))
+
+    assert atom1 == atom2
+
+    atom3 = TypeAtom(create_type("const int32"))
+    assert atom1 != atom3
+
+    atom4 = TypeAtom(DynamicType.INDEX_TYPE)
+    atom5 = TypeAtom(DynamicType.NUMERIC_TYPE)
+
+    assert atom3 != atom4
+    assert atom4 != atom5
+
+    dump = pickle.dumps(atom1)
+    atom1_reconst = pickle.loads(dump)
+
+    assert atom1_reconst == atom1
+
+
+def test_typed_symbol():
+    x = TypedSymbol("x", "uint32")
+    x2 = TypedSymbol("x", "uint64 *")
+    z = TypedSymbol("z", "float32")
+
+    assert x == TypedSymbol("x", np.uint32)
+    assert x != x2
+
+    assert x.dtype == UInt(32)
+    assert x2.dtype == Ptr(UInt(64))
+
+    assert x.is_integer
+    assert x.is_nonnegative
+
+    assert not x2.is_integer
+
+    assert z.is_real
+    assert not z.is_nonnegative
+
+
+def test_casts():
+    x, y = sp.symbols("x, y")
+
+    #   Pickling
+    expr = tcast(x, "int32")
+    dump = pickle.dumps(expr)
+    expr_reconst = pickle.loads(dump)
+    assert expr_reconst == expr
+
+    #   Boolean Casts
+    bool_expr = tcast(x, "bool")
+    assert isinstance(bool_expr, boolalg.Boolean)
+    
+    #   Check that we can construct boolean expressions with cast results
+    _ = boolalg.Or(bool_expr, y)
+    
+    #   Assumptions
+    expr = tcast(x, "int32")
+    assert expr.is_integer
+    assert expr.is_real
+    assert expr.is_nonnegative is None
+
+    expr = tcast(x, "uint32")
+    assert expr.is_integer
+    assert expr.is_real
+    assert expr.is_nonnegative
+
+    expr = tcast(x, "float32")
+    assert expr.is_integer is None
+    assert expr.is_real
+    assert expr.is_nonnegative is None

tests/jit/test_cpujit.py

+import pytest
+
+import sympy as sp
+import numpy as np
+from pystencils import create_kernel, Assignment, fields, Field
+from pystencils.jit import CpuJit
+
+
+@pytest.fixture
+def cpu_jit(tmp_path) -> CpuJit:
+    return CpuJit.create(objcache=tmp_path)
+
+
+def test_basic_cpu_kernel(cpu_jit):
+    f, g = fields("f, g: [2D]")
+    asm = Assignment(f.center(), 2.0 * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    rng = np.random.default_rng()
+    f_arr = rng.random(size=(34, 26), dtype="float64")
+    g_arr = np.zeros_like(f_arr)
+
+    kfunc(f=f_arr, g=g_arr)
+
+    np.testing.assert_almost_equal(g_arr, 2.0 * f_arr)
+
+
+def test_argument_type_error(cpu_jit):
+    f, g = fields("f, g: [2D]")
+    c = sp.Symbol("c")
+    asm = Assignment(f.center(), c * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    arr_fp16 = np.zeros((23, 12), dtype="float16")
+    arr_fp32 = np.zeros((23, 12), dtype="float32")
+    arr_fp64 = np.zeros((23, 12), dtype="float64")
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp32, g=arr_fp64, c=2.0)
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp64, g=arr_fp32, c=2.0)
+
+    with pytest.raises(TypeError):
+        kfunc(f=arr_fp16, g=arr_fp16, c=2.0)
+
+    #   Wrong scalar types are OK, though
+    kfunc(f=arr_fp64, g=arr_fp64, c=np.float16(1.0))
+
+
+def test_fixed_shape(cpu_jit):
+    a = np.zeros((12, 23), dtype="float64")
+    b = np.zeros((13, 21), dtype="float64")
+    
+    f = Field.create_from_numpy_array("f", a)
+    g = Field.create_from_numpy_array("g", a)
+
+    asm = Assignment(f.center(), 2.0 * g.center())
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    kfunc(f=a, g=a)
+
+    with pytest.raises(ValueError):
+        kfunc(f=b, g=a)
+
+    with pytest.raises(ValueError):
+        kfunc(f=a, g=b)
+
+
+def test_fixed_index_shape(cpu_jit):
+    f, g = fields("f(3), g(2, 2): [2D]")
+
+    asm = Assignment(f.center(1), g.center(0, 0) + g.center(0, 1) + g.center(1, 0) + g.center(1, 1))
+    ker = create_kernel(asm)
+    kfunc = cpu_jit.compile(ker)
+
+    f_arr = np.zeros((12, 14, 3))
+    g_arr = np.zeros((12, 14, 2, 2))
+    kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 2))
+        g_arr = np.zeros((12, 14, 2, 2))
+        kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 3))
+        g_arr = np.zeros((12, 14, 4))
+        kfunc(f=f_arr, g=g_arr)
+
+    with pytest.raises(ValueError):
+        f_arr = np.zeros((12, 14, 3))
+        g_arr = np.zeros((12, 14, 1, 3))
+        kfunc(f=f_arr, g=g_arr)

tests/test_blocking.py → tests/kernelcreation/test_blocking.py

 import numpy as np
 import sympy as sp
+import pytest
 
 import pystencils as ps
 
@@ -39,6 +40,7 @@ def check_equivalence(assignments, src_arr):
             np.testing.assert_almost_equal(ref_arr, dst3_arr)
 
 
+@pytest.mark.xfail(reason="CPU blocking is not yet implemented in the new backend")
 def test_jacobi3d_var_size():
     src, dst = ps.fields("src, dst: double[3D]", layout='c')
 
@@ -55,6 +57,7 @@ def test_jacobi3d_var_size():
     check_equivalence(jacobi(dst, src), arr)
 
 
+@pytest.mark.xfail(reason="CPU blocking is not yet implemented in the new backend")
 def test_jacobi3d_fixed_size():
     print("Fixed Size: Large non divisible sizes")
     arr = np.empty([10, 10, 9])
@@ -72,6 +75,7 @@ def test_jacobi3d_fixed_size():
     check_equivalence(jacobi(dst, src), arr)
 
 
+@pytest.mark.xfail(reason="CPU blocking is not yet implemented in the new backend")
 def test_jacobi3d_fixed_field_size():
     src, dst = ps.fields("src, dst: double[3, 5, 6]", layout='c')
 

tests/test_blocking_staggered.py → tests/kernelcreation/test_blocking_staggered.py

 import numpy as np
+import pytest
 
 import pystencils as ps
 
 
+@pytest.mark.xfail(reason="Blocking and staggered accesses are not yet implemented in the new backend")
 def test_blocking_staggered():
     f = ps.fields("f: double[3D]")
     stag = ps.fields("stag(3): double[3D]", field_type=ps.FieldType.STAGGERED)

tests/test_buffer.py → tests/kernelcreation/test_buffer.py

 """Tests  (un)packing (from)to buffers."""
 
+import pytest
 import numpy as np
 
 import pystencils as ps
@@ -41,9 +42,7 @@ def test_full_scalar_field():
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = [Assignment(buffer.center(), src_field.center())]
-        config = ps.CreateKernelConfig(data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
-        pack_code = create_kernel(pack_eqs, config=config)
-        code = ps.get_code_str(pack_code)
+        pack_code = create_kernel(pack_eqs)
         ps.show_code(pack_code)
 
         pack_kernel = pack_code.compile()
@@ -51,8 +50,7 @@ def test_full_scalar_field():
 
         unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
 
-        config = ps.CreateKernelConfig(data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
-        unpack_code = create_kernel(unpack_eqs, config=config)
+        unpack_code = create_kernel(unpack_eqs)
 
         unpack_kernel = unpack_code.compile()
         unpack_kernel(dst_field=dst_arr, buffer=buffer_arr)
@@ -77,8 +75,7 @@ def test_field_slice():
 
             pack_eqs = [Assignment(buffer.center(), src_field.center())]
 
-            config = ps.CreateKernelConfig(data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
-            pack_code = create_kernel(pack_eqs, config=config)
+            pack_code = create_kernel(pack_eqs)
 
             pack_kernel = pack_code.compile()
             pack_kernel(buffer=bufferArr, src_field=src_arr[pack_slice])
@@ -86,8 +83,7 @@ def test_field_slice():
             # Unpack into ghost layer of dst_field in N direction
             unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
 
-            config = ps.CreateKernelConfig(data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
-            unpack_code = create_kernel(unpack_eqs, config=config)
+            unpack_code = create_kernel(unpack_eqs)
 
             unpack_kernel = unpack_code.compile()
             unpack_kernel(buffer=bufferArr, dst_field=dst_arr[unpack_slice])
@@ -102,7 +98,7 @@ def test_all_cell_values():
     for (src_arr, dst_arr, bufferArr) in fields:
         src_field = Field.create_from_numpy_array("src_field", src_arr, index_dimensions=1)
         dst_field = Field.create_from_numpy_array("dst_field", dst_arr, index_dimensions=1)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = []
@@ -112,8 +108,7 @@ def test_all_cell_values():
             eq = Assignment(buffer(idx), src_field(idx))
             pack_eqs.append(eq)
 
-        config = ps.CreateKernelConfig(data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
-        pack_code = create_kernel(pack_eqs, config=config)
+        pack_code = create_kernel(pack_eqs)
         pack_kernel = pack_code.compile()
         pack_kernel(buffer=bufferArr, src_field=src_arr)
 
@@ -123,8 +118,7 @@ def test_all_cell_values():
             eq = Assignment(dst_field(idx), buffer(idx))
             unpack_eqs.append(eq)
 
-        config = ps.CreateKernelConfig(data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
-        unpack_code = create_kernel(unpack_eqs, config=config)
+        unpack_code = create_kernel(unpack_eqs)
         unpack_kernel = unpack_code.compile()
         unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
 
@@ -140,7 +134,7 @@ def test_subset_cell_values():
     for (src_arr, dst_arr, bufferArr) in fields:
         src_field = Field.create_from_numpy_array("src_field", src_arr, index_dimensions=1)
         dst_field = Field.create_from_numpy_array("dst_field", dst_arr, index_dimensions=1)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(len(cell_indices),),
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = []
@@ -150,8 +144,7 @@ def test_subset_cell_values():
             eq = Assignment(buffer(buffer_idx), src_field(cell_idx))
             pack_eqs.append(eq)
 
-        config = ps.CreateKernelConfig(data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
-        pack_code = create_kernel(pack_eqs, config=config)
+        pack_code = create_kernel(pack_eqs)
         pack_kernel = pack_code.compile()
         pack_kernel(buffer=bufferArr, src_field=src_arr)
 
@@ -161,8 +154,7 @@ def test_subset_cell_values():
             eq = Assignment(dst_field(cell_idx), buffer(buffer_idx))
             unpack_eqs.append(eq)
 
-        config = ps.CreateKernelConfig(data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
-        unpack_code = create_kernel(unpack_eqs, config=config)
+        unpack_code = create_kernel(unpack_eqs)
         unpack_kernel = unpack_code.compile()
         unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
 
@@ -177,7 +169,7 @@ def test_field_layouts():
         for (src_arr, dst_arr, bufferArr) in fields:
             src_field = Field.create_from_numpy_array("src_field", src_arr, index_dimensions=1)
             dst_field = Field.create_from_numpy_array("dst_field", dst_arr, index_dimensions=1)
-            buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+            buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                           field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
             pack_eqs = []
@@ -187,8 +179,7 @@ def test_field_layouts():
                 eq = Assignment(buffer(idx), src_field(idx))
                 pack_eqs.append(eq)
 
-            config = ps.CreateKernelConfig(data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
-            pack_code = create_kernel(pack_eqs, config=config)
+            pack_code = create_kernel(pack_eqs)
             pack_kernel = pack_code.compile()
             pack_kernel(buffer=bufferArr, src_field=src_arr)
 
@@ -198,8 +189,7 @@ def test_field_layouts():
                 eq = Assignment(dst_field(idx), buffer(idx))
                 unpack_eqs.append(eq)
 
-            config = ps.CreateKernelConfig(data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
-            unpack_code = create_kernel(unpack_eqs, config=config)
+            unpack_code = create_kernel(unpack_eqs)
             unpack_kernel = unpack_code.compile()
             unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
 
@@ -214,7 +204,7 @@ def test_iteration_slices():
         # dst_field = Field.create_from_numpy_array("dst_field", dst_arr, index_dimensions=1)
         src_field = Field.create_generic("src_field", spatial_dimensions, index_shape=(num_cell_values,), dtype=dt)
         dst_field = Field.create_generic("dst_field", spatial_dimensions, index_shape=(num_cell_values,), dtype=dt)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = []
@@ -227,14 +217,13 @@ def test_iteration_slices():
         dim = src_field.spatial_dimensions
 
         #   Pack only the leftmost slice, only every second cell
-        pack_slice = (slice(None, None, 2),) * (dim - 1) + (0,)
+        pack_slice = (slice(None, None, 2),) * (dim - 1) + (slice(0, 1, None),)
 
         #   Fill the entire array with data
         src_arr[(slice(None, None, 1),) * dim] = np.arange(num_cell_values)
         dst_arr.fill(0)
 
-        config = ps.CreateKernelConfig(iteration_slice=pack_slice,
-                                       data_type={'src_field': src_arr.dtype, 'buffer': buffer.dtype})
+        config = ps.CreateKernelConfig(iteration_slice=pack_slice)
 
         pack_code = create_kernel(pack_eqs, config=config)
         pack_kernel = pack_code.compile()
@@ -246,8 +235,7 @@ def test_iteration_slices():
             eq = Assignment(dst_field(idx), buffer(idx))
             unpack_eqs.append(eq)
 
-        config = ps.CreateKernelConfig(iteration_slice=pack_slice,
-                                       data_type={'dst_field': dst_arr.dtype, 'buffer': buffer.dtype})
+        config = ps.CreateKernelConfig(iteration_slice=pack_slice)
 
         unpack_code = create_kernel(unpack_eqs, config=config)
         unpack_kernel = unpack_code.compile()

tests/test_buffer_gpu.py → tests/kernelcreation/test_buffer_gpu.py

@@ -6,7 +6,7 @@ import pytest
 
 import pystencils
 from pystencils import Assignment, Field, FieldType, Target, CreateKernelConfig, create_kernel, fields
-from pystencils.bit_masks import flag_cond
+from pystencils.sympyextensions.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)
@@ -57,18 +57,16 @@ def test_full_scalar_field():
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = [Assignment(buffer.center(), src_field.center())]
-        pack_types = {'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
 
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=pack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         pack_ast = create_kernel(pack_eqs, config=config)
 
         pack_kernel = pack_ast.compile()
         pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
 
         unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
-        unpack_types = {'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
 
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=unpack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         unpack_ast = create_kernel(unpack_eqs, config=config)
 
         unpack_kernel = unpack_ast.compile()
@@ -95,9 +93,8 @@ def test_field_slice():
                                           field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
             pack_eqs = [Assignment(buffer.center(), src_field.center())]
-            pack_types = {'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
 
-            config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=pack_types)
+            config = CreateKernelConfig(target=pystencils.Target.GPU)
             pack_ast = create_kernel(pack_eqs, config=config)
 
             pack_kernel = pack_ast.compile()
@@ -105,9 +102,8 @@ def test_field_slice():
 
             # Unpack into ghost layer of dst_field in N direction
             unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
-            unpack_types = {'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
 
-            config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=unpack_types)
+            config = CreateKernelConfig(target=pystencils.Target.GPU)
             unpack_ast = create_kernel(unpack_eqs, config=config)
 
             unpack_kernel = unpack_ast.compile()
@@ -125,7 +121,7 @@ def test_all_cell_values():
     for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
         src_field = Field.create_from_numpy_array("src_field", gpu_src_arr, index_dimensions=1)
         dst_field = Field.create_from_numpy_array("dst_field", gpu_src_arr, index_dimensions=1)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                       field_type=FieldType.BUFFER, dtype=gpu_src_arr.dtype)
 
         pack_eqs = []
@@ -135,9 +131,7 @@ def test_all_cell_values():
             eq = Assignment(buffer(idx), src_field(idx))
             pack_eqs.append(eq)
 
-        pack_types = {'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=pack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         pack_code = create_kernel(pack_eqs, config=config)
         pack_kernel = pack_code.compile()
 
@@ -149,9 +143,7 @@ def test_all_cell_values():
             eq = Assignment(dst_field(idx), buffer(idx))
             unpack_eqs.append(eq)
 
-        unpack_types = {'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=unpack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         unpack_ast = create_kernel(unpack_eqs, config=config)
         unpack_kernel = unpack_ast.compile()
         unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
@@ -170,7 +162,7 @@ def test_subset_cell_values():
     for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
         src_field = Field.create_from_numpy_array("src_field", gpu_src_arr, index_dimensions=1)
         dst_field = Field.create_from_numpy_array("dst_field", gpu_src_arr, index_dimensions=1)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(len(cell_indices),),
                                       field_type=FieldType.BUFFER, dtype=gpu_src_arr.dtype)
 
         pack_eqs = []
@@ -181,7 +173,7 @@ def test_subset_cell_values():
             pack_eqs.append(eq)
 
         pack_types = {'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=pack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         pack_ast = create_kernel(pack_eqs, config=config)
         pack_kernel = pack_ast.compile()
         pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
@@ -193,7 +185,7 @@ def test_subset_cell_values():
             unpack_eqs.append(eq)
 
         unpack_types = {'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-        config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=unpack_types)
+        config = CreateKernelConfig(target=pystencils.Target.GPU)
         unpack_ast = create_kernel(unpack_eqs, config=config)
         unpack_kernel = unpack_ast.compile()
 
@@ -212,7 +204,7 @@ def test_field_layouts():
         for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
             src_field = Field.create_from_numpy_array("src_field", gpu_src_arr, index_dimensions=1)
             dst_field = Field.create_from_numpy_array("dst_field", gpu_src_arr, index_dimensions=1)
-            buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+            buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                           field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
             pack_eqs = []
@@ -223,7 +215,7 @@ def test_field_layouts():
                 pack_eqs.append(eq)
 
             pack_types = {'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-            config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=pack_types)
+            config = CreateKernelConfig(target=pystencils.Target.GPU)
             pack_ast = create_kernel(pack_eqs, config=config)
             pack_kernel = pack_ast.compile()
 
@@ -236,13 +228,14 @@ def test_field_layouts():
                 unpack_eqs.append(eq)
 
             unpack_types = {'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype}
-            config = CreateKernelConfig(target=pystencils.Target.GPU, data_type=unpack_types)
+            config = CreateKernelConfig(target=pystencils.Target.GPU)
             unpack_ast = create_kernel(unpack_eqs, config=config)
             unpack_kernel = unpack_ast.compile()
 
             unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
 
 
+@pytest.mark.xfail(reason="flag_cond is not available yet")
 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]')
@@ -256,7 +249,7 @@ def test_buffer_indexing():
     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)
+    config = replace(config, iteration_slice=iteration_slice)
 
     ast = create_kernel(up, config=config)
     parameters = ast.get_parameters()
@@ -277,13 +270,16 @@ def test_buffer_indexing():
 
 @pytest.mark.parametrize('gpu_indexing', ("block", "line"))
 def test_iteration_slices(gpu_indexing):
+    if gpu_indexing == "line":
+        pytest.xfail("Line indexing not available yet")
+
     num_cell_values = 19
     dt = np.uint64
     fields = _generate_fields(dt=dt, stencil_directions=num_cell_values)
     for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
         src_field = Field.create_from_numpy_array("src_field", gpu_src_arr, index_dimensions=1)
         dst_field = Field.create_from_numpy_array("dst_field", gpu_src_arr, index_dimensions=1)
-        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_dimensions=1,
+        buffer = Field.create_generic("buffer", spatial_dimensions=1, index_shape=(num_cell_values,),
                                       field_type=FieldType.BUFFER, dtype=src_arr.dtype)
 
         pack_eqs = []
@@ -303,9 +299,7 @@ def test_iteration_slices(gpu_indexing):
         gpu_src_arr.set(src_arr)
         gpu_dst_arr.fill(0)
 
-        config = CreateKernelConfig(target=Target.GPU, iteration_slice=pack_slice,
-                                    data_type={'src_field': gpu_src_arr.dtype, 'buffer': gpu_buffer_arr.dtype},
-                                    gpu_indexing=gpu_indexing)
+        config = CreateKernelConfig(target=Target.GPU, iteration_slice=pack_slice)
 
         pack_code = create_kernel(pack_eqs, config=config)
         pack_kernel = pack_code.compile()
@@ -317,9 +311,7 @@ def test_iteration_slices(gpu_indexing):
             eq = Assignment(dst_field(idx), buffer(idx))
             unpack_eqs.append(eq)
 
-        config = CreateKernelConfig(target=Target.GPU, iteration_slice=pack_slice,
-                                    data_type={'dst_field': gpu_dst_arr.dtype, 'buffer': gpu_buffer_arr.dtype},
-                                    gpu_indexing=gpu_indexing)
+        config = CreateKernelConfig(target=Target.GPU, iteration_slice=pack_slice)
 
         unpack_code = create_kernel(unpack_eqs, config=config)
         unpack_kernel = unpack_code.compile()

tests/kernelcreation/test_domain_kernels.py

+import pytest
+import sympy as sp
+import numpy as np
+
+from dataclasses import replace
+
+from pystencils import (
+    fields,
+    Field,
+    AssignmentCollection,
+    Target,
+    CreateKernelConfig,
+)
+from pystencils.assignment import assignment_from_stencil
+
+from pystencils import create_kernel, Kernel
+from pystencils.backend.emission import emit_code
+
+
+def inspect_dp_kernel(kernel: Kernel, gen_config: CreateKernelConfig):
+    code = emit_code(kernel)
+
+    match gen_config.target:
+        case Target.X86_SSE:
+            assert "_mm_loadu_pd" in code
+            assert "_mm_storeu_pd" in code
+        case Target.X86_AVX:
+            assert "_mm256_loadu_pd" in code
+            assert "_mm256_storeu_pd" in code
+        case Target.X86_AVX512:
+            assert "_mm512_loadu_pd" in code
+            assert "_mm512_storeu_pd" in code
+
+
+def test_filter_kernel(gen_config):
+    if gen_config.target == Target.CUDA:
+        import cupy as cp
+
+        xp = cp
+    else:
+        xp = np
+
+    weight = sp.Symbol("weight")
+    stencil = [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
+
+    src, dst = fields("src, dst: [2D]")
+    asm = assignment_from_stencil(stencil, src, dst, normalization_factor=weight)
+    asms = AssignmentCollection([asm])
+
+    ast = create_kernel(asms, gen_config)
+    inspect_dp_kernel(ast, gen_config)
+    kernel = ast.compile()
+
+    src_arr = xp.ones((42, 31))
+    dst_arr = xp.zeros_like(src_arr)
+
+    kernel(src=src_arr, dst=dst_arr, weight=2.0)
+
+    expected = xp.zeros_like(src_arr)
+    expected[1:-1, 1:-1].fill(18.0)
+
+    xp.testing.assert_allclose(dst_arr, expected)
+
+
+def test_filter_kernel_fixedsize(gen_config):
+    if gen_config.target == Target.CUDA:
+        import cupy as cp
+
+        xp = cp
+    else:
+        xp = np
+
+    weight = sp.Symbol("weight")
+    stencil = [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
+
+    src_arr = xp.ones((42, 31))
+    dst_arr = xp.zeros_like(src_arr)
+
+    src = Field.create_from_numpy_array("src", src_arr)
+    dst = Field.create_from_numpy_array("dst", dst_arr)
+
+    asm = assignment_from_stencil(stencil, src, dst, normalization_factor=weight)
+    asms = AssignmentCollection([asm])
+
+    ast = create_kernel(asms, gen_config)
+    inspect_dp_kernel(ast, gen_config)
+    kernel = ast.compile()
+
+    kernel(src=src_arr, dst=dst_arr, weight=2.0)
+
+    expected = xp.zeros_like(src_arr)
+    expected[1:-1, 1:-1].fill(18.0)
+
+    xp.testing.assert_allclose(dst_arr, expected)

tests/kernelcreation/test_functions.py

+import sympy as sp
+import numpy as np
+import pytest
+from dataclasses import replace
+from itertools import product
+
+from pystencils import (
+    fields,
+    create_kernel,
+    CreateKernelConfig,
+    Target,
+    Assignment,
+    Field,
+)
+from pystencils.backend.ast import dfs_preorder
+from pystencils.backend.ast.expressions import PsCall
+
+
+def unary_function(name, xp):
+    return {
+        "exp": (sp.exp, xp.exp),
+        "log": (sp.log, xp.log),
+        "sin": (sp.sin, xp.sin),
+        "cos": (sp.cos, xp.cos),
+        "tan": (sp.tan, xp.tan),
+        "sinh": (sp.sinh, xp.sinh),
+        "cosh": (sp.cosh, xp.cosh),
+        "asin": (sp.asin, xp.arcsin),
+        "acos": (sp.acos, xp.arccos),
+        "atan": (sp.atan, xp.arctan),
+        "sqrt": (sp.sqrt, xp.sqrt),
+        "abs": (sp.Abs, xp.abs),
+        "floor": (sp.floor, xp.floor),
+        "ceil": (sp.ceiling, xp.ceil),
+    }[name]
+
+
+def binary_function(name, xp):
+    return {
+        "min": (sp.Min, xp.fmin),
+        "max": (sp.Max, xp.fmax),
+        "pow": (sp.Pow, xp.power),
+        "atan2": (sp.atan2, xp.arctan2),
+    }[name]
+
+
+AVAIL_TARGETS = Target.available_targets()
+
+
+@pytest.fixture
+def function_domain(function_name, dtype):
+    eps = 1e-6
+    rng = np.random.default_rng()
+
+    match function_name:
+        case (
+            "exp" | "sin" | "cos" | "sinh" | "cosh" | "atan" | "abs" | "floor" | "ceil"
+        ):
+            return np.concatenate(
+                [
+                    [0.0, -1.0, 1.0],
+                    rng.uniform(-0.1, 0.1, 5),
+                    rng.uniform(-5.0, 5.0, 8),
+                    rng.uniform(-10, 10, 8),
+                ]
+            ).astype(dtype)
+        case "tan":
+            return np.concatenate(
+                [
+                    [0.0, -1.0, 1.0],
+                    rng.uniform(-np.pi / 2.0 + eps, np.pi / 2.0 - eps, 13),
+                ]
+            ).astype(dtype)
+        case "asin" | "acos":
+            return np.concatenate(
+                [
+                    [0.0, 1.0, -1.0],
+                    rng.uniform(-1.0, 1.0, 13),
+                ]
+            ).astype(dtype)
+        case "log" | "sqrt":
+            return np.concatenate(
+                [
+                    [1.0],
+                    rng.uniform(eps, 0.1, 7),
+                    rng.uniform(eps, 1.0, 8),
+                    rng.uniform(eps, 1e6, 8),
+                ]
+            ).astype(dtype)
+        case "min" | "max" | "atan2":
+            return np.concatenate(
+                [
+                    rng.uniform(-0.1, 0.1, 8),
+                    rng.uniform(-5.0, 5.0, 8),
+                    rng.uniform(-10, 10, 8),
+                ]
+            ).astype(dtype), np.concatenate(
+                [
+                    rng.uniform(-0.1, 0.1, 8),
+                    rng.uniform(-5.0, 5.0, 8),
+                    rng.uniform(-10, 10, 8),
+                ]
+            ).astype(
+                dtype
+            )
+        case "pow":
+            return np.concatenate(
+                [
+                    [0., 1., 1.],
+                    rng.uniform(-1., 1., 8),
+                    rng.uniform(0., 5., 8),
+                ]
+            ).astype(dtype), np.concatenate(
+                [
+                    [1., 0., 2.],
+                    np.arange(2., 10., 1.),
+                    rng.uniform(-2.0, 2.0, 8),
+                ]
+            ).astype(
+                dtype
+            )
+        case _:
+            assert False, "I don't know the domain of that function"
+
+
+@pytest.mark.parametrize(
+    "function_name, target",
+    list(
+        product(
+            (
+                "exp",
+                "log",
+                "sin",
+                "cos",
+                "tan",
+                "sinh",
+                "cosh",
+                "asin",
+                "acos",
+                "atan",
+            ),
+            [t for t in AVAIL_TARGETS if Target._X86 not in t],
+        )
+    )
+    + list(
+        product(
+            ["floor", "ceil"], [t for t in AVAIL_TARGETS if Target._AVX512 not in t]
+        )
+    )
+    + list(product(["sqrt", "abs"], AVAIL_TARGETS)),
+)
+@pytest.mark.parametrize("dtype", (np.float32, np.float64))
+def test_unary_functions(gen_config, xp, function_name, dtype, function_domain):
+    sp_func, xp_func = unary_function(function_name, xp)
+    resolution = np.finfo(dtype).resolution
+
+    #   Array size should be larger than eight, such that vectorized kernels don't just run their remainder loop
+    inp = xp.array(function_domain)
+    outp = xp.zeros_like(inp)
+
+    reference = xp_func(inp)
+
+    inp_field = Field.create_from_numpy_array("inp", inp)
+    outp_field = inp_field.new_field_with_different_name("outp")
+
+    asms = [Assignment(outp_field.center(), sp_func(inp_field.center()))]
+    gen_config = replace(gen_config, default_dtype=dtype)
+
+    kernel = create_kernel(asms, gen_config)
+    kfunc = kernel.compile()
+    kfunc(inp=inp, outp=outp)
+
+    xp.testing.assert_allclose(outp, reference, rtol=resolution)
+
+
+@pytest.mark.parametrize(
+    "function_name,target",
+    list(product(["min", "max"], AVAIL_TARGETS))
+    + list(
+        product(["pow", "atan2"], [t for t in AVAIL_TARGETS if Target._X86 not in t])
+    ),
+)
+@pytest.mark.parametrize("dtype", (np.float32, np.float64))
+def test_binary_functions(gen_config, xp, function_name, dtype, function_domain):
+    sp_func, xp_func = binary_function(function_name, xp)
+    resolution: dtype = np.finfo(dtype).resolution
+
+    inp = xp.array(function_domain[0])
+    inp2 = xp.array(function_domain[1])
+    outp = xp.zeros_like(inp)
+
+    reference = xp_func(inp, inp2)
+
+    inp_field = Field.create_from_numpy_array("inp", inp)
+    inp2_field = Field.create_from_numpy_array("inp2", inp)
+    outp_field = inp_field.new_field_with_different_name("outp")
+
+    asms = [
+        Assignment(
+            outp_field.center(), sp_func(inp_field.center(), inp2_field.center())
+        )
+    ]
+    gen_config = replace(gen_config, default_dtype=dtype)
+
+    kernel = create_kernel(asms, gen_config)
+    kfunc = kernel.compile()
+    kfunc(inp=inp, inp2=inp2, outp=outp)
+
+    xp.testing.assert_allclose(outp, reference, rtol=resolution)
+
+
+dtype_and_target_for_integer_funcs = pytest.mark.parametrize(
+    "dtype, target",
+    list(product([np.int32], [t for t in AVAIL_TARGETS if t is not Target.CUDA]))
+    + list(
+        product(
+            [np.int64],
+            [
+                t
+                for t in AVAIL_TARGETS
+                if t not in (Target.X86_SSE, Target.X86_AVX, Target.CUDA)
+            ],
+        )
+    ),
+)
+
+
+@dtype_and_target_for_integer_funcs
+def test_integer_abs(gen_config, xp, dtype):
+    sp_func, xp_func = unary_function("abs", xp)
+
+    smallest = np.iinfo(dtype).min
+    largest = np.iinfo(dtype).max
+
+    inp = xp.array([-1, 0, 1, 3, -5, -312, smallest + 1, largest], dtype=dtype)
+
+    outp = xp.zeros_like(inp)
+    reference = xp_func(inp)
+
+    inp_field = Field.create_from_numpy_array("inp", inp)
+    outp_field = inp_field.new_field_with_different_name("outp")
+
+    asms = [Assignment(outp_field.center(), sp_func(inp_field.center()))]
+    gen_config = replace(gen_config, default_dtype=dtype)
+
+    kernel = create_kernel(asms, gen_config)
+    kfunc = kernel.compile()
+    kfunc(inp=inp, outp=outp)
+
+    xp.testing.assert_array_equal(outp, reference)
+
+
+@pytest.mark.parametrize("function_name", ("min", "max"))
+@dtype_and_target_for_integer_funcs
+def test_integer_binary_functions(gen_config, xp, function_name, dtype):
+    sp_func, xp_func = binary_function(function_name, xp)
+
+    smallest = np.iinfo(dtype).min
+    largest = np.iinfo(dtype).max
+
+    inp1 = xp.array([-1, 0, 1, 3, -5, -312, smallest + 1, largest], dtype=dtype)
+    inp2 = xp.array([3, -5, 1, 12, 1, 11, smallest + 42, largest - 3], dtype=dtype)
+
+    outp = xp.zeros_like(inp1)
+    reference = xp_func(inp1, inp2)
+
+    inp_field = Field.create_from_numpy_array("inp1", inp1)
+    inp2_field = Field.create_from_numpy_array("inp2", inp2)
+    outp_field = inp_field.new_field_with_different_name("outp")
+
+    asms = [
+        Assignment(
+            outp_field.center(), sp_func(inp_field.center(), inp2_field.center())
+        )
+    ]
+    gen_config = replace(gen_config, default_dtype=dtype)
+
+    kernel = create_kernel(asms, gen_config)
+    kfunc = kernel.compile()
+    kfunc(inp1=inp1, inp2=inp2, outp=outp)
+
+    xp.testing.assert_array_equal(outp, reference)
+
+
+@pytest.mark.parametrize("a", [sp.Symbol("a"), fields("a:  float64[2d]").center])
+def test_avoid_pow(a):
+    x = fields("x:  float64[2d]")
+
+    up = Assignment(x.center_vector[0], 2 * a**2 / 3)
+    func = create_kernel(up)
+
+    powers = list(
+        dfs_preorder(
+            func.body, lambda n: isinstance(n, PsCall) and "pow" in n.function.name
+        )
+    )
+    assert not powers
+
+
+@pytest.mark.xfail(reason="fast_div not available yet")
+def test_avoid_pow_fast_div():
+    x = fields("x:  float64[2d]")
+    a = fields("a:  float64[2d]").center
+
+    up = Assignment(x.center_vector[0], fast_division(1, (a**2)))
+    func = create_kernel(up, config=CreateKernelConfig(target=Target.GPU))
+
+    powers = list(
+        dfs_preorder(
+            func.body, lambda n: isinstance(n, PsCall) and "pow" in n.function.name
+        )
+    )
+    assert not powers
+
+
+def test_avoid_pow_move_constants():
+    # At the end of the kernel creation the function move_constants_before_loop will be called
+    # This function additionally contains substitutions for symbols with the same value
+    # Thus it simplifies the equations again
+    x = fields("x:  float64[2d]")
+    a, b, c = sp.symbols("a, b, c")
+
+    up = [
+        Assignment(a, 0.0),
+        Assignment(b, 0.0),
+        Assignment(c, 0.0),
+        Assignment(
+            x.center_vector[0],
+            a**2 / 18 - a * b / 6 - a / 18 + b**2 / 18 + b / 18 - c**2 / 36,
+        ),
+    ]
+    func = create_kernel(up)
+
+    powers = list(
+        dfs_preorder(
+            func.body, lambda n: isinstance(n, PsCall) and "pow" in n.function.name
+        )
+    )
+    assert not powers

tests/test_gpu.py → tests/kernelcreation/test_gpu.py

 import pytest
 
 import numpy as np
-import cupy as cp
 import sympy as sp
 from scipy.ndimage import convolve
 
 from pystencils import Assignment, Field, fields, CreateKernelConfig, create_kernel, Target
-from pystencils.gpu import BlockIndexing
+# from pystencils.gpu import BlockIndexing
 from pystencils.simp import sympy_cse_on_assignment_list
 from pystencils.slicing import add_ghost_layers, make_slice, remove_ghost_layers, normalize_slice
 
 try:
-    import cupy
-    device_numbers = range(cupy.cuda.runtime.getDeviceCount())
+    import cupy as cp
+    device_numbers = range(cp.cuda.runtime.getDeviceCount())
 except ImportError:
     device_numbers = []
 
@@ -74,7 +73,7 @@ def test_multiple_index_dimensions():
     """Sums along the last axis of a numpy array"""
     src_size = (7, 6, 4)
     dst_size = src_size[:2]
-    src_arr = np.asfortranarray(np.random.rand(*src_size))
+    src_arr = np.array(np.random.rand(*src_size))
     dst_arr = np.zeros(dst_size)
 
     src_field = Field.create_from_numpy_array('src', src_arr, index_dimensions=1)
@@ -102,6 +101,7 @@ def test_multiple_index_dimensions():
     np.testing.assert_almost_equal(reference, dst_arr)
 
 
+@pytest.mark.xfail(reason="Line indexing not available yet")
 def test_ghost_layer():
     size = (6, 5)
     src_arr = np.ones(size)
@@ -112,7 +112,7 @@ def test_ghost_layer():
     update_rule = Assignment(dst_field[0, 0], src_field[0, 0])
     ghost_layers = [(1, 2), (2, 1)]
 
-    config = CreateKernelConfig(target=Target.GPU, ghost_layers=ghost_layers, gpu_indexing="line")
+    config = CreateKernelConfig(target=Target.GPU, ghost_layers=ghost_layers, gpu="line")
     ast = create_kernel(sympy_cse_on_assignment_list([update_rule]), config=config)
     kernel = ast.compile()
 
@@ -126,6 +126,7 @@ def test_ghost_layer():
     np.testing.assert_equal(reference, dst_arr)
 
 
+@pytest.mark.xfail(reason="Line indexing not available yet")
 def test_setting_value():
     arr_cpu = np.arange(25, dtype=np.float64).reshape(5, 5)
     arr_gpu = cp.asarray(arr_cpu)
@@ -134,7 +135,7 @@ def test_setting_value():
     f = Field.create_generic("f", 2)
     update_rule = [Assignment(f(0), sp.Symbol("value"))]
 
-    config = CreateKernelConfig(target=Target.GPU, gpu_indexing="line", iteration_slice=iteration_slice)
+    config = CreateKernelConfig(target=Target.GPU, gpu="line", iteration_slice=iteration_slice)
     ast = create_kernel(sympy_cse_on_assignment_list(update_rule), config=config)
     kernel = ast.compile()
 
@@ -156,12 +157,13 @@ def test_periodicity():
     cpu_result = np.copy(arr_cpu)
     periodic_cpu_kernel(cpu_result)
 
-    periodic_gpu_kernel(pdfs=arr_gpu)
+    periodic_gpu_kernel(arr_gpu)
     gpu_result = arr_gpu.get()
     np.testing.assert_equal(cpu_result, gpu_result)
 
 
 @pytest.mark.parametrize("device_number", device_numbers)
+@pytest.mark.xfail(reason="Block indexing specification is not available yet")
 def test_block_indexing(device_number):
     f = fields("f: [3D]")
     s = normalize_slice(make_slice[:, :, :], f.spatial_shape)
@@ -194,6 +196,7 @@ def test_block_indexing(device_number):
 @pytest.mark.parametrize('gpu_indexing', ("block", "line"))
 @pytest.mark.parametrize('layout', ("C", "F"))
 @pytest.mark.parametrize('shape', ((5, 5, 5, 5), (3, 17, 387, 4), (23, 44, 21, 11)))
+@pytest.mark.xfail(reason="4D kernels not available yet")
 def test_four_dimensional_kernel(gpu_indexing, layout, shape):
     n_elements = np.prod(shape)
 
@@ -204,7 +207,7 @@ def test_four_dimensional_kernel(gpu_indexing, layout, shape):
     f = Field.create_from_numpy_array("f", arr_cpu)
     update_rule = [Assignment(f.center, sp.Symbol("value"))]
 
-    config = CreateKernelConfig(target=Target.GPU, gpu_indexing=gpu_indexing, iteration_slice=iteration_slice)
+    config = CreateKernelConfig(target=Target.GPU, gpu=gpu_indexing, iteration_slice=iteration_slice)
     ast = create_kernel(update_rule, config=config)
     kernel = ast.compile()
 

tests/test_half_precision.py → tests/kernelcreation/test_half_precision.py

@@ -29,7 +29,7 @@ def test_half_precison(target):
 
     up = ps.Assignment(f3.center, f1.center + 2.1 * f2.center)
 
-    config = ps.CreateKernelConfig(target=dh.default_target, default_number_float=np.float32)
+    config = ps.CreateKernelConfig(target=dh.default_target, default_dtype=np.float32)
     ast = ps.create_kernel(up, config=config)
 
     kernel = ast.compile()

tests/kernelcreation/test_index_kernels.py

+import numpy as np
+import pytest
+
+from pystencils import Assignment, Field, FieldType, AssignmentCollection, Target
+from pystencils import create_kernel, CreateKernelConfig
+
+
+@pytest.mark.parametrize("target", [Target.CPU, Target.GPU])
+def test_indexed_kernel(target):
+    if target == Target.GPU:
+        cp = pytest.importorskip("cupy")
+        xp = cp
+    else:
+        xp = np
+
+    arr = xp.zeros((3, 4))
+    dtype = np.dtype([('x', int), ('y', int), ('value', arr.dtype)], align=True)
+    
+    cpu_index_arr = np.zeros((3,), dtype=dtype)
+    cpu_index_arr[0] = (0, 2, 3.0)
+    cpu_index_arr[1] = (1, 3, 42.0)
+    cpu_index_arr[2] = (2, 1, 5.0)
+
+    if target == Target.GPU:
+        gpu_index_arr = cp.empty(cpu_index_arr.shape, cpu_index_arr.dtype)
+        gpu_index_arr.set(cpu_index_arr)
+        index_arr = gpu_index_arr
+    else:
+        index_arr = cpu_index_arr
+
+    index_field = Field.create_from_numpy_array('index', index_arr, field_type=FieldType.INDEXED)
+    normal_field = Field.create_from_numpy_array('f', arr)
+    update_rule = AssignmentCollection([
+        Assignment(normal_field[0, 0], index_field('value'))
+    ])
+
+    options = CreateKernelConfig(index_field=index_field, target=target)
+    ast = create_kernel(update_rule, options)
+    kernel = ast.compile()
+
+    kernel(f=arr, index=index_arr)
+
+    if target == Target.GPU:
+        arr = cp.asnumpy(arr)
+
+    for i in range(cpu_index_arr.shape[0]):
+        np.testing.assert_allclose(arr[cpu_index_arr[i]['x'], cpu_index_arr[i]['y']], cpu_index_arr[i]['value'], atol=1e-13)

tests/kernelcreation/test_iteration_slices.py

+import numpy as np
+import sympy as sp
+import pytest
+
+from dataclasses import replace
+
+from pystencils import (
+    DEFAULTS,
+    Assignment,
+    Field,
+    TypedSymbol,
+    create_kernel,
+    make_slice,
+    Target,
+    CreateKernelConfig,
+    DynamicType,
+)
+from pystencils.sympyextensions.integer_functions import int_rem
+from pystencils.simp import sympy_cse_on_assignment_list
+from pystencils.slicing import normalize_slice
+from pystencils.jit.gpu_cupy import CupyKernelWrapper
+
+
+def test_sliced_iteration():
+    size = (4, 4)
+    src_arr = np.ones(size)
+    dst_arr = np.zeros_like(src_arr)
+    src_field = Field.create_from_numpy_array("src", src_arr)
+    dst_field = Field.create_from_numpy_array("dst", dst_arr)
+
+    a, b = sp.symbols("a b")
+    update_rule = Assignment(
+        dst_field[0, 0],
+        (
+            a * src_field[0, 1]
+            + a * src_field[0, -1]
+            + b * src_field[1, 0]
+            + b * src_field[-1, 0]
+        )
+        / 4,
+    )
+
+    x_end = TypedSymbol("x_end", "int")
+    s = make_slice[1:x_end, 1]
+    x_end_value = size[1] - 1
+    kernel = create_kernel(
+        sympy_cse_on_assignment_list([update_rule]), iteration_slice=s
+    ).compile()
+
+    kernel(src=src_arr, dst=dst_arr, a=1.0, b=1.0, x_end=x_end_value)
+
+    expected_result = np.zeros(size)
+    expected_result[1:x_end_value, 1] = 1
+    np.testing.assert_almost_equal(expected_result, dst_arr)
+
+
+@pytest.mark.parametrize(
+    "islice",
+    [
+        make_slice[1:-1, 1:-1],
+        make_slice[3, 2:-2],
+        make_slice[2:-2:2, ::3],
+        make_slice[10:, :-5:2],
+        make_slice[-5:-1, -1],
+        make_slice[-3, -1]
+    ],
+)
+def test_numerical_slices(gen_config: CreateKernelConfig, xp, islice):
+    shape = (16, 16)
+
+    f_arr = xp.zeros(shape)
+    expected = xp.zeros_like(f_arr)
+    expected[islice] = 1.0
+
+    f = Field.create_from_numpy_array("f", f_arr)
+
+    update = Assignment(f.center(), 1)
+    gen_config = replace(gen_config, iteration_slice=islice)
+
+    try:
+        kernel = create_kernel(update, gen_config).compile()
+    except NotImplementedError:
+        if gen_config.get_target().is_vector_cpu():
+            #   TODO Gather/Scatter not implemented yet
+            pytest.xfail("Gather/Scatter not available yet")
+
+    kernel(f=f_arr)
+
+    xp.testing.assert_array_equal(f_arr, expected)
+
+
+def test_symbolic_slice(gen_config: CreateKernelConfig, xp):
+    shape = (16, 16)
+
+    sx, sy, ex, ey = [
+        TypedSymbol(n, DynamicType.INDEX_TYPE) for n in ("sx", "sy", "ex", "ey")
+    ]
+
+    f_arr = xp.zeros(shape)
+
+    f = Field.create_from_numpy_array("f", f_arr)
+
+    update = Assignment(f.center(), 1)
+    islice = make_slice[sy:ey, sx:ex]
+    gen_config = replace(gen_config, iteration_slice=islice)
+
+    print(repr(gen_config))
+
+    kernel = create_kernel(update, gen_config).compile()
+
+    for slic in [make_slice[:, :], make_slice[1:-1, 2:-2], make_slice[8:14, 7:11]]:
+        slic = normalize_slice(slic, shape)
+        expected = xp.zeros_like(f_arr)
+        expected[slic] = 1.0
+
+        f_arr[:] = 0.0
+
+        kernel(
+            f=f_arr,
+            sy=slic[0].start,
+            ey=slic[0].stop,
+            sx=slic[1].start,
+            ex=slic[1].stop,
+        )
+
+        xp.testing.assert_array_equal(f_arr, expected)
+
+
+def test_triangle_pattern(gen_config: CreateKernelConfig, xp):
+    shape = (16, 16)
+
+    f_arr = xp.zeros(shape)
+    f = Field.create_from_numpy_array("f", f_arr)
+
+    expected = xp.zeros_like(f_arr)
+    for r in range(shape[0]):
+        expected[r, r:] = 1.0
+
+    update = Assignment(f.center(), 1)
+    outer_counter = DEFAULTS.spatial_counters[0]
+    islice = make_slice[:, outer_counter:]
+    gen_config = replace(gen_config, iteration_slice=islice)
+
+    if gen_config.target == Target.CUDA:
+        gen_config.gpu.manual_launch_grid = True
+
+    kernel = create_kernel(update, gen_config).compile()
+
+    if isinstance(kernel, CupyKernelWrapper):
+        kernel.block_size = shape + (1,)
+        kernel.num_blocks = (1, 1, 1)
+
+    kernel(f=f_arr)
+
+    xp.testing.assert_array_equal(f_arr, expected)
+
+
+def test_red_black_pattern(gen_config: CreateKernelConfig, xp):
+    shape = (16, 16)
+
+    f_arr = xp.zeros(shape)
+    f = Field.create_from_numpy_array("f", f_arr)
+
+    expected = xp.zeros_like(f_arr)
+    for r in range(shape[0]):
+        start = 0 if r % 2 == 0 else 1
+        expected[r, start::2] = 1.0
+
+    update = Assignment(f.center(), 1)
+    outer_counter = DEFAULTS.spatial_counters[0]
+    start = sp.Piecewise((0, sp.Eq(int_rem(outer_counter, 2), 0)), (1, True))
+    islice = make_slice[:, start::2]
+    gen_config.iteration_slice = islice
+
+    if gen_config.target == Target.CUDA:
+        gen_config.gpu.manual_launch_grid = True
+
+    try:
+        kernel = create_kernel(update, gen_config).compile()
+    except NotImplementedError:
+        if gen_config.get_target().is_vector_cpu():
+            pytest.xfail("Gather/Scatter not implemented yet")
+
+    if isinstance(kernel, CupyKernelWrapper):
+        kernel.block_size = (8, 16, 1)
+        kernel.num_blocks = (1, 1, 1)
+
+    kernel(f=f_arr)
+
+    xp.testing.assert_array_equal(f_arr, expected)

tests/kernelcreation/test_kernel_decorator.py

+import numpy as np
+import sympy as sp
+import pystencils as ps
+
+
+def test_kernel_decorator_config():
+    config = ps.CreateKernelConfig()
+    a, b, c = ps.fields(a=np.ones(100), b=np.ones(100), c=np.ones(100))
+
+    @ps.kernel_config(config)
+    def test():
+        a[0] @= b[0] + c[0]
+
+    ps.create_kernel(**test)
+
+
+def test_kernel_decorator2():
+    h = sp.symbols("h")
+    dtype = "float64"
+
+    src, dst = ps.fields(f"src, src_tmp: {dtype}[3D]")
+
+    @ps.kernel
+    def kernel_func():
+        dst[0, 0, 0] @= (src[1, 0, 0] + src[-1, 0, 0]
+                         + src[0, 1, 0] + src[0, -1, 0]
+                         + src[0, 0, 1] + src[0, 0, -1]) / (6 * h ** 2)
+
+    # assignments = ps.assignment_from_stencil(stencil, src, dst, normalization_factor=2)
+    ast = ps.create_kernel(kernel_func)
+
+    _ = ps.get_code_str(ast)

tests/test_logarithm.py → tests/kernelcreation/test_logarithm.py

@@ -12,7 +12,7 @@ def test_log(dtype):
 
     assignments = ps.AssignmentCollection({x.center(): sp.log(a)})
 
-    ast = ps.create_kernel(assignments)
+    ast = ps.create_kernel(assignments, default_dtype=dtype)
     code = ps.get_code_str(ast)
     kernel = ast.compile()
 

tests/test_size_and_layout_checks.py → tests/kernelcreation/test_size_and_layout_checks.py

@@ -5,10 +5,13 @@ import pystencils
 import sympy as sp
 
 from pystencils import Assignment, Field, create_kernel, fields
+from pystencils.backend.exceptions import KernelConstraintsError
 
 
+@pytest.mark.xfail(reason="Size checks not yet reimplemented in CPU JIT")
 def test_size_check():
     """Kernel with two fixed-sized fields creating with same size but calling with wrong size"""
+    
     src = np.zeros((20, 21, 9))
     dst = np.zeros_like(src)
 
@@ -23,6 +26,8 @@ def test_size_check():
     src = np.zeros(new_shape)
     dst = np.zeros(new_shape)
 
+    assert False  #  TODO: The function call below is not valid and will likely segfault
+
     with pytest.raises(ValueError) as e:
         func(src=src, dst=dst)
     assert 'Wrong shape' in str(e.value)
@@ -38,9 +43,9 @@ def test_fixed_size_mismatch_check():
     update_rule = Assignment(sym_dst(0),
                              sym_src[-1, 1](1) + sym_src[1, -1](2))
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([update_rule])
-    assert 'Differently sized field accesses' in str(e.value)
+    assert 'Fixed-shape fields of different sizes encountered' in str(e.value)
 
 
 def test_fixed_and_variable_field_check():
@@ -53,11 +58,12 @@ def test_fixed_and_variable_field_check():
     update_rule = Assignment(sym_dst(0),
                              sym_src[-1, 1](1) + sym_src[1, -1](2))
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel(update_rule)
-    assert 'Mixing fixed-shaped and variable-shape fields' in str(e.value)
+    assert 'Cannot mix fixed- and variable-shape fields' in str(e.value)
 
 
+@pytest.mark.xfail(reason="Shape checks not yet reimplemented in CPU JIT")
 def test_two_variable_shaped_fields():
     src = np.zeros((20, 21, 9))
     dst = np.zeros((22, 21, 9))
@@ -70,6 +76,8 @@ def test_two_variable_shaped_fields():
     ast = create_kernel([update_rule])
     func = ast.compile()
 
+    assert False  #  TODO: The function call below is not valid
+
     with pytest.raises(TypeError) as e:
         func(src=src, dst=dst)
     assert 'must have same' in str(e.value)
@@ -79,19 +87,19 @@ def test_ssa_checks():
     f, g = fields("f, g : double[2D]")
     a, b, c = sp.symbols("a b c")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(c, f[0, 1]),
                        Assignment(c, f[1, 0]),
                        Assignment(g[0, 0], c)])
     assert 'Assignments not in SSA form' in str(e.value)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(c, a + 3),
                        Assignment(a, 42),
                        Assignment(g[0, 0], c)])
     assert 'Symbol a is written, after it has been read' in str(e.value)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(c, c + 1),
                        Assignment(g[0, 0], c)])
     assert 'Symbol c is written, after it has been read' in str(e.value)
@@ -101,13 +109,13 @@ def test_loop_independence_checks():
     f, g = fields("f, g : double[2D]")
     v = fields("v(2) : double[2D]")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(g[0, 1], f[0, 1]),
                        Assignment(g[0, 0], f[1, 0])])
     assert 'Field g is written at two different locations' in str(e.value)
 
     # This is not allowed - because this is not SSA (it can be overwritten with allow_double_writes)
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(g[0, 2], f[0, 1]),
                        Assignment(g[0, 2], 2 * g[0, 2])])
 
@@ -116,12 +124,12 @@ def test_loop_independence_checks():
                    Assignment(g[0, 2], 2 * g[0, 2])],
                   config=pystencils.CreateKernelConfig(allow_double_writes=True))
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(v[0, 2](1), f[0, 1]),
                        Assignment(v[0, 1](0), 4),
                        Assignment(v[0, 2](1), 2 * v[0, 2](1))])
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(KernelConstraintsError) as e:
         create_kernel([Assignment(g[0, 1], 3),
                        Assignment(f[0, 1], 2 * g[0, 2])])
     assert 'Field g is read at (0, 2) and written at (0, 1)' in str(e.value)

tests/kernelcreation/test_spatial_counters.py

+import pytest
+import numpy as np
+
+from pystencils import (
+    Field,
+    Assignment,
+    create_kernel,
+    CreateKernelConfig,
+    DEFAULTS,
+    FieldType,
+)
+from pystencils.sympyextensions import tcast
+
+
+@pytest.mark.parametrize("index_dtype", ["int16", "int32", "uint32", "int64"])
+def test_spatial_counters_dense(index_dtype):
+    #   Parametrized over index_dtype to make sure the `DynamicType.INDEX` in the
+    #   DEFAULTS works validly
+    x, y, z = DEFAULTS.spatial_counters
+
+    f = Field.create_generic("f", 3, "float64", index_shape=(3,), layout="fzyx")
+
+    asms = [
+        Assignment(f(0), tcast.as_numeric(z)),
+        Assignment(f(1), tcast.as_numeric(y)),
+        Assignment(f(2), tcast.as_numeric(x)),
+    ]
+
+    cfg = CreateKernelConfig(index_dtype=index_dtype)
+    kernel = create_kernel(asms, cfg).compile()
+
+    f_arr = np.zeros((16, 16, 16, 3))
+    kernel(f=f_arr)
+
+    expected = np.mgrid[0:16, 0:16, 0:16].astype(np.float64).transpose()
+
+    np.testing.assert_equal(f_arr, expected)
+
+
+@pytest.mark.parametrize("index_dtype", ["int16", "int32", "uint32", "int64"])
+def test_spatial_counters_sparse(index_dtype):
+    x, y, z = DEFAULTS.spatial_counters
+
+    f = Field.create_generic("f", 3, "float64", index_shape=(3,), layout="fzyx")
+
+    asms = [
+        Assignment(f(0), tcast.as_numeric(x)),
+        Assignment(f(1), tcast.as_numeric(y)),
+        Assignment(f(2), tcast.as_numeric(z)),
+    ]
+
+    idx_struct = DEFAULTS.index_struct(index_dtype, 3)
+    idx_field = Field.create_generic(
+        "index", 1, idx_struct, field_type=FieldType.INDEXED
+    )
+
+    cfg = CreateKernelConfig(index_dtype=index_dtype, index_field=idx_field)
+    kernel = create_kernel(asms, cfg).compile()
+
+    f_arr = np.zeros((16, 16, 16, 3))
+    idx_arr = np.array(
+        [(1, 4, 3), (5, 1, 6), (9, 5, 1), (3, 13, 7)], dtype=idx_struct.numpy_dtype
+    )
+
+    kernel(f=f_arr, index=idx_arr)
+
+    for t in idx_arr:
+        assert f_arr[t[0], t[1], t[2], 0] == t[0].astype(np.float64)
+        assert f_arr[t[0], t[1], t[2], 1] == t[1].astype(np.float64)
+        assert f_arr[t[0], t[1], t[2], 2] == t[2].astype(np.float64)

tests/test_staggered_kernel.py → tests/kernelcreation/test_staggered_kernel.py

@@ -5,11 +5,12 @@ import pytest
 
 import pystencils as ps
 from pystencils import x_staggered_vector, TypedSymbol
-from pystencils.enums import Target
+from pystencils import Target
 
 
 class TestStaggeredDiffusion:
     def _run(self, num_neighbors, target=ps.Target.CPU, openmp=False):
+        pytest.xfail("Staggered kernels not available yet")
         L = (40, 40)
         D = 0.066
         dt = 1
@@ -76,6 +77,7 @@ class TestStaggeredDiffusion:
         self._run(4, openmp=True)
 
 
+@pytest.mark.xfail(reason="Staggered kernels not available yet")
 def test_staggered_subexpressions():
     dh = ps.create_data_handling((10, 10), periodicity=True, default_target=Target.CPU)
     j = dh.add_array('j', values_per_cell=2, field_type=ps.FieldType.STAGGERED)
@@ -85,6 +87,7 @@ def test_staggered_subexpressions():
     ps.create_staggered_kernel(assignments, target=dh.default_target).compile()
 
 
+@pytest.mark.xfail(reason="Staggered kernels not available yet")
 def test_staggered_loop_cutting():
     pytest.importorskip('islpy')
     dh = ps.create_data_handling((4, 4), periodicity=True, default_target=Target.CPU)
@@ -94,6 +97,7 @@ def test_staggered_loop_cutting():
     assert not ast.atoms(ps.astnodes.Conditional)
 
 
+@pytest.mark.xfail(reason="Staggered kernels not available yet")
 def test_staggered_vector():
     dim = 2
     v = x_staggered_vector(dim)

tests/kernelcreation/test_struct_types.py

+import numpy as np
+import pytest
+
+from pystencils import Assignment, Field, create_kernel
+
+
+@pytest.mark.parametrize("order", ['c', 'f'])
+@pytest.mark.parametrize("align", [True, False])
+def test_fixed_sized_field(order, align):
+    if not align:
+        pytest.xfail("Non-Aligned structs not supported")
+    dt = np.dtype([('e1', np.float32), ('e2', np.double), ('e3', np.double)], align=align)
+    arr = np.zeros((3, 2), dtype=dt, order=order)
+
+    f = Field.create_from_numpy_array("f", arr)
+    d = Field.create_from_numpy_array("d", arr)
+    update_rules = [Assignment(d[0, 0]['e2'], f[0, 0]['e3'])]
+    result = arr.copy(order=order)
+    assert result.strides == arr.strides
+    arr['e2'] = 0
+    arr['e3'] = np.random.rand(3, 2)
+
+    kernel = create_kernel(update_rules).compile()
+    kernel(f=arr, d=result)
+    np.testing.assert_almost_equal(result['e2'], arr['e3'])
+    np.testing.assert_equal(arr['e2'], np.zeros((3, 2)))
+
+
+@pytest.mark.parametrize("order", ['c', 'f'])
+@pytest.mark.parametrize("align", [True, False])
+def test_variable_sized_field(order, align):
+    if not align:
+        pytest.xfail("Non-Aligned structs not supported")
+    dt = np.dtype([('e1', np.float32), ('e2', np.double), ('e3', np.double)], align=align)
+
+    f = Field.create_generic("f", 2, dt, layout=order)
+    d = Field.create_generic("d", 2, dt, layout=order)
+    update_rules = [Assignment(d[0, 0]['e2'], f[0, 0]['e3'])]
+
+    arr = np.zeros((3, 2), dtype=dt, order=order)
+    result = arr.copy(order=order)
+
+    arr['e2'] = 0
+    arr['e3'] = np.random.rand(3, 2)
+
+    kernel = create_kernel(update_rules).compile()
+    kernel(f=arr, d=result)
+    np.testing.assert_almost_equal(result['e2'], arr['e3'])
+    np.testing.assert_equal(arr['e2'], np.zeros((3, 2)))

tests/test_sum_prod.py → tests/kernelcreation/test_sum_prod.py

@@ -10,17 +10,16 @@
 import pytest
 import numpy as np
 
-import pystencils.config
 import sympy as sp
 import sympy.abc
 
 import pystencils as ps
-from pystencils.typing import create_type
+from pystencils import create_type
 
 
 @pytest.mark.parametrize('dtype', ["float64", "float32"])
+@pytest.mark.xfail(reason="Sum and product reductions are not available yet")
 def test_sum(dtype):
-
     sum = sp.Sum(sp.abc.k, (sp.abc.k, 1, 100))
     expanded_sum = sum.doit()
 
@@ -46,8 +45,8 @@ def test_sum(dtype):
 
 
 @pytest.mark.parametrize('dtype', ["int32", "int64", "float64", "float32"])
+@pytest.mark.xfail(reason="Sum and product reductions are not available yet")
 def test_product(dtype):
-
     k = ps.TypedSymbol('k', create_type(dtype))
 
     sum = sympy.Product(k, (k, 1, 10))
@@ -60,7 +59,7 @@ def test_product(dtype):
 
     assignments = ps.AssignmentCollection({x.center(): sum})
 
-    config = pystencils.config.CreateKernelConfig()
+    config = ps.CreateKernelConfig()
 
     ast = ps.create_kernel(assignments, config=config)
     code = ps.get_code_str(ast)