fix conversion of structs from GPU to CPU (#137)

PTsolvers · Apr 29, 2024 · f9c943d · f9c943d
1 parent 8afbd4b
commit f9c943d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/src/array_conversions.jl b/src/array_conversions.jl
@@ -5,24 +5,24 @@ struct CPUDeviceTrait <: DeviceTrait end
 struct NonCPUDeviceTrait <: DeviceTrait end
 
 @inline iscpu(::Array) = CPUDeviceTrait()
-@inline iscpu(::AbstractArray) = NonCPUDeviceTrait()
+@inline iscpu(::T) where {T<:AbstractArray} = NonCPUDeviceTrait()
 @inline iscpu(::T) where {T} = throw(ArgumentError("Unknown device"))
 
 @inline iscpu(::Velocity{Array{T,N}}) where {T,N} = CPUDeviceTrait()
-@inline iscpu(::Velocity{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
+@inline iscpu(::Velocity{<:AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
 
 @inline iscpu(::SymmetricTensor{Array{T,N}}) where {T,N} = CPUDeviceTrait()
-@inline iscpu(::SymmetricTensor{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
+@inline iscpu(::SymmetricTensor{<:AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
 
 @inline iscpu(::Residual{Array{T,N}}) where {T,N} = CPUDeviceTrait()
-@inline iscpu(::Residual{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
+@inline iscpu(::Residual{<:AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
 
 @inline iscpu(::ThermalArrays{Array{T,N}}) where {T,N} = CPUDeviceTrait()
-@inline iscpu(::ThermalArrays{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
+@inline iscpu(::ThermalArrays{<:AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait()
 
 @inline iscpu(::StokesArrays{M,A,B,C,Array{T,N},nDim}) where {M,A,B,C,T,N,nDim} =
     CPUDeviceTrait()
-@inline iscpu(::StokesArrays{M,A,B,C,AbstractArray{T,N},nDim}) where {M,A,B,C,T,N,nDim} =
+@inline iscpu(::StokesArrays{M,A,B,C,<:AbstractArray{T,N},nDim}) where {M,A,B,C,T,N,nDim} =
     NonCPUDeviceTrait()
 
 ## Conversion of structs to CPU

diff --git a/src/stokes/MetaStokes.jl b/src/stokes/MetaStokes.jl
@@ -20,7 +20,7 @@ function make_velocity_struct!(ndim::Integer; name::Symbol=:Velocity)
                 return new{$PTArray}(@zeros(ni[1]...), @zeros(ni[2]...), @zeros(ni[3]...))
             end
 
-            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{$PTArray}(args...)
+            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...)
         end
     end
 end
@@ -83,7 +83,7 @@ function make_symmetrictensor_struct!(nDim::Integer; name::Symbol=:SymmetricTens
                 )
             end
 
-            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{$PTArray}(args...)
+            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...)
         end
     end
 end
@@ -111,7 +111,7 @@ function make_residual_struct!(ndim; name::Symbol=:Residual)
                 return new{typeof(Rx)}(Rx, Ry, Rz, RP)
             end
 
-            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{$PTArray}(args...)
+            $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...)
         end
     end
 end

diff --git a/src/thermal_diffusion/MetaDiffusion.jl b/src/thermal_diffusion/MetaDiffusion.jl
@@ -19,7 +19,7 @@ function make_thermal_arrays!(ndim)
             ResT::_T
 
             function ThermalArrays(args::Vararg{T,N}) where {T<:AbstractArray,N}
-                return new{$PTArray}(args...)
+                return new{T}(args...)
             end
 
             function ThermalArrays(ni::NTuple{1,Integer})