xref: /aosp_15_r20/external/eigen/Eigen/src/Core/products/GeneralMatrixMatrixTriangular.h (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2009-2010 Gael Guennebaud <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #ifndef EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H
11 #define EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H
12 
13 namespace Eigen {
14 
15 template<typename Scalar, typename Index, int StorageOrder, int UpLo, bool ConjLhs, bool ConjRhs>
16 struct selfadjoint_rank1_update;
17 
18 namespace internal {
19 
20 /**********************************************************************
21 * This file implements a general A * B product while
22 * evaluating only one triangular part of the product.
23 * This is a more general version of self adjoint product (C += A A^T)
24 * as the level 3 SYRK Blas routine.
25 **********************************************************************/
26 
27 // forward declarations (defined at the end of this file)
28 template<typename LhsScalar, typename RhsScalar, typename Index, int mr, int nr, bool ConjLhs, bool ConjRhs, int ResInnerStride, int UpLo>
29 struct tribb_kernel;
30 
31 /* Optimized matrix-matrix product evaluating only one triangular half */
32 template <typename Index,
33           typename LhsScalar, int LhsStorageOrder, bool ConjugateLhs,
34           typename RhsScalar, int RhsStorageOrder, bool ConjugateRhs,
35                               int ResStorageOrder, int ResInnerStride, int  UpLo, int Version = Specialized>
36 struct general_matrix_matrix_triangular_product;
37 
38 // as usual if the result is row major => we transpose the product
39 template <typename Index, typename LhsScalar, int LhsStorageOrder, bool ConjugateLhs,
40                           typename RhsScalar, int RhsStorageOrder, bool ConjugateRhs,
41                           int ResInnerStride, int  UpLo, int Version>
42 struct general_matrix_matrix_triangular_product<Index,LhsScalar,LhsStorageOrder,ConjugateLhs,RhsScalar,RhsStorageOrder,ConjugateRhs,RowMajor,ResInnerStride,UpLo,Version>
43 {
44   typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;
45   static EIGEN_STRONG_INLINE void run(Index size, Index depth,const LhsScalar* lhs, Index lhsStride,
46                                       const RhsScalar* rhs, Index rhsStride, ResScalar* res, Index resIncr, Index resStride,
47                                       const ResScalar& alpha, level3_blocking<RhsScalar,LhsScalar>& blocking)
48   {
49     general_matrix_matrix_triangular_product<Index,
50         RhsScalar, RhsStorageOrder==RowMajor ? ColMajor : RowMajor, ConjugateRhs,
51         LhsScalar, LhsStorageOrder==RowMajor ? ColMajor : RowMajor, ConjugateLhs,
52         ColMajor, ResInnerStride, UpLo==Lower?Upper:Lower>
53       ::run(size,depth,rhs,rhsStride,lhs,lhsStride,res,resIncr,resStride,alpha,blocking);
54   }
55 };
56 
57 template <typename Index, typename LhsScalar, int LhsStorageOrder, bool ConjugateLhs,
58                           typename RhsScalar, int RhsStorageOrder, bool ConjugateRhs,
59                           int ResInnerStride, int  UpLo, int Version>
60 struct general_matrix_matrix_triangular_product<Index,LhsScalar,LhsStorageOrder,ConjugateLhs,RhsScalar,RhsStorageOrder,ConjugateRhs,ColMajor,ResInnerStride,UpLo,Version>
61 {
62   typedef typename ScalarBinaryOpTraits<LhsScalar, RhsScalar>::ReturnType ResScalar;
63   static EIGEN_STRONG_INLINE void run(Index size, Index depth,const LhsScalar* _lhs, Index lhsStride,
64                                       const RhsScalar* _rhs, Index rhsStride,
65                                       ResScalar* _res, Index resIncr, Index resStride,
66                                       const ResScalar& alpha, level3_blocking<LhsScalar,RhsScalar>& blocking)
67   {
68     typedef gebp_traits<LhsScalar,RhsScalar> Traits;
69 
70     typedef const_blas_data_mapper<LhsScalar, Index, LhsStorageOrder> LhsMapper;
71     typedef const_blas_data_mapper<RhsScalar, Index, RhsStorageOrder> RhsMapper;
72     typedef blas_data_mapper<typename Traits::ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;
73     LhsMapper lhs(_lhs,lhsStride);
74     RhsMapper rhs(_rhs,rhsStride);
75     ResMapper res(_res, resStride, resIncr);
76 
77     Index kc = blocking.kc();
78     Index mc = (std::min)(size,blocking.mc());
79 
80     // !!! mc must be a multiple of nr:
81     if(mc > Traits::nr)
82       mc = (mc/Traits::nr)*Traits::nr;
83 
84     std::size_t sizeA = kc*mc;
85     std::size_t sizeB = kc*size;
86 
87     ei_declare_aligned_stack_constructed_variable(LhsScalar, blockA, sizeA, blocking.blockA());
88     ei_declare_aligned_stack_constructed_variable(RhsScalar, blockB, sizeB, blocking.blockB());
89 
90     gemm_pack_lhs<LhsScalar, Index, LhsMapper, Traits::mr, Traits::LhsProgress, typename Traits::LhsPacket4Packing, LhsStorageOrder> pack_lhs;
91     gemm_pack_rhs<RhsScalar, Index, RhsMapper, Traits::nr, RhsStorageOrder> pack_rhs;
92     gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs> gebp;
93     tribb_kernel<LhsScalar, RhsScalar, Index, Traits::mr, Traits::nr, ConjugateLhs, ConjugateRhs, ResInnerStride, UpLo> sybb;
94 
95     for(Index k2=0; k2<depth; k2+=kc)
96     {
97       const Index actual_kc = (std::min)(k2+kc,depth)-k2;
98 
99       // note that the actual rhs is the transpose/adjoint of mat
100       pack_rhs(blockB, rhs.getSubMapper(k2,0), actual_kc, size);
101 
102       for(Index i2=0; i2<size; i2+=mc)
103       {
104         const Index actual_mc = (std::min)(i2+mc,size)-i2;
105 
106         pack_lhs(blockA, lhs.getSubMapper(i2, k2), actual_kc, actual_mc);
107 
108         // the selected actual_mc * size panel of res is split into three different part:
109         //  1 - before the diagonal => processed with gebp or skipped
110         //  2 - the actual_mc x actual_mc symmetric block => processed with a special kernel
111         //  3 - after the diagonal => processed with gebp or skipped
112         if (UpLo==Lower)
113           gebp(res.getSubMapper(i2, 0), blockA, blockB, actual_mc, actual_kc,
114                (std::min)(size,i2), alpha, -1, -1, 0, 0);
115 
116         sybb(_res+resStride*i2 + resIncr*i2, resIncr, resStride, blockA, blockB + actual_kc*i2, actual_mc, actual_kc, alpha);
117 
118         if (UpLo==Upper)
119         {
120           Index j2 = i2+actual_mc;
121           gebp(res.getSubMapper(i2, j2), blockA, blockB+actual_kc*j2, actual_mc,
122                actual_kc, (std::max)(Index(0), size-j2), alpha, -1, -1, 0, 0);
123         }
124       }
125     }
126   }
127 };
128 
129 // Optimized packed Block * packed Block product kernel evaluating only one given triangular part
130 // This kernel is built on top of the gebp kernel:
131 // - the current destination block is processed per panel of actual_mc x BlockSize
132 //   where BlockSize is set to the minimal value allowing gebp to be as fast as possible
133 // - then, as usual, each panel is split into three parts along the diagonal,
134 //   the sub blocks above and below the diagonal are processed as usual,
135 //   while the triangular block overlapping the diagonal is evaluated into a
136 //   small temporary buffer which is then accumulated into the result using a
137 //   triangular traversal.
138 template<typename LhsScalar, typename RhsScalar, typename Index, int mr, int nr, bool ConjLhs, bool ConjRhs, int ResInnerStride, int UpLo>
139 struct tribb_kernel
140 {
141   typedef gebp_traits<LhsScalar,RhsScalar,ConjLhs,ConjRhs> Traits;
142   typedef typename Traits::ResScalar ResScalar;
143 
144   enum {
145     BlockSize  = meta_least_common_multiple<EIGEN_PLAIN_ENUM_MAX(mr,nr),EIGEN_PLAIN_ENUM_MIN(mr,nr)>::ret
146   };
147   void operator()(ResScalar* _res, Index resIncr, Index resStride, const LhsScalar* blockA, const RhsScalar* blockB, Index size, Index depth, const ResScalar& alpha)
148   {
149     typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned, ResInnerStride> ResMapper;
150     typedef blas_data_mapper<ResScalar, Index, ColMajor, Unaligned> BufferMapper;
151     ResMapper res(_res, resStride, resIncr);
152     gebp_kernel<LhsScalar, RhsScalar, Index, ResMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel1;
153     gebp_kernel<LhsScalar, RhsScalar, Index, BufferMapper, mr, nr, ConjLhs, ConjRhs> gebp_kernel2;
154 
155     Matrix<ResScalar,BlockSize,BlockSize,ColMajor> buffer((internal::constructor_without_unaligned_array_assert()));
156 
157     // let's process the block per panel of actual_mc x BlockSize,
158     // again, each is split into three parts, etc.
159     for (Index j=0; j<size; j+=BlockSize)
160     {
161       Index actualBlockSize = std::min<Index>(BlockSize,size - j);
162       const RhsScalar* actual_b = blockB+j*depth;
163 
164       if(UpLo==Upper)
165         gebp_kernel1(res.getSubMapper(0, j), blockA, actual_b, j, depth, actualBlockSize, alpha,
166                      -1, -1, 0, 0);
167 
168       // selfadjoint micro block
169       {
170         Index i = j;
171         buffer.setZero();
172         // 1 - apply the kernel on the temporary buffer
173         gebp_kernel2(BufferMapper(buffer.data(), BlockSize), blockA+depth*i, actual_b, actualBlockSize, depth, actualBlockSize, alpha,
174                      -1, -1, 0, 0);
175 
176         // 2 - triangular accumulation
177         for(Index j1=0; j1<actualBlockSize; ++j1)
178         {
179           typename ResMapper::LinearMapper r = res.getLinearMapper(i,j+j1);
180           for(Index i1=UpLo==Lower ? j1 : 0;
181               UpLo==Lower ? i1<actualBlockSize : i1<=j1; ++i1)
182             r(i1) += buffer(i1,j1);
183         }
184       }
185 
186       if(UpLo==Lower)
187       {
188         Index i = j+actualBlockSize;
189         gebp_kernel1(res.getSubMapper(i, j), blockA+depth*i, actual_b, size-i,
190                      depth, actualBlockSize, alpha, -1, -1, 0, 0);
191       }
192     }
193   }
194 };
195 
196 } // end namespace internal
197 
198 // high level API
199 
200 template<typename MatrixType, typename ProductType, int UpLo, bool IsOuterProduct>
201 struct general_product_to_triangular_selector;
202 
203 
204 template<typename MatrixType, typename ProductType, int UpLo>
205 struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,true>
206 {
207   static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)
208   {
209     typedef typename MatrixType::Scalar Scalar;
210 
211     typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;
212     typedef internal::blas_traits<Lhs> LhsBlasTraits;
213     typedef typename LhsBlasTraits::DirectLinearAccessType ActualLhs;
214     typedef typename internal::remove_all<ActualLhs>::type _ActualLhs;
215     typename internal::add_const_on_value_type<ActualLhs>::type actualLhs = LhsBlasTraits::extract(prod.lhs());
216 
217     typedef typename internal::remove_all<typename ProductType::RhsNested>::type Rhs;
218     typedef internal::blas_traits<Rhs> RhsBlasTraits;
219     typedef typename RhsBlasTraits::DirectLinearAccessType ActualRhs;
220     typedef typename internal::remove_all<ActualRhs>::type _ActualRhs;
221     typename internal::add_const_on_value_type<ActualRhs>::type actualRhs = RhsBlasTraits::extract(prod.rhs());
222 
223     Scalar actualAlpha = alpha * LhsBlasTraits::extractScalarFactor(prod.lhs().derived()) * RhsBlasTraits::extractScalarFactor(prod.rhs().derived());
224 
225     if(!beta)
226       mat.template triangularView<UpLo>().setZero();
227 
228     enum {
229       StorageOrder = (internal::traits<MatrixType>::Flags&RowMajorBit) ? RowMajor : ColMajor,
230       UseLhsDirectly = _ActualLhs::InnerStrideAtCompileTime==1,
231       UseRhsDirectly = _ActualRhs::InnerStrideAtCompileTime==1
232     };
233 
234     internal::gemv_static_vector_if<Scalar,Lhs::SizeAtCompileTime,Lhs::MaxSizeAtCompileTime,!UseLhsDirectly> static_lhs;
235     ei_declare_aligned_stack_constructed_variable(Scalar, actualLhsPtr, actualLhs.size(),
236       (UseLhsDirectly ? const_cast<Scalar*>(actualLhs.data()) : static_lhs.data()));
237     if(!UseLhsDirectly) Map<typename _ActualLhs::PlainObject>(actualLhsPtr, actualLhs.size()) = actualLhs;
238 
239     internal::gemv_static_vector_if<Scalar,Rhs::SizeAtCompileTime,Rhs::MaxSizeAtCompileTime,!UseRhsDirectly> static_rhs;
240     ei_declare_aligned_stack_constructed_variable(Scalar, actualRhsPtr, actualRhs.size(),
241       (UseRhsDirectly ? const_cast<Scalar*>(actualRhs.data()) : static_rhs.data()));
242     if(!UseRhsDirectly) Map<typename _ActualRhs::PlainObject>(actualRhsPtr, actualRhs.size()) = actualRhs;
243 
244 
245     selfadjoint_rank1_update<Scalar,Index,StorageOrder,UpLo,
246                               LhsBlasTraits::NeedToConjugate && NumTraits<Scalar>::IsComplex,
247                               RhsBlasTraits::NeedToConjugate && NumTraits<Scalar>::IsComplex>
248           ::run(actualLhs.size(), mat.data(), mat.outerStride(), actualLhsPtr, actualRhsPtr, actualAlpha);
249   }
250 };
251 
252 template<typename MatrixType, typename ProductType, int UpLo>
253 struct general_product_to_triangular_selector<MatrixType,ProductType,UpLo,false>
254 {
255   static void run(MatrixType& mat, const ProductType& prod, const typename MatrixType::Scalar& alpha, bool beta)
256   {
257     typedef typename internal::remove_all<typename ProductType::LhsNested>::type Lhs;
258     typedef internal::blas_traits<Lhs> LhsBlasTraits;
259     typedef typename LhsBlasTraits::DirectLinearAccessType ActualLhs;
260     typedef typename internal::remove_all<ActualLhs>::type _ActualLhs;
261     typename internal::add_const_on_value_type<ActualLhs>::type actualLhs = LhsBlasTraits::extract(prod.lhs());
262 
263     typedef typename internal::remove_all<typename ProductType::RhsNested>::type Rhs;
264     typedef internal::blas_traits<Rhs> RhsBlasTraits;
265     typedef typename RhsBlasTraits::DirectLinearAccessType ActualRhs;
266     typedef typename internal::remove_all<ActualRhs>::type _ActualRhs;
267     typename internal::add_const_on_value_type<ActualRhs>::type actualRhs = RhsBlasTraits::extract(prod.rhs());
268 
269     typename ProductType::Scalar actualAlpha = alpha * LhsBlasTraits::extractScalarFactor(prod.lhs().derived()) * RhsBlasTraits::extractScalarFactor(prod.rhs().derived());
270 
271     if(!beta)
272       mat.template triangularView<UpLo>().setZero();
273 
274     enum {
275       IsRowMajor = (internal::traits<MatrixType>::Flags&RowMajorBit) ? 1 : 0,
276       LhsIsRowMajor = _ActualLhs::Flags&RowMajorBit ? 1 : 0,
277       RhsIsRowMajor = _ActualRhs::Flags&RowMajorBit ? 1 : 0,
278       SkipDiag = (UpLo&(UnitDiag|ZeroDiag))!=0
279     };
280 
281     Index size = mat.cols();
282     if(SkipDiag)
283       size--;
284     Index depth = actualLhs.cols();
285 
286     typedef internal::gemm_blocking_space<IsRowMajor ? RowMajor : ColMajor,typename Lhs::Scalar,typename Rhs::Scalar,
287           MatrixType::MaxColsAtCompileTime, MatrixType::MaxColsAtCompileTime, _ActualRhs::MaxColsAtCompileTime> BlockingType;
288 
289     BlockingType blocking(size, size, depth, 1, false);
290 
291     internal::general_matrix_matrix_triangular_product<Index,
292       typename Lhs::Scalar, LhsIsRowMajor ? RowMajor : ColMajor, LhsBlasTraits::NeedToConjugate,
293       typename Rhs::Scalar, RhsIsRowMajor ? RowMajor : ColMajor, RhsBlasTraits::NeedToConjugate,
294       IsRowMajor ? RowMajor : ColMajor, MatrixType::InnerStrideAtCompileTime, UpLo&(Lower|Upper)>
295       ::run(size, depth,
296             &actualLhs.coeffRef(SkipDiag&&(UpLo&Lower)==Lower ? 1 : 0,0), actualLhs.outerStride(),
297             &actualRhs.coeffRef(0,SkipDiag&&(UpLo&Upper)==Upper ? 1 : 0), actualRhs.outerStride(),
298             mat.data() + (SkipDiag ? (bool(IsRowMajor) != ((UpLo&Lower)==Lower) ? mat.innerStride() : mat.outerStride() ) : 0),
299             mat.innerStride(), mat.outerStride(), actualAlpha, blocking);
300   }
301 };
302 
303 template<typename MatrixType, unsigned int UpLo>
304 template<typename ProductType>
305 EIGEN_DEVICE_FUNC TriangularView<MatrixType,UpLo>& TriangularViewImpl<MatrixType,UpLo,Dense>::_assignProduct(const ProductType& prod, const Scalar& alpha, bool beta)
306 {
307   EIGEN_STATIC_ASSERT((UpLo&UnitDiag)==0, WRITING_TO_TRIANGULAR_PART_WITH_UNIT_DIAGONAL_IS_NOT_SUPPORTED);
308   eigen_assert(derived().nestedExpression().rows() == prod.rows() && derived().cols() == prod.cols());
309 
310   general_product_to_triangular_selector<MatrixType, ProductType, UpLo, internal::traits<ProductType>::InnerSize==1>::run(derived().nestedExpression().const_cast_derived(), prod, alpha, beta);
311 
312   return derived();
313 }
314 
315 } // end namespace Eigen
316 
317 #endif // EIGEN_GENERAL_MATRIX_MATRIX_TRIANGULAR_H
318