devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add support for custom packed Lhs/Rhs blocks in tensor contractions

Browse Source
This commit is contained in:
Eugene Zhulenev
2019-04-01 11:47:31 -07:00
parent 45e65fbb77
commit 4e2f6de1a8
3 changed files with 247 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
unsupported/Eigen/CXX11/src/Tensor/TensorContractionMapper.h
View File

@@ -24,12 +24,17 @@ enum {
*/
/// The make pointer class is used by sycl in order to build the mapper class on the device. For other platform the default make pointer is used which
/// is scalar * for CoeffLoader.
template <typename Tensor, bool HasRawAccess, template <class> class MakePointer_ = MakePointer> struct CoeffLoader;
template<typename Scalar, typename Index, int side, typename Tensor, typename nocontract_t, typename contract_t,
int packet_size, bool inner_dim_contiguous, bool inner_dim_reordered, int Alignment,
template <class> class MakePointer_ = MakePointer> class BaseTensorContractionMapper;
template <typename Tensor, bool HasRawAccess, template <class> class MakePointer_ = MakePointer>
struct CoeffLoader;
template <typename Tensor, bool HasRawAccess, template <class> class MakePointer_> struct CoeffLoader {
template <typename Scalar, typename Index, int side, typename Tensor,
typename nocontract_t, typename contract_t, int packet_size,
bool inner_dim_contiguous, bool inner_dim_reordered, int Alignment,
template <class> class MakePointer_ = MakePointer>
class BaseTensorContractionMapper;
template <typename Tensor, bool HasRawAccess, template <class> class MakePointer_>
struct CoeffLoader {
enum {
DirectOffsets = false
};
@@ -40,6 +45,12 @@ template <typename Tensor, bool HasRawAccess, template <class> class MakePointer
eigen_assert(false && "unsupported");
}
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const typename MakePointer_<const typename Tensor::Scalar>::Type
data() const {
eigen_assert(false && "unsupported");
return NULL;
}
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE typename Tensor::Scalar coeff(typename Tensor::Index index) const { return m_tensor.coeff(index); }
template<int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
@@ -48,12 +59,12 @@ template <typename Tensor, bool HasRawAccess, template <class> class MakePointer
return m_tensor.template packet<LoadMode>(index);
}
private:
const Tensor m_tensor;
};
template <typename Tensor, template <class> class MakePointer_> struct CoeffLoader<Tensor, true, MakePointer_> {
template <typename Tensor, template <class> class MakePointer_>
struct CoeffLoader<Tensor, true, MakePointer_> {
enum {
DirectOffsets = true
};
@@ -64,6 +75,11 @@ template <typename Tensor, template <class> class MakePointer_> struct CoeffLoad
m_data += offset;
}
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const typename MakePointer_<const typename Tensor::Scalar>::Type
data() const {
return m_data;
}
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE typename Tensor::Scalar coeff(typename Tensor::Index index) const { return loadConstant(m_data+index); }
template<int LoadMode> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
@@ -214,6 +230,17 @@ class SimpleTensorContractionMapper {
return ((side == Lhs) && inner_dim_contiguous && array_size<contract_t>::value > 0) ? m_contract_strides[0] : 1;
}
const CoeffLoader<Tensor, Tensor::RawAccess, MakePointer_>& tensor() const {
return m_tensor;
}
const nocontract_t& nocontract_strides() const {
return m_nocontract_strides;
}
const nocontract_t& ij_strides() const { return m_ij_strides; }
const contract_t& contract_strides() const { return m_contract_strides; }
const contract_t& k_strides() const { return m_k_strides; }
protected:
CoeffLoader<Tensor, Tensor::RawAccess, MakePointer_> m_tensor;
const nocontract_t m_nocontract_strides;
@@ -445,6 +472,10 @@ class TensorContractionSubMapper {
return false;
}
const ParentMapper& base_mapper() const { return m_base_mapper; }
Index vert_offset() const { return m_vert_offset; }
Index horiz_offset() const { return m_horiz_offset; }
private:
ParentMapper m_base_mapper;
const Index m_vert_offset;