Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

SmallMatrix: Support 1-based indexing #4188

Merged
merged 1 commit into from
Oct 11, 2024
Merged

SmallMatrix: Support 1-based indexing #4188

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 93 additions & 62 deletions Src/Base/AMReX_SmallMatrix.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,23 +20,24 @@ namespace amrex {
/**
* \brief Matrix class with compile-time size
*
* The starting index for both rows and columns is always zero. Also
* note that column vectors and row vectors are special cases of a
* Note that column vectors and row vectors are special cases of a
* Matrix.
*
* \tparam T Matrix element data type.
* \tparam NRows Number of rows.
* \tparam NCols Number of columns.
* \tparam ORDER Memory layout order. Order::F (i.e., column-major) by default.
* \tparam StartIndex Starting index. Either 0 or 1.
*/
template <class T, int NRows, int NCols, Order ORDER = Order::F>
template <class T, int NRows, int NCols, Order ORDER = Order::F, int StartIndex = 0>
struct SmallMatrix
{
using value_type = T;
using reference_type = T&;
static constexpr int row_size = NRows;
static constexpr int column_size = NCols;
static constexpr Order ordering = ORDER;
static constexpr int starting_index = StartIndex;

/**
* \brief Default constructor
Expand Down Expand Up @@ -78,10 +79,10 @@ namespace amrex {
explicit SmallMatrix (std::initializer_list<std::initializer_list<T>> const& init)
{
AMREX_ASSERT(NRows == init.size());
int i = 0;
int i = StartIndex;
for (auto const& row : init) {
AMREX_ASSERT(NCols == row.size());
int j = 0;
int j = StartIndex;
for (auto const& x : row) {
(*this)(i,j) = x;
++j;
Expand All @@ -93,6 +94,11 @@ namespace amrex {
//! Returns a const reference to the element at row i and column j.
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
const T& operator() (int i, int j) const noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
--j;
}
AMREX_ASSERT(i < NRows && j < NCols);
if constexpr (ORDER == Order::F) {
return m_mat[i+j*NRows];
Expand All @@ -104,6 +110,11 @@ namespace amrex {
//! Returns a reference to the element at row i and column j.
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
T& operator() (int i, int j) noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
--j;
}
AMREX_ASSERT(i < NRows && j < NCols);
if constexpr (ORDER == Order::F) {
return m_mat[i+j*NRows];
Expand All @@ -116,6 +127,10 @@ namespace amrex {
template <int MM=NRows, int NN=NCols, std::enable_if_t<(MM==1 || NN==1), int> = 0>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
const T& operator() (int i) const noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
}
AMREX_ASSERT(i < NRows*NCols);
return m_mat[i];
}
Expand All @@ -124,6 +139,10 @@ namespace amrex {
template <int MM=NRows, int NN=NCols, std::enable_if_t<(MM==1 || NN==1), int> = 0>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
T& operator() (int i) noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
}
AMREX_ASSERT(i < NRows*NCols);
return m_mat[i];
}
Expand All @@ -132,6 +151,10 @@ namespace amrex {
template <int MM=NRows, int NN=NCols, std::enable_if_t<(MM==1 || NN==1), int> = 0>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
const T& operator[] (int i) const noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
}
AMREX_ASSERT(i < NRows*NCols);
return m_mat[i];
}
Expand All @@ -140,6 +163,10 @@ namespace amrex {
template <int MM=NRows, int NN=NCols, std::enable_if_t<(MM==1 || NN==1), int> = 0>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
T& operator[] (int i) noexcept {
static_assert(StartIndex == 0 || StartIndex == 1);
if constexpr (StartIndex == 1) {
--i;
}
AMREX_ASSERT(i < NRows*NCols);
return m_mat[i];
}
Expand Down Expand Up @@ -174,7 +201,7 @@ namespace amrex {

//! Set all elements in the matrix to the given value
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>&
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>&
setVal (T val)
{
for (auto& x : m_mat) { x = val; }
Expand All @@ -185,30 +212,32 @@ namespace amrex {
template <int MM=NRows, int NN=NCols, std::enable_if_t<MM==NN, int> = 0>
static constexpr
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
Identity () noexcept {
SmallMatrix<T,NRows,NCols,ORDER> I{};
constexpr_for<0,NRows>([&] (int i) { I(i,i) = T(1); });
static_assert(StartIndex == 0 || StartIndex == 1);
SmallMatrix<T,NRows,NCols,ORDER,StartIndex> I{};
constexpr_for<StartIndex,NRows+StartIndex>(
[&] (int i) { I(i,i) = T(1); });
return I;
}

//! Returns a matrix initialized with zeros
static constexpr
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
Zero () noexcept {
SmallMatrix<T,NRows,NCols,ORDER> Z{};
SmallMatrix<T,NRows,NCols,ORDER,StartIndex> Z{};
return Z;
}

//! Returns transposed matrix
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NCols,NRows,ORDER>
SmallMatrix<T,NCols,NRows,ORDER,StartIndex>
transpose () const
{
SmallMatrix<T,NCols,NRows,ORDER> r;
for (int j = 0; j < NRows; ++j) {
for (int i = 0; i < NCols; ++i) {
SmallMatrix<T,NCols,NRows,ORDER,StartIndex> r;
for (int j = StartIndex; j < NRows+StartIndex; ++j) {
for (int i = StartIndex; i < NCols+StartIndex; ++i) {
r(i,j) = (*this)(j,i);
}
}
Expand All @@ -218,11 +247,12 @@ namespace amrex {
//! Transposes a square matrix in-place.
template <int MM=NRows, int NN=NCols, std::enable_if_t<MM==NN,int> = 0>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>&
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>&
transposeInPlace ()
{
for (int j = 1; j < NCols; ++j) {
for (int i = 0; i < j; ++i) {
static_assert(StartIndex == 0 || StartIndex == 1);
for (int j = 1+StartIndex; j < NCols+StartIndex; ++j) {
for (int i = StartIndex; i < j; ++i) {
amrex::Swap((*this)(i,j), (*this)(j,i));
}
}
Expand Down Expand Up @@ -257,14 +287,14 @@ namespace amrex {
T trace () const
{
T t = 0;
constexpr_for<0,MM>([&] (int i) { t += (*this)(i,i); });
constexpr_for<StartIndex,MM+StartIndex>([&] (int i) { t += (*this)(i,i); });
return t;
}

//! Operator += performing matrix addition as in (*this) += rhs
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>&
operator += (SmallMatrix<T,NRows,NCols,ORDER> const& rhs)
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>&
operator += (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> const& rhs)
{
for (int n = 0; n < NRows*NCols; ++n) {
m_mat[n] += rhs.m_mat[n];
Expand All @@ -274,18 +304,18 @@ namespace amrex {

//! Binary operator + returning the result of maxtrix addition, lhs+rhs
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
friend SmallMatrix<T,NRows,NCols,ORDER>
operator+ (SmallMatrix<T,NRows,NCols,ORDER> lhs,
SmallMatrix<T,NRows,NCols,ORDER> const& rhs)
friend SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
operator+ (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> lhs,
SmallMatrix<T,NRows,NCols,ORDER,StartIndex> const& rhs)
{
lhs += rhs;
return lhs;
}

//! Operator -= performing matrix subtraction as in (*this) -= rhs
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>&
operator -= (SmallMatrix<T,NRows,NCols,ORDER> const& rhs)
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>&
operator -= (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> const& rhs)
{
for (int n = 0; n < NRows*NCols; ++n) {
m_mat[n] -= rhs.m_mat[n];
Expand All @@ -295,25 +325,25 @@ namespace amrex {

//! Binary operator - returning the result of maxtrix subtraction, lhs-rhs
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
friend SmallMatrix<T,NRows,NCols,ORDER>
operator- (SmallMatrix<T,NRows,NCols,ORDER> lhs,
SmallMatrix<T,NRows,NCols,ORDER> const& rhs)
friend SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
operator- (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> lhs,
SmallMatrix<T,NRows,NCols,ORDER,StartIndex> const& rhs)
{
lhs -= rhs;
return lhs;
}

//! Unary minus operator
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
operator- () const
{
return (*this) * T(-1);
}

//! Operator *= that scales this matrix in place by a scalar.
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<T,NRows,NCols,ORDER>&
SmallMatrix<T,NRows,NCols,ORDER,StartIndex>&
operator *= (T a)
{
for (auto& x : m_mat) {
Expand All @@ -324,32 +354,32 @@ namespace amrex {

//! Returns the product of a matrix and a scalar
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
friend SmallMatrix<T,NRows,NCols,ORDER>
operator* (SmallMatrix<T,NRows,NCols,ORDER> m, T a)
friend SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
operator* (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> m, T a)
{
m *= a;
return m;
}

//! Returns the product of a scalar and a matrix
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
friend SmallMatrix<T,NRows,NCols,ORDER>
operator* (T a, SmallMatrix<T,NRows,NCols,ORDER> m)
friend SmallMatrix<T,NRows,NCols,ORDER,StartIndex>
operator* (T a, SmallMatrix<T,NRows,NCols,ORDER,StartIndex> m)
{
m *= a;
return m;
}

//! Returns matrix product of two matrices
template <class U, int N1, int N2, int N3, Order Ord>
template <class U, int N1, int N2, int N3, Order Ord, int SI>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
friend SmallMatrix<U,N1,N3,Ord>
operator* (SmallMatrix<U,N1,N2,Ord> const& lhs,
SmallMatrix<U,N2,N3,Ord> const& rhs);
friend SmallMatrix<U,N1,N3,Ord,SI>
operator* (SmallMatrix<U,N1,N2,Ord,SI> const& lhs,
SmallMatrix<U,N2,N3,Ord,SI> const& rhs);

//! Returns the dot product of two vectors
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
T dot (SmallMatrix<T,NRows,NCols,ORDER> const& rhs) const
T dot (SmallMatrix<T,NRows,NCols,ORDER,StartIndex> const& rhs) const
{
T r = 0;
for (int n = 0; n < NRows*NCols; ++n) {
Expand All @@ -362,30 +392,31 @@ namespace amrex {
T m_mat[NRows*NCols];
};

template <class U, int N1, int N2, int N3, Order Ord>
template <class U, int N1, int N2, int N3, Order Ord, int SI>
[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
SmallMatrix<U,N1,N3,Ord>
operator* (SmallMatrix<U,N1,N2,Ord> const& lhs,
SmallMatrix<U,N2,N3,Ord> const& rhs)
SmallMatrix<U,N1,N3,Ord,SI>
operator* (SmallMatrix<U,N1,N2,Ord,SI> const& lhs,
SmallMatrix<U,N2,N3,Ord,SI> const& rhs)
{
SmallMatrix<U,N1,N3,Ord> r;
static_assert(SI == 0 || SI == 1);
SmallMatrix<U,N1,N3,Ord,SI> r;
if constexpr (Ord == Order::F) {
for (int j = 0; j < N3; ++j) {
constexpr_for<0,N1>([&] (int i) { r(i,j) = U(0); });
for (int k = 0; k < N2; ++k) {
for (int j = SI; j < N3+SI; ++j) {
constexpr_for<SI,N1+SI>([&] (int i) { r(i,j) = U(0); });
for (int k = SI; k < N2+SI; ++k) {
auto b = rhs(k,j);
constexpr_for<0,N1>([&] (int i)
constexpr_for<SI,N1+SI>([&] (int i)
{
r(i,j) += lhs(i,k) * b;
});
}
}
} else {
for (int i = 0; i < N1; ++i) {
constexpr_for<0,N3>([&] (int j) { r(i,j) = U(0); });
for (int k = 0; k < N2; ++k) {
for (int i = SI; i < N1+SI; ++i) {
constexpr_for<SI,N3+SI>([&] (int j) { r(i,j) = U(0); });
for (int k = SI; k < N2+SI; ++k) {
auto a = lhs(i,k);
constexpr_for<0,N3>([&] (int j)
constexpr_for<SI,N3+SI>([&] (int j)
{
r(i,j) += a * rhs(k,j);
});
Expand All @@ -395,25 +426,25 @@ namespace amrex {
return r;
}

template <class T, int NRows, int NCols, Order ORDER>
template <class T, int NRows, int NCols, Order ORDER, int SI>
std::ostream& operator<< (std::ostream& os,
SmallMatrix<T,NRows,NCols,ORDER> const& mat)
SmallMatrix<T,NRows,NCols,ORDER,SI> const& mat)
{
for (int i = 0; i < NRows; ++i) {
os << mat(i,0);
for (int j = 1; j < NCols; ++j) {
for (int i = SI; i < NRows+SI; ++i) {
os << mat(i,SI);
for (int j = 1+SI; j < NCols+SI; ++j) {
os << " " << mat(i,j);
}
os << "\n";
}
return os;
}

template <class T, int N>
using SmallVector = SmallMatrix<T,N,1>;
template <class T, int N, int StartIndex = 0>
using SmallVector = SmallMatrix<T,N,1,Order::F,StartIndex>;

template <class T, int N>
using SmallRowVector = SmallMatrix<T,1,N>;
template <class T, int N, int StartIndex = 0>
using SmallRowVector = SmallMatrix<T,1,N,Order::F,StartIndex>;
}

#endif
Expand Down
Loading
Loading