Files
mfem/linalg/ttensor.hpp
T
Veselin Dobrev b28b5370d1 Remove the MFEM_HOST_DEVICE specifiers from the operator[] methods
in class TVector -- this was causing compilation errors when CUDA is
enabled.
2020-05-19 20:32:50 -07:00

680 lines
22 KiB
C++

// Copyright (c) 2010-2020, Lawrence Livermore National Security, LLC. Produced
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
// LICENSE and NOTICE for details. LLNL-CODE-806117.
//
// This file is part of the MFEM library. For more information and source code
// availability visit https://mfem.org.
//
// MFEM is free software; you can redistribute it and/or modify it under the
// terms of the BSD-3 license. We welcome feedback and contributions, see file
// CONTRIBUTING.md for details.
#ifndef MFEM_TEMPLATE_TENSOR
#define MFEM_TEMPLATE_TENSOR
#include "../config/tconfig.hpp"
#include "../linalg/simd.hpp"
#include "../general/tassign.hpp"
#include "tlayout.hpp"
#include "tmatrix.hpp"
#include "../general/cuda.hpp"
#include "../general/hip.hpp"
// Templated tensor implementation (up to order 4)
namespace mfem
{
// Element-wise tensor operations
namespace internal
{
template <int Rank>
struct TensorOps;
template <>
struct TensorOps<1> // rank = 1
{
// Assign: A {=,+=,*=} scalar_value
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
const scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 1, "invalid rank");
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1)], value);
}
}
// Assign: A {=,+=,*=} scalar_value, host+device version
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
MFEM_HOST_DEVICE
static void AssignHD(const A_layout_t &A_layout, A_data_t &A_data,
const scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 1, "invalid rank");
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::AssignHD<Op>(A_data[A_layout.ind(i1)], value);
}
}
// Assign: A {=,+=,*=} B
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 1 && B_layout_t::rank == 1,
"invalid ranks");
MFEM_STATIC_ASSERT(A_layout_t::dim_1 == B_layout_t::dim_1,
"invalid dimensions");
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1)], B_data[B_layout.ind(i1)]);
}
}
};
template <>
struct TensorOps<2> // rank = 2
{
// Assign: A {=,+=,*=} scalar_value
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2, "invalid rank");
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2)], value);
}
}
}
// Assign: A {=,+=,*=} scalar_value, host+device version
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
MFEM_HOST_DEVICE
static void AssignHD(const A_layout_t &A_layout, A_data_t &A_data,
scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2, "invalid rank");
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::AssignHD<Op>(A_data[A_layout.ind(i1,i2)], value);
}
}
}
// Assign: A {=,+=,*=} B
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2 && B_layout_t::rank == 2,
"invalid ranks");
MFEM_STATIC_ASSERT(A_layout_t::dim_1 == B_layout_t::dim_1 &&
A_layout_t::dim_2 == B_layout_t::dim_2,
"invalid dimensions");
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2)],
B_data[B_layout.ind(i1,i2)]);
}
}
}
};
template <>
struct TensorOps<3> // rank = 3
{
// Assign: A {=,+=,*=} scalar_value
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 3, "invalid rank");
for (int i3 = 0; i3 < A_layout_t::dim_3; i3++)
{
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2,i3)], value);
}
}
}
}
// Assign: A {=,+=,*=} B
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 3 && B_layout_t::rank == 3,
"invalid ranks");
MFEM_STATIC_ASSERT(A_layout_t::dim_1 == B_layout_t::dim_1 &&
A_layout_t::dim_2 == B_layout_t::dim_2 &&
A_layout_t::dim_3 == B_layout_t::dim_3,
"invalid dimensions");
for (int i3 = 0; i3 < A_layout_t::dim_3; i3++)
{
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2,i3)],
B_data[B_layout.ind(i1,i2,i3)]);
}
}
}
}
};
template <>
struct TensorOps<4> // rank = 4
{
// Assign: A {=,+=,*=} scalar_value
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
scalar_t value)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 4, "invalid rank");
for (int i4 = 0; i4 < A_layout_t::dim_4; i4++)
{
for (int i3 = 0; i3 < A_layout_t::dim_3; i3++)
{
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2,i3,i4)], value);
}
}
}
}
}
// Assign: A {=,+=,*=} B
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t>
static void Assign(const A_layout_t &A_layout, A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 4 && B_layout_t::rank == 4,
"invalid ranks");
MFEM_STATIC_ASSERT(A_layout_t::dim_1 == B_layout_t::dim_1 &&
A_layout_t::dim_2 == B_layout_t::dim_2 &&
A_layout_t::dim_3 == B_layout_t::dim_3 &&
A_layout_t::dim_4 == B_layout_t::dim_4,
"invalid dimensions");
for (int i4 = 0; i4 < A_layout_t::dim_4; i4++)
{
for (int i3 = 0; i3 < A_layout_t::dim_3; i3++)
{
for (int i2 = 0; i2 < A_layout_t::dim_2; i2++)
{
for (int i1 = 0; i1 < A_layout_t::dim_1; i1++)
{
mfem::Assign<Op>(A_data[A_layout.ind(i1,i2,i3,i4)],
B_data[B_layout.ind(i1,i2,i3,i4)]);
}
}
}
}
}
};
} // namespace mfem::internal
// Tensor or sub-tensor assign function: A {=,+=,*=} scalar_value.
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
inline void TAssign(const A_layout_t &A_layout, A_data_t &A_data,
const scalar_t value)
{
internal::TensorOps<A_layout_t::rank>::
template Assign<Op>(A_layout, A_data, value);
}
// Tensor or sub-tensor assign function: A {=,+=,*=} scalar_value.
// Host+device version.
template <AssignOp::Type Op, typename A_layout_t, typename A_data_t,
typename scalar_t>
MFEM_HOST_DEVICE
inline void TAssignHD(const A_layout_t &A_layout, A_data_t &A_data,
const scalar_t value)
{
internal::TensorOps<A_layout_t::rank>::
template AssignHD<Op>(A_layout, A_data, value);
}
// Tensor assign function: A {=,+=,*=} B that allows different input and output
// layouts. With suitable layouts this function can be used to permute
// (transpose) tensors, extract sub-tensors, etc.
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t>
inline void TAssign(const A_layout_t &A_layout, A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data)
{
internal::TensorOps<A_layout_t::rank>::
template Assign<Op>(A_layout, A_data, B_layout, B_data);
}
// classes TVector, TMatrix, TTensor3, TTensor4
template <int S, typename data_t = double, bool align = false>
struct TVector
{
public:
static const int size = S;
static const int aligned_size = align ? MFEM_ALIGN_SIZE(S,data_t) : size;
typedef data_t data_type;
data_t data[aligned_size>0?aligned_size:1];
typedef StridedLayout1D<S,1> layout_type;
static const layout_type layout;
data_t &operator[](int i) { return data[i]; }
const data_t &operator[](int i) const { return data[i]; }
template <AssignOp::Type Op>
void Assign(const data_t d)
{
TAssign<Op>(layout, data, d);
}
template <AssignOp::Type Op, typename src_data_t>
void Assign(const src_data_t &src)
{
TAssign<Op>(layout, data, layout, src);
}
template <AssignOp::Type Op, typename dest_data_t>
void AssignTo(dest_data_t &dest)
{
TAssign<Op>(layout, dest, layout, data);
}
void Set(const data_t d)
{
Assign<AssignOp::Set>(d);
}
template <typename src_data_t>
void Set(const src_data_t &src)
{
Assign<AssignOp::Set>(src);
}
template <typename dest_data_t>
void Assemble(dest_data_t &dest) const
{
AssignTo<AssignOp::Add>(dest);
}
void Scale(const data_t scale)
{
Assign<AssignOp::Mult>(scale);
}
};
template <int S, typename data_t, bool align>
const typename TVector<S,data_t,align>::layout_type
TVector<S,data_t,align>::layout = layout_type();
template <int N1, int N2, typename data_t = double, bool align = false>
struct TMatrix : public TVector<N1*N2,data_t,align>
{
typedef TVector<N1*N2,data_t,align> base_class;
using base_class::size;
using base_class::data;
typedef ColumnMajorLayout2D<N1,N2> layout_type;
static const layout_type layout;
static inline int ind(int i1, int i2) { return layout.ind(i1,i2); }
data_t &operator()(int i, int j) { return data[ind(i,j)]; }
const data_t &operator()(int i, int j) const { return data[ind(i,j)]; }
inline data_t Det() const
{
return TDet<data_t>(layout, data);
}
inline void Adjugate(TMatrix<N1,N2,data_t> &adj) const
{
TAdjugate<data_t>(layout, data, layout, adj.data);
}
// Compute the adjugate and the determinant of a (small) matrix.
inline data_t AdjDet(TMatrix<N2,N1,data_t> &adj) const
{
return TAdjDet<data_t>(layout, data, layout, adj.data);
}
};
template <int N1, int N2, typename data_t, bool align>
const typename TMatrix<N1,N2,data_t,align>::layout_type
TMatrix<N1,N2,data_t,align>::layout = layout_type();
template <int N1, int N2, int N3, typename data_t = double, bool align = false>
struct TTensor3 : TVector<N1*N2*N3,data_t,align>
{
typedef TVector<N1*N2*N3,data_t,align> base_class;
using base_class::size;
using base_class::data;
typedef ColumnMajorLayout3D<N1,N2,N3> layout_type;
static const layout_type layout;
static inline int ind(int i1, int i2, int i3)
{ return layout.ind(i1,i2,i3); }
data_t &operator()(int i, int j, int k) { return data[ind(i,j,k)]; }
const data_t &operator()(int i, int j, int k) const
{ return data[ind(i,j,k)]; }
};
template <int N1, int N2, int N3, typename data_t, bool align>
const typename TTensor3<N1,N2,N3,data_t,align>::layout_type
TTensor3<N1,N2,N3,data_t,align>::layout = layout_type();
template <int N1, int N2, int N3, int N4, typename data_t = double,
bool align = false>
struct TTensor4 : TVector<N1*N2*N3*N4,data_t,align>
{
typedef TVector<N1*N2*N3*N4,data_t,align> base_class;
using base_class::size;
using base_class::data;
typedef ColumnMajorLayout4D<N1,N2,N3,N4> layout_type;
static const layout_type layout;
static inline int ind(int i1, int i2, int i3, int i4)
{ return layout.ind(i1,i2,i3,i4); }
data_t &operator()(int i, int j, int k, int l)
{ return data[ind(i,j,k,l)]; }
const data_t &operator()(int i, int j, int k, int l) const
{ return data[ind(i,j,k,l)]; }
};
template <int N1, int N2, int N3, int N4, typename data_t, bool align>
const typename TTensor4<N1,N2,N3,N4,data_t,align>::layout_type
TTensor4<N1,N2,N3,N4,data_t,align>::layout = layout_type();
// Tensor products
// C_{i,j,k} {=|+=} \sum_s A_{s,j} B_{i,s,k}
template <bool Add,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t,
typename C_layout_t, typename C_data_t>
MFEM_ALWAYS_INLINE inline
void Mult_1_2(const A_layout_t &A_layout, const A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data,
const C_layout_t &C_layout, C_data_t &C_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2 && B_layout_t::rank == 3 &&
C_layout_t::rank == 3, "invalid ranks");
const int B3 = B_layout_t::dim_3;
const int C3 = C_layout_t::dim_3;
MFEM_STATIC_ASSERT(B3 == C3, "invalid dimentions");
for (int k = 0; k < B3; k++)
{
Mult_AB<Add>(B_layout.ind3(k), B_data,
A_layout, A_data,
C_layout.ind3(k), C_data);
}
}
// C_{i,j,k} {=|+=} \sum_s A_{i,s} B_{s,j,k}
template <bool Add,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t,
typename C_layout_t, typename C_data_t>
MFEM_ALWAYS_INLINE inline
void Mult_2_1(const A_layout_t &A_layout, const A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data,
const C_layout_t &C_layout, C_data_t &C_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2 && B_layout_t::rank == 3 &&
C_layout_t::rank == 3, "invalid ranks");
Mult_AB<Add>(A_layout, A_data,
B_layout.merge_23(), B_data,
C_layout.merge_23(), C_data);
}
// C_{i,k,j,l} {=|+=} \sum_s A_{s,i} A_{s,j} B_{k,s,l}
template <bool Add,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t,
typename C_layout_t, typename C_data_t>
MFEM_ALWAYS_INLINE inline
void TensorAssemble(const A_layout_t &A_layout, const A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data,
const C_layout_t &C_layout, C_data_t &C_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2 && B_layout_t::rank == 3 &&
C_layout_t::rank == 4, "invalid ranks");
const int A1 = A_layout_t::dim_1;
const int A2 = A_layout_t::dim_2;
const int B1 = B_layout_t::dim_1;
const int B2 = B_layout_t::dim_2;
const int B3 = B_layout_t::dim_3;
const int C1 = C_layout_t::dim_1;
const int C2 = C_layout_t::dim_2;
const int C3 = C_layout_t::dim_3;
const int C4 = C_layout_t::dim_4;
MFEM_STATIC_ASSERT(A1 == B2 && A2 == C1 && A2 == C3 && B1 == C2 && B3 == C4,
"invalid dimensions");
#if 1
// Impl == 3
MFEM_FLOPS_ADD(3*A1*A2*A2*B1*B3);
if (!Add) { TAssign<AssignOp::Set>(C_layout, C_data, 0.0); }
for (int j = 0; j < A2; j++)
{
for (int i = 0; i < A2; i++)
{
for (int l = 0; l < B3; l++)
{
for (int k = 0; k < B1; k++)
{
for (int s = 0; s < A1; s++)
{
// C(i,k,j,l) += A(s,i) * A(s,j) * B(k,s,l);
C_data[C_layout.ind(i,k,j,l)] +=
A_data[A_layout.ind(s,i)] *
A_data[A_layout.ind(s,j)] *
B_data[B_layout.ind(k,s,l)];
}
}
}
}
}
#else
// Impl == 1
if (!Add) { TAssign<AssignOp::Set>(C_layout, C_data, 0.0); }
for (int s = 0; s < A1; s++)
{
for (int i = 0; i < A2; i++)
{
for (int k = 0; k < B1; k++)
{
for (int j = 0; j < A2; j++)
{
for (int l = 0; l < B3; l++)
{
// C(i,k,j,l) += A(s,i) * A(s,j) * B(k,s,l);
C_data[C_layout.ind(i,k,j,l)] +=
A_data[A_layout.ind(s,i)] *
A_data[A_layout.ind(s,j)] *
B_data[B_layout.ind(k,s,l)];
}
}
}
}
}
#endif
}
// D_{i,k,j,l} {=|+=} \sum_s A_{i,s} B_{s,j} C_{k,s,l}
template <bool Add,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t,
typename C_layout_t, typename C_data_t,
typename D_layout_t, typename D_data_t>
MFEM_ALWAYS_INLINE inline
void TensorAssemble(const A_layout_t &A_layout, const A_data_t &A_data,
const B_layout_t &B_layout, const B_data_t &B_data,
const C_layout_t &C_layout, const C_data_t &C_data,
const D_layout_t &D_layout, D_data_t &D_data)
{
MFEM_STATIC_ASSERT(A_layout_t::rank == 2 && B_layout_t::rank == 2 &&
C_layout_t::rank == 3 && D_layout_t::rank == 4,
"invalid ranks");
const int A1 = A_layout_t::dim_1;
const int A2 = A_layout_t::dim_2;
const int B1 = B_layout_t::dim_1;
const int B2 = B_layout_t::dim_2;
const int C1 = C_layout_t::dim_1;
const int C2 = C_layout_t::dim_2;
const int C3 = C_layout_t::dim_3;
const int D1 = D_layout_t::dim_1;
const int D2 = D_layout_t::dim_2;
const int D3 = D_layout_t::dim_3;
const int D4 = D_layout_t::dim_4;
MFEM_STATIC_ASSERT(A2 == B1 && A2 == C2 && A1 == D1 && B2 == D3 &&
C1 == D2 && C3 == D4, "invalid dimensions");
#if 0
TTensor4<A1,C1,A2,C3> H;
// H_{i,k,s,l} = A_{i,s} C_{k,s,l}
for (int l = 0; l < C3; l++)
{
for (int s = 0; s < B1; s++)
{
for (int k = 0; k < C1; k++)
{
for (int i = 0; i < A1; i++)
{
H(i,k,s,l) = A_data[A_layout.ind(i,s)]*
C_data[C_layout.ind(k,s,l)];
}
}
}
}
// D_{(i,k),j,l} = \sum_s B_{s,j} H_{(i,k),s,l}
Mult_1_2<Add>(B_layout, B_data, H.layout.merge_12(), H,
D_layout.merge_12(), D_data);
#elif 1
MFEM_FLOPS_ADD(A1*B1*C1*C3); // computation of H(l)
for (int l = 0; l < C3; l++)
{
TTensor3<A1,C1,A2,typename C_data_t::data_type> H;
// H(l)_{i,k,s} = A_{i,s} C_{k,s,l}
for (int s = 0; s < B1; s++)
{
for (int k = 0; k < C1; k++)
{
for (int i = 0; i < A1; i++)
{
H(i,k,s) = A_data[A_layout.ind(i,s)]*
C_data[C_layout.ind(k,s,l)];
}
}
}
// D_{(i,k),j,l} = \sum_s H(l)_{(i,k),s} B_{s,j}
Mult_AB<Add>(H.layout.merge_12(), H, B_layout, B_data,
D_layout.merge_12().ind3(l), D_data);
}
#else
TTensor4<B1,C1,B2,C3> F;
for (int l = 0; l < C3; l++)
{
for (int j = 0; j < B2; j++)
{
for (int k = 0; k < C1; k++)
{
for (int s = 0; s < B1; s++)
{
F(s,k,j,l) = B_data[B_layout.ind(s,j)]*
C_data[C_layout.ind(k,s,l)];
}
}
}
}
Mult_AB<Add>(A_layout, A_data, F.layout.merge_34().merge_23(), F,
D_layout.merge_34().merge_23(), D_data);
#endif
}
// C_{i,j,k,l} {=|+=} A_{i,j,k} B_{j,l}
template <AssignOp::Type Op,
typename A_layout_t, typename A_data_t,
typename B_layout_t, typename B_data_t,
typename C_layout_t, typename C_data_t>
MFEM_ALWAYS_INLINE inline
void TensorProduct(const A_layout_t &a, const A_data_t &A,
const B_layout_t &b, const B_data_t &B,
const C_layout_t &c, C_data_t &C)
{
const int A1 = A_layout_t::dim_1;
const int A2 = A_layout_t::dim_2;
const int A3 = A_layout_t::dim_3;
const int B1 = B_layout_t::dim_1;
const int B2 = B_layout_t::dim_2;
const int C1 = C_layout_t::dim_1;
const int C2 = C_layout_t::dim_2;
const int C3 = C_layout_t::dim_3;
const int C4 = C_layout_t::dim_4;
MFEM_STATIC_ASSERT(A1 == C1 && A2 == B1 && A2 == C2 && A3 == C3 && B2 == C4,
"invalid dimensions");
MFEM_FLOPS_ADD(A1*A2*A3*B2);
for (int l = 0; l < B2; l++)
{
for (int k = 0; k < A3; k++)
{
for (int j = 0; j < A2; j++)
{
for (int i = 0; i < A1; i++)
{
mfem::Assign<Op>(C[c.ind(i,j,k,l)],
A[a.ind(i,j,k)]*B[b.ind(j,l)]);
}
}
}
}
}
} // namespace mfem
#endif // MFEM_TEMPLATE_TENSOR