Merge pull request #178 from slimgroup/cuda-ctx

Init openacc ctx at init to avoid cuda conflicts

Project.toml

@@ -1,7 +1,7 @@
 name = "JUDI"
 uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 authors = ["Philipp Witte, Mathias Louboutin"]
-version = "3.2.2"
+version = "3.2.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/JUDI.jl

@@ -119,6 +119,11 @@ function __init__()
     # Initialize lock at session start
     PYLOCK[] = ReentrantLock()
 
+    if get(ENV, "DEVITO_PLATFORM", "") == "nvidiaX"
+        @info "Initializing openacc/openmp offloading"
+        devito_model(Model((21, 21, 21), (10., 10., 10.), (0., 0., 0.), randn(Float32, 21, 21, 21)), Options())
+    end
+
     @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
         Zygote.unbroadcast(x::AbstractArray, x̄::LazyPropagation) = Zygote.unbroadcast(x, eval_prop(x̄))
     end

src/TimeModeling/LinearOperators/basics.jl

@@ -25,8 +25,10 @@ setindex!(S::AbstractSize, v, I::Symbol) = setindex!(S.dims, v, I)
 iterate(S::AbstractSize) = iterate(S.dims)
 iterate(S::AbstractSize, state) = iterate(S.dims, state)
 
+Base.isless(i::Int64, a::AbstractSize) = isless(i, nsamples(a))
+
 convert(::Type{T}, S::AbstractSize) where T<:Number = convert(T, nsamples(S))
-(::Type{T})(S::AbstractSize) where T<:Number = convert(T, nsamples(S))
+(::Type{T})(S::AbstractSize) where T<:Union{Integer, AbstractFloat} = convert(T, nsamples(S))
 (c::Colon)(i::T, S::AbstractSize) where T = c(i, T(S))
 
 Base.keys(S::AbstractSize) = keys(S.dims)

src/TimeModeling/Preconditioners/ModelPreconditioners.jl

@@ -77,6 +77,7 @@ end
 matvec(D::TopMute{T, N}, x::PhysicalParameter{T}) where {T, N} = PhysicalParameter(x, matvec(D, x.data))
 matvec(D::TopMute{T, N}, x::judiWeights{T}) where {T, N} = judiWeights{T}(x.nsrc, [matvec(D, x.data[s]) for s=1:x.nsrc])
 matvec(D::TopMute{T, N}, x::Vector{T}) where {T, N} = vec(matvec(D, reshape(x, size(D.wb)..., :)))
+matvec_T(D::TopMute{T, N}, x) where {T, N} = matvec(D, x)
 
 # Real diagonal operator
 conj(I::TopMute{T, N}) where {T, N} = I

src/TimeModeling/Types/ModelStructure.jl

@@ -203,6 +203,8 @@ function similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{PhysicalParamete
     PhysicalParameter(Ad, A.d, A.o)
 end
 
+similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{PhysicalParameter}}) = similar(bc, eltype(find_pm(bc)))
+
 "`A = find_pm(As)` returns the first PhysicalParameter among the arguments."
 find_pm(bc::Base.Broadcast.Broadcasted) = find_pm(bc.args)
 find_pm(args::Tuple) = find_pm(find_pm(args[1]), Base.tail(args))