def test_origin_unchanged(backend):
@gtscript.stencil(backend=backend)
def calc_damp(outp: Field[float], inp: Field[K, float]):
with computation(FORWARD), interval(...):
outp = inp
outp = gt_storage.ones(
backend=backend,
default_origin=(1, 1, 1),
shape=(4, 4, 4),
dtype=float,
mask=[True, True, True],
)
inp = gt_storage.ones(
backend=backend,
default_origin=(1,),
shape=(4, 4, 4),
dtype=float,
mask=[False, False, True],
)
origin = {"_all_": (1, 1, 1), "inp": (1,)}
origin_ref = dict(origin)
calc_damp(outp, inp, origin=origin, domain=(3, 3, 3))
assert all(k in origin_ref for k in origin.keys())
assert all(k in origin for k in origin_ref.keys())
assert all(v is origin_ref[k] for k, v in origin.items())
def test_default_arguments(backend):
branch_true = gtscript.stencil(backend=backend,
definition=a_stencil,
externals={"BRANCH": True},
rebuild=True)
branch_false = gtscript.stencil(backend=backend,
definition=a_stencil,
externals={"BRANCH": False},
rebuild=True)
arg1 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg2 = gt_storage.zeros(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg3 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
tmp = np.asarray(arg3)
tmp *= 2
branch_true(arg1, None, arg3, par1=2.0)
np.testing.assert_equal(arg1, 14 * np.ones((3, 3, 3)))
branch_true(arg1, None, par1=2.0)
np.testing.assert_equal(arg1, 196 * np.ones((3, 3, 3)))
branch_false(arg1, arg2, arg3, par1=2.0, par3=2.0)
np.testing.assert_equal(arg1, 56 * np.ones((3, 3, 3)))
with pytest.raises((ValueError, AssertionError)):
branch_false(arg1, arg2, par1=2.0, par3=2.0)
arg1 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg2 = gt_storage.zeros(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
arg3 = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
tmp = np.asarray(arg3)
tmp *= 2
branch_true(arg1, arg2=None, par1=2.0, par2=5.0, par3=3.0)
np.testing.assert_equal(arg1, 10 * np.ones((3, 3, 3)))
branch_true(arg1, arg2=None, par1=2.0, par2=5.0)
np.testing.assert_equal(arg1, 100 * np.ones((3, 3, 3)))
branch_false(arg1, arg2, arg3, par1=2.0, par2=5.0, par3=3.0)
np.testing.assert_equal(arg1, 60 * np.ones((3, 3, 3)))
with pytest.raises((TypeError, AssertionError)):
branch_false(arg1, arg2, arg3, par1=2.0, par2=5.0)
def test_halo_checks(backend):
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test default works
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field, out_field=out_field)
assert (out_field[1:-1, 1:-1, :] == 1).all()
# test setting arbitrary, small domain works
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(10, 10, 10))
assert (out_field[2:12, 2:12, :] == 1).all()
# test setting domain+origin too large raises
in_field = gt_storage.ones(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(22, 22, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
with pytest.raises(ValueError):
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10))
# test 2*origin+domain does not raise if still fits (c.f. previous bug in c++ check.)
in_field = gt_storage.ones(backend=backend,
shape=(23, 23, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
out_field = gt_storage.zeros(backend=backend,
shape=(23, 23, 10),
default_origin=(1, 1, 0),
dtype=np.float64)
stencil(in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10))
def test_assert_same_shape(backend):
stencil_call = gtscript.stencil(definition=stencil, backend=backend)
A = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(0, 0, 0))
B = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(3, 3, 3),
default_origin=(2, 2, 2))
C = gt_storage.ones(backend=backend,
dtype=np.float32,
shape=(3, 3, 3),
default_origin=(0, 1, 0))
stencil_call(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
A = gt_storage.ones(backend=backend,
dtype=np.float64,
shape=(5, 5, 5),
default_origin=(0, 0, 0))
A = A[1:-1, 1:-1, 1:-1]
A.is_stencil_view = True
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
C = gt_storage.ones(backend=backend,
dtype=np.float32,
shape=(5, 5, 5),
default_origin=(0, 1, 0))
with pytest.raises(ValueError):
stencil_call(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
with pytest.raises(ValueError):
stencil_call(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
def test_higher_dimensional_fields(backend):
FLOAT64_VEC2 = (np.float64, (2, ))
FLOAT64_MAT22 = (np.float64, (2, 2))
@gtscript.stencil(backend=backend)
def stencil(
field: gtscript.Field[np.float64],
vec_field: gtscript.Field[FLOAT64_VEC2],
mat_field: gtscript.Field[FLOAT64_MAT22],
):
with computation(PARALLEL), interval(...):
tmp = vec_field[0, 0, 0][0] + vec_field[0, 0, 0][1]
with computation(FORWARD):
with interval(0, 1):
vec_field[0, 0, 0][0] = field[1, 0, 0]
vec_field[0, 0, 0][1] = field[0, 1, 0]
with interval(1, -1):
vec_field[0, 0, 0][0] = 2 * field[1, 0, -1]
vec_field[0, 0, 0][1] = 2 * field[0, 1, -1]
with interval(-1, None):
vec_field[0, 0, 0][0] = field[1, 0, 0]
vec_field[0, 0, 0][1] = field[0, 1, 0]
with computation(PARALLEL), interval(...):
mat_field[0, 0, 0][0, 0] = vec_field[0, 0, 0][0] + tmp[0, 0, 0]
mat_field[0, 0, 0][1, 1] = vec_field[0, 0, 0][1] + tmp[1, 1, 0]
full_shape = (6, 6, 6)
default_origin = (1, 1, 0)
field = gt_storage.ones(backend,
default_origin,
full_shape,
dtype=np.float64)
assert field.shape == full_shape[:]
vec_field = 2.0 * gt_storage.ones(
backend, default_origin, full_shape, dtype=FLOAT64_VEC2)
assert vec_field.shape[:-1] == full_shape
mat_field = gt_storage.ones(backend,
default_origin,
full_shape,
dtype=FLOAT64_MAT22)
assert mat_field.shape[:-2] == full_shape
stencil(field, vec_field, mat_field, origin=(1, 1, 0), domain=(4, 4, 6))
mat_field.device_to_host()
np.testing.assert_allclose(
mat_field.view(np.ndarray)[1:-1, 1:-1, 1:1], 2.0 + 5.0)
stencil(field, vec_field, mat_field)
def test_origin_selection():
stencil = gtscript.stencil(definition=base_stencil, backend="numpy")
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
assert A[1, 1, 1] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin={"_all_": (1, 1, 1), "field1": (2, 2, 2)}, domain=(1, 1, 1))
assert A[2, 2, 2] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(
backend="gt:cpu_ifirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gt:cpu_kfirst", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin={"field1": (2, 2, 2)}, domain=(1, 1, 1))
assert A[2, 2, 2] == 4
assert B[2, 2, 2] == 7
assert C[0, 1, 0] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
def test_exec_info(backend):
"""test that proper warnings are raised depending on field type."""
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
exec_info = {}
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(1, 1, 0),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend, shape=(23, 23, 10), default_origin=(1, 1, 0), dtype=np.float64
)
stencil(
in_field=in_field,
out_field=out_field,
origin=(2, 2, 0),
domain=(20, 20, 10),
exec_info=exec_info,
)
timings = ["call", "call_run", "run"]
assert all([k + "_start_time" in exec_info for k in timings])
assert all([k + "_end_time" in exec_info for k in timings])
assert all([exec_info[k + "_end_time"] > exec_info[k + "_start_time"] for k in timings])
if backend.startswith("gt:"):
assert "run_cpp_start_time" in exec_info
assert "run_cpp_end_time" in exec_info
assert exec_info["run_cpp_end_time"] > exec_info["run_cpp_start_time"]
def test_ndarray_warning():
"""test that proper warnings are raised depending on field type."""
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
with pytest.warns(RuntimeWarning):
stencil(
in_field=in_field.view(np.ndarray),
out_field=out_field.view(np.ndarray),
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
with pytest.warns(None) as record:
stencil(
in_field=in_field,
out_field=out_field,
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
assert len(record) == 0
def test_start_offset(self, backend):
backend_cls = gt_backend.from_name(backend)
default_origin = (1, 2, 3)
stor = gt_store.ones(
backend=backend,
managed_memory=False,
shape=(3, 7, 13),
default_origin=default_origin,
dtype=np.float64,
)
descriptor: dace.data.Array = stor.__descriptor__()
raveled = TestDescriptor.ravel_with_padding(stor)[descriptor.start_offset :]
if backend_cls.storage_info["device"] == "gpu":
assert raveled.data.ptr % (backend_cls.storage_info["alignment"] * stor.itemsize) == 0
assert (
backend_cls.storage_info["alignment"] == 1
or cp.asarray(stor).data.ptr
% (backend_cls.storage_info["alignment"] * stor.itemsize)
!= 0
)
else:
assert (
raveled.ctypes.data % (backend_cls.storage_info["alignment"] * stor.itemsize) == 0
)
assert (
backend_cls.storage_info["alignment"] == 1
or stor.ctypes.data % (backend_cls.storage_info["alignment"] * stor.itemsize) != 0
)
def test_stencil_without_effect(backend):
def definition1(field_in: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
tmp = 0.0
def definition2(f_in: gtscript.Field[np.float_]):
from __externals__ import flag
with computation(PARALLEL), interval(...):
if __INLINED(flag):
B = f_in
stencil1 = gtscript.stencil(backend, definition1)
stencil2 = gtscript.stencil(backend, definition2, externals={"flag": False})
field_in = gt_storage.ones(
dtype=np.float_, backend=backend, shape=(23, 23, 23), default_origin=(0, 0, 0)
)
# test with explicit domain specified
stencil1(field_in, domain=(3, 3, 3))
stencil2(field_in, domain=(3, 3, 3))
# test without domain specified
stencil1(field_in)
def test_stage_merger_induced_interval_block_reordering(backend):
field_in = gt_storage.ones(
dtype=np.float_, backend=backend, shape=(23, 23, 23), default_origin=(0, 0, 0)
)
field_out = gt_storage.zeros(
dtype=np.float_, backend=backend, shape=(23, 23, 23), default_origin=(0, 0, 0)
)
@gtscript.stencil(backend=backend)
def stencil(field_in: gtscript.Field[np.float_], field_out: gtscript.Field[np.float_]):
with computation(BACKWARD):
with interval(-2, -1): # block 1
field_out = field_in
with interval(0, -2): # block 2
field_out = field_in
with computation(BACKWARD):
with interval(-1, None): # block 3
field_out = 2 * field_in
with interval(0, -1): # block 4
field_out = 3 * field_in
stencil(field_in, field_out)
np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 0:-1], 3)
np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, -1], 2)
def test_negative_origin(backend):
def stencil_i(
input_field: gtscript.Field[gtscript.IJK, np.int32],
output_field: gtscript.Field[gtscript.IJK, np.int32],
):
with computation(PARALLEL), interval(...):
output_field = input_field[1, 0, 0]
def stencil_k(
input_field: gtscript.Field[gtscript.IJK, np.int32],
output_field: gtscript.Field[gtscript.IJK, np.int32],
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 1]
input_field = gt_storage.ones(
backend, default_origin=(0, 0, 0), shape=(1, 1, 1), dtype=np.int32
)
output_field = gt_storage.zeros(
backend, default_origin=(0, 0, 0), shape=(1, 1, 1), dtype=np.int32
)
for origin, stencil in (((-1, 0, 0), stencil_i), ((0, 0, -1), stencil_k)):
gtscript.stencil(definition=stencil, backend=backend)(
input_field, output_field, origin={"input_field": origin}
)
assert output_field[0, 0, 0] == 1
def test_stencil_object_cache(backend: str):
@gtscript.stencil(backend=backend)
def stencil(in_field: Field[float], out_field: Field[float], *,
offset: float):
with computation(PARALLEL), interval(...):
out_field = ( # noqa: F841 # local variable 'out_field' is assigned to but never used
in_field + offset)
shape = (4, 4, 4)
in_storage = gt_storage.ones(backend=backend,
default_origin=(0, 0, 0),
shape=shape,
dtype=float)
out_storage = gt_storage.ones(backend=backend,
default_origin=(0, 0, 0),
shape=shape,
dtype=float)
def runit(*args, **kwargs) -> float:
exec_info: Dict[str, Any] = {}
stencil(*args, **kwargs, exec_info=exec_info)
run_time: float = exec_info["run_end_time"] - exec_info[
"run_start_time"]
call_time: float = exec_info["call_run_end_time"] - exec_info[
"call_run_start_time"]
return call_time - run_time
base_time = runit(in_storage, out_storage, offset=1.0)
fast_time = runit(in_storage, out_storage, offset=1.0)
assert fast_time < base_time
# When an origin changes, it needs to recompute more, so the time should increase
other_out_storage = gt_storage.ones(backend=backend,
default_origin=(1, 0, 0),
shape=shape,
dtype=float)
other_origin_time = runit(in_storage, other_out_storage, offset=1.0)
assert other_origin_time > fast_time
# When the cache is cleared, everything is recomputed and the time will increase
assert len(stencil._domain_origin_cache) > 0
stencil.clean_call_args_cache()
assert len(stencil._domain_origin_cache) == 0
cleaned_cache_time = runit(in_storage, out_storage, offset=1.0)
assert cleaned_cache_time > fast_time
def test_lower_dimensional_inputs(backend):
@gtscript.stencil(backend=backend)
def stencil(
field_3d: gtscript.Field[np.float_, gtscript.IJK],
field_2d: gtscript.Field[np.float_, gtscript.IJ],
field_1d: gtscript.Field[np.float_, gtscript.K],
):
with computation(FORWARD):
with interval(0, 1):
field_d = field_1d[1] + field_3d[0, 1, 0]
with computation(PARALLEL):
with interval(0, 1):
tmp = field_2d[0, 1] + field_1d[1]
field_3d = tmp[1, 0, 0] + field_1d[1]
with interval(1, None):
field_3d = tmp[-1, 0, 0]
full_shape = (6, 6, 6)
default_origin = (1, 1, 0)
dtype = float
field_3d = gt_storage.ones(backend,
default_origin,
full_shape,
dtype,
mask=None)
assert field_3d.shape == full_shape[:]
field_2d = gt_storage.ones(backend,
default_origin[:-1],
full_shape[:-1],
dtype,
mask=[True, True, False])
assert field_2d.shape == full_shape[:-1]
field_1d = gt_storage.zeros(backend, (default_origin[-1], ),
(full_shape[-1], ),
dtype,
mask=[False, False, True])
assert field_1d.shape == (full_shape[-1], )
stencil(field_3d, field_2d, field_1d, origin=(1, 1, 0))
stencil(field_3d, field_2d, field_1d)
def test_lower_dimensional_inputs(backend):
@gtscript.stencil(backend=backend)
def stencil(
field_3d: gtscript.Field[gtscript.IJK, np.float_],
field_2d: gtscript.Field[gtscript.IJ, np.float_],
field_1d: gtscript.Field[gtscript.K, np.float_],
):
with computation(FORWARD):
with interval(0, 1):
field_2d = field_1d[1]
with computation(PARALLEL):
with interval(0, -1):
tmp = field_2d[0, 1] + field_1d[1]
with interval(-1, None):
tmp = field_2d[0, 1] + field_1d[0]
with computation(PARALLEL):
with interval(0, 1):
field_3d = tmp[1, 0, 0] + field_1d[1]
with interval(1, None):
field_3d = tmp[-1, 0, 0]
full_shape = (6, 6, 6)
default_origin = (1, 1, 0)
dtype = float
field_3d = gt_storage.zeros(backend,
default_origin,
full_shape,
dtype,
mask=None)
assert field_3d.shape == full_shape[:]
field_2d = gt_storage.zeros(backend,
default_origin[:-1],
full_shape[:-1],
dtype,
mask=[True, True, False])
assert field_2d.shape == full_shape[:-1]
field_1d = gt_storage.ones(backend, (default_origin[-1], ),
(full_shape[-1], ),
dtype,
mask=[False, False, True])
assert field_1d.shape == (full_shape[-1], )
stencil(field_3d, field_2d, field_1d, origin=(1, 1, 0), domain=(4, 3, 6))
field_3d.device_to_host()
np.testing.assert_allclose(field_3d.view(np.ndarray)[1:-1, 1:-2, :1], 3)
np.testing.assert_allclose(field_3d.view(np.ndarray)[1:-1, 1:-2, 1:], 2)
stencil(field_3d, field_2d, field_1d)
def test_domain_selection():
A = gt_storage.ones(
backend="gtmc", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gtx86", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
assert A[1, 1, 1] == 4
assert B[1, 1, 1] == 7
assert C[1, 1, 1] == 21
assert np.sum(np.asarray(A)) == 30
assert np.sum(np.asarray(B)) == 33
assert np.sum(np.asarray(C)) == 47
A = gt_storage.ones(
backend="gtmc", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gtx86", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(0, 0, 0))
assert np.all(A == 4)
assert np.all(B == 7)
assert np.all(C == 21)
def test_device(self, backend):
backend_cls = gt_backend.from_name(backend)
stor = gt_store.ones(
backend=backend,
managed_memory=False,
shape=(3, 7, 13),
default_origin=(1, 2, 3),
dtype=np.float64,
)
descriptor: dace.data.Array = stor.__descriptor__()
if backend_cls.storage_info["device"] == "gpu":
assert descriptor.storage == dace.StorageType.GPU_Global
else:
assert descriptor.storage == dace.StorageType.CPU_Heap
def test_read_data_dim_indirect_addressing(backend):
INT32_VEC2 = (np.int32, (2,))
def stencil(
input_field: gtscript.Field[gtscript.IJK, INT32_VEC2],
output_field: gtscript.Field[gtscript.IJK, np.int32],
index: int,
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 0][index]
default_origin = (0, 0, 0)
full_shape = (1, 1, 2)
input_field = gt_storage.ones(backend, default_origin, full_shape, dtype=INT32_VEC2)
output_field = gt_storage.zeros(backend, default_origin, full_shape, dtype=np.int32)
gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, 1)
assert output_field[0, 0, 0] == 1
def test_generation(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend, stencil_definition, externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=(v.dtype, v.data_dims) if v.data_dims else v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
# vertical domain size >= 16 required for test_large_k_interval
stencil(**args, origin=(10, 10, 5), domain=(3, 3, 16))
def test_generation_gpu(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend,
stencil_definition,
externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
stencil(**args, origin=(10, 10, 10), domain=(3, 3, 3))
def test_slice_gpu():
stor = gt_store.ones(
backend="gtcuda",
managed_memory=False,
shape=(10, 10, 10),
default_origin=(0, 0, 0),
dtype=np.float64,
)
stor.synchronize()
view = stor[1:-1, 1:-1, 1:-1]
gpu_stor = stor._device_field
gpu_view = view._device_field
view_start = gpu_view.data.ptr
storage_start = gpu_stor.data.ptr
view_end = gpu_view[-1:, -1:, -1:].data.ptr
storage_end = gpu_stor[-1:, -1:, -1:].data.ptr
assert view_start > storage_start
assert view_end < storage_end
def test_read_data_dim_indirect_addressing(backend):
INT32_VEC2 = (np.int32, (2,))
def stencil(
input_field: gtscript.Field[gtscript.IJK, INT32_VEC2],
output_field: gtscript.Field[gtscript.IJK, np.int32],
index: int,
):
with computation(PARALLEL), interval(...):
output_field = input_field[0, 0, 0][index]
default_origin = (0, 0, 0)
full_shape = (1, 1, 2)
input_field = gt_storage.ones(backend, default_origin, full_shape, dtype=INT32_VEC2)
output_field = gt_storage.zeros(backend, default_origin, full_shape, dtype=np.int32)
if backend in (backend.values[0] for backend in LEGACY_GRIDTOOLS_BACKENDS):
with pytest.raises(ValueError):
gtscript.stencil(definition=stencil, backend=backend)
else:
gtscript.stencil(definition=stencil, backend=backend)(input_field, output_field, 1)
assert output_field[0, 0, 0] == 1
def test_np_array_int_types():
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=np.asarray((23, 23, 10), dtype=np.int64),
default_origin=np.asarray((1, 1, 0), dtype=np.int64),
dtype=np.float64,
)
stencil(
in_field=in_field,
out_field=out_field,
origin=np.asarray((2, 2, 0), dtype=np.int64),
domain=np.asarray((20, 20, 10), dtype=np.int64),
)
def test_np_int_types():
backend = "numpy"
stencil = gtscript.stencil(definition=avg_stencil, backend=backend)
# test numpy int types are accepted
in_field = gt_storage.ones(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(np.int64(1), int(1), 0),
dtype=np.float64,
)
out_field = gt_storage.zeros(
backend=backend,
shape=(np.int8(23), np.int16(23), np.int32(10)),
default_origin=(np.int64(1), int(1), 0),
dtype=np.float64,
)
stencil(
in_field=in_field,
out_field=out_field,
origin=(np.int8(2), np.int16(2), np.int32(0)),
domain=(np.int64(20), int(20), 10),
)
def test_masked_vector_assignment():
from gt4py.gtscript import FORWARD, IJ, Field, computation, interval, stencil
from gt4py.storage import ones
BACKEND = "gtc:numpy"
dtype = np.float64
@stencil(BACKEND)
def masked_vector_assignment(fld2D: Field[IJ, dtype]):
with computation(FORWARD), interval(0, None):
fld2D += fld2D
if fld2D >= 1.0:
fld2D = 0.0
origin = (0, 0, 0)
fld2D = ones(shape=(2, 3),
dtype=dtype,
backend=BACKEND,
default_origin=origin)
masked_vector_assignment(fld2D)
assert np.allclose(fld2D, np.zeros((2, 3)))
def ones(self, *args, **kwargs):
"""Create a storage with all fields initialized to 1."""
keywords = self._kwargs.copy()
keywords.update(kwargs)
return storage.ones(*args, **keywords)
def test_assert_same_shape():
A = gt_storage.ones(
backend="gtmc", dtype=np.float64, shape=(3, 3, 3), default_origin=(0, 0, 0)
)
B = gt_storage.ones(
backend="gtx86", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
B = gt_storage.ones(
backend="gtx86", dtype=np.float64, shape=(3, 3, 3), default_origin=(2, 2, 2)
)
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(3, 3, 3), default_origin=(0, 1, 0)
)
stencil(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
A = gt_storage.ones(
backend="gtmc", dtype=np.float64, shape=(5, 5, 5), default_origin=(0, 0, 0)
)
A = A[1:-1, 1:-1, 1:-1]
A.is_stencil_view = True
stencil(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
try:
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(5, 5, 5), default_origin=(0, 1, 0)
)
stencil(A, B, C, param=3.0, origin=(1, 1, 1), domain=(1, 1, 1))
except ValueError:
pass
else:
assert False
try:
C = gt_storage.ones(
backend="numpy", dtype=np.float32, shape=(5, 5, 5), default_origin=(0, 1, 0)
)
stencil(
A,
B,
C,
param=3.0,
origin=dict(field1=(2, 2, 2), field2=(0, 0, 0), field3=(1, 1, 1)),
domain=(1, 1, 1),
)
except ValueError:
pass
else:
assert False
