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