Skip to content

Commit 0c43fba

Browse files
zhwesky2010zhwesky2010
committed
add new API/OP: paddle.linalg.triangular_solve (#36714)
* add new API: paddle.linalg.triangular_solve * add new API/OP: paddle.linalg.triangular_solve * add new API/OP: paddle.linalg.triangular_solve * fix comment
1 parent 824c4ef commit 0c43fba

File tree

17 files changed

+1245
-113
lines changed

17 files changed

+1245
-113
lines changed

‎paddle/fluid/operators/math/blas.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,12 @@ class Blas {
253253
void BatchedGETRS(CBLAS_TRANSPOSE trans, int n, int nrhs, const T** a,
254254
int lda, int* ipiv, T** b, int ldb, int* info,
255255
int batch_size) const;
256+
257+
// cuBlas triangular_solve
258+
template <typename T>
259+
void BatchedTRSM(CBLAS_SIDE side, CBLAS_UPLO uplo, CBLAS_TRANSPOSE transA,
260+
CBLAS_DIAG diag, int M, int N, T alpha, const T** a, int lda,
261+
T** b, int ldb, int batch_size) const;
256262
#endif
257263

258264
private:
@@ -414,6 +420,12 @@ class BlasT : private Blas<DeviceContext> {
414420
void BatchedGETRS(ARGS... args) const {
415421
Base()->template BatchedGETRS<T>(args...);
416422
}
423+
424+
// triangular_solve
425+
template <typename... ARGS>
426+
void BatchedTRSM(ARGS... args) const {
427+
Base()->template BatchedTRSM<T>(args...);
428+
}
417429
#endif
418430

419431
private:

‎paddle/fluid/operators/math/blas_impl.cu.h‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ struct CUBlas<float> {
120120
PADDLE_ENFORCE_CUDA_SUCCESS(
121121
platform::dynload::cublasSgetrsBatched(args...));
122122
}
123+
124+
template <typename... ARGS>
125+
static void TRSM_BATCH(ARGS... args) {
126+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasStrsmBatched(args...));
127+
}
123128
};
124129

125130
template <>
@@ -194,6 +199,11 @@ struct CUBlas<double> {
194199
PADDLE_ENFORCE_CUDA_SUCCESS(
195200
platform::dynload::cublasDgetrsBatched(args...));
196201
}
202+
203+
template <typename... ARGS>
204+
static void TRSM_BATCH(ARGS... args) {
205+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasDtrsmBatched(args...));
206+
}
197207
};
198208

199209
template <>
@@ -339,6 +349,19 @@ struct CUBlas<platform::complex<float>> {
339349
reinterpret_cast<cuFloatComplex *>(C), ldc));
340350
}
341351

352+
static void TRSM(cublasHandle_t handle, cublasSideMode_t side,
353+
cublasFillMode_t uplo, cublasOperation_t transa,
354+
cublasDiagType_t diag, int m, int n,
355+
const paddle::platform::complex<float> *alpha,
356+
const paddle::platform::complex<float> *A, int lda,
357+
paddle::platform::complex<float> *B, int ldb) {
358+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasCtrsm(
359+
handle, side, uplo, transa, diag, m, n,
360+
reinterpret_cast<const cuFloatComplex *>(alpha),
361+
reinterpret_cast<const cuFloatComplex *>(A), lda,
362+
reinterpret_cast<cuFloatComplex *>(B), ldb));
363+
}
364+
342365
// NOTES: GEMM_EX can use Tensor Core to accelerate matrix multiply.
343366
// https://docs.nvidia.com/cuda/cublas/index.html#cublassetmathmode
344367
template <typename... ARGS>
@@ -370,6 +393,20 @@ struct CUBlas<platform::complex<float>> {
370393
"cublasGemmEx is not supported on cuda <= 7.5"));
371394
#endif
372395
}
396+
397+
static void TRSM_BATCH(cublasHandle_t handle, cublasSideMode_t side,
398+
cublasFillMode_t uplo, cublasOperation_t transa,
399+
cublasDiagType_t diag, int m, int n,
400+
const paddle::platform::complex<float> *alpha,
401+
const paddle::platform::complex<float> **A, int lda,
402+
paddle::platform::complex<float> **B, int ldb,
403+
int batch_size) {
404+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasCtrsmBatched(
405+
handle, side, uplo, transa, diag, m, n,
406+
reinterpret_cast<const cuFloatComplex *>(alpha),
407+
reinterpret_cast<const cuFloatComplex **>(A), lda,
408+
reinterpret_cast<cuFloatComplex **>(B), ldb, batch_size));
409+
}
373410
};
374411

375412
template <>
@@ -440,6 +477,33 @@ struct CUBlas<platform::complex<double>> {
440477
reinterpret_cast<cuDoubleComplex *>(C), ldc));
441478
}
442479

480+
static void TRSM(cublasHandle_t handle, cublasSideMode_t side,
481+
cublasFillMode_t uplo, cublasOperation_t transa,
482+
cublasDiagType_t diag, int m, int n,
483+
const paddle::platform::complex<double> *alpha,
484+
const paddle::platform::complex<double> *A, int lda,
485+
paddle::platform::complex<double> *B, int ldb) {
486+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasZtrsm(
487+
handle, side, uplo, transa, diag, m, n,
488+
reinterpret_cast<const cuDoubleComplex *>(alpha),
489+
reinterpret_cast<const cuDoubleComplex *>(A), lda,
490+
reinterpret_cast<cuDoubleComplex *>(B), ldb));
491+
}
492+
493+
static void TRSM_BATCH(cublasHandle_t handle, cublasSideMode_t side,
494+
cublasFillMode_t uplo, cublasOperation_t transa,
495+
cublasDiagType_t diag, int m, int n,
496+
const paddle::platform::complex<double> *alpha,
497+
const paddle::platform::complex<double> **A, int lda,
498+
paddle::platform::complex<double> **B, int ldb,
499+
int batch_size) {
500+
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cublasZtrsmBatched(
501+
handle, side, uplo, transa, diag, m, n,
502+
reinterpret_cast<const cuDoubleComplex *>(alpha),
503+
reinterpret_cast<const cuDoubleComplex **>(A), lda,
504+
reinterpret_cast<cuDoubleComplex **>(B), ldb, batch_size));
505+
}
506+
443507
// NOTES: GEMM_EX can use Tensor Core to accelerate matrix multiply.
444508
// https://docs.nvidia.com/cuda/cublas/index.html#cublassetmathmode
445509
template <typename... ARGS>
@@ -897,6 +961,30 @@ void Blas<platform::CUDADeviceContext>::BatchedGETRS(
897961
});
898962
}
899963

964+
template <>
965+
template <typename T>
966+
void Blas<platform::CUDADeviceContext>::BatchedTRSM(
967+
CBLAS_SIDE side, CBLAS_UPLO uplo, CBLAS_TRANSPOSE transA, CBLAS_DIAG diag,
968+
int M, int N, T alpha, const T **A, int lda, T **B, int ldb,
969+
int batch_size) const {
970+
// solve row major `op ( A ) X = α B` by taking it as `X' op ( A' ) = α B'`
971+
// where ' stands for transpose
972+
cublasSideMode_t cuSide =
973+
(side == CblasLeft) ? CUBLAS_SIDE_RIGHT : CUBLAS_SIDE_LEFT;
974+
cublasFillMode_t cuUplo =
975+
(uplo == CblasLower) ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER;
976+
// use CUBLAS_OP_C (conjugate transpose) for complex
977+
cublasOperation_t cuTransA =
978+
(transA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
979+
cublasDiagType_t cuDiag =
980+
(diag == CblasUnit) ? CUBLAS_DIAG_UNIT : CUBLAS_DIAG_NON_UNIT;
981+
982+
context_.CublasCall([&](cublasHandle_t handle) {
983+
CUBlas<T>::TRSM_BATCH(handle, cuSide, cuUplo, cuTransA, cuDiag, N, M,
984+
&alpha, A, lda, B, ldb, batch_size);
985+
});
986+
}
987+
900988
} // namespace math
901989
} // namespace operators
902990
} // namespace paddle

‎paddle/fluid/operators/math/blas_impl.h‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,6 +434,17 @@ struct CBlas<platform::complex<float>> {
434434
a_, lda, b_, ldb, &beta, c_, ldc);
435435
}
436436

437+
static void TRSM(CBLAS_LAYOUT layout, CBLAS_SIDE side, CBLAS_UPLO uplo,
438+
CBLAS_TRANSPOSE trans_a, CBLAS_DIAG diag, int M, int N,
439+
paddle::platform::complex<float> alpha,
440+
const paddle::platform::complex<float> *A, int lda,
441+
paddle::platform::complex<float> *B, int ldb) {
442+
const void *a_ = (const void *)(A);
443+
void *b_ = static_cast<void *>(B);
444+
platform::dynload::cblas_ctrsm(layout, side, uplo, trans_a, diag, M, N,
445+
&alpha, a_, lda, b_, ldb);
446+
}
447+
437448
template <typename... ARGS>
438449
static void GEMM_BATCH(CBLAS_LAYOUT layout, CBLAS_TRANSPOSE *trans_a,
439450
CBLAS_TRANSPOSE *trans_b, int *M, int *N, int *K,
@@ -562,6 +573,17 @@ struct CBlas<platform::complex<double>> {
562573
a_, lda, b_, ldb, &beta, c_, ldc);
563574
}
564575

576+
static void TRSM(CBLAS_LAYOUT layout, CBLAS_SIDE side, CBLAS_UPLO uplo,
577+
CBLAS_TRANSPOSE trans_a, CBLAS_DIAG diag, int M, int N,
578+
paddle::platform::complex<double> alpha,
579+
const paddle::platform::complex<double> *A, int lda,
580+
paddle::platform::complex<double> *B, int ldb) {
581+
const void *a_ = (const void *)(A);
582+
void *b_ = static_cast<void *>(B);
583+
platform::dynload::cblas_ztrsm(layout, side, uplo, trans_a, diag, M, N,
584+
&alpha, a_, lda, b_, ldb);
585+
}
586+
565587
template <typename... ARGS>
566588
static void GEMM_BATCH(CBLAS_LAYOUT layout, CBLAS_TRANSPOSE *trans_a,
567589
CBLAS_TRANSPOSE *trans_b, int *M, int *N, int *K,
@@ -682,6 +704,15 @@ struct CBlas<platform::complex<float>> {
682704
cblas_cgemm(layout, TransA, TransB, M, N, K, &alpha, A, lda, B, ldb, &beta,
683705
C, ldc);
684706
}
707+
708+
static void TRSM(const CBLAS_LAYOUT layout, const CBLAS_SIDE side,
709+
const CBLAS_UPLO uplo, const CBLAS_TRANSPOSE transA,
710+
const CBLAS_DIAG diag, const int M, const int N,
711+
const paddle::platform::complex<float> alpha,
712+
const paddle::platform::complex<float> *A, const int lda,
713+
paddle::platform::complex<double> *B, const int ldb) {
714+
cblas_ctrsm(layout, side, uplo, transA, diag, M, N, &alpha, A, lda, B, ldb);
715+
}
685716
};
686717

687718
template <>
@@ -720,6 +751,15 @@ struct CBlas<platform::complex<double>> {
720751
cblas_zgemm(layout, TransA, TransB, M, N, K, &alpha, A, lda, B, ldb, &beta,
721752
C, ldc);
722753
}
754+
755+
static void TRSM(const CBLAS_LAYOUT layout, const CBLAS_SIDE side,
756+
const CBLAS_UPLO uplo, const CBLAS_TRANSPOSE transA,
757+
const CBLAS_DIAG diag, const int M, const int N,
758+
const paddle::platform::complex<double> alpha,
759+
const paddle::platform::complex<double> *A, const int lda,
760+
paddle::platform::complex<double> *B, const int ldb) {
761+
cblas_ztrsm(layout, side, uplo, transA, diag, M, N, &alpha, A, lda, B, ldb);
762+
}
723763
};
724764

725765
#endif

‎paddle/fluid/operators/math/blas_impl.hip.h‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,12 @@ struct CUBlas<float> {
9090
PADDLE_THROW(platform::errors::Unimplemented(
9191
"cublasSmatinvBatched is not supported on HIP platform."));
9292
}
93+
94+
template <typename... ARGS>
95+
static void TRSM_BATCH(ARGS... args) {
96+
PADDLE_THROW(platform::errors::Unimplemented(
97+
"cublasStrsmBatched is not supported on HIP platform."));
98+
}
9399
};
94100

95101
template <>
@@ -153,6 +159,12 @@ struct CUBlas<double> {
153159
PADDLE_THROW(platform::errors::Unimplemented(
154160
"cublasDmatinvBatched is not supported on HIP platform."));
155161
}
162+
163+
template <typename... ARGS>
164+
static void TRSM_BATCH(ARGS... args) {
165+
PADDLE_THROW(platform::errors::Unimplemented(
166+
"cublasDtrsmBatched is not supported on HIP platform."));
167+
}
156168
};
157169

158170
template <>
@@ -730,6 +742,32 @@ void Blas<platform::CUDADeviceContext>::BatchedGETRS(
730742
batch_size);
731743
});
732744
}
745+
746+
template <>
747+
template <typename T>
748+
void Blas<platform::CUDADeviceContext>::BatchedTRSM(
749+
CBLAS_SIDE side, CBLAS_UPLO uplo, CBLAS_TRANSPOSE transA, CBLAS_DIAG diag,
750+
int M, int N, T alpha, const T **A, int lda, T **B, int ldb,
751+
int batch_size) const {
752+
// solve row major `op ( A ) X = α B` by taking it as `X' op ( A' ) = α B'`
753+
// where ' stands for transpose
754+
rocblas_side cuSide =
755+
(side == CblasLeft) ? rocblas_side_right : rocblas_side_left;
756+
rocblas_fill cuUplo =
757+
(uplo == CblasLower) ? rocblas_fill_upper : rocblas_fill_lower;
758+
// use CUBLAS_OP_C (conjugate transpose) for complex
759+
rocblas_operation cuTransA = (transA == CblasNoTrans)
760+
? rocblas_operation_none
761+
: rocblas_operation_transpose;
762+
rocblas_diagonal cuDiag =
763+
(diag == CblasUnit) ? rocblas_diagonal_unit : rocblas_diagonal_non_unit;
764+
765+
context_.CublasCall([&](rocblas_handle handle) {
766+
CUBlas<T>::TRSM_BATCH(handle, cuSide, cuUplo, cuTransA, cuDiag, N, M,
767+
&alpha, A, lda, B, ldb, batch_size);
768+
});
769+
}
770+
733771
} // namespace math
734772
} // namespace operators
735773
} // namespace paddle

‎paddle/fluid/operators/math/matrix_solve.cc‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,45 @@ class MatrixSolveFunctor<platform::CPUDeviceContext, T> {
3434
template class MatrixSolveFunctor<platform::CPUDeviceContext, float>;
3535
template class MatrixSolveFunctor<platform::CPUDeviceContext, double>;
3636

37+
template <typename T>
38+
class TriangularSolveFunctor<platform::CPUDeviceContext, T> {
39+
public:
40+
void operator()(const platform::CPUDeviceContext& context,
41+
const framework::Tensor* a, framework::Tensor* b, bool left,
42+
bool upper, bool transpose, bool unitriangular) {
43+
CBLAS_SIDE side = left ? CblasLeft : CblasRight;
44+
CBLAS_UPLO uplo = upper ? CblasUpper : CblasLower;
45+
CBLAS_TRANSPOSE transA = transpose ? CblasTrans : CblasNoTrans;
46+
CBLAS_DIAG diag = unitriangular ? CblasUnit : CblasNonUnit;
47+
48+
const T* a_data = a->data<T>();
49+
T* b_data = b->mutable_data<T>(context.GetPlace());
50+
51+
int a_dim_size = a->dims().size();
52+
int b_dim_size = b->dims().size();
53+
54+
int M = static_cast<int>(b->dims()[b_dim_size - 2]);
55+
int N = static_cast<int>(b->dims()[b_dim_size - 1]);
56+
auto lda = left ? std::max(1, M) : std::max(1, N);
57+
auto ldb = std::max(1, N);
58+
59+
int batch_size = 1;
60+
auto& a_dim = a->dims();
61+
for (int i = 0; i < a_dim_size - 2; i++) {
62+
batch_size *= a_dim[i];
63+
}
64+
65+
auto blas = math::GetBlas<platform::CPUDeviceContext, T>(context);
66+
for (int i = 0; i < batch_size; i++) {
67+
blas.TRSM(side, uplo, transA, diag, M, N, T(1), a_data + i * M * M, lda,
68+
b_data + i * N * M, ldb);
69+
}
70+
}
71+
};
72+
73+
template class TriangularSolveFunctor<platform::CPUDeviceContext, float>;
74+
template class TriangularSolveFunctor<platform::CPUDeviceContext, double>;
75+
3776
} // namespace math
3877
} // namespace operators
3978
} // namespace paddle

‎paddle/fluid/operators/math/matrix_solve.cu.cc‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,68 @@ class MatrixSolveFunctor<platform::CUDADeviceContext, T> {
163163
template class MatrixSolveFunctor<platform::CUDADeviceContext, float>;
164164
template class MatrixSolveFunctor<platform::CUDADeviceContext, double>;
165165

166+
template <typename T>
167+
class TriangularSolveFunctor<platform::CUDADeviceContext, T> {
168+
public:
169+
void operator()(const platform::CUDADeviceContext& context, const Tensor* a,
170+
Tensor* b, bool left, bool upper, bool transpose,
171+
bool unitriangular) {
172+
CBLAS_SIDE side = left ? CblasLeft : CblasRight;
173+
CBLAS_UPLO uplo = upper ? CblasUpper : CblasLower;
174+
CBLAS_TRANSPOSE transA = transpose ? CblasTrans : CblasNoTrans;
175+
CBLAS_DIAG diag = unitriangular ? CblasUnit : CblasNonUnit;
176+
177+
const T* a_data = a->data<T>();
178+
T* b_data = b->mutable_data<T>(context.GetPlace());
179+
180+
int a_dim_size = a->dims().size();
181+
int b_dim_size = b->dims().size();
182+
183+
int M = static_cast<int>(b->dims()[b_dim_size - 2]);
184+
int N = static_cast<int>(b->dims()[b_dim_size - 1]);
185+
auto lda = left ? std::max(1, M) : std::max(1, N);
186+
auto ldb = std::max(1, N);
187+
188+
int batch_size = 1;
189+
auto& a_dim = a->dims();
190+
for (int i = 0; i < a_dim_size - 2; i++) {
191+
batch_size *= a_dim[i];
192+
}
193+
194+
auto blas = math::GetBlas<platform::CUDADeviceContext, T>(context);
195+
if (batch_size <= 8 && M >= 64) {
196+
for (auto i = 0; i < batch_size; i++) {
197+
blas.TRSM(side, uplo, transA, diag, M, N, static_cast<T>(1.0),
198+
a_data + i * M * M, lda, b_data + i * N * M, ldb);
199+
}
200+
} else {
201+
std::vector<const T*> cpu_ptrs(batch_size * 2);
202+
for (int i = 0; i < batch_size; ++i) {
203+
cpu_ptrs[i] = a_data + i * M * M;
204+
cpu_ptrs[i + batch_size] = b_data + i * M * N;
205+
}
206+
207+
// Copy the addresses of A and tmp_b from host to device.
208+
memory::allocation::AllocationPtr tmp_gpu_ptrs_data =
209+
memory::Alloc(context, cpu_ptrs.size() * sizeof(T*));
210+
memory::Copy(BOOST_GET_CONST(platform::CUDAPlace, context.GetPlace()),
211+
tmp_gpu_ptrs_data->ptr(), platform::CPUPlace(),
212+
static_cast<void*>(cpu_ptrs.data()),
213+
cpu_ptrs.size() * sizeof(T*), context.stream());
214+
215+
const T** gpu_a_ptrs =
216+
reinterpret_cast<const T**>(tmp_gpu_ptrs_data->ptr());
217+
T** gpu_b_ptrs =
218+
reinterpret_cast<T**>(tmp_gpu_ptrs_data->ptr()) + batch_size;
219+
blas.BatchedTRSM(side, uplo, transA, diag, M, N, static_cast<T>(1.0),
220+
gpu_a_ptrs, lda, gpu_b_ptrs, ldb, batch_size);
221+
}
222+
}
223+
};
224+
225+
template class TriangularSolveFunctor<platform::CUDADeviceContext, float>;
226+
template class TriangularSolveFunctor<platform::CUDADeviceContext, double>;
227+
166228
} // namespace math
167229
} // namespace operators
168230
} // namespace paddle

0 commit comments

Comments
 (0)