Support disabling implicit synchronization (#2662)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: vchuravy

src/array.jl

@@ -472,10 +472,35 @@ function Base.unsafe_convert(::Type{CuDeviceArray{T,N,AS.Global}}, a::DenseCuArr
 end
 
 
-## memory copying
+## synchronization
 
 synchronize(x::CuArray) = synchronize(x.data[])
 
+"""
+    enable_synchronization!(arr::CuArray, enable::Bool)
+
+By default `CuArray`s are implicitly synchronized when they are accessed on different CUDA
+devices or streams. This may be unwanted when e.g. using disjoint slices of memory across
+different tasks. This function allows to enable or disable this behavior.
+
+!!! warning
+
+    Disabling implicit synchronization affects _all_ `CuArray`s that are referring to the
+    same underlying memory. Unsafe use of this API _will_ result in data corruption.
+
+    This API is only provided as an escape hatch, and should not be used without careful
+    consideration. If automatic synchronization is generally problematic for your use case,
+    it is recommended to figure out a better model instead and file an issue or pull request.
+    For more details see [this discussion](https://github.com/JuliaGPU/CUDA.jl/issues/2617).
+"""
+function enable_synchronization!(arr::CuArray, enable::Bool=true)
+    arr.data[].synchronizing = enable
+    return arr
+end
+
+
+## memory copying
+
 if VERSION >= v"1.11.0-DEV.753"
 function typetagdata(a::Array, i=1)
   ptr_or_offset = Int(a.ref.ptr_or_offset)

src/memory.jl

@@ -503,16 +503,20 @@ mutable struct Managed{M}
   # which stream is currently using the memory.
   stream::CuStream
 
+  # whether accessing this memory can cause implicit synchronization
+  synchronizing::Bool
+
   # whether there are outstanding operations that haven't been synchronized
   dirty::Bool
 
   # whether the memory has been captured in a way that would make the dirty bit unreliable
   captured::Bool
 
-  function Managed(mem::AbstractMemory; stream=CUDA.stream(), dirty=true, captured=false)
+  function Managed(mem::AbstractMemory; stream = CUDA.stream(), synchronizing = true,
+                   dirty = true, captured = false)
     # NOTE: memory starts as dirty, because stream-ordered allocations are only
     #       guaranteed to be physically allocated at a synchronization event.
-    new{typeof(mem)}(mem, stream, dirty, captured)
+    new{typeof(mem)}(mem, stream, synchronizing, dirty, captured)
   end
 end
 
@@ -524,7 +528,7 @@ function synchronize(managed::Managed)
   managed.dirty = false
 end
 function maybe_synchronize(managed::Managed)
-  if managed.dirty || managed.captured
+  if managed.synchronizing && (managed.dirty || managed.captured)
     synchronize(managed)
   end
 end

test/base/array.jl

@@ -51,6 +51,13 @@ using ChainRulesCore: add!!, is_inplaceable_destination
   end
 end
 
+@testset "synchronization" begin
+  a = CUDA.zeros(2, 2)
+  synchronize(a)
+  CUDA.enable_synchronization!(a, false)
+  CUDA.enable_synchronization!(a)
+end
+
 @testset "unsafe_wrap" begin
     # managed memory -> CuArray
     for a in [cu([1]; device=true), cu([1]; unified=true)]