fixing lots of Knet tests

Author: SimonDanisch

src/blas.jl

@@ -62,13 +62,28 @@ for elty in (Float32, Float64, Complex64, Complex128)
             elseif trans == 'T' && (length(X) != m || length(Y) != n)
                 throw(DimensionMismatch("A.' has dimensions $n, $m, X has length $(length(X)) and Y has length $(length(Y))"))
             end
-            ctx = context(A)
-            blasmod = blas_module(ctx)
+            blasmod = blas_module(A)
             blasmod.gemv!(
                 trans, alpha,
-                blasbuffer(ctx, A), blasbuffer(ctx, X), beta, blasbuffer(ctx, Y)
+                blasbuffer(A), blasbuffer(X), beta, blasbuffer(Y)
             )
             Y
         end
     end
 end
+
+
+for elty in (Float32, Float64, Complex64, Complex128)
+    @eval begin
+        function Base.BLAS.axpy!(
+                alpha::Number, x::GPUArray{$elty}, y::GPUArray{$elty}
+            )
+            if length(x) != length(y)
+                throw(DimensionMismatch("x has length $(length(x)), but y has length $(length(y))"))
+            end
+            blasmod = blas_module(A)
+            blasmod.axpy!($elty(alpha), blasbuffer(dx), blasbuffer(dx))
+            y
+        end
+    end
+end

src/jlbackend.jl

@@ -78,6 +78,10 @@ function gpu_call(f, A::JLArray, args::Tuple, blocks = nothing, threads = C_NULL
 end
 
 # "intrinsics"
+struct JLDevice end
+device(x::JLArray) = JLDevice()
+threads(dev::JLDevice) = 256
+
 
 @inline synchronize_threads(::JLState) = nothing
 

src/mapreduce.jl

@@ -1,5 +1,9 @@
 #############################
 # 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), (u)))
+count(pred, A::GPUArray) = Int(mapreduce(pred, +, Cuint(0), A))
+
 
 # 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
@@ -28,8 +32,7 @@ function Base.mapreduce{T, N}(f::Function, op::Function, A::GPUArray{T, N})
     v0 = startvalue(op, OT) # TODO do this better
     mapreduce(f, op, v0, A)
 end
-
-
+function acc_mapreduce end
 function Base.mapreduce(f, op, v0, A::GPUArray, B::GPUArray, C::Number)
     acc_mapreduce(f, op, v0, A, (B, C))
 end
@@ -65,7 +68,3 @@ function Base._mapreducedim!(f, op, R::GPUArray, A::GPUArray)
     gpu_call(mapreducedim_kernel, R, (f, op, R, A, Cuint(slice_size), Cuint.(size(A)), Cuint(dim)))
     return R
 end
-
-
-any(pred, A::GPUArray) = Bool(mapreduce(isnan, |, Cint(0), (u)))
-count(pred, A::GPUArray) = Int(mapreduce(pred, +, Cuint(0), A))

src/random.jl

@@ -0,0 +1,56 @@
+using GPUArrays
+function TausStep(z::Unsigned, S1::Integer, S2::Integer, S3::Integer, M::Unsigned)
+    b = (((z << S1) ⊻ z) >> S2)
+    return (((z & M) << S3) ⊻ b)
+end
+
+LCGStep(z::Unsigned, A::Unsigned, C::Unsigned) = A * z + C
+
+make_rand_num(::Type{Float64}, tmp) = 2.3283064365387e-10 * Float64(tmp)
+make_rand_num(::Type{Float32}, tmp) = 2.3283064f-10 * Float32(tmp)
+
+
+function next_rand(::Type{FT}, state::NTuple{4, T}) where {FT, T <: Unsigned}
+    state = (
+        TausStep(state[1], Cint(13), Cint(19), Cint(12), T(4294967294)),
+        TausStep(state[2], Cint(2), Cint(25), Cint(4), T(4294967288)),
+        TausStep(state[3], Cint(3), Cint(11), Cint(17), T(4294967280)),
+        LCGStep(state[4], T(1664525), T(1013904223))
+    )
+    tmp = (state[1] ⊻ state[2] ⊻ state[3] ⊻ state[4])
+    return (
+        state,
+        make_rand_num(FT, tmp)
+    )
+end
+
+function gpu_rand(::Type{T}, state, randstate::AbstractVector{NTuple{4, Cuint}}) where T
+    threadid = GPUArrays.threadidx_x(state)
+    stateful_rand = next_rand(T, randstate[threadid])
+    randstate[threadid] = stateful_rand[1]
+    return stateful_rand[2]
+end
+
+global cached_state, clear_cache
+let rand_state_dict = Dict()
+    clear_cache() = (empty!(rand_state_dict); return)
+    function cached_state(x)
+        dev = GPUArrays.device(x)
+        get!(rand_state_dict, dev) do
+            N = GPUArrays.threads(dev)
+            res = similar(x, NTuple{4, Cuint}, N)
+            copy!(res, [ntuple(i-> rand(Cuint), 4) for i=1:N])
+            res
+        end
+    end
+end
+function Base.rand!(A::GPUArray{T}) where T <: AbstractFloat
+    rstates = cached_state(A)
+    gpu_call(A, (rstates, A,)) do state, randstates, a
+        idx = linear_index(state)
+        idx > length(a) && return
+        a[idx] = gpu_rand(T, state, randstates)
+        return
+    end
+    A
+end