Add direct BLAS calls for trmm! and trsm!.

JuliaGPU · May 5, 2020 · 7ee3747 · 7ee3747
1 parent 9b2ec43
commit 7ee3747
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/blas/linalg.jl b/src/blas/linalg.jl
@@ -301,3 +301,11 @@ LinearAlgebra.rdiv!(A::CuMatrix{T}, B::Transpose{T,<:LowerTriangular{T, <:CuMatr
     CUBLAS.trsm!('R', 'L', 'T', 'N', one(T), parent(parent(B)), A)
 LinearAlgebra.rdiv!(A::CuMatrix{T}, B::Transpose{T,<:UnitLowerTriangular{T, <:CuMatrix{T}}}) where T<:CublasFloat =
     CUBLAS.trsm!('R', 'L', 'T', 'U', one(T), parent(parent(B)), A)
+
+# Direct BLAS calls
+for T in Base.uniontypes(CublasFloat)
+    @eval LinearAlgebra.BLAS.trmm!(side::AbstractChar, uplo::AbstractChar, transa::AbstractChar, diag::AbstractChar, alpha::$T, A::CuMatrix{$T}, B::CuMatrix{$T}) =
+        CuArrays.CUBLAS.trmm!(side, uplo, transa, diag, alpha, A, B, B)
+    @eval LinearAlgebra.BLAS.trsm!(side::AbstractChar, uplo::AbstractChar, transa::AbstractChar, diag::AbstractChar, alpha::$T, A::CuMatrix{$T}, B::CuMatrix{$T}) =
+        CuArrays.CUBLAS.trsm!(side, uplo, transa, diag, alpha, A, B)
+end