Enable allowscalar(false) testing for JLArray.

Author: maleadt

src/array.jl

@@ -1,5 +1,8 @@
-# Very simple Julia backend which is just for testing the implementation
-# and can be used as a reference implementation
+# Very simple Julia back-end which is just for testing the implementation and can be used as
+# a reference implementation
+
+
+## construction
 
 struct JLArray{T, N} <: GPUArray{T, N}
     data::Array{T, N}
@@ -12,6 +15,22 @@ end
 
 JLArray(data::AbstractArray{T, N}, size::Dims{N}) where {T,N} = JLArray{T,N}(data, size)
 
+(::Type{<: JLArray{T}})(x::AbstractArray) where T = JLArray(convert(Array{T}, x), size(x))
+
+function JLArray{T, N}(size::NTuple{N, Integer}) where {T, N}
+    JLArray{T, N}(Array{T, N}(undef, size), size)
+end
+
+
+## getters
+
+size(x::JLArray) = x.size
+
+pointer(x::JLArray) = pointer(x.data)
+
+
+## I/O
+
 Base.show(io::IO, x::JLArray) = show(io, collect(x))
 Base.show(io::IO, x::LinearAlgebra.Adjoint{<:Any,<:JLArray}) = show(io, LinearAlgebra.adjoint(collect(x.parent)))
 Base.show(io::IO, x::LinearAlgebra.Transpose{<:Any,<:JLArray}) = show(io, LinearAlgebra.transpose(collect(x.parent)))
@@ -20,15 +39,16 @@ Base.show(io::IO, ::MIME"text/plain", x::JLArray) = show(io, MIME"text/plain"(),
 Base.show(io::IO, ::MIME"text/plain", x::LinearAlgebra.Adjoint{<:Any,<:JLArray}) = show(io, MIME"text/plain"(), LinearAlgebra.adjoint(collect(x.parent)))
 Base.show(io::IO, ::MIME"text/plain", x::LinearAlgebra.Transpose{<:Any,<:JLArray}) = show(io, MIME"text/plain"(), LinearAlgebra.transpose(collect(x.parent)))
 
+
+## other
+
 """
 Thread group local memory
 """
 struct LocalMem{N, T}
     x::NTuple{N, Vector{T}}
 end
 
-size(x::JLArray) = x.size
-pointer(x::JLArray) = pointer(x.data)
 to_device(state, x::JLArray) = x.data
 to_device(state, x::Tuple) = to_device.(Ref(state), x)
 to_device(state, x::RefValue{<: JLArray}) = RefValue(to_device(state, x[]))
@@ -40,12 +60,6 @@ to_blocks(state, x) = x
 # unpacks local memory for each block
 to_blocks(state, x::LocalMem) = x.x[blockidx_x(state)]
 
-(::Type{<: JLArray{T}})(x::AbstractArray) where T = JLArray(convert(Array{T}, x), size(x))
-
-function JLArray{T, N}(size::NTuple{N, Integer}) where {T, N}
-    JLArray{T, N}(Array{T, N}(undef, size), size)
-end
-
 similar(::Type{<: JLArray}, ::Type{T}, size::Base.Dims{N}) where {T, N} = JLArray{T, N}(size)
 
 function unsafe_reinterpret(::Type{T}, A::JLArray{ET}, size::NTuple{N, Integer}) where {T, ET, N}
@@ -131,7 +145,8 @@ function _gpu_call(f, A::JLArray, args::Tuple, blocks_threads::Tuple{T, T}) wher
         block_args = to_blocks.(Ref(state), device_args)
         for threadidx in CartesianIndices(threads)
             thread_state = JLState(state, threadidx.I)
-            tasks[threadidx] = @async f(thread_state, block_args...)
+            tasks[threadidx] = @async @allowscalar f(thread_state, block_args...)
+            # TODO: @async obfuscates the trace to any exception which happens during f
         end
         for t in tasks
             fetch(t)
@@ -146,7 +161,6 @@ device(x::JLArray) = JLDevice()
 threads(dev::JLDevice) = 256
 blocks(dev::JLDevice) = (256, 256, 256)
 
-
 @inline function synchronize_threads(::JLState)
     #=
     All threads are getting started asynchronously,so a yield will
@@ -168,8 +182,9 @@ end
 blas_module(::JLArray) = LinearAlgebra.BLAS
 blasbuffer(A::JLArray) = A.data
 
+# defining our own plan type is the easiest way to pass around the plans in Base interface
+# without ambiguities
 
-# defining our own plan type is the easiest way to pass around the plans in Base interface without ambiguities
 struct FFTPlan{T}
     p::T
 end
@@ -192,7 +207,6 @@ function plan_ifft(A::JLArray; kw_args...)
     FFTPlan(plan_ifft(A.data; kw_args...))
 end
 
-
 function *(plan::FFTPlan, A::JLArray)
     x = plan.p * A.data
     JLArray(x)

src/indexing.jl

@@ -1,10 +1,20 @@
-const _allowslow = Ref(true)
+const _allowscalar = Ref(true)
 
-allowslow(flag = true) = (_allowslow[] = flag)
+allowscalar(flag = true) = (_allowscalar[] = flag)
 
-function assertslow(op = "Operation")
-    # _allowslow[] || error("$op is disabled")
-    return
+function assertscalar(op = "Operation")
+  _allowscalar[] || error("$op is disabled")
+  return
+end
+
+macro allowscalar(ex)
+    quote
+        local prev = _allowscalar[]
+        _allowscalar[] = true
+        local ret = $(esc(ex))
+        _allowscalar[] = prev
+        ret
+    end
 end
 
 Base.IndexStyle(::Type{<:GPUArray}) = Base.IndexLinear()
@@ -16,7 +26,7 @@ function _getindex(xs::GPUArray{T}, i::Integer) where T
 end
 
 function Base.getindex(xs::GPUArray{T}, i::Integer) where T
-    assertslow("getindex")
+  ndims(xs) > 0 && assertscalar("scalar getindex")
     _getindex(xs, i)
 end
 
@@ -27,12 +37,13 @@ function _setindex!(xs::GPUArray{T}, v::T, i::Integer) where T
 end
 
 function Base.setindex!(xs::GPUArray{T}, v::T, i::Integer) where T
-    assertslow("setindex!")
+  assertscalar("scalar setindex!")
     _setindex!(xs, v, i)
 end
 
 Base.setindex!(xs::GPUArray, v, i::Integer) = xs[i] = convert(eltype(xs), v)
 
+
 # Vector indexing
 
 to_index(a, x) = x

src/testsuite.jl

@@ -60,7 +60,11 @@ end
 Runs the GPUArrays test suite on array type `Typ`
 """
 function test(Typ)
-    GPUArrays.allowslow(false)
+    # TODO: more fine-grained allowscalar within test_indexing
+    GPUArrays.allowscalar(true)
+    TestSuite.test_indexing(Typ)
+
+    GPUArrays.allowscalar(false)
     TestSuite.test_gpuinterface(Typ)
     TestSuite.test_base(Typ)
     TestSuite.test_blas(Typ)
@@ -69,6 +73,5 @@ function test(Typ)
     TestSuite.test_fft(Typ)
     TestSuite.test_linalg(Typ)
     TestSuite.test_mapreduce(Typ)
-    TestSuite.test_indexing(Typ)
     TestSuite.test_random(Typ)
 end

src/testsuite/base.jl

@@ -68,9 +68,6 @@ function test_base(Typ)
             @test Array(a) == x
         end
 
-        GPUArrays.allowslow(true)
-        # right now in CLArrays we fallback to geindex since on some hardware
-        # somehow the vcat kernel segfaults -.-
         @testset "vcat + hcat" begin
             x = fill(0f0, (10, 10))
             y = rand(Float32, 20, 10)
@@ -84,7 +81,6 @@ function test_base(Typ)
             against_base(hcat, Typ{Float32}, (3, 3), (3, 3))
             against_base(vcat, Typ{Float32}, (3, 3), (3, 3))
         end
-        GPUArrays.allowslow(false)
 
         @testset "reinterpret" begin
             a = rand(ComplexF32, 22)

src/testsuite/indexing.jl

@@ -4,11 +4,9 @@ function test_indexing(Typ)
             @testset "Indexing with $T" begin
                 x = rand(T, 32)
                 src = Typ(x)
-                GPUArrays.allowslow(true)
                 for (i, xi) in enumerate(x)
                     @test src[i] == xi
                 end
-                GPUArrays.allowslow(false)
                 @test Array(src[1:3]) == x[1:3]
                 @test Array(src[3:end]) == x[3:end]
             end
@@ -25,7 +23,6 @@ function test_indexing(Typ)
             @testset "Indexing with $T" begin
                 x = fill(zero(T), 7)
                 src = Typ(x)
-                GPUArrays.allowslow(true)
                 for i = 1:7
                     src[i] = i
                 end
@@ -34,7 +31,6 @@ function test_indexing(Typ)
                 @test Array(src[1:3]) == T[77, 22, 11]
                 src[1] = T(0)
                 src[2:end] = T(77)
-                GPUArrays.allowslow(false)
                 @test Array(src) == T[0, 77, 77, 77, 77, 77, 77]
             end
         end
@@ -43,18 +39,15 @@ function test_indexing(Typ)
             @testset "issue #42 with $T" begin
                 Ac = rand(Float32, 2, 2)
                 A = Typ(Ac)
-                GPUArrays.allowslow(true)
                 @test A[1] == Ac[1]
                 @test A[end] == Ac[end]
                 @test A[1, 1] == Ac[1, 1]
-                GPUArrays.allowslow(false)
             end
         end
         for T in (Float32, Int32)
             @testset "Colon() $T" begin
                 Ac = rand(T, 10)
                 A = Typ(Ac)
-                GPUArrays.allowslow(false)
                 A[:] = T(1)
                 @test all(x-> x == 1, A)
                 A[:] = Typ(Ac)