fixup! datadeps: Allow disabling in-place moves

Author: jpsamaroo

src/datadeps.jl

@@ -147,18 +147,22 @@ struct DataDepsState{State<:Union{DataDepsAliasingState,DataDepsNonAliasingState
     # The mapping of memory space to remote argument copies
     remote_args::Dict{MemorySpace,IdDict{Any,Any}}
 
+    # Cache of whether arguments supports in-place move
+    supports_inplace_cache::IdDict{Any,Bool}
+
     # The aliasing analysis state
     alias_state::State
 
     function DataDepsState(aliasing::Bool)
         dependencies = Pair{DTask,Vector{Tuple{Bool,Bool,<:AbstractAliasing,<:Any,<:Any}}}[]
         remote_args = Dict{MemorySpace,IdDict{Any,Any}}()
+        supports_inplace_cache = IdDict{Any,Bool}()
         if aliasing
             state = DataDepsAliasingState()
         else
             state = DataDepsNonAliasingState()
         end
-        return new{typeof(state)}(aliasing, dependencies, remote_args, state)
+        return new{typeof(state)}(aliasing, dependencies, remote_args, supports_inplace_cache, state)
     end
 end
 
@@ -168,6 +172,12 @@ function aliasing(astate::DataDepsAliasingState, arg, dep_mod)
     end
 end
 
+function supports_inplace_move(state::DataDepsState, arg)
+    return get!(state.supports_inplace_cache, arg) do
+        return supports_inplace_move(arg)
+    end
+end
+
 # Determine which arguments could be written to, and thus need tracking
 
 "Whether `arg` has any writedep in this datadeps region."
@@ -708,7 +718,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
             end
 
             # Is the data writeable?
-            if !supports_inplace_move(arg)
+            if !supports_inplace_move(state, arg)
                 @dagdebug nothing :spawn_datadeps "($(repr(spec.f)))[$idx] Skipped copy-to (non-writeable)"
                 spec.args[idx] = pos => arg
                 continue
@@ -781,6 +791,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
             arg, deps = unwrap_inout(arg)
             arg = arg isa DTask ? fetch(arg; raw=true) : arg
             type_may_alias(typeof(arg)) || continue
+            supports_inplace_move(state, arg) || continue
             if queue.aliasing
                 for (dep_mod, _, writedep) in deps
                     ainfo = aliasing(astate, arg, dep_mod)
@@ -862,7 +873,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
                 end
 
                 # Skip non-writeable arguments
-                if !supports_inplace_move(arg)
+                if !supports_inplace_move(state, arg)
                     @dagdebug nothing :spawn_datadeps "Skipped copy-from (non-writeable)"
                     continue
                 end
@@ -919,7 +930,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
             end
 
             # Can the data be written back to?
-            if !supports_inplace_move(arg)
+            if !supports_inplace_move(state, arg)
                 @dagdebug nothing :spawn_datadeps "Skipped copy-from (non-writeable)"
             end