Fix cache retrieval (#718)

JuliaGPU · Jan 9, 2025 · 8c937e3 · 8c937e3 · pxl-th · Jan 9, 2025
1 parent b6f3376
commit 8c937e3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <[email protected]>", "Valentin Churavy <[email protected]>", "Anton Smirnov <[email protected]>"]
-version = "1.2.0"
+version = "1.2.1"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

diff --git a/src/array.jl b/src/array.jl
@@ -6,10 +6,15 @@ mutable struct ROCArray{T, N, B} <: AbstractGPUArray{T, N}
     function ROCArray{T, N, B}(::UndefInitializer, dims::Dims{N}) where {T, N, B <: Mem.AbstractAMDBuffer}
         @assert isbitstype(T) "ROCArray only supports bits types"
         sz::Int64 = prod(dims) * sizeof(T)
-        return GPUArrays.cached_alloc((ROCArray, AMDGPU.device(), T, B, sz)) do
+        x = GPUArrays.cached_alloc((ROCArray, AMDGPU.device(), T, B, sz)) do
             @debug "Allocate `T=$T`, `dims=$dims`: $(Base.format_bytes(sz))"
             data = DataRef(pool_free, pool_alloc(B, sz))
             return finalizer(unsafe_free!, new{T, N, B}(data, dims, 0))
+        end
+        return if size(x) != dims
+            reshape(x, dims)
+        else
+            x
         end::ROCArray{T, N, B}
     end