Add reshape for CuDeviceArray (#1561)

Author: omlins

src/device/array.jl

@@ -265,3 +265,22 @@ function Base.reinterpret(::Type{T}, a::CuDeviceArray{S,N,A}) where {T,S,N,A}
   osize = tuple(size1, Base.tail(isize)...)
   return CuDeviceArray{T,N,A}(osize, reinterpret(LLVMPtr{T,A}, a.ptr), a.maxsize)
 end
+
+
+## reshape
+
+function Base.reshape(a::CuDeviceArray{T,M}, dims::NTuple{N,Int}) where {T,N,M}
+  if prod(dims) != length(a)
+      throw(DimensionMismatch("new dimensions (argument `dims`) must be consistent with array size (`size(a)`)"))
+  end
+  if N == M && dims == size(a)
+      return a
+  end
+  _derived_array(T, N, a, dims)
+end
+
+# create a derived device array (reinterpreted or reshaped) that's still a CuDeviceArray
+@inline function _derived_array(::Type{T}, N::Int, a::CuDeviceArray{T,M,A}, osize::Dims) where {T, M, A}
+  return CuDeviceArray{T,N,A}(osize, a.ptr, a.maxsize)
+end
+

test/device/array.jl

@@ -137,6 +137,29 @@ end
     @test array == Array(array_dev)
 end
 
+@testset "reshape" begin
+    function kernel(array)
+        i = (blockIdx().x-1i32) * blockDim().x + threadIdx().x
+        j = (blockIdx().y-1i32) * blockDim().y + threadIdx().y
+
+        _array2d = reshape(array, 10, 10)
+        _array2d[i,j] = i + (j-1)*size(_array2d,1)
+
+        return
+    end
+
+    array = zeros(Int64, 100)
+    array_dev = CuArray(array)
+
+    array2d = reshape(array, 10, 10)
+    for i in 1:size(array2d,1), j in 1:size(array2d,2)
+        array2d[i,j] = i + (j-1)*size(array2d,1)
+    end
+
+    @cuda threads=(10, 10) kernel(array_dev)
+    @test array == Array(array_dev)
+end
+
 @testset "non-Int index to unsafe_load" begin
     function kernel(a)
         a[UInt64(1)] = 1