diff --git a/src/array/copy.jl b/src/array/copy.jl index d1a229ba..db5493f9 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -109,3 +109,26 @@ function copyto_view!(Bpart, Brange, Apart, Arange) copyto!(view(Bpart, Brange), view(Apart, Arange)) return end + +function Base.copyto!(B::DArray{T,N}, A::Array{T,N}) where {T,N} + if size(B) != size(A) + # Fallback to the default implementation + return Base.invoke(copyto!, Tuple{AbstractArray, AbstractArray}, B, A) + end + + A_view = view(A, B.partitioning) + copyto!(B, A_view) + + return B +end +function Base.copyto!(B::Array{T,N}, A::DArray{T,N}) where {T,N} + if size(B) != size(A) + # Fallback to the default implementation + return Base.invoke(copyto!, Tuple{AbstractArray, AbstractArray}, B, A) + end + + B_view = view(B, A.partitioning) + copyto!(B_view, A) + + return B +end diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 0ef0d120..2f0f9370 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -269,11 +269,13 @@ function _memory_spans(a::StridedAliasing{T,N,S}, spans, ptr, dim) where {T,N,S} return spans end +memory_space(x::SubArray) = memory_space(parent(x)) function aliasing(x::SubArray{T,N,A}) where {T,N,A<:Array} if isbitstype(T) - S = CPURAMMemorySpace - return StridedAliasing{T,ndims(x),S}(RemotePtr{Cvoid}(pointer(parent(x))), - RemotePtr{Cvoid}(pointer(x)), + space = memory_space(x) + S = typeof(space) + return StridedAliasing{T,ndims(x),S}(RemotePtr{Cvoid}(pointer(parent(x)), space), + RemotePtr{Cvoid}(pointer(x), space), parentindices(x), size(x), strides(parent(x))) else diff --git a/test/array/fft.jl b/test/array/fft.jl index 54a4afe2..f55b11f7 100644 --- a/test/array/fft.jl +++ b/test/array/fft.jl @@ -1,141 +1,156 @@ using FFTW -@testset "FFT" begin - @testset for T in (ComplexF64, ComplexF32) - @testset "1D" begin - # Out-of-place - A = rand(T, 100) - DA = DArray(A) - B = fft(A) - DB = fft(DA) - @test DB isa DVector{T} - @test B ≈ collect(DB) - - # In-place - A = rand(T, 100) - DA = DArray(A) - fft!(A) - fft!(DA) - @test A ≈ collect(DA) - end - - @testset "2D" begin - # Out-of-place - A = rand(T, 100, 100) - DA = DArray(A) - B = fft(A) - DB = fft(DA) - @test DB isa DMatrix{T} - @test B ≈ collect(DB) - - # In-place - A = rand(T, 100, 100) - DA = DArray(A) - fft!(A) - fft!(DA) - @test A ≈ collect(DA) +function test_fft() + @testset "FFT" begin + @testset for T in (ComplexF64, ComplexF32) + @testset "1D" begin + # Out-of-place + A = rand(T, 100) + DA = DArray(A) + B = fft(A) + DB = fft(DA) + @test DB isa DVector{T} + @test B ≈ collect(DB) + + # In-place + A = rand(T, 100) + DA = DArray(A) + fft!(A) + fft!(DA) + @test A ≈ collect(DA) + end + + @testset "2D" begin + # Out-of-place + A = rand(T, 100, 100) + DA = DArray(A) + B = fft(A) + DB = fft(DA) + @test DB isa DMatrix{T} + @test B ≈ collect(DB) + + # In-place + A = rand(T, 100, 100) + DA = DArray(A) + fft!(A) + fft!(DA) + @test A ≈ collect(DA) + end + + @testset "3D" begin + # Out-of-place (Pencil) + A = rand(T, 100, 100, 100) + DA = DArray(A) + B = fft(A) + DB = fft(DA; decomp=:pencil) + @test DB isa DArray{T, 3} + @test B ≈ collect(DB) + + # Out-of-place (Slab) + A = rand(T, 100, 100, 100) + DA = DArray(A) + B = fft(A) + DB = fft(DA; decomp=:slab) + @test DB isa DArray{T, 3} + @test B ≈ collect(DB) + + # In-place (Pencil) + A = rand(T, 100, 100, 100) + DA = DArray(A) + fft!(A) + fft!(DA; decomp=:pencil) + @test A ≈ collect(DA) + + # In-place (Slab) + A = rand(T, 100, 100, 100) + DA = DArray(A) + fft!(A) + fft!(DA; decomp=:slab) + @test A ≈ collect(DA) + end end + end - @testset "3D" begin - # Out-of-place (Pencil) - A = rand(T, 100, 100, 100) - DA = DArray(A) - B = fft(A) - DB = fft(DA; decomp=:pencil) - @test DB isa DArray{T, 3} - @test B ≈ collect(DB) - - # Out-of-place (Slab) - A = rand(T, 100, 100, 100) - DA = DArray(A) - B = fft(A) - DB = fft(DA; decomp=:slab) - @test DB isa DArray{T, 3} - @test B ≈ collect(DB) - - # In-place (Pencil) - A = rand(T, 100, 100, 100) - DA = DArray(A) - fft!(A) - fft!(DA; decomp=:pencil) - @test A ≈ collect(DA) - - # In-place (Slab) - A = rand(T, 100, 100, 100) - DA = DArray(A) - fft!(A) - fft!(DA; decomp=:slab) - @test A ≈ collect(DA) + @testset "IFFT" begin + for T in (ComplexF64, ComplexF32) + @testset "1D" begin + # Out-of-place + A = rand(T, 100) + DA = DArray(A) + B = ifft(A) + DB = ifft(DA) + @test DB isa DVector{T} + @test B ≈ collect(DB) + + # In-place + A = rand(T, 100) + DA = DArray(A) + ifft!(A) + ifft!(DA) + @test A ≈ collect(DA) + end + + @testset "2D" begin + # Out-of-place + A = rand(T, 100, 100) + DA = DArray(A) + B = ifft(A) + DB = ifft(DA) + @test DB isa DMatrix{T} + @test B ≈ collect(DB) + + # In-place + A = rand(T, 100, 100) + DA = DArray(A) + ifft!(A) + ifft!(DA) + @test A ≈ collect(DA) + end + + @testset "3D" begin + # Out-of-place (Pencil) + A = rand(T, 100, 100, 100) + DA = DArray(A) + B = ifft(A) + DB = ifft(DA; decomp=:pencil) + @test DB isa DArray{T, 3} + @test B ≈ collect(DB) + + # Out-of-place (Slab) + A = rand(T, 100, 100, 100) + DA = DArray(A) + B = ifft(A) + DB = ifft(DA; decomp=:slab) + @test DB isa DArray{T, 3} + @test B ≈ collect(DB) + + # In-place (Pencil) + A = rand(T, 100, 100, 100) + DA = DArray(A) + ifft!(A) + ifft!(DA; decomp=:pencil) + @test A ≈ collect(DA) + + # In-place (Slab) + A = rand(T, 100, 100, 100) + DA = DArray(A) + ifft!(A) + ifft!(DA; decomp=:slab) + @test A ≈ collect(DA) + end end end end -@testset "IFFT" begin - for T in (ComplexF64, ComplexF32) - @testset "1D" begin - # Out-of-place - A = rand(T, 100) - DA = DArray(A) - B = ifft(A) - DB = ifft(DA) - @test DB isa DVector{T} - @test B ≈ collect(DB) - - # In-place - A = rand(T, 100) - DA = DArray(A) - ifft!(A) - ifft!(DA) - @test A ≈ collect(DA) - end - - @testset "2D" begin - # Out-of-place - A = rand(T, 100, 100) - DA = DArray(A) - B = ifft(A) - DB = ifft(DA) - @test DB isa DMatrix{T} - @test B ≈ collect(DB) - - # In-place - A = rand(T, 100, 100) - DA = DArray(A) - ifft!(A) - ifft!(DA) - @test A ≈ collect(DA) - end +@testset "CPU" begin + test_fft() +end - @testset "3D" begin - # Out-of-place (Pencil) - A = rand(T, 100, 100, 100) - DA = DArray(A) - B = ifft(A) - DB = ifft(DA; decomp=:pencil) - @test DB isa DArray{T, 3} - @test B ≈ collect(DB) - - # Out-of-place (Slab) - A = rand(T, 100, 100, 100) - DA = DArray(A) - B = ifft(A) - DB = ifft(DA; decomp=:slab) - @test DB isa DArray{T, 3} - @test B ≈ collect(DB) - - # In-place (Pencil) - A = rand(T, 100, 100, 100) - DA = DArray(A) - ifft!(A) - ifft!(DA; decomp=:pencil) - @test A ≈ collect(DA) - - # In-place (Slab) - A = rand(T, 100, 100, 100) - DA = DArray(A) - ifft!(A) - ifft!(DA; decomp=:slab) - @test A ≈ collect(DA) +for (kind, scope) in GPU_SCOPES + kind == :CUDA || kind == :ROCm || continue + @testset "$kind" begin + Dagger.with_options(;scope) do + test_fft() end end -end \ No newline at end of file +end diff --git a/test/runtests.jl b/test/runtests.jl index 8ef8cb3a..06fadfa7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,7 @@ if USE_GPU tests = [ ("GPU", "gpu.jl"), ("Array - Stencils", "array/stencil.jl"), + ("Array - FFT", "array/fft.jl"), ] end all_test_names = map(test -> replace(last(test), ".jl"=>""), tests)