Skip to content

Commit

Permalink
Switch from LocalArray -> ArrayND (#1339)
Browse files Browse the repository at this point in the history
  • Loading branch information
@rcknaus
rcknaus authored Jan 17, 2025
1 parent 5db7851 commit da1f193
Show file tree
Hide file tree
Showing 48 changed files with 448 additions and 414 deletions.
180 changes: 180 additions & 0 deletions include/ArrayND.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,180 @@
#ifndef ARRAY_ND_H
#define ARRAY_ND_H

#include "Kokkos_Macros.hpp"
#include <type_traits>

namespace sierra::nalu {

// stack array set to interoperate with Kokkos views
template <typename ArrayType, typename = void>
struct ArrayND
{
};

template <typename ArrayType, decltype(std::rank_v<ArrayType>) r>
using enable_if_rank = std::enable_if_t<std::rank_v<ArrayType> == r>;

template <typename ArrayType>
struct ArrayND<ArrayType, enable_if_rank<ArrayType, 1>>
{
KOKKOS_DEFAULTED_FUNCTION constexpr ArrayND() = default;

static constexpr int rank = 1;
using value_type = std::remove_all_extents_t<ArrayType>;
value_type internal_data_[std::extent_v<ArrayType, 0>];

[[nodiscard]] static constexpr int extent_int(int /*unused*/)
{
return int(std::extent_v<ArrayType, 0>);
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type
operator()(int i) const noexcept
{
return internal_data_[i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type
operator[](int i) const noexcept
{
return internal_data_[i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type&
operator()(int i) noexcept
{
return internal_data_[i];
}
};

template <typename ArrayType>
struct ArrayND<ArrayType, enable_if_rank<ArrayType, 2>>
{
KOKKOS_DEFAULTED_FUNCTION constexpr ArrayND() = default;

static constexpr int rank = 2;

using value_type = std::remove_all_extents_t<ArrayType>;
value_type internal_data_[std::extent_v<ArrayType, 0>]
[std::extent_v<ArrayType, 1>];

[[nodiscard]] static constexpr int extent_int(int n)
{
return (n == 0) ? int(std::extent_v<ArrayType, 0>)
: int(std::extent_v<ArrayType, 1>);
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type
operator()(int j, int i) const noexcept
{
return internal_data_[j][i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type&
operator()(int j, int i) noexcept
{
return internal_data_[j][i];
}
};

template <typename ArrayType>
struct ArrayND<ArrayType, enable_if_rank<ArrayType, 3>>
{
KOKKOS_DEFAULTED_FUNCTION constexpr ArrayND() = default;

static constexpr int rank = 3;
[[nodiscard]] static constexpr int extent_int(int n)
{
return (n == 0) ? int(std::extent_v<ArrayType, 0>)
: (n == 1) ? int(std::extent_v<ArrayType, 1>)
: int(std::extent_v<ArrayType, 2>);
}

using value_type = std::remove_all_extents_t<ArrayType>;
value_type internal_data_[std::extent_v<ArrayType, 0>]
[std::extent_v<ArrayType, 1>]
[std::extent_v<ArrayType, 2>];

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr auto
operator()(int k, int j, int i) const noexcept
{
return internal_data_[k][j][i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr auto&
operator()(int k, int j, int i) noexcept
{
return internal_data_[k][j][i];
}
};

template <typename ArrayType>
struct ArrayND<ArrayType, enable_if_rank<ArrayType, 4>>
{
KOKKOS_DEFAULTED_FUNCTION constexpr ArrayND() = default;

static constexpr int rank = 4;
[[nodiscard]] static constexpr int extent_int(int n)
{
return (n == 0) ? int(std::extent_v<ArrayType, 0>)
: (n == 1) ? int(std::extent_v<ArrayType, 1>)
: (n == 2) ? int(std::extent_v<ArrayType, 2>)
: int(std::extent_v<ArrayType, 3>);
}

using value_type = std::remove_all_extents_t<ArrayType>;
value_type
internal_data_[std::extent_v<ArrayType, 0>][std::extent_v<ArrayType, 1>]
[std::extent_v<ArrayType, 2>][std::extent_v<ArrayType, 3>];

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type
operator()(int l, int k, int j, int i) const noexcept
{
return internal_data_[l][k][j][i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type&
operator()(int l, int k, int j, int i) noexcept
{
return internal_data_[l][k][j][i];
}
};

template <typename ArrayType>
struct ArrayND<ArrayType, enable_if_rank<ArrayType, 5>>
{
KOKKOS_DEFAULTED_FUNCTION constexpr ArrayND() = default;

static constexpr int rank = 5;
[[nodiscard]] static constexpr int extent_int(int n)
{
return (n == 0) ? int(std::extent_v<ArrayType, 0>)
: (n == 1) ? int(std::extent_v<ArrayType, 1>)
: (n == 2) ? int(std::extent_v<ArrayType, 2>)
: (n == 3) ? int(std::extent_v<ArrayType, 3>)
: int(std::extent_v<ArrayType, 4>);
}

using value_type = std::remove_all_extents_t<ArrayType>;
value_type
internal_data_[std::extent_v<ArrayType, 0>][std::extent_v<ArrayType, 1>]
[std::extent_v<ArrayType, 2>][std::extent_v<ArrayType, 3>]
[std::extent_v<ArrayType, 4>];

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type
operator()(int m, int l, int k, int j, int i) const noexcept
{
return internal_data_[m][l][k][j][i];
}

[[nodiscard]] KOKKOS_FORCEINLINE_FUNCTION constexpr value_type&
operator()(int m, int l, int k, int j, int i) noexcept
{
return internal_data_[m][l][k][j][i];
}
};

} // namespace sierra::nalu

#endif
52 changes: 26 additions & 26 deletions include/matrix_free/Coefficients.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
#ifndef CVFEM_COEFFICIENTS_H
#define CVFEM_COEFFICIENTS_H

#include "matrix_free/LocalArray.h"
#include "ArrayND.h"

// An option to use a symmetric version of
// the P=2 CVFEM operator. This loses one to two orders
Expand All @@ -32,10 +32,10 @@ struct Coeffs
template <>
struct Coeffs<1>
{
using nodal_matrix_type = LocalArray<double[2][2]>;
using scs_matrix_type = LocalArray<double[1][2]>;
using linear_nodal_matrix_type = LocalArray<double[2][2]>;
using linear_scs_matrix_type = LocalArray<double[2][1]>;
using nodal_matrix_type = ArrayND<double[2][2]>;
using scs_matrix_type = ArrayND<double[1][2]>;
using linear_nodal_matrix_type = ArrayND<double[2][2]>;
using linear_scs_matrix_type = ArrayND<double[2][1]>;
static constexpr nodal_matrix_type W = {{{0.75, 0.25}, {0.25, 0.75}}};

static constexpr nodal_matrix_type D = {{{-0.5, +0.5}, {-0.5, 0.5}}};
Expand All @@ -46,17 +46,17 @@ struct Coeffs<1>
static constexpr linear_nodal_matrix_type Nlin = {{{1, 0}, {0, 1}}};
static constexpr linear_scs_matrix_type Ntlin = {{{0.5}, {0.5}}};

static constexpr LocalArray<double[2]> Wl = {{1, 1}};
static constexpr ArrayND<double[2]> Wl = {{1, 1}};
};

#if USE_SYMMETRIC
template <>
struct Coeffs<2>
{
using nodal_matrix_type = LocalArray<double[3][3]>;
using scs_matrix_type = LocalArray<double[2][3]>;
using linear_nodal_matrix_type = LocalArray<double[2][3]>;
using linear_scs_matrix_type = LocalArray<double[2][2]>;
using nodal_matrix_type = ArrayND<double[3][3]>;
using scs_matrix_type = ArrayND<double[2][3]>;
using linear_nodal_matrix_type = ArrayND<double[2][3]>;
using linear_scs_matrix_type = ArrayND<double[2][2]>;

static constexpr nodal_matrix_type W = {
{{0.2561728395061728395062, 0.09876543209876543209877,
Expand Down Expand Up @@ -87,18 +87,18 @@ struct Coeffs<2>
{{0.8333333333333333333333, 0.1666666666666666666667},
{0.1666666666666666666667, 0.8333333333333333333333}}};

static constexpr LocalArray<double[3]> Wl = {
static constexpr ArrayND<double[3]> Wl = {
{0.3333333333333333333333, 1.333333333333333333333,
0.3333333333333333333333}};
};
#else
template <>
struct Coeffs<2>
{
using nodal_matrix_type = LocalArray<double[3][3]>;
using scs_matrix_type = LocalArray<double[2][3]>;
using linear_nodal_matrix_type = LocalArray<double[2][3]>;
using linear_scs_matrix_type = LocalArray<double[2][2]>;
using nodal_matrix_type = ArrayND<double[3][3]>;
using scs_matrix_type = ArrayND<double[2][3]>;
using linear_nodal_matrix_type = ArrayND<double[2][3]>;
using linear_scs_matrix_type = ArrayND<double[2][2]>;

static constexpr nodal_matrix_type W = {{
{+0.301258318378354124194, +0.15346642738699932044,
Expand Down Expand Up @@ -129,18 +129,18 @@ struct Coeffs<2>
{{+0.7886751345948129, +0.21132486540518708},
{+0.21132486540518708, +0.7886751345948129}}};

static constexpr LocalArray<double[3]> Wl = {
static constexpr ArrayND<double[3]> Wl = {
{0.422649730810374235491, 1.154700538379251529018,
0.422649730810374235491}};
};
#endif
template <>
struct Coeffs<3>
{
using nodal_matrix_type = LocalArray<double[4][4]>;
using scs_matrix_type = LocalArray<double[3][4]>;
using linear_nodal_matrix_type = LocalArray<double[2][4]>;
using linear_scs_matrix_type = LocalArray<double[2][3]>;
using nodal_matrix_type = ArrayND<double[4][4]>;
using scs_matrix_type = ArrayND<double[3][4]>;
using linear_nodal_matrix_type = ArrayND<double[2][4]>;
using linear_scs_matrix_type = ArrayND<double[2][3]>;

static constexpr nodal_matrix_type W = {
{{0.1583333333333333333333, 0.08527003148341972055897,
Expand Down Expand Up @@ -190,18 +190,18 @@ struct Coeffs<3>
{0.1127016653792582978610, 0.5000000000000000000000,
0.8872983346207417021390}}};

static constexpr LocalArray<double[4]> Wl = {
static constexpr ArrayND<double[4]> Wl = {
{0.225403330758516622964, 0.774596669241483377036, 0.774596669241483377036,
0.225403330758516622964}};
};

template <>
struct Coeffs<4>
{
using nodal_matrix_type = LocalArray<double[5][5]>;
using scs_matrix_type = LocalArray<double[4][5]>;
using linear_nodal_matrix_type = LocalArray<double[2][5]>;
using linear_scs_matrix_type = LocalArray<double[2][4]>;
using nodal_matrix_type = ArrayND<double[5][5]>;
using scs_matrix_type = ArrayND<double[4][5]>;
using linear_nodal_matrix_type = ArrayND<double[2][5]>;
using linear_scs_matrix_type = ArrayND<double[2][4]>;

static constexpr nodal_matrix_type W = {
{{0.09695253015044793809126, 0.05331486033726615325831,
Expand Down Expand Up @@ -268,7 +268,7 @@ struct Coeffs<4>
{0.06943184420297371373110, 0.3300094782075718713443,
0.6699905217924281286557, 0.9305681557970262307578}}};

static constexpr LocalArray<double[5]> Wl = {
static constexpr ArrayND<double[5]> Wl = {
{0.13886368840594742478, 0.52115526800919631042, 0.679962087169712529605,
0.52115526800919631042, 0.13886368840594742478}};
};
Expand Down
6 changes: 3 additions & 3 deletions include/matrix_free/ConductionInterior.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

#include "matrix_free/PolynomialOrders.h"
#include "matrix_free/KokkosViewTypes.h"
#include "matrix_free/LocalArray.h"
#include "ArrayND.h"

#include "Kokkos_Array.hpp"
#include "Tpetra_MultiVector.hpp"
Expand All @@ -29,7 +29,7 @@ namespace impl {
template <int p>
struct conduction_residual_t
{
using narray = LocalArray<ftype[p + 1][p + 1][p + 1]>;
using narray = ArrayND<ftype[p + 1][p + 1][p + 1]>;

static void invoke(
Kokkos::Array<double, 3> gammas,
Expand All @@ -47,7 +47,7 @@ namespace impl {
template <int p>
struct conduction_linearized_residual_t
{
using narray = LocalArray<ftype[p + 1][p + 1][p + 1]>;
using narray = ArrayND<ftype[p + 1][p + 1][p + 1]>;

static void invoke(
double gamma,
Expand Down
6 changes: 3 additions & 3 deletions include/matrix_free/ContinuityInterior.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#define CONTINUITY_INTERIOR_H

#include "matrix_free/KokkosViewTypes.h"
#include "matrix_free/LocalArray.h"
#include "ArrayND.h"
#include "matrix_free/PolynomialOrders.h"

#include "Teuchos_RCP.hpp"
Expand All @@ -29,7 +29,7 @@ namespace impl {
template <int p>
struct continuity_residual_t
{
using narray = LocalArray<ftype[p + 1][p + 1][p + 1]>;
using narray = ArrayND<ftype[p + 1][p + 1][p + 1]>;

static void invoke(
double scaling,
Expand All @@ -44,7 +44,7 @@ namespace impl {
template <int p>
struct continuity_linearized_residual_t
{
using narray = LocalArray<ftype[p + 1][p + 1][p + 1]>;
using narray = ArrayND<ftype[p + 1][p + 1][p + 1]>;

static void invoke(
const_elem_offset_view<p> offsets,
Expand Down
Loading

0 comments on commit da1f193

Please sign in to comment.