You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
280 lines
10 KiB
280 lines
10 KiB
// This file is part of Eigen, a lightweight C++ template library |
|
// for linear algebra. |
|
// |
|
// Copyright (C) 2012 Gael Guennebaud <gael.guennebaud@inria.fr> |
|
// |
|
// This Source Code Form is subject to the terms of the Mozilla |
|
// Public License v. 2.0. If a copy of the MPL was not distributed |
|
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/. |
|
|
|
#ifndef EIGEN_SPARSELU_GEMM_KERNEL_H |
|
#define EIGEN_SPARSELU_GEMM_KERNEL_H |
|
|
|
namespace Eigen { |
|
|
|
namespace internal { |
|
|
|
|
|
/** \internal |
|
* A general matrix-matrix product kernel optimized for the SparseLU factorization. |
|
* - A, B, and C must be column major |
|
* - lda and ldc must be multiples of the respective packet size |
|
* - C must have the same alignment as A |
|
*/ |
|
template<typename Scalar> |
|
EIGEN_DONT_INLINE |
|
void sparselu_gemm(Index m, Index n, Index d, const Scalar* A, Index lda, const Scalar* B, Index ldb, Scalar* C, Index ldc) |
|
{ |
|
using namespace Eigen::internal; |
|
|
|
typedef typename packet_traits<Scalar>::type Packet; |
|
enum { |
|
NumberOfRegisters = EIGEN_ARCH_DEFAULT_NUMBER_OF_REGISTERS, |
|
PacketSize = packet_traits<Scalar>::size, |
|
PM = 8, // peeling in M |
|
RN = 2, // register blocking |
|
RK = NumberOfRegisters>=16 ? 4 : 2, // register blocking |
|
BM = 4096/sizeof(Scalar), // number of rows of A-C per chunk |
|
SM = PM*PacketSize // step along M |
|
}; |
|
Index d_end = (d/RK)*RK; // number of columns of A (rows of B) suitable for full register blocking |
|
Index n_end = (n/RN)*RN; // number of columns of B-C suitable for processing RN columns at once |
|
Index i0 = internal::first_default_aligned(A,m); |
|
|
|
eigen_internal_assert(((lda%PacketSize)==0) && ((ldc%PacketSize)==0) && (i0==internal::first_default_aligned(C,m))); |
|
|
|
// handle the non aligned rows of A and C without any optimization: |
|
for(Index i=0; i<i0; ++i) |
|
{ |
|
for(Index j=0; j<n; ++j) |
|
{ |
|
Scalar c = C[i+j*ldc]; |
|
for(Index k=0; k<d; ++k) |
|
c += B[k+j*ldb] * A[i+k*lda]; |
|
C[i+j*ldc] = c; |
|
} |
|
} |
|
// process the remaining rows per chunk of BM rows |
|
for(Index ib=i0; ib<m; ib+=BM) |
|
{ |
|
Index actual_b = std::min<Index>(BM, m-ib); // actual number of rows |
|
Index actual_b_end1 = (actual_b/SM)*SM; // actual number of rows suitable for peeling |
|
Index actual_b_end2 = (actual_b/PacketSize)*PacketSize; // actual number of rows suitable for vectorization |
|
|
|
// Let's process two columns of B-C at once |
|
for(Index j=0; j<n_end; j+=RN) |
|
{ |
|
const Scalar* Bc0 = B+(j+0)*ldb; |
|
const Scalar* Bc1 = B+(j+1)*ldb; |
|
|
|
for(Index k=0; k<d_end; k+=RK) |
|
{ |
|
|
|
// load and expand a RN x RK block of B |
|
Packet b00, b10, b20, b30, b01, b11, b21, b31; |
|
{ b00 = pset1<Packet>(Bc0[0]); } |
|
{ b10 = pset1<Packet>(Bc0[1]); } |
|
if(RK==4) { b20 = pset1<Packet>(Bc0[2]); } |
|
if(RK==4) { b30 = pset1<Packet>(Bc0[3]); } |
|
{ b01 = pset1<Packet>(Bc1[0]); } |
|
{ b11 = pset1<Packet>(Bc1[1]); } |
|
if(RK==4) { b21 = pset1<Packet>(Bc1[2]); } |
|
if(RK==4) { b31 = pset1<Packet>(Bc1[3]); } |
|
|
|
Packet a0, a1, a2, a3, c0, c1, t0, t1; |
|
|
|
const Scalar* A0 = A+ib+(k+0)*lda; |
|
const Scalar* A1 = A+ib+(k+1)*lda; |
|
const Scalar* A2 = A+ib+(k+2)*lda; |
|
const Scalar* A3 = A+ib+(k+3)*lda; |
|
|
|
Scalar* C0 = C+ib+(j+0)*ldc; |
|
Scalar* C1 = C+ib+(j+1)*ldc; |
|
|
|
a0 = pload<Packet>(A0); |
|
a1 = pload<Packet>(A1); |
|
if(RK==4) |
|
{ |
|
a2 = pload<Packet>(A2); |
|
a3 = pload<Packet>(A3); |
|
} |
|
else |
|
{ |
|
// workaround "may be used uninitialized in this function" warning |
|
a2 = a3 = a0; |
|
} |
|
|
|
#define KMADD(c, a, b, tmp) {tmp = b; tmp = pmul(a,tmp); c = padd(c,tmp);} |
|
#define WORK(I) \ |
|
c0 = pload<Packet>(C0+i+(I)*PacketSize); \ |
|
c1 = pload<Packet>(C1+i+(I)*PacketSize); \ |
|
KMADD(c0, a0, b00, t0) \ |
|
KMADD(c1, a0, b01, t1) \ |
|
a0 = pload<Packet>(A0+i+(I+1)*PacketSize); \ |
|
KMADD(c0, a1, b10, t0) \ |
|
KMADD(c1, a1, b11, t1) \ |
|
a1 = pload<Packet>(A1+i+(I+1)*PacketSize); \ |
|
if(RK==4){ KMADD(c0, a2, b20, t0) }\ |
|
if(RK==4){ KMADD(c1, a2, b21, t1) }\ |
|
if(RK==4){ a2 = pload<Packet>(A2+i+(I+1)*PacketSize); }\ |
|
if(RK==4){ KMADD(c0, a3, b30, t0) }\ |
|
if(RK==4){ KMADD(c1, a3, b31, t1) }\ |
|
if(RK==4){ a3 = pload<Packet>(A3+i+(I+1)*PacketSize); }\ |
|
pstore(C0+i+(I)*PacketSize, c0); \ |
|
pstore(C1+i+(I)*PacketSize, c1) |
|
|
|
// process rows of A' - C' with aggressive vectorization and peeling |
|
for(Index i=0; i<actual_b_end1; i+=PacketSize*8) |
|
{ |
|
EIGEN_ASM_COMMENT("SPARSELU_GEMML_KERNEL1"); |
|
prefetch((A0+i+(5)*PacketSize)); |
|
prefetch((A1+i+(5)*PacketSize)); |
|
if(RK==4) prefetch((A2+i+(5)*PacketSize)); |
|
if(RK==4) prefetch((A3+i+(5)*PacketSize)); |
|
|
|
WORK(0); |
|
WORK(1); |
|
WORK(2); |
|
WORK(3); |
|
WORK(4); |
|
WORK(5); |
|
WORK(6); |
|
WORK(7); |
|
} |
|
// process the remaining rows with vectorization only |
|
for(Index i=actual_b_end1; i<actual_b_end2; i+=PacketSize) |
|
{ |
|
WORK(0); |
|
} |
|
#undef WORK |
|
// process the remaining rows without vectorization |
|
for(Index i=actual_b_end2; i<actual_b; ++i) |
|
{ |
|
if(RK==4) |
|
{ |
|
C0[i] += A0[i]*Bc0[0]+A1[i]*Bc0[1]+A2[i]*Bc0[2]+A3[i]*Bc0[3]; |
|
C1[i] += A0[i]*Bc1[0]+A1[i]*Bc1[1]+A2[i]*Bc1[2]+A3[i]*Bc1[3]; |
|
} |
|
else |
|
{ |
|
C0[i] += A0[i]*Bc0[0]+A1[i]*Bc0[1]; |
|
C1[i] += A0[i]*Bc1[0]+A1[i]*Bc1[1]; |
|
} |
|
} |
|
|
|
Bc0 += RK; |
|
Bc1 += RK; |
|
} // peeled loop on k |
|
} // peeled loop on the columns j |
|
// process the last column (we now perform a matrix-vector product) |
|
if((n-n_end)>0) |
|
{ |
|
const Scalar* Bc0 = B+(n-1)*ldb; |
|
|
|
for(Index k=0; k<d_end; k+=RK) |
|
{ |
|
|
|
// load and expand a 1 x RK block of B |
|
Packet b00, b10, b20, b30; |
|
b00 = pset1<Packet>(Bc0[0]); |
|
b10 = pset1<Packet>(Bc0[1]); |
|
if(RK==4) b20 = pset1<Packet>(Bc0[2]); |
|
if(RK==4) b30 = pset1<Packet>(Bc0[3]); |
|
|
|
Packet a0, a1, a2, a3, c0, t0/*, t1*/; |
|
|
|
const Scalar* A0 = A+ib+(k+0)*lda; |
|
const Scalar* A1 = A+ib+(k+1)*lda; |
|
const Scalar* A2 = A+ib+(k+2)*lda; |
|
const Scalar* A3 = A+ib+(k+3)*lda; |
|
|
|
Scalar* C0 = C+ib+(n_end)*ldc; |
|
|
|
a0 = pload<Packet>(A0); |
|
a1 = pload<Packet>(A1); |
|
if(RK==4) |
|
{ |
|
a2 = pload<Packet>(A2); |
|
a3 = pload<Packet>(A3); |
|
} |
|
else |
|
{ |
|
// workaround "may be used uninitialized in this function" warning |
|
a2 = a3 = a0; |
|
} |
|
|
|
#define WORK(I) \ |
|
c0 = pload<Packet>(C0+i+(I)*PacketSize); \ |
|
KMADD(c0, a0, b00, t0) \ |
|
a0 = pload<Packet>(A0+i+(I+1)*PacketSize); \ |
|
KMADD(c0, a1, b10, t0) \ |
|
a1 = pload<Packet>(A1+i+(I+1)*PacketSize); \ |
|
if(RK==4){ KMADD(c0, a2, b20, t0) }\ |
|
if(RK==4){ a2 = pload<Packet>(A2+i+(I+1)*PacketSize); }\ |
|
if(RK==4){ KMADD(c0, a3, b30, t0) }\ |
|
if(RK==4){ a3 = pload<Packet>(A3+i+(I+1)*PacketSize); }\ |
|
pstore(C0+i+(I)*PacketSize, c0); |
|
|
|
// agressive vectorization and peeling |
|
for(Index i=0; i<actual_b_end1; i+=PacketSize*8) |
|
{ |
|
EIGEN_ASM_COMMENT("SPARSELU_GEMML_KERNEL2"); |
|
WORK(0); |
|
WORK(1); |
|
WORK(2); |
|
WORK(3); |
|
WORK(4); |
|
WORK(5); |
|
WORK(6); |
|
WORK(7); |
|
} |
|
// vectorization only |
|
for(Index i=actual_b_end1; i<actual_b_end2; i+=PacketSize) |
|
{ |
|
WORK(0); |
|
} |
|
// remaining scalars |
|
for(Index i=actual_b_end2; i<actual_b; ++i) |
|
{ |
|
if(RK==4) |
|
C0[i] += A0[i]*Bc0[0]+A1[i]*Bc0[1]+A2[i]*Bc0[2]+A3[i]*Bc0[3]; |
|
else |
|
C0[i] += A0[i]*Bc0[0]+A1[i]*Bc0[1]; |
|
} |
|
|
|
Bc0 += RK; |
|
#undef WORK |
|
} |
|
} |
|
|
|
// process the last columns of A, corresponding to the last rows of B |
|
Index rd = d-d_end; |
|
if(rd>0) |
|
{ |
|
for(Index j=0; j<n; ++j) |
|
{ |
|
enum { |
|
Alignment = PacketSize>1 ? Aligned : 0 |
|
}; |
|
typedef Map<Matrix<Scalar,Dynamic,1>, Alignment > MapVector; |
|
typedef Map<const Matrix<Scalar,Dynamic,1>, Alignment > ConstMapVector; |
|
if(rd==1) MapVector(C+j*ldc+ib,actual_b) += B[0+d_end+j*ldb] * ConstMapVector(A+(d_end+0)*lda+ib, actual_b); |
|
|
|
else if(rd==2) MapVector(C+j*ldc+ib,actual_b) += B[0+d_end+j*ldb] * ConstMapVector(A+(d_end+0)*lda+ib, actual_b) |
|
+ B[1+d_end+j*ldb] * ConstMapVector(A+(d_end+1)*lda+ib, actual_b); |
|
|
|
else MapVector(C+j*ldc+ib,actual_b) += B[0+d_end+j*ldb] * ConstMapVector(A+(d_end+0)*lda+ib, actual_b) |
|
+ B[1+d_end+j*ldb] * ConstMapVector(A+(d_end+1)*lda+ib, actual_b) |
|
+ B[2+d_end+j*ldb] * ConstMapVector(A+(d_end+2)*lda+ib, actual_b); |
|
} |
|
} |
|
|
|
} // blocking on the rows of A and C |
|
} |
|
#undef KMADD |
|
|
|
} // namespace internal |
|
|
|
} // namespace Eigen |
|
|
|
#endif // EIGEN_SPARSELU_GEMM_KERNEL_H
|
|
|