GitHub-Proxy/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2025-09-12 09:23:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

Allow custom TENSOR_CONTRACTION_DISPATCH macro.

Browse Source 
Currently TF lite needs to hack around with the Tensor headers in order
to customize the contraction dispatch method. Here we add simple `#ifndef`
guards to allow them to provide their own dispatch prior to inclusion.
This commit is contained in:
Antonio Sanchez 2021-06-11 08:30:41 -07:00 committed by Rasmus Munk Larsen
parent fc87e2cbaa
commit 6aec83263d
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h
View File

@ -633,6 +633,7 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator
} }
#endif // EIGEN_USE_THREADS #endif // EIGEN_USE_THREADS
#ifndef TENSOR_CONTRACTION_DISPATCH
#define TENSOR_CONTRACTION_DISPATCH(METHOD, ALIGNMENT, ARGS) \ #define TENSOR_CONTRACTION_DISPATCH(METHOD, ALIGNMENT, ARGS) \
if (this->m_lhs_inner_dim_contiguous) { \ if (this->m_lhs_inner_dim_contiguous) { \
if (this->m_rhs_inner_dim_contiguous) { \ if (this->m_rhs_inner_dim_contiguous) { \
@ -663,7 +664,9 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator
} \ } \
} \ } \
} }
#endif
#ifndef TENSOR_CONTRACTION_ASYNC_DISPATCH
#define TENSOR_CONTRACTION_ASYNC_DISPATCH(METHOD, DONE, ALIGNMENT, ARGS, FN) \ #define TENSOR_CONTRACTION_ASYNC_DISPATCH(METHOD, DONE, ALIGNMENT, ARGS, FN) \
if (this->m_lhs_inner_dim_contiguous) { \ if (this->m_lhs_inner_dim_contiguous) { \
if (this->m_rhs_inner_dim_contiguous) { \ if (this->m_rhs_inner_dim_contiguous) { \
@ -694,6 +697,7 @@ struct TensorContractionEvaluatorBase : internal::no_assignment_operator
} \ } \
} \ } \
} }
#endif
EIGEN_DEVICE_FUNC void evalTo(Scalar* buffer) const { EIGEN_DEVICE_FUNC void evalTo(Scalar* buffer) const {
static_cast<const Derived*>(this)->template evalProduct<Unaligned>(buffer); static_cast<const Derived*>(this)->template evalProduct<Unaligned>(buffer);