diff --git a/ext/JustPICAMDGPUExt.jl b/ext/JustPICAMDGPUExt.jl index 576eb059..af61dc51 100644 --- a/ext/JustPICAMDGPUExt.jl +++ b/ext/JustPICAMDGPUExt.jl @@ -7,6 +7,17 @@ import JustPIC: AbstractBackend, AMDGPUBackend JustPIC.TA(::Type{AMDGPUBackend}) = ROCArray +function AMDGPU.ROCArray(particles::JustPIC.Particles{JustPIC.AMDGPUBackend}) + (; coords, index, nxcell, max_xcell, min_xcell, np) = particles + coords_gpu = CuArray.(coords); + return Particles(CUDABackend, coords_gpu, CuArray(index), nxcell, max_xcell, min_xcell, np) +end + +function AMDGPU.ROCArray(phase_ratios::JustPIC.PhaseRatios{JustPIC.AMDGPUBackend}) + (; vertex, center) = phase_ratios + return JustPIC.PhaseRatios(CUDABackend, CuArray(vertex), CuArray(center)) +end + module _2D using AMDGPU using ImplicitGlobalGrid @@ -36,6 +47,7 @@ module _2D include(joinpath(@__DIR__, "../src/common.jl")) include(joinpath(@__DIR__, "../src/AMDGPUExt/CellArrays.jl")) + function JustPIC._2D.Particles( coords, index::CellArray{StaticArraysCore.SVector{N1,Bool},2,0,ROCArray{Bool,N2}}, @@ -232,7 +244,7 @@ module _2D function JustPIC._2D.PhaseRatios( ::Type{AMDGPUBackend}, nphases::Integer, ni::NTuple{N,Integer} ) where {N} - return PhaseRatios(Float64, AMDGPUBackend, nphases, ni) + return JustPIC._2D.PhaseRatios(Float64, AMDGPUBackend, nphases, ni) end function JustPIC._2D.PhaseRatios( @@ -241,7 +253,7 @@ module _2D center = cell_array(0.0, (nphases,), ni) vertex = cell_array(0.0, (nphases,), ni .+ 1) - return PhaseRatios{B,typeof(center)}(center, vertex) + return JustPIC.PhaseRatios(AMDGPUBackend, center, vertex) end function JustPIC._2D.phase_ratios_center!( @@ -297,6 +309,7 @@ module _3D include(joinpath(@__DIR__, "../src/common.jl")) include(joinpath(@__DIR__, "../src/AMDGPUExt/CellArrays.jl")) + function JustPIC._3D.Particles( coords, index::CellArray{StaticArraysCore.SVector{N1,Bool},3,0,ROCArray{Bool,N2}}, @@ -470,7 +483,7 @@ module _3D function JustPIC._3D.PhaseRatios( ::Type{AMDGPUBackend}, nphases::Integer, ni::NTuple{N,Integer} ) where {N} - return PhaseRatios(Float64, AMDGPUBackend, nphases, ni) + return JustPIC._3D.PhaseRatios(Float64, AMDGPUBackend, nphases, ni) end function JustPIC._3D.PhaseRatios( @@ -479,7 +492,7 @@ module _3D center = cell_array(0.0, (nphases,), ni) vertex = cell_array(0.0, (nphases,), ni .+ 1) - return PhaseRatios{B,typeof(center)}(center, vertex) + return JustPIC.PhaseRatios(AMDGPUBackend, center, vertex) end function JustPIC._3D.phase_ratios_center!( diff --git a/ext/JustPICCUDAExt.jl b/ext/JustPICCUDAExt.jl index cd68fc4f..c290ef86 100644 --- a/ext/JustPICCUDAExt.jl +++ b/ext/JustPICCUDAExt.jl @@ -5,6 +5,17 @@ using JustPIC JustPIC.TA(::Type{CUDABackend}) = CuArray +function CUDA.CuArray(particles::JustPIC.Particles{JustPIC.CPUBackend}) + (; coords, index, nxcell, max_xcell, min_xcell, np) = particles + coords_gpu = CuArray.(coords); + return Particles(CUDABackend, coords_gpu, CuArray(index), nxcell, max_xcell, min_xcell, np) +end + +function CUDA.CuArray(phase_ratios::JustPIC.PhaseRatios{JustPIC.CPUBackend}) + (; vertex, center) = phase_ratios + return JustPIC.PhaseRatios(CUDABackend, CuArray(vertex), CuArray(center)) +end + module _2D using CUDA using ImplicitGlobalGrid @@ -36,6 +47,8 @@ module _2D include(joinpath(@__DIR__, "../src/common.jl")) include(joinpath(@__DIR__, "../src/CUDAExt/CellArrays.jl")) + # Conversions + function JustPIC._2D.Particles( coords, index::CellArray{StaticArraysCore.SVector{N1,Bool},2,0,CuArray{Bool,N2}}, @@ -226,7 +239,7 @@ module _2D function JustPIC._2D.PhaseRatios( ::Type{CUDABackend}, nphases::Integer, ni::NTuple{N,Integer} ) where {N} - return PhaseRatios(Float64, CUDABackend, nphases, ni) + return JustPIC._2D.PhaseRatios(Float64, CUDABackend, nphases, ni) end function JustPIC._2D.PhaseRatios( @@ -235,7 +248,7 @@ module _2D center = cell_array(0.0, (nphases,), ni) vertex = cell_array(0.0, (nphases,), ni .+ 1) - return PhaseRatios{B,typeof(center)}(center, vertex) + return JustPIC.PhaseRatios(CUDABackend, center, vertex) end function JustPIC._2D.phase_ratios_center!( @@ -291,6 +304,8 @@ module _3D include(joinpath(@__DIR__, "../src/common.jl")) include(joinpath(@__DIR__, "../src/CUDAExt/CellArrays.jl")) + + # Conversions function JustPIC._3D.Particles( coords, @@ -465,7 +480,7 @@ module _3D function JustPIC._3D.PhaseRatios( ::Type{CUDABackend}, nphases::Integer, ni::NTuple{N,Integer} ) where {N} - return PhaseRatios(Float64, CUDABackend, nphases, ni) + return JustPIC._3D.PhaseRatios(Float64, CUDABackend, nphases, ni) end function JustPIC._3D.PhaseRatios( @@ -474,7 +489,7 @@ module _3D center = cell_array(0.0, (nphases,), ni) vertex = cell_array(0.0, (nphases,), ni .+ 1) - return PhaseRatios{B,typeof(center)}(center, vertex) + return JustPIC.PhaseRatios(CUDABackend, center, vertex) end function JustPIC._3D.phase_ratios_center!( diff --git a/src/CellArrays/conversion.jl b/src/CellArrays/conversion.jl index 259f06e7..d0015312 100644 --- a/src/CellArrays/conversion.jl +++ b/src/CellArrays/conversion.jl @@ -3,7 +3,7 @@ import Base: Array, copy @inline remove_parameters(::T) where {T} = Base.typename(T).wrapper # detect if we are on the CPU (`Val{false}`) or GPU (`Val{true}`) -@inline isdevice(::Type{Array}) = Val(false) +@inline isdevice(::Type{Array{T, N}}) where {T,N} = Val(false) @inline isdevice(::Type{T}) where {T<:AbstractArray} = Val(true) # this is a big assumption but still @inline isdevice(::T) where {T} = throw(ArgumentError("$(T) is not a supported CellArray type.")) @@ -33,8 +33,8 @@ function Array(x::T) where {T<:AbstractParticles} return T_clean(CPUBackend, cpu_fields...) end +_Array(x) = x _Array(::Nothing) = nothing -_Array(::T) where {T} = T _Array(x::AbstractArray) = Array(x) _Array(x::NTuple{N,T}) where {N,T} = ntuple(i -> _Array(x[i]), Val(N)) diff --git a/src/JustPIC.jl b/src/JustPIC.jl index d4c5516b..ff2e6883 100644 --- a/src/JustPIC.jl +++ b/src/JustPIC.jl @@ -17,12 +17,12 @@ function CA end TA() = Array TA(::Type{CPUBackend}) = Array -include("PhaseRatios/PhaseRatios.jl") -export nphases, numphases - include("particles.jl") export AbstractParticles, Particles, MarkerChain, PassiveMarkers, cell_index, cell_length +include("PhaseRatios/PhaseRatios.jl") +export nphases, numphases + include("Advection/types.jl") export AbstractAdvectionIntegrator, Euler, RungeKutta2 diff --git a/src/PhaseRatios/PhaseRatios.jl b/src/PhaseRatios/PhaseRatios.jl index 13ff05ec..4647f61a 100644 --- a/src/PhaseRatios/PhaseRatios.jl +++ b/src/PhaseRatios/PhaseRatios.jl @@ -1,8 +1,13 @@ -struct PhaseRatios{Backend,T} +struct PhaseRatios{Backend,T} <: AbstractParticles center::T vertex::T + + function PhaseRatios(::Type{B}, center::T, vertex::T) where {B, T} + return new{B, T}(center, vertex) + end end + """ nphases(x::PhaseRatios) diff --git a/src/PhaseRatios/constructors.jl b/src/PhaseRatios/constructors.jl index c8a2cae3..f94397d4 100644 --- a/src/PhaseRatios/constructors.jl +++ b/src/PhaseRatios/constructors.jl @@ -4,9 +4,13 @@ function PhaseRatios( center = cell_array(0.0, (nphases,), ni) vertex = cell_array(0.0, (nphases,), ni .+ 1) - return JustPIC.PhaseRatios{B,typeof(center)}(center, vertex) + return JustPIC.PhaseRatios(B, center, vertex) end +# function PhaseRatios(::Type{B}, center, vertex) where {B} +# return JustPIC.PhaseRatios{B, typeof(center)}(center, vertex) +# end + function PhaseRatios(nphases::Integer, ni::NTuple{N,Integer}) where {N} return PhaseRatios(Float64, CPUBackend, nphases, ni) end diff --git a/test/Project.toml b/test/Project.toml index 3ed7dcae..b53ffc52 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,6 @@ [deps] CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ParallelStencil = "94395366-693c-11ea-3b26-d9b7aac5d958" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" diff --git a/test/test_save_load.jl b/test/test_save_load.jl new file mode 100644 index 00000000..905b3fa9 --- /dev/null +++ b/test/test_save_load.jl @@ -0,0 +1,56 @@ +using JLD2, JustPIC, JustPIC._2D + +@testset "Save and load" begin + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 6, 6, 6 + n = 128 + nx = ny = n-1 + ni = nx, ny + Lx = Ly = 1.0 + # nodal vertices + xvi = xv, yv = range(0, Lx, length=n), range(0, Ly, length=n) + + particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi...,); + phases, = init_cell_arrays(particles, Val(1)); + phase_ratios = PhaseRatios(backend, 2, ni); + phase_ratios = PhaseRatios(JustPIC.CPUBackend, 2, ni); + + jldsave( + "particles.jld2"; + particles = Array(particles), + phases = Array(phases), + phase_ratios = Array(phase_ratios) + ) + + data = load("particles.jld2") + particles2 = data["particles"] + phases2 = data["phases"] + phase_ratios2 = data["phase_ratios"] + + @test Array(particles.coords[1].data) == particles2.coords[1].data + @test Array(particles.coords[2].data) == particles2.coords[2].data + @test Array(particles.index.data) == particles2.index.data + @test Array(phase_ratios.center.data) == phase_ratios2.center.data + @test Array(phase_ratios.vertex.data) == phase_ratios2.vertex.data + @test Array(phases.data) == phases2.data + + if isdefined(Main, :CUDA) + particles_cuda = CuArray(particles2) + phase_ratios_cuda = CuArray(phase_ratios2) + phases_cuda = CuArray(phases2) + + @test particles_cuda isa JustPIC.Particles{CUDABackend} + @test phase_ratios_cuda isa JustPIC.PhaseRatios{CUDABackend} + @test phases_cuda isa CuArray + + elseif isdefined(Main, :AMDGPU) + particles_amdgpu = ROCArray(particles2) + phase_ratios_amdgpu = ROCArray(phase_ratios2) + phases_amdgpu = ROCArray(phases2) + + @test particles_cuda isa JustPIC.Particles{AMDGPUBackend} + @test phase_ratios_cuda isa JustPIC.PhaseRatios{AMDGPUBackend} + @test phases_cuda isa ROCArray + + end +end