Extend support for self-adjoint matrices on CUDA devices

fwyzard · smuzaffar · commit f7f365a553ff · 2021-07-15T19:35:23.000+02:00
diff --git a/Eigen/src/Core/ProductEvaluators.h b/Eigen/src/Core/ProductEvaluators.h
@@ -813,6 +813,7 @@ struct generic_product_impl<Lhs,Rhs,DenseShape,SelfAdjointShape,ProductTag>
   typedef typename Product<Lhs,Rhs>::Scalar Scalar;
 
   template<typename Dest>
+  EIGEN_DEVICE_FUNC
   static void scaleAndAddTo(Dest& dst, const Lhs& lhs, const Rhs& rhs, const Scalar& alpha)
   {
     selfadjoint_product_impl<Lhs,0,Lhs::IsVectorAtCompileTime,typename Rhs::MatrixType,Rhs::Mode,false>::run(dst, lhs, rhs.nestedExpression(), alpha);
diff --git a/Eigen/src/Core/TriangularMatrix.h b/Eigen/src/Core/TriangularMatrix.h
@@ -961,6 +961,8 @@ template< typename DstXprType, typename Lhs, typename Rhs, typename Scalar>
 struct Assignment<DstXprType, Product<Lhs,Rhs,DefaultProduct>, internal::assign_op<Scalar,typename Product<Lhs,Rhs,DefaultProduct>::Scalar>, Dense2Triangular>
 {
   typedef Product<Lhs,Rhs,DefaultProduct> SrcXprType;
+
+  EIGEN_DEVICE_FUNC
   static void run(DstXprType &dst, const SrcXprType &src, const internal::assign_op<Scalar,typename SrcXprType::Scalar> &)
   {
     Index dstRows = src.rows();
diff --git a/Eigen/src/Core/products/GeneralBlockPanelKernel.h b/Eigen/src/Core/products/GeneralBlockPanelKernel.h
@@ -2361,11 +2361,11 @@ template<typename Scalar, typename Index, typename DataMapper, int Pack1, int Pa
 struct gemm_pack_lhs<Scalar, Index, DataMapper, Pack1, Pack2, Packet, RowMajor, Conjugate, PanelMode>
 {
   typedef typename DataMapper::LinearMapper LinearMapper;
-  EIGEN_DONT_INLINE void operator()(Scalar* blockA, const DataMapper& lhs, Index depth, Index rows, Index stride=0, Index offset=0);
+  EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE void operator()(Scalar* blockA, const DataMapper& lhs, Index depth, Index rows, Index stride=0, Index offset=0);
 };
 
 template<typename Scalar, typename Index, typename DataMapper, int Pack1, int Pack2, typename Packet, bool Conjugate, bool PanelMode>
-EIGEN_DONT_INLINE void gemm_pack_lhs<Scalar, Index, DataMapper, Pack1, Pack2, Packet, RowMajor, Conjugate, PanelMode>
+EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE void gemm_pack_lhs<Scalar, Index, DataMapper, Pack1, Pack2, Packet, RowMajor, Conjugate, PanelMode>
   ::operator()(Scalar* blockA, const DataMapper& lhs, Index depth, Index rows, Index stride, Index offset)
 {
   typedef typename unpacket_traits<Packet>::half HalfPacket;
@@ -2624,8 +2624,8 @@ struct gemm_pack_rhs<Scalar, Index, DataMapper, nr, RowMajor, Conjugate, PanelMo
   typedef typename DataMapper::LinearMapper LinearMapper;
   enum { PacketSize = packet_traits<Scalar>::size,
          HalfPacketSize = unpacket_traits<HalfPacket>::size,
-		 QuarterPacketSize = unpacket_traits<QuarterPacket>::size};
-  EIGEN_DONT_INLINE void operator()(Scalar* blockB, const DataMapper& rhs, Index depth, Index cols, Index stride=0, Index offset=0)
+         QuarterPacketSize = unpacket_traits<QuarterPacket>::size};
+  EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE void operator()(Scalar* blockB, const DataMapper& rhs, Index depth, Index cols, Index stride=0, Index offset=0)
   {
     EIGEN_ASM_COMMENT("EIGEN PRODUCT PACK RHS ROWMAJOR");
     EIGEN_UNUSED_VARIABLE(stride);
diff --git a/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h b/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h
@@ -60,6 +60,8 @@ template <typename Index, typename LhsScalar, int LhsStorageOrder, bool Conjugat
 struct general_matrix_matrix_triangular_product<Index,LhsScalar,LhsStorageOrder,ConjugateLhs,RhsScalar,RhsStorageOrder,ConjugateRhs,ColMajor,ResInnerStride,UpLo,Version>
 {
   typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;
+
+  EIGEN_DEVICE_FUNC
   static EIGEN_STRONG_INLINE void run(Index size, Index depth,const LhsScalar* _lhs, Index lhsStride,
                                       const RhsScalar* _rhs, Index rhsStride,
                                       ResScalar* _res, Index resIncr, Index resStride,
@@ -144,6 +146,7 @@ struct tribb_kernel
   enum {
     BlockSize  = meta_least_common_multiple<EIGEN_PLAIN_ENUM_MAX(mr,nr),EIGEN_PLAIN_ENUM_MIN(mr,nr)>::ret
   };
+  EIGEN_DEVICE_FUNC
   void operator()(ResScalar* _res, Index resIncr, Index resStride, const LhsScalar* blockA, const RhsScalar* blockB, Index size, Index depth, const ResScalar& alpha)
   {
     typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;
@@ -252,6 +255,7 @@ struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,true>
 template<typename MatrixType, typename ProductType, int UpLo>
 struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,false>
 {
+  EIGEN_DEVICE_FUNC
   static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)
   {
     typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;
diff --git a/Eigen/src/Core/products/SelfadjointMatrixMatrix.h b/Eigen/src/Core/products/SelfadjointMatrixMatrix.h
@@ -101,6 +101,8 @@ template<typename Scalar, typename Index, int nr, int StorageOrder>
 struct symm_pack_rhs
 {
   enum { PacketSize = packet_traits<Scalar>::size };
+
+  EIGEN_DEVICE_FUNC
   void operator()(Scalar* blockB, const Scalar* _rhs, Index rhsStride, Index rows, Index cols, Index k2)
   {
     Index end_k = k2 + rows;
@@ -423,6 +425,7 @@ template <typename Scalar, typename Index,
 struct product_selfadjoint_matrix<Scalar,Index,LhsStorageOrder,false,ConjugateLhs, RhsStorageOrder,true,ConjugateRhs,ColMajor,ResInnerStride>
 {
 
+  EIGEN_DEVICE_FUNC
   static EIGEN_DONT_INLINE void run(
     Index rows, Index cols,
     const Scalar* _lhs, Index lhsStride,
@@ -435,7 +438,7 @@ template <typename Scalar, typename Index,
           int LhsStorageOrder, bool ConjugateLhs,
           int RhsStorageOrder, bool ConjugateRhs,
           int ResInnerStride>
-EIGEN_DONT_INLINE void product_selfadjoint_matrix<Scalar,Index,LhsStorageOrder,false,ConjugateLhs, RhsStorageOrder,true,ConjugateRhs,ColMajor,ResInnerStride>::run(
+EIGEN_DEVICE_FUNC EIGEN_DONT_INLINE void product_selfadjoint_matrix<Scalar,Index,LhsStorageOrder,false,ConjugateLhs, RhsStorageOrder,true,ConjugateRhs,ColMajor,ResInnerStride>::run(
     Index rows, Index cols,
     const Scalar* _lhs, Index lhsStride,
     const Scalar* _rhs, Index rhsStride,
@@ -505,6 +508,7 @@ struct selfadjoint_product_impl<Lhs,LhsMode,false,Rhs,RhsMode,false>
   };
   
   template<typename Dest>
+  EIGEN_DEVICE_FUNC
   static void run(Dest &dst, const Lhs &a_lhs, const Rhs &a_rhs, const Scalar& alpha)
   {
     eigen_assert(dst.rows()==a_lhs.rows() && dst.cols()==a_rhs.cols());
diff --git a/Eigen/src/Core/products/SelfadjointMatrixVector.h b/Eigen/src/Core/products/SelfadjointMatrixVector.h
@@ -246,7 +246,8 @@ struct selfadjoint_product_impl<Lhs,0,true,Rhs,RhsMode,false>
   enum { RhsUpLo = RhsMode&(Upper|Lower)  };
 
   template<typename Dest>
-  static void run(Dest& dest, const Lhs &a_lhs, const Rhs &a_rhs, const Scalar& alpha)
+  static EIGEN_DEVICE_FUNC
+  void run(Dest& dest, const Lhs &a_lhs, const Rhs &a_rhs, const Scalar& alpha)
   {
     // let's simply transpose the product
     Transpose<Dest> destT(dest);
diff --git a/Eigen/src/Core/util/BlasUtil.h b/Eigen/src/Core/util/BlasUtil.h
@@ -594,8 +594,8 @@ struct blas_traits<Transpose<NestedXpr> >
   enum {
     IsTransposed = Base::IsTransposed ? 0 : 1
   };
-  static inline ExtractType extract(const XprType& x) { return ExtractType(Base::extract(x.nestedExpression())); }
-  static inline Scalar extractScalarFactor(const XprType& x) { return Base::extractScalarFactor(x.nestedExpression()); }
+  EIGEN_DEVICE_FUNC static inline ExtractType extract(const XprType& x) { return ExtractType(Base::extract(x.nestedExpression())); }
+  EIGEN_DEVICE_FUNC static inline Scalar extractScalarFactor(const XprType& x) { return Base::extractScalarFactor(x.nestedExpression()); }
 };
 
 template<typename T>

Original file line number	Diff line number	Diff line change
`@@ -813,6 +813,7 @@ struct generic_product_impl<Lhs,Rhs,DenseShape,SelfAdjointShape,ProductTag>`
`813`	`813`	`typedef typename Product<Lhs,Rhs>::Scalar Scalar;`
`814`	`814`
`815`	`815`	`template<typename Dest>`
	`816`	`+ EIGEN_DEVICE_FUNC`
`816`	`817`	`static void scaleAndAddTo(Dest& dst, const Lhs& lhs, const Rhs& rhs, const Scalar& alpha)`
`817`	`818`	`{`
`818`	`819`	`selfadjoint_product_impl<Lhs,0,Lhs::IsVectorAtCompileTime,typename Rhs::MatrixType,Rhs::Mode,false>::run(dst, lhs, rhs.nestedExpression(), alpha);`
Original file line number	Diff line number	Diff line change
`@@ -961,6 +961,8 @@ template< typename DstXprType, typename Lhs, typename Rhs, typename Scalar>`
`961`	`961`	`struct Assignment<DstXprType, Product<Lhs,Rhs,DefaultProduct>, internal::assign_op<Scalar,typename Product<Lhs,Rhs,DefaultProduct>::Scalar>, Dense2Triangular>`
`962`	`962`	`{`
`963`	`963`	`typedef Product<Lhs,Rhs,DefaultProduct> SrcXprType;`
	`964`	`+`
	`965`	`+ EIGEN_DEVICE_FUNC`
`964`	`966`	`static void run(DstXprType &dst, const SrcXprType &src, const internal::assign_op<Scalar,typename SrcXprType::Scalar> &)`
`965`	`967`	`{`
`966`	`968`	`Index dstRows = src.rows();`
Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,8 @@ template <typename Index, typename LhsScalar, int LhsStorageOrder, bool Conjugat`
`60`	`60`	`struct general_matrix_matrix_triangular_product<Index,LhsScalar,LhsStorageOrder,ConjugateLhs,RhsScalar,RhsStorageOrder,ConjugateRhs,ColMajor,ResInnerStride,UpLo,Version>`
`61`	`61`	`{`
`62`	`62`	`typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;`
	`63`	`+`
	`64`	`+ EIGEN_DEVICE_FUNC`
`63`	`65`	`static EIGEN_STRONG_INLINE void run(Index size, Index depth,const LhsScalar* _lhs, Index lhsStride,`
`64`	`66`	`const RhsScalar* _rhs, Index rhsStride,`
`65`	`67`	`ResScalar* _res, Index resIncr, Index resStride,`
`@@ -144,6 +146,7 @@ struct tribb_kernel`
`144`	`146`	`enum {`
`145`	`147`	`BlockSize = meta_least_common_multiple<EIGEN_PLAIN_ENUM_MAX(mr,nr),EIGEN_PLAIN_ENUM_MIN(mr,nr)>::ret`
`146`	`148`	`};`
	`149`	`+ EIGEN_DEVICE_FUNC`
`147`	`150`	`void operator()(ResScalar* _res, Index resIncr, Index resStride, const LhsScalar* blockA, const RhsScalar* blockB, Index size, Index depth, const ResScalar& alpha)`
`148`	`151`	`{`
`149`	`152`	`typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;`
`@@ -252,6 +255,7 @@ struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,true>`
`252`	`255`	`template<typename MatrixType, typename ProductType, int UpLo>`
`253`	`256`	`struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,false>`
`254`	`257`	`{`
	`258`	`+ EIGEN_DEVICE_FUNC`
`255`	`259`	`static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)`
`256`	`260`	`{`
`257`	`261`	`typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;`
Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,8 @@ struct selfadjoint_product_impl<Lhs,0,true,Rhs,RhsMode,false>`
`246`	`246`	`enum { RhsUpLo = RhsMode&(Upper\|Lower) };`
`247`	`247`
`248`	`248`	`template<typename Dest>`
`249`		`- static void run(Dest& dest, const Lhs &a_lhs, const Rhs &a_rhs, const Scalar& alpha)`
	`249`	`+ static EIGEN_DEVICE_FUNC`
	`250`	`+ void run(Dest& dest, const Lhs &a_lhs, const Rhs &a_rhs, const Scalar& alpha)`
`250`	`251`	`{`
`251`	`252`	`// let's simply transpose the product`
`252`	`253`	`Transpose<Dest> destT(dest);`