Esempio n. 1
0
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)
Esempio n. 2
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))
Esempio n. 3
0
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)
Esempio n. 4
0
def test_optional_arg_provide(backend):
    @gtscript.stencil(backend=backend)
    def stencil(
        inp: gtscript.Field[np.float64],
        unused_field: gtscript.Field[np.float64],
        outp: gtscript.Field[np.float64],
        unused_par: float,
    ):
        with computation(PARALLEL), interval(...):
            outp = inp  # noqa F841: local variable 'outp' is assigned to but never used

    frozen_stencil = stencil.freeze(
        domain=(3, 3, 10),
        origin={
            "inp": (2, 2, 0),
            "outp": (2, 2, 0),
            "unused_field": (0, 0, 0)
        },
    )

    inp = gt_storage.from_array(data=7.0,
                                dtype=np.float64,
                                shape=(10, 10, 10),
                                default_origin=(0, 0, 0),
                                backend=backend)
    outp = gt_storage.zeros(dtype=np.float64,
                            shape=(10, 10, 10),
                            default_origin=(0, 0, 0),
                            backend=backend)
    unused_field = gt_storage.zeros(dtype=np.float64,
                                    shape=(10, 10, 10),
                                    default_origin=(0, 0, 0),
                                    backend=backend)
    inp.host_to_device()
    outp.host_to_device()

    @dace.program(
        device=dace.DeviceType.GPU if "gpu" in backend else dace.DeviceType.CPU
    )
    def call_frozen_stencil():
        frozen_stencil(inp=inp,
                       unused_field=unused_field,
                       outp=outp,
                       unused_par=7.0)

    call_frozen_stencil()

    outp.device_to_host(force=True)

    assert np.allclose(inp, 7.0)
    assert np.allclose(np.asarray(outp)[2:5, 2:5, :], 7.0)
    assert np.sum(np.asarray(outp), axis=(0, 1, 2)) == 90 * 7.0
Esempio n. 5
0
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
Esempio n. 6
0
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 run_rank_adder_test(backend: str, rebuild: bool):
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    set_backend(backend)
    StencilTable.clear()

    origin = (0, 0, 0)
    domain = (1, 1, size)

    out_field = gt_storage.zeros(shape=domain,
                                 default_origin=origin,
                                 dtype=np.int64,
                                 backend=backend)
    ref_field = np.arange(1, size + 1, dtype=np.int64)

    for rank in range(0, size):
        stencil_object = future_stencil(
            definition=add_rank,
            backend=backend,
            rebuild=rebuild,
            externals={"rank": rank},
        )
        stencil_object(out_field, domain=domain, origin=origin)

    assert np.array_equal(out_field[0, 0, :], ref_field)
Esempio n. 8
0
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"]
Esempio n. 9
0
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
Esempio n. 10
0
def test_origin_k_fields(backend):
    @gtscript.stencil(backend=backend, rebuild=True)
    def k_to_ijk(outp: Field[np.float64], inp: Field[gtscript.K, np.float64]):
        with computation(PARALLEL), interval(...):
            outp = inp

    origin = {"outp": (0, 0, 1), "inp": (2,)}
    domain = (2, 2, 8)

    data = np.arange(10, dtype=np.float64)
    inp = gt_storage.from_array(
        data=data,
        shape=(10,),
        default_origin=(0,),
        dtype=np.float64,
        mask=[False, False, True],
        backend=backend,
    )
    outp = gt_storage.zeros(
        shape=(2, 2, 10), default_origin=(0, 0, 0), dtype=np.float64, backend=backend
    )

    k_to_ijk(outp, inp, origin=origin, domain=domain)

    inp.device_to_host()
    outp.device_to_host()
    np.testing.assert_allclose(data, np.asarray(inp))
    np.testing.assert_allclose(
        np.broadcast_to(data[2:], shape=(2, 2, 8)), np.asarray(outp)[:, :, 1:-1]
    )
    np.testing.assert_allclose(0.0, np.asarray(outp)[:, :, 0])
    np.testing.assert_allclose(0.0, np.asarray(outp)[:, :, -1])
Esempio n. 11
0
def test_origin_offsetting_nofrozen(dace_stencil, domain, outp_origin):
    backend = dace_stencil.backend
    inp = gt_storage.from_array(data=7.0,
                                dtype=np.float64,
                                shape=(10, 10, 10),
                                default_origin=(0, 0, 0),
                                backend=backend)
    outp = gt_storage.zeros(dtype=np.float64,
                            shape=(10, 10, 10),
                            default_origin=(0, 0, 0),
                            backend=backend)

    origin = {"inp": (0, 0, 0), "outp": outp_origin}

    inp.host_to_device()
    outp.host_to_device()

    @dace.program(
        device=dace.DeviceType.GPU if "gpu" in backend else dace.DeviceType.CPU
    )
    def call_stencil_object():
        dace_stencil(inp=inp, outp=outp, domain=domain, origin=origin)

    call_stencil_object()

    outp.device_to_host(force=True)

    assert np.allclose(inp, 7.0)
    assert np.allclose(
        np.asarray(outp)[outp_origin[0]:outp_origin[0] + domain[0],
                         outp_origin[1]:outp_origin[1] + domain[1],
                         outp_origin[2]:outp_origin[2] + domain[2], ],
        7.0,
    )
    assert np.sum(np.asarray(outp), axis=(0, 1, 2)) == np.prod(domain) * 7.0
Esempio n. 12
0
    def setup_and_run(backend, **kwargs):
        field_a = gt_storage.zeros(
            dtype=np.float_, backend=backend, shape=(3, 3, 1), default_origin=(0, 0, 0)
        )

        @gtscript.stencil(backend=backend, **kwargs)
        def divide_by_zero(field_a: gtscript.Field[np.float_]):
            with computation(PARALLEL), interval(...):
                field_a = 1.0 / field_a

        divide_by_zero(field_a)
Esempio n. 13
0
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
Esempio n. 14
0
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)
Esempio n. 15
0
def test_mask_with_offset_written_in_conditional(backend):
    @gtscript.stencil(backend, externals={"mord": 5})
    def stencil(
        outp: gtscript.Field[np.float_],
    ):

        with computation(PARALLEL), interval(...):
            cond = True
            if cond[0, -1, 0] or cond[0, 0, 0]:
                outp = 1.0
            else:
                outp = 0.0

    outp = gt_storage.zeros(
        shape=(10, 10, 10), backend=backend, default_origin=(0, 0, 0), dtype=float
    )

    stencil(outp)

    outp.device_to_host()
    assert np.allclose(1.0, np.asarray(outp))
Esempio n. 16
0
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
Esempio n. 17
0
def test_nondace_raises():
    @gtscript.stencil(backend="numpy")
    def numpy_stencil(inp: gtscript.Field[np.float64],
                      outp: gtscript.Field[np.float64]):
        with computation(PARALLEL), interval(...):
            outp = inp  # noqa F841: local variable 'outp' is assigned to but never used

    frozen_stencil = numpy_stencil.freeze(domain=(3, 3, 3),
                                          origin={
                                              "inp": (0, 0, 0),
                                              "outp": (0, 0, 0)
                                          })

    inp = gt_storage.from_array(
        data=7.0,
        dtype=np.float64,
        shape=(10, 10, 10),
        default_origin=(0, 0, 0),
        backend="numpy",
    )
    outp = gt_storage.zeros(
        dtype=np.float64,
        shape=(10, 10, 10),
        default_origin=(0, 0, 0),
        backend="numpy",
    )

    @dace.program
    def call_frozen_stencil():
        frozen_stencil(inp=inp, outp=outp)

    with pytest.raises(
            TypeError,
            match=re.escape(
                "Only dace backends are supported in DaCe-orchestrated programs."
                ' (found "numpy")'),
    ):
        call_frozen_stencil()
Esempio n. 18
0
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),
    )
Esempio n. 19
0
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),
    )
Esempio n. 20
0
def test_basic(backend):
    @gtscript.stencil(backend=backend)
    def defn(outp: gtscript.Field[np.float64], par: np.float64):
        with computation(PARALLEL), interval(...):
            outp = par  # noqa F841: local variable 'outp' is assigned to but never used

    outp = gt_storage.zeros(dtype=np.float64,
                            shape=(10, 10, 10),
                            default_origin=(0, 0, 0),
                            backend=backend)

    inp = 7.0

    outp.host_to_device()

    @dace.program(
        device=dace.DeviceType.GPU if "gpu" in backend else dace.DeviceType.CPU
    )
    def call_stencil_object(locoutp, locinp):
        defn(locoutp, par=locinp)

    call_stencil_object(locoutp=outp, locinp=inp)
    outp.device_to_host(force=True)
    assert np.allclose(outp, 7.0)
Esempio n. 21
0
 def zeros(self, *args, **kwargs):
     """Create a storage with all fields initialized to 0."""
     keywords = self._kwargs.copy()
     keywords.update(kwargs)
     return storage.zeros(*args, **keywords)
Esempio n. 22
0
def test_optional_arg_provide_aot(backend):
    @gtscript.stencil(backend=backend)
    def stencil(
        inp: gtscript.Field[np.float64],
        unused_field: gtscript.Field[np.float64],
        outp: gtscript.Field[np.float64],
        unused_par: float,
    ):
        with computation(PARALLEL), interval(...):
            outp = inp  # noqa F841: local variable 'outp' is assigned to but never used

    frozen_stencil = stencil.freeze(
        domain=(3, 3, 10),
        origin={
            "inp": (2, 2, 0),
            "outp": (2, 2, 0),
            "unused_field": (0, 0, 0)
        },
    )

    inp = gt_storage.from_array(data=7.0,
                                dtype=np.float64,
                                shape=(10, 10, 10),
                                default_origin=(0, 0, 0),
                                backend=backend)
    outp = gt_storage.zeros(dtype=np.float64,
                            shape=(10, 10, 10),
                            default_origin=(0, 0, 0),
                            backend=backend)
    unused_field = gt_storage.zeros(dtype=np.float64,
                                    shape=(10, 10, 10),
                                    default_origin=(0, 0, 0),
                                    backend=backend)
    inp.host_to_device()
    outp.host_to_device()

    storage = dace.StorageType.GPU_Global if "gpu" in backend else dace.StorageType.CPU_Heap

    @dace.program(
        device=dace.DeviceType.GPU if "gpu" in backend else dace.DeviceType.CPU
    )
    def call_frozen_stencil(
            inp: dace.data.Array(
                shape=inp.shape,
                strides=tuple(s // inp.itemsize for s in inp.strides),
                dtype=dace.float64,
                storage=storage,
            ),
            outp: dace.data.Array(
                shape=outp.shape,
                strides=tuple(s // outp.itemsize for s in outp.strides),
                dtype=dace.float64,
                storage=storage,
            ),
            unused_field: dace.data.Array(
                shape=unused_field.shape,
                strides=tuple(s // unused_field.itemsize
                              for s in unused_field.strides),
                dtype=dace.float64,
                storage=storage,
            ),  # type: ignore
            unused_par: dace.float64,  # type: ignore
    ):
        frozen_stencil(inp=inp,
                       unused_field=unused_field,
                       outp=outp,
                       unused_par=unused_par)

    csdfg = call_frozen_stencil.compile()
    csdfg(inp=inp, outp=outp, unused_field=unused_field, unused_par=7.0)

    outp.device_to_host(force=True)

    assert np.allclose(inp, 7.0)
    assert np.allclose(np.asarray(outp)[2:5, 2:5, :], 7.0)
    assert np.sum(np.asarray(outp), axis=(0, 1, 2)) == 90 * 7.0