allow construction from isbits iterators

Author: SimonDanisch

src/construction.jl

@@ -41,6 +41,45 @@ similar(x::X, ::Type{T}, size::Base.Dims{N}) where {X <: GPUArray, T, N} = simil
 
 convert(AT::Type{<: GPUArray{T, N}}, A::GPUArray{T, N}) where {T, N} = A
 
+function indexstyle(x::T) where T
+    style = try
+        Base.IndexStyle(x)
+    catch
+        nothing
+    end
+    style
+end
+
+function collect_kernel(state, A, iter, ::IndexCartesian)
+    idx = @cartesianidx(A, state)
+    @inbounds A[idx...] = iter[idx...]
+    return
+end
+
+function collect_kernel(state, A, iter, ::IndexLinear)
+    idx = linear_index(state)
+    @inbounds A[idx] = iter[idx]
+    return
+end
+
+eltype_or(::Type{<: GPUArray}, or) = or
+eltype_or(::Type{<: GPUArray{T}}, or) where T = T
+eltype_or(::Type{<: GPUArray{T, N}}, or) where {T, N} = T
+
+function convert(AT::Type{<: GPUArray}, iter)
+    isize = Base.iteratorsize(iter)
+    style = indexstyle(iter)
+    ettrait = Base.iteratoreltype(iter)
+    if isbits(iter) && isize == Base.HasShape() && style != nothing && ettrait == Base.HasEltype()
+        # We can collect on the GPU
+        A = similar(AT, eltype_or(AT, eltype(iter)), size(iter))
+        gpu_call(collect_kernel, A, (A, iter, style))
+        A
+    else
+        convert(AT, collect(iter))
+    end
+end
+
 function convert(AT::Type{<: GPUArray{T, N}}, A::DenseArray{T, N}) where {T, N}
     copy!(AT(Base.size(A)), A)
 end

src/mapreduce.jl

@@ -3,11 +3,12 @@ import Base: any, count, countnz
 #############################
 # reduce
 # functions in base implemented with a direct loop need to be overloaded to use mapreduce
-any(pred, A::GPUArray) = Bool(mapreduce(pred, |, Cint(0), A))
+any(pred, A::GPUArray) = Bool(mapreduce(pred, |, Int32(0), A))
 count(pred, A::GPUArray) = Int(mapreduce(pred, +, UInt32(0), A))
 countnz(A::GPUArray) = Int(mapreduce(x-> x != 0, +, UInt32(0), A))
 countnz(A::GPUArray, dim) = Int(mapreducedim(x-> x != 0, +, UInt32(0), A, dim))
 
+Base.:(==)(A::GPUArray, B::GPUArray) = Bool(mapreduce(==, &, Int32(1), A, B))
 
 # hack to get around of fetching the first element of the GPUArray
 # as a startvalue, which is a bit complicated with the current reduce implementation

src/testsuite/construction.jl

@@ -1,11 +1,44 @@
 using GPUArrays
 using Base.Test, GPUArrays.TestSuite
 
+
+
+# It's kind of annoying to make FillArrays only a test dependency
+# so for texting the conversion to GPUArrays of shaped iterators,
+# I just copied the core types from FillArrays:s
+
+abstract type AbstractFill{T, N} <: AbstractArray{T, N} end
+@inline function Base.getindex(F::AbstractFill, k::Integer)
+    @boundscheck checkbounds(F, k)
+    getindex_value(F)
+end
+@inline function Base.getindex(F::AbstractFill{T, N}, kj::Vararg{<:Integer, N}) where {T, N}
+    @boundscheck checkbounds(F, kj...)
+    getindex_value(F)
+end
+Base.IndexStyle(F::AbstractFill) = IndexLinear()
+struct Fill{T, N} <: AbstractFill{T, N}
+    value::T
+    size::NTuple{N, Int}
+end
+getindex_value(x::Fill) = x.value
+@inline Base.size(F::Fill) = F.size
+struct Eye{T} <: AbstractMatrix{T}
+    size::NTuple{2, Int}
+end
+Base.size(E::Eye) = E.size
+@inline function Base.getindex(E::Eye{T}, k::Integer, j::Integer) where T
+    @boundscheck checkbounds(E, k, j)
+    ifelse(k == j, one(T), zero(T))
+end
+
+
 function run_construction(Typ)
     @testset "Construction" begin
         constructors(Typ)
         conversion(Typ)
         value_constructor(Typ)
+        iterator_constructors(Typ)
     end
 end
 
@@ -124,7 +157,7 @@ function value_constructor(Typ)
             @test all(x-> x == Int32(77), Array(x2))
 
             x = eye(T, 2, 2)
-            
+
             x1 = eye(Typ{T, 2}, 2, 2)
             x2 = eye(Typ{T}, (2, 2))
             x3 = eye(Typ{T, 2}, (2, 2))
@@ -135,3 +168,17 @@ function value_constructor(Typ)
         end
     end
 end
+function iterator_constructors(Typ)
+    @testset "iterator constructors" begin
+        for T in supported_eltypes()
+            @test Typ(Fill(T(0), (10,))) == zeros(Typ{T}, 10)
+            @test Typ(Fill(T(0), (10, 10))) == zeros(Typ{T}, 10, 10)
+            x = Typ{Float32}(Fill(T(0), (10, 10)))
+            eltype(x) == Float32
+
+            @test Typ(Eye{T}((10, 10))) == eye(Typ{T}, 10, 10)
+            x = Typ{Float32}(Eye{T}((10, 10)))
+            @test eltype(x) == Float32
+        end
+    end
+end