1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
/*
* GridTools
*
* Copyright (c) 2014-2023, ETH Zurich
* All rights reserved.
*
* Please, refer to the LICENSE file in the root directory.
* SPDX-License-Identifier: BSD-3-Clause
*/
/**
@file This file shows an implementation of the "horizontal diffusion" stencil, similar to the one used in COSMO since
it implements flux-limiting
*/
#include <cstdlib>
#include <iostream>
#include <gridtools/common/defs.hpp>
#include <gridtools/stencil/cartesian.hpp>
#include <gridtools/storage/builder.hpp>
#include <gridtools/storage/sid.hpp>
#ifdef GT_CUDACC
#include <gridtools/stencil/gpu.hpp>
#include <gridtools/storage/gpu.hpp>
using stencil_backend_t = gridtools::stencil::gpu<>;
using storage_traits_t = gridtools::storage::gpu;
#else
#include <gridtools/stencil/cpu_ifirst.hpp>
#include <gridtools/storage/cpu_ifirst.hpp>
using stencil_backend_t = gridtools::stencil::cpu_ifirst<>;
using storage_traits_t = gridtools::storage::cpu_ifirst;
#endif
namespace gt = gridtools;
namespace st = gt::stencil;
// These are the stencil operators that compose the multistage stencil in this test
struct lap_function {
using out = st::cartesian::inout_accessor<0>;
using in = st::cartesian::in_accessor<1, st::extent<-1, 1, -1, 1>>;
using param_list = st::make_param_list<out, in>;
template <class Eval>
GT_FUNCTION static void apply(Eval &&eval) {
eval(out()) = 4. * eval(in()) - (eval(in(1, 0)) + eval(in(0, 1)) + eval(in(-1, 0)) + eval(in(0, -1)));
}
};
struct flx_function {
using out = st::cartesian::inout_accessor<0>;
using in = st::cartesian::in_accessor<1, st::extent<0, 1, 0, 0>>;
using lap = st::cartesian::in_accessor<2, st::extent<0, 1, 0, 0>>;
using param_list = st::make_param_list<out, in, lap>;
template <class Eval>
GT_FUNCTION static void apply(Eval &&eval) {
auto res = eval(lap(1, 0)) - eval(lap(0, 0));
eval(out()) = res * (eval(in(1, 0)) - eval(in(0, 0))) > 0 ? 0 : res;
}
};
struct fly_function {
using out = st::cartesian::inout_accessor<0>;
using in = st::cartesian::in_accessor<1, st::extent<0, 0, 0, 1>>;
using lap = st::cartesian::in_accessor<2, st::extent<0, 0, 0, 1>>;
using param_list = st::make_param_list<out, in, lap>;
template <class Eval>
GT_FUNCTION static void apply(Eval &&eval) {
auto res = eval(lap(0, 1)) - eval(lap(0, 0));
eval(out()) = res * (eval(in(0, 1)) - eval(in(0, 0))) > 0 ? 0 : res;
}
};
struct out_function {
using out = st::cartesian::inout_accessor<0>;
using in = st::cartesian::in_accessor<1>;
using flx = st::cartesian::in_accessor<2, st::extent<-1, 0, 0, 0>>;
using fly = st::cartesian::in_accessor<3, st::extent<0, 0, -1, 0>>;
using coeff = st::cartesian::in_accessor<4>;
using param_list = st::make_param_list<out, in, flx, fly, coeff>;
template <class Eval>
GT_FUNCTION static void apply(Eval &&eval) {
eval(out()) = eval(in()) - eval(coeff()) * (eval(flx()) - eval(flx(-1, 0)) + eval(fly()) - eval(fly(0, -1)));
}
};
int main(int argc, char **argv) {
constexpr unsigned halo = 2;
unsigned d1, d2, d3;
if (argc != 4) {
std::cerr << "Usage: " << argv[0] << " dimx dimy dimz\n";
return 1;
} else {
d1 = atoi(argv[1]);
d2 = atoi(argv[2]);
d3 = atoi(argv[3]);
}
// Add dimensions and halos to the storage builder
auto storage_builder = gt::storage::builder<storage_traits_t>.dimensions(d1, d2, d3).halos(halo, halo, 0);
// Definition of the actual data fields that are used for input/output
auto in = storage_builder.type<double const>().value(42).build();
auto coeff = storage_builder.type<double const>().value(42).build();
auto out = storage_builder.type<double>().build();
// Here we specify the stencil composition of the computation.
auto spec = [](auto coeff, auto in, auto out) {
// temporary data (the library will take care of that and it is not observable by the user)
GT_DECLARE_TMP(double, lap, flx, fly);
return st::execute_parallel()
.ij_cached(lap, flx, fly)
.stage(lap_function(), lap, in)
.stage(flx_function(), flx, in, lap)
.stage(fly_function(), fly, in, lap)
.stage(out_function(), out, in, flx, fly, coeff);
};
// Now we describe the iteration space. The first two dimensions are described with a tuple of values (minus, plus,
// begin, end, length) begin and end, for each dimension represent the space where the output data will be located
// in the data_stores, while minus and plus indicate the number of halo points in the indices before begin and after
// end, respectively. The length, is not needed, and will be removed in future versions, but we keep it for now
// since the data structure used is the same used in the communication library and there the length is used.
gt::halo_descriptor di{halo, halo, halo, d1 - halo - 1, d1};
gt::halo_descriptor dj{halo, halo, halo, d2 - halo - 1, d2};
// The grid represent the iteration space. The third dimension is indicated here as a size and the iteration space
// is deduced by the fact that there is not an axis definition. More complex third dimensions are possible but not
// described in this example.
auto grid = st::make_grid(di, dj, d3);
// Here we perform the computation, specifying the backend, the grid (iteration space), binding spec arguments to
// the data fields
st::run(spec, stencil_backend_t(), grid, coeff, in, out);
}
|
