Add specializations for istriu/istril to speed up isdiag. (#502)

Author: maleadt

src/host/linalg.jl

@@ -165,6 +165,38 @@ function LinearAlgebra.triu!(A::AbstractGPUMatrix{T}, d::Integer = 0) where T
   return A
 end
 
+# check if upper triangular starting from the kth superdiagonal.
+function LinearAlgebra.istriu(A::AbstractGPUMatrix, k::Integer = 0)
+    function mapper(a, I)
+        row, col = Tuple(I)
+        if col < row + k
+            return iszero(a)
+        else
+            true
+        end
+    end
+    function reducer(a, b)
+        a && b
+    end
+    mapreduce(mapper, reducer, A, eachindex(IndexCartesian(), A); init=true)
+end
+
+# check if lower triangular starting from the kth subdiagonal.
+function LinearAlgebra.istril(A::AbstractGPUMatrix, k::Integer = 0)
+    function mapper(a, I)
+        row, col = Tuple(I)
+        if col > row + k
+            return iszero(a)
+        else
+            true
+        end
+    end
+    function reducer(a, b)
+        a && b
+    end
+    mapreduce(mapper, reducer, A, eachindex(IndexCartesian(), A); init=true)
+end
+
 
 ## diagonal
 

test/testsuite/linalg.jl

@@ -105,6 +105,22 @@
                 @test gpu_c isa TR
             end
         end
+
+        @testset "istril" begin
+            for rows in 3:4, cols in 3:4, diag in -4:4
+                A = tril(rand(Float32, rows,cols), diag)
+                B = AT(A)
+                @test istril(A) == istril(B)
+            end
+        end
+
+        @testset "istriu" begin
+            for rows in 3:4, cols in 3:4, diag in -4:4
+                A = triu(rand(Float32, rows,cols), diag)
+                B = AT(A)
+                @test istriu(A) == istriu(B)
+            end
+        end
     end
 
     @testset "diagonal" begin