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
144
145
146
147
148
149
150
151
152
153
154
155
156
|
/*
* 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
*/
#include <gtest/gtest.h>
#include <gridtools/fn/cartesian.hpp>
#include <fn_select.hpp>
#include <test_environment.hpp>
#include "../horizontal_diffusion_repository.hpp"
namespace {
using namespace gridtools;
using namespace fn;
using namespace cartesian;
using namespace literals;
struct laplacian {
GT_FUNCTION constexpr auto operator()() const {
return [](auto const &in) {
constexpr auto i = cartesian::dim::i();
constexpr auto j = cartesian::dim::j();
return 4 * deref(in) - (deref(shift(in, i, 1)) + deref(shift(in, i, -1)) + deref(shift(in, j, 1)) +
deref(shift(in, j, -1)));
};
}
};
template <class D>
struct flux {
GT_FUNCTION constexpr auto operator()() const {
return [](auto const &in, auto const &lap) {
auto tmp = deref(shift(lap, D(), 1)) - deref(lap);
return tmp * (deref(shift(in, D(), 1)) - deref(in)) > 0 ? 0 : tmp;
};
}
};
struct hdiff {
GT_FUNCTION constexpr auto operator()() const {
return [](auto const &in, auto const &coeff, auto const &flx, auto const &fly) {
constexpr auto i = cartesian::dim::i();
constexpr auto j = cartesian::dim::j();
return deref(in) -
deref(coeff) * (deref(flx) - deref(shift(flx, i, -1)) + deref(fly) - deref(shift(fly, j, -1)));
};
}
};
struct hdiff_fused {
GT_FUNCTION constexpr auto operator()() const {
return [](auto const &in, auto const &coeff) {
constexpr auto i = cartesian::dim::i();
constexpr auto j = cartesian::dim::j();
auto lap = 4 * deref(in) - (deref(shift(in, i, 1)) + deref(shift(in, i, -1)) + deref(shift(in, j, 1)) +
deref(shift(in, j, -1)));
auto lap_ip1 =
4 * deref(shift(in, i, 1)) -
(deref(shift(in, i, 2)) + deref(in) + deref(shift(in, i, 1, j, 1)) + deref(shift(in, i, 1, j, -1)));
auto lap_im1 =
4 * deref(shift(in, i, -1)) - (deref(in) + deref(shift(in, i, -2)) + deref(shift(in, i, -1, j, 1)) +
deref(shift(in, i, -1, j, -1)));
auto lap_jp1 =
4 * deref(shift(in, j, 1)) -
(deref(shift(in, i, 1, j, 1)) + deref(shift(in, i, -1, j, 1)) + deref(shift(in, j, 2)) + deref(in));
auto lap_jm1 =
4 * deref(shift(in, j, -1)) - (deref(shift(in, i, 1, j, -1)) + deref(shift(in, i, -1, j, -1)) +
deref(in) + deref(shift(in, j, -2)));
auto tmp0 = lap_ip1 - lap;
auto flx = tmp0 * (deref(shift(in, i, 1)) - deref(in)) > 0 ? 0 : tmp0;
auto tmp1 = lap - lap_im1;
auto flx_im1 = tmp1 * (deref(in) - deref(shift(in, i, -1))) > 0 ? 0 : tmp1;
auto tmp2 = lap_jp1 - lap;
auto fly = tmp2 * (deref(shift(in, j, 1)) - deref(in)) > 0 ? 0 : tmp2;
auto tmp3 = lap - lap_jm1;
auto fly_jm1 = tmp3 * (deref(in) - deref(shift(in, j, -1))) > 0 ? 0 : tmp3;
return deref(in) - deref(coeff) * (flx - flx_im1 + fly - fly_jm1);
};
}
};
GT_REGRESSION_TEST(fn_cartesian_horizontal_diffusion, test_environment<2>, fn_backend_t) {
using float_t = typename TypeParam::float_t;
horizontal_diffusion_repository repo(TypeParam::d(0), TypeParam::d(1), TypeParam::d(2));
auto out = TypeParam::make_storage();
auto fencil = [&](int i, int j, int k, auto &out, auto const &in, auto const &coeff) {
using sizes_t = hymap::keys<dim::i, dim::j, dim::k>::values<int, int, int>;
auto be = fn_backend_t();
auto full_domain = cartesian_domain(sizes_t{i, j, k});
auto full_domain_backend = make_backend(be, full_domain);
auto alloc = tmp_allocator(be);
auto lap = allocate_global_tmp<float_t>(alloc, sizes_t{i, j, k});
auto flx = allocate_global_tmp<float_t>(alloc, sizes_t{i, j, k});
auto fly = allocate_global_tmp<float_t>(alloc, sizes_t{i, j, k});
auto domain = cartesian_domain(sizes_t{i - 2, j - 2, k}, sizes_t{1, 1, 0});
auto backend = make_backend(fn_backend_t(), domain);
backend.stencil_executor()().arg(lap).arg(in).assign(0_c, laplacian(), 1_c).execute();
backend.stencil_executor()()
.arg(flx)
.arg(fly)
.arg(in)
.arg(lap)
.assign(0_c, flux<dim::i>(), 2_c, 3_c)
.assign(1_c, flux<dim::j>(), 2_c, 3_c)
.execute();
backend.stencil_executor()()
.arg(out)
.arg(in)
.arg(coeff)
.arg(flx)
.arg(fly)
.assign(0_c, hdiff(), 1_c, 2_c, 3_c, 4_c)
.execute();
};
auto comp =
[&, coeff = TypeParam::make_const_storage(repo.coeff), in = TypeParam::make_const_storage(repo.in)] {
fencil(TypeParam::d(0), TypeParam::d(1), TypeParam::d(2), out, in, coeff);
};
comp();
TypeParam::verify(repo.out, out);
TypeParam::benchmark("fn_cartesian_horizontal_diffusion", comp);
}
GT_REGRESSION_TEST(fn_cartesian_horizontal_diffusion_fused, test_environment<2>, fn_backend_t) {
horizontal_diffusion_repository repo(TypeParam::d(0), TypeParam::d(1), TypeParam::d(2));
auto out = TypeParam::make_storage();
auto fencil = [&](int i, int j, int k, auto &out, auto const &in, auto const &coeff) {
using sizes_t = hymap::keys<dim::i, dim::j, dim::k>::values<int, int, int>;
auto domain = cartesian_domain(sizes_t{i - 4, j - 4, k}, sizes_t{2, 2, 0});
auto backend = make_backend(fn_backend_t(), domain);
backend.stencil_executor()().arg(out).arg(in).arg(coeff).assign(0_c, hdiff_fused(), 1_c, 2_c).execute();
};
auto comp =
[&, coeff = TypeParam::make_const_storage(repo.coeff), in = TypeParam::make_const_storage(repo.in)] {
fencil(TypeParam::d(0), TypeParam::d(1), TypeParam::d(2), out, in, coeff);
};
comp();
TypeParam::verify(repo.out, out);
TypeParam::benchmark("fn_cartesian_horizontal_diffusion_fused", comp);
}
} // namespace
|
