Add iszero and isone for AbstractGPUMatrix (#419)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: GVigne

src/host/linalg.jl

@@ -418,3 +418,21 @@ LinearAlgebra.axpby!(alpha::Number, x::AbstractGPUArray,
                      beta::Number,  y::AbstractGPUArray) = y .= x.*alpha .+ y.*beta
 
 LinearAlgebra.axpy!(alpha::Number, x::AbstractGPUArray, y::AbstractGPUArray) = y .+= x.*alpha
+
+## identity and zero equality check
+
+Base.iszero(x::AbstractGPUMatrix{T}) where {T} = all(iszero, x)
+function Base.isone(x::AbstractGPUMatrix{T}) where {T}
+    n,m = size(x)
+    m != n && return false
+
+    # lazily perform `x-I`
+    bc = Broadcast.broadcasted(x, CartesianIndices(x)) do _x, inds
+        _x - (inds[1] == inds[2] ? one(_x) : zero(_x))
+    end
+    # call `GPUArrays.mapreducedim!` directly, which supports Broadcasted inputs
+    y = similar(x, Bool, 1)
+    GPUArrays.mapreducedim!(iszero, &, y, Broadcast.instantiate(bc); init=true)
+
+    Array(y)[]
+end

test/testsuite/linalg.jl

@@ -222,6 +222,16 @@
         @test axpby!(alpha,x,beta,y) ≈ T.([7,10,13])
         @test axpy!(alpha,x,y) ≈ T.([8,12,16])
     end
+
+    @testset "iszero and isone" for T in eltypes
+        A = one(AT(rand(T, 2, 2)))
+        @test isone(A)
+        @test iszero(A) == false
+
+        A = zero(AT(rand(T, 2, 2)))
+        @test iszero(A)
+        @test isone(A) == false
+    end
 end
 
 @testsuite "linalg/mul!/vector-matrix" (AT, eltypes)->begin