Use optimized findall on GPU

Requires JuliaGPU/CuArrays.jl#446
SciML · Oct 10, 2019 · df1671b · df1671b
1 parent 803bc9c
commit df1671b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/src/callbacks.jl b/src/callbacks.jl
@@ -301,7 +301,7 @@ end
   next_condition = get_condition(integrator, callback, abst)
   @. next_sign = sign(next_condition)
 
-  event_idx = findall(x-> ((prev_sign[x] < 0 && callback.affect! !== nothing) || (prev_sign[x] > 0 && callback.affect_neg! !== nothing)) && prev_sign[x]*next_sign[x]<=0, keys(prev_sign))
+  event_idx = findall_events(affect!,affect_neg!,prev_sign,next_sign)
   if length(event_idx) != 0
     event_occurred = true
     interp_index = callback.interp_points
@@ -310,7 +310,7 @@ end
     for i in 2:length(Θs)
       abst = integrator.tprev+integrator.dt*Θs[i]
       new_sign = get_condition(integrator, callback, abst)
-      _event_idx = findall(x -> ((prev_sign[x] < 0 && callback.affect! !== nothing) || (prev_sign[x] > 0 && callback.affect_neg! !== nothing)) && prev_sign[x]*new_sign[x]<0, keys(prev_sign))
+      _event_idx = findall_events(affect!,affect_neg!,prev_sign,new_sign)
       if length(_event_idx) != 0
         event_occurred = true
         event_idx = _event_idx
@@ -398,6 +398,11 @@ end
   event_occurred,interp_index,Θs,prev_sign,prev_sign_index,event_idx
 end
 
+## Different definition for GPUs
+function findall_events(affect!,affect_neg,prev_sign,next_sign)
+  findall(x-> ((prev_sign[x] < 0 && affect! !== nothing) || (prev_sign[x] > 0 && affect_neg! !== nothing)) && prev_sign[x]*next_sign[x]<=0, keys(prev_sign))
+end
+
 function find_callback_time(integrator,callback::ContinuousCallback,counter)
   event_occurred,interp_index,Θs,prev_sign,prev_sign_index,event_idx = determine_event_occurance(integrator,callback,counter)
   if event_occurred

diff --git a/src/init.jl b/src/init.jl
@@ -136,6 +136,14 @@ function __init__()
     function LinearAlgebra.ldiv!(x::CuArrays.CuArray,_qr::CuArrays.CUSOLVER.CuQR,b::CuArrays.CuArray)
       _x = UpperTriangular(_qr.R) \ (_qr.Q' * reshape(b,length(b),1))
       x .= vec(_x)
+      unsafe_free!(_x)
+    end
+    function findall_events(affect!,affect_neg,prev_sign::CuArrays.CuArray,next_sign::CuArrays.CuArray)
+      f = (prev_sign,next_sign)-> ((prev_sign < 0 && affect! !== nothing) || (prev_sign > 0 && affect_neg! !== nothing)) && prev_sign*next_sign<=0
+      A = map(f,prev_sign,next_sign)
+      out = findall(A)
+      unsafe_free!(A)
+      out
     end
   end
 end