Add specializations for some triangular-triangular multiplications (#1538)

Author: Red-Portal

lib/cublas/linalg.jl

@@ -443,6 +443,68 @@ for (t, uploc, isunitc) in ((:LowerTriangular, 'U', 'N'),
     end
 end
 
+function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                            A::LowerTriangular{T,<:DenseCuMatrix},
+                            B::UpperTriangular{T,<:DenseCuMatrix},
+                            ) where {T<:CublasFloat}
+    triu!(parent(B))
+    trmm!('L', 'L', 'N', 'N', one(T), parent(A), parent(B), parent(X))
+    X
+end
+
+function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                            A::UpperTriangular{T,<:DenseCuMatrix},
+                            B::LowerTriangular{T,<:DenseCuMatrix},
+                            ) where {T<:CublasFloat}
+    tril!(parent(B))
+    trmm!('L', 'U', 'N', 'N', one(T), parent(A), parent(B), parent(X))
+    X
+end
+
+for (trtype, valtype) in ((:Transpose, :CublasFloat),
+                          (:Adjoint,   :CublasReal),
+                          (:Adjoint,   :CublasComplex))
+    @eval begin
+        function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                                    A::UpperTriangular{T,<:DenseCuMatrix},
+                                    B::LowerTriangular{<:Any,<:$trtype{T,<:DenseCuMatrix}},
+                                    ) where {T<:$valtype}
+            # operation is reversed to avoid executing the tranpose
+            triu!(parent(A))
+            CUBLAS.trmm!('R', 'U', 'T', 'N', one(T), parent(parent(B)), parent(A), parent(X))
+            X
+        end
+
+        function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                                    A::UpperTriangular{<:Any,<:$trtype{T,<:DenseCuMatrix}},
+                                    B::LowerTriangular{T,<:DenseCuMatrix},
+                                    ) where {T<:$valtype}
+            tril!(parent(B))
+            CUBLAS.trmm!('L', 'L', 'T', 'N', one(T), parent(parent(A)), parent(B), parent(X))
+            X
+        end
+
+        function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                                    A::LowerTriangular{<:Any,<:$trtype{T,<:DenseCuMatrix}},
+                                    B::UpperTriangular{T,<:DenseCuMatrix},
+                                    ) where {T<:$valtype}
+            triu!(parent(B))
+            CUBLAS.trmm!('L', 'U', 'T', 'N', one(T), parent(parent(A)), parent(B), parent(X))
+            X
+        end
+
+        function LinearAlgebra.mul!(X::DenseCuMatrix{T},
+                                    A::LowerTriangular{T,<:DenseCuMatrix},
+                                    B::UpperTriangular{<:Any,<:$trtype{T,<:DenseCuMatrix}},
+                                    ) where {T<:$valtype}
+            # operation is reversed to avoid executing the tranpose
+            tril!(parent(A))
+            CUBLAS.trmm!('R', 'L', 'T', 'N', one(T), parent(parent(B)), parent(A), parent(X))
+            X
+        end
+    end
+end
+
 # symmetric mul!
 # level 2
 @inline function LinearAlgebra.mul!(y::CuVector{T}, A::Hermitian{T,<:CuMatrix}, x::CuVector{T},

test/cublas.jl

@@ -1019,7 +1019,7 @@ end
             end
         end
 
-        @testset "triangular mul!" begin
+        @testset "triangular-dense mul!" begin
             A = triu(rand(elty, m, m))
             B = rand(elty,m,n)
             C = zeros(elty,m,n)
@@ -1052,6 +1052,44 @@ end
             end
         end
 
+        @testset "triangular-triangular mul!" begin
+            A  = triu(rand(elty, m, m))
+            B  = triu(rand(elty, m, m))
+            C0 = zeros(elty,m,m)
+
+            sA = rand(elty,m,m)
+            sA = sA + transpose(sA)
+            sB = rand(elty,m,m)
+            sB = sB + transpose(sB)
+
+            for (TRa, ta, TRb, tb, TRc) in (
+                (UpperTriangular, identity,  LowerTriangular, identity,  Matrix),
+                (LowerTriangular, identity,  UpperTriangular, identity,  Matrix),
+                (UpperTriangular, identity,  UpperTriangular, transpose, Matrix),
+                (UpperTriangular, transpose, UpperTriangular, identity,  Matrix),
+                (LowerTriangular, identity,  LowerTriangular, transpose, Matrix),
+                (LowerTriangular, transpose, LowerTriangular, identity,  Matrix),
+                )
+
+                A = copy(sA) |> TRa
+                B = copy(sB) |> TRb
+                C = copy(C0) |> TRc
+                dA = CuArray(parent(sA)) |> TRa
+                dB = CuArray(parent(sB)) |> TRb
+                dC = if TRc == Matrix
+                    CuArray(C0) |> DenseCuMatrix
+                else
+                    CuArray(C0) |> TRc
+                end
+
+                D = mul!(C, ta(A), tb(B))
+                dD = mul!(dC, ta(dA), tb(dB))
+
+                @test C ≈ Array(dC)
+                @test D ≈ Array(dD)
+            end
+        end
+
         B = rand(elty,m,n)
         C = rand(elty,m,n)
         d_B = CuArray(B)