add device array reshape (#713)

* add device array reshape

* add device array reshape unit test

* fix error handling

src/device/gcn/array.jl

@@ -133,3 +133,20 @@ end
         nothing
     end
 end
+
+## reshape
+
+function Base.reshape(a::ROCDeviceArray{T,M,A}, dims::NTuple{N,Int}) where {T,N,M,A}
+    if prod(dims) != length(a)
+        throw(nothing) # standing for: 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(a, T, dims)
+end
+
+# create a derived device array (reinterpreted or reshaped) that's still a ROCDeviceArray
+@inline function _derived_array(a::ROCDeviceArray{<:Any,<:Any,A}, ::Type{T}, dims::Dims{N}) where {T, N, A}
+    return ROCDeviceArray{T,N,A}(dims, a.ptr)
+end

test/device/array.jl

@@ -44,3 +44,26 @@ end
     o2 = occursin("4×4 device array wrapper Adjoint", adj_repr)
     @test o1 || o2
 end
+
+@testset "reshape" begin
+    function kernel(array)
+        i = (workgroupIdx().x-1) * workgroupDim().x + workitemIdx().x
+        j = (workgroupIdx().y-1) * workgroupDim().y + workitemIdx().y
+
+        _array2d = reshape(array, 10, 10)
+        _array2d[i,j] = i + (j-1)*size(_array2d,1)
+
+        return
+    end
+
+    array = zeros(Int64, 100)
+    array_dev = ROCArray(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
+
+    @roc groupsize=(10, 10) kernel(array_dev)
+    @test array == Array(array_dev)
+end