Skip to content

Commit e67aded

Browse files
tcclevengertcclevenger
authored
Merge branch 'tcclevenger/eamxx/fix_tracer_layout_in_shoc_funcs' (PR #7143)
This was a bug where I was passing around StridedLayout views to functions expecting LayoutLeft. Only caught in the PM-GPU nightlies and by PR #6916, so I will ensure those pass before merging. [BFB]
2 parents 1e40ce2 + 8d6d8e0 commit e67aded

File tree

4 files changed

+27
-17
lines changed

4 files changed

+27
-17
lines changed

‎components/eamxx/src/physics/shoc/disp/shoc_update_prognostics_implicit_disp.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ ::update_prognostics_implicit_disp(
2828
const WorkspaceMgr& workspace_mgr,
2929
const view_2d<Spack>& thetal,
3030
const view_2d<Spack>& qw,
31-
const view_3d<Spack>& tracer,
31+
const view_3d_strided<Spack>& tracer,
3232
const view_2d<Spack>& tke,
3333
const view_2d<Spack>& u_wind,
3434
const view_2d<Spack>& v_wind)
@@ -58,7 +58,7 @@ ::update_prognostics_implicit_disp(
5858
workspace,
5959
ekat::subview(thetal, i),
6060
ekat::subview(qw, i),
61-
ekat::subview(tracer, i),
61+
Kokkos::subview(tracer, i, Kokkos::ALL(), Kokkos::ALL()),
6262
ekat::subview(tke, i),
6363
ekat::subview(u_wind, i),
6464
ekat::subview(v_wind, i));

‎components/eamxx/src/physics/shoc/impl/shoc_main_impl.hpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ void Functions<S,D>::shoc_main_internal(
112112
const uview_1d<Spack>& u_wind,
113113
const uview_1d<Spack>& v_wind,
114114
const uview_1d<Spack>& wthv_sec,
115-
const uview_2d<Spack>& qtracers,
115+
const uview_2d_strided<Spack>& qtracers,
116116
const uview_1d<Spack>& tk,
117117
const uview_1d<Spack>& shoc_cldfrac,
118118
const uview_1d<Spack>& shoc_ql,
@@ -367,7 +367,7 @@ void Functions<S,D>::shoc_main_internal(
367367
const uview_2d<Spack>& u_wind,
368368
const uview_2d<Spack>& v_wind,
369369
const view_2d<Spack>& wthv_sec,
370-
const view_3d<Spack>& qtracers,
370+
const view_3d_strided<Spack>& qtracers,
371371
const view_2d<Spack>& tk,
372372
const view_2d<Spack>& shoc_cldfrac,
373373
const view_2d<Spack>& shoc_ql,
@@ -595,10 +595,10 @@ Int Functions<S,D>::shoc_main(
595595
auto start = std::chrono::steady_clock::now();
596596

597597
// Runtime options
598-
const Scalar lambda_low = shoc_runtime.lambda_low;
599-
const Scalar lambda_high = shoc_runtime.lambda_high;
600-
const Scalar lambda_slope = shoc_runtime.lambda_slope;
601-
const Scalar lambda_thresh = shoc_runtime.lambda_thresh;
598+
const Scalar lambda_low = shoc_runtime.lambda_low;
599+
const Scalar lambda_high = shoc_runtime.lambda_high;
600+
const Scalar lambda_slope = shoc_runtime.lambda_slope;
601+
const Scalar lambda_thresh = shoc_runtime.lambda_thresh;
602602
const Scalar thl2tune = shoc_runtime.thl2tune;
603603
const Scalar qw2tune = shoc_runtime.qw2tune;
604604
const Scalar qwthl2tune = shoc_runtime.qwthl2tune;

‎components/eamxx/src/physics/shoc/impl/shoc_update_prognostics_implicit_impl.hpp‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ void Functions<S,D>::update_prognostics_implicit(
3434
const Workspace& workspace,
3535
const uview_1d<Spack>& thetal,
3636
const uview_1d<Spack>& qw,
37-
const uview_2d<Spack>& qtracers,
37+
const uview_2d_strided<Spack>& qtracers,
3838
const uview_1d<Spack>& tke,
3939
const uview_1d<Spack>& u_wind,
4040
const uview_1d<Spack>& v_wind)
@@ -83,7 +83,6 @@ void Functions<S,D>::update_prognostics_implicit(
8383
const auto thetal_s = ekat::scalarize(thetal);
8484
const auto qw_s = ekat::scalarize(qw);
8585
const auto tke_s = ekat::scalarize(tke);
86-
const auto qtracers_s = ekat::scalarize(qtracers);
8786
const auto qtracers_rhs_s = ekat::scalarize(qtracers_rhs);
8887
const auto wtracer_sfc_s = ekat::scalarize(wtracer_sfc);
8988

@@ -140,8 +139,10 @@ void Functions<S,D>::update_prognostics_implicit(
140139
tke_s(nlev-1) += cmnfac*wtke_sfc;
141140
});
142141

142+
const auto sfc_lev_idx = (nlev-1)/Spack::n;
143+
const auto sfc_pack_idx = (nlev-1)%Spack::n;
143144
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, num_qtracers), [&] (const Int& q) {
144-
qtracers_s(q, nlev-1) += cmnfac*wtracer_sfc_s(q);
145+
qtracers(q, sfc_lev_idx)[sfc_pack_idx] += cmnfac*wtracer_sfc_s(q);
145146
});
146147
}
147148

@@ -152,8 +153,10 @@ void Functions<S,D>::update_prognostics_implicit(
152153
wind_rhs_s(k,1) = v_wind_s(k);
153154

154155
// The rhs version of the tracers is the transpose of the input/output layout
156+
const auto lev_idx = k/Spack::n;
157+
const auto pack_idx = k%Spack::n;
155158
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, num_qtracers), [&] (const Int& q) {
156-
qtracers_rhs_s(k, q) = qtracers_s(q, k);
159+
qtracers_rhs_s(k, q) = qtracers(q, lev_idx)[pack_idx];
157160
});
158161
qtracers_rhs_s(k, num_qtracers) = thetal_s(k);
159162
qtracers_rhs_s(k, num_qtracers+1) = qw_s(k);
@@ -189,8 +192,10 @@ void Functions<S,D>::update_prognostics_implicit(
189192
v_wind_s(k) = wind_rhs_s(k, 1);
190193

191194
// Transpose tracers back to input/output layout
195+
const auto lev_idx = k/Spack::n;
196+
const auto pack_idx = k%Spack::n;
192197
Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, num_qtracers), [&] (const Int& q) {
193-
qtracers_s(q, k) = qtracers_rhs_s(k, q);
198+
qtracers(q, lev_idx)[pack_idx] = qtracers_rhs_s(k, q);
194199
});
195200
thetal_s(k) = qtracers_rhs_s(k, num_qtracers);
196201
qw_s(k) = qtracers_rhs_s(k, num_qtracers+1);

‎components/eamxx/src/physics/shoc/shoc_functions.hpp‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ struct Functions
5555
template <typename S>
5656
using view_3d = typename KT::template view_3d<S>;
5757

58+
template <typename S>
59+
using view_2d_strided = typename KT::template sview<S**>;
5860
template <typename S>
5961
using view_3d_strided = typename KT::template sview<S***>;
6062

@@ -67,6 +69,9 @@ struct Functions
6769
template <typename S>
6870
using uview_2d = typename ekat::template Unmanaged<view_2d<S> >;
6971

72+
template <typename S>
73+
using uview_2d_strided = typename ekat::template Unmanaged<view_2d_strided<S> >;
74+
7075
using MemberType = typename KT::MemberType;
7176

7277
using WorkspaceMgr = typename ekat::WorkspaceManager<Spack, Device>;
@@ -649,7 +654,7 @@ struct Functions
649654
const Workspace& workspace,
650655
const uview_1d<Spack>& thetal,
651656
const uview_1d<Spack>& qw,
652-
const uview_2d<Spack>& tracer,
657+
const uview_2d_strided<Spack>& tracer,
653658
const uview_1d<Spack>& tke,
654659
const uview_1d<Spack>& u_wind,
655660
const uview_1d<Spack>& v_wind);
@@ -675,7 +680,7 @@ struct Functions
675680
const WorkspaceMgr& workspace_mgr,
676681
const view_2d<Spack>& thetal,
677682
const view_2d<Spack>& qw,
678-
const view_3d<Spack>& tracer,
683+
const view_3d_strided<Spack>& tracer,
679684
const view_2d<Spack>& tke,
680685
const view_2d<Spack>& u_wind,
681686
const view_2d<Spack>& v_wind);
@@ -1036,7 +1041,7 @@ struct Functions
10361041
const uview_1d<Spack>& u_wind,
10371042
const uview_1d<Spack>& v_wind,
10381043
const uview_1d<Spack>& wthv_sec,
1039-
const uview_2d<Spack>& qtracers,
1044+
const uview_2d_strided<Spack>& qtracers,
10401045
const uview_1d<Spack>& tk,
10411046
const uview_1d<Spack>& shoc_cldfrac,
10421047
const uview_1d<Spack>& shoc_ql,
@@ -1110,7 +1115,7 @@ struct Functions
11101115
const uview_2d<Spack>& u_wind,
11111116
const uview_2d<Spack>& v_wind,
11121117
const view_2d<Spack>& wthv_sec,
1113-
const view_3d<Spack>& qtracers,
1118+
const view_3d_strided<Spack>& qtracers,
11141119
const view_2d<Spack>& tk,
11151120
const view_2d<Spack>& shoc_cldfrac,
11161121
const view_2d<Spack>& shoc_ql,

0 commit comments

Comments
 (0)