Repository: PTsolvers/Chmy.jl
Pull request #65: Disable task sync in CUDA
Author: login: luraess
Created: 2025-02-17T22:56:28Z | Merged: 2025-04-07T15:58:40Z

This PR implements the changes needed to support opting out of syncing on task switching for the CUDA backend https://github.com/JuliaGPU/CUDA.jl/pull/2662.

It thus avoids serialisation of async operations and brings back async execution
![Screenshot 2025-02-17 at 23 47 36](https://github.com/user-attachments/assets/cc08ce76-a3b2-4742-a55b-fe5eef06b496)

I am unsure whether this is the best way of addressing things tho...

Requires `CUDA#vc/unsafe_stream_switching` for testing


Diff:
diff --git a/Project.toml b/Project.toml
index 4c9448eb..120ed1c9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chmy"
 uuid = "33a72cf0-4690-46d7-b987-06506c2248b9"
 authors = ["Ivan Utkin <iutkin@ethz.ch>, Ludovic Raess <ludovic.rass@gmail.com>, and contributors"]
-version = "0.1.22"
+version = "0.1.23"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -23,9 +23,9 @@ ChmyMetalExt = "Metal"
 [compat]
 AMDGPU = "0.8, 0.9, 1"
 Adapt = "4"
-CUDA = "5"
+CUDA = "5.7.2"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 MacroTools = "0.5"
 Metal = "1"
-julia = "1.9"
+julia = "1.10"
diff --git a/ext/ChmyCUDAExt/ChmyCUDAExt.jl b/ext/ChmyCUDAExt/ChmyCUDAExt.jl
index 4c58e7dd..d09abcc2 100644
--- a/ext/ChmyCUDAExt/ChmyCUDAExt.jl
+++ b/ext/ChmyCUDAExt/ChmyCUDAExt.jl
@@ -2,12 +2,16 @@ module ChmyCUDAExt
 
 using CUDA, KernelAbstractions
 
-import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype
+import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype, disable_task_sync!, enable_task_sync!
 
 Base.unsafe_wrap(::CUDABackend, ptr::CuPtr, dims) = unsafe_wrap(CuArray, ptr, dims)
 
 pointertype(::CUDABackend, T::DataType) = CuPtr{T}
 
+# because of https://github.com/JuliaGPU/CUDA.jl/pull/2335
+disable_task_sync!(x::CuArray) = CUDA.enable_synchronization!(x, false)
+enable_task_sync!(x::CuArray) = CUDA.enable_synchronization!(x, true)
+
 set_device!(dev::CuDevice) = CUDA.device!(dev)
 
 get_device(::CUDABackend, id::Integer) = CuDevice(id - 1)
diff --git a/src/Architectures.jl b/src/Architectures.jl
index c054abfc..fcb7adda 100644
--- a/src/Architectures.jl
+++ b/src/Architectures.jl
@@ -65,7 +65,7 @@ get_device(arch::SingleDeviceArchitecture) = arch.device
 """
     activate!(arch::SingleDeviceArchitecture; priority=:normal)
 
-Activate the given architecture on the specified device and set the priority of the 
+Activate the given architecture on the specified device and set the priority of the
 backend. For the priority accepted values are `:normal`, `:low` and `:high`.
 """
 function activate!(arch::SingleDeviceArchitecture; priority=:normal)
@@ -82,4 +82,8 @@ Base.unsafe_wrap(::CPU, ptr::Ptr, dims) = unsafe_wrap(Array, ptr, dims)
 
 pointertype(::CPU, T::DataType) = Ptr{T}
 
+# because of https://github.com/JuliaGPU/CUDA.jl/pull/2335
+disable_task_sync!(::Any) = nothing
+enable_task_sync!(::Any)  = nothing
+
 end
diff --git a/src/Chmy.jl b/src/Chmy.jl
index b107271c..eacdfdaa 100644
--- a/src/Chmy.jl
+++ b/src/Chmy.jl
@@ -3,13 +3,14 @@ module Chmy
 using MacroTools
 using KernelAbstractions
 
-export 
+export
     # utils
     Dim, Side, Left, Right, remove_dim, insert_dim, Offset,
 
     # Architectures
-    Architecture, SingleDeviceArchitecture, Arch, get_backend, get_device, activate!, set_device!, heuristic_groupsize, pointertype,
-    
+    Architecture, SingleDeviceArchitecture, Arch, get_backend, get_device, activate!, set_device!,
+    heuristic_groupsize, pointertype,
+
     # BoundaryConditions
     FieldBoundaryCondition, FirstOrderBC, Dirichlet, Neumann, bc!,
     BoundaryFunction,
@@ -29,7 +30,7 @@ export
     AbstractField, Field, VectorField, TensorField, ConstantField, ZeroField, OneField, ValueField, FunctionField,
     location, halo, interior, set!,
     divg,
-    
+
     # Grids
     Location, Center, Vertex, flip,
     Connectivity, Bounded, Connected, Periodic, Flat,
@@ -97,4 +98,4 @@ using .GridOperators
 using .KernelLaunch
 using .Workers
 
-end # module Chmy
\ No newline at end of file
+end # module Chmy
diff --git a/src/Distributed/exchange_halo.jl b/src/Distributed/exchange_halo.jl
index 2daf81d9..a957163a 100644
--- a/src/Distributed/exchange_halo.jl
+++ b/src/Distributed/exchange_halo.jl
@@ -1,5 +1,5 @@
 """
-    exchange_halo!(side::Side, dim::Dim, arch, grid, fields...; async=false)
+    exchange_halo!(side::Side, dim::Dim, arch, grid, fields...)
 
 Perform halo exchange communication between neighboring processes in a distributed architecture.
 
@@ -9,14 +9,11 @@ Perform halo exchange communication between neighboring processes in a distribut
 - `arch`: The distributed architecture used for communication.
 - `grid`: The structured grid on which the halo exchange is performed.
 - `fields...`: The fields to be exchanged.
-
-# Optional Arguments
-- `async=false`: Whether to perform the halo exchange asynchronously.
 """
 function exchange_halo!(side::Side{S}, dim::Dim{D},
                         arch::DistributedArchitecture,
                         ::StructuredGrid,
-                        fields::Vararg{Field,K}; async=false) where {S,D,K}
+                        fields::Vararg{Field,K}) where {S,D,K}
     comm = cart_comm(topology(arch))
     nbrank = neighbor(topology(arch), D, S)
     @assert nbrank != MPI.PROC_NULL "no neighbor to communicate"
@@ -56,7 +53,7 @@ function exchange_halo!(side::Side{S}, dim::Dim{D},
         yield()
     end
 
-    async || KernelAbstractions.synchronize(Architectures.get_backend(arch))
+    KernelAbstractions.synchronize(Architectures.get_backend(arch))
 
     return
 end
@@ -88,7 +85,7 @@ end
     BoundaryConditions.bc!(side::Side, dim::Dim,
                                 arch::DistributedArchitecture,
                                 grid::StructuredGrid,
-                                batch::ExchangeBatch; kwargs...)
+                                batch::ExchangeBatch)
 
 Apply boundary conditions on a distributed grid with halo exchange performed internally.
 
@@ -99,11 +96,11 @@ Apply boundary conditions on a distributed grid with halo exchange performed int
 - `grid`: The structured grid on which the halo exchange is performed.
 - `batch`: The batch set to apply boundary conditions to.
 """
-function BoundaryConditions.bc!(side::Side, 
+function BoundaryConditions.bc!(side::Side,
                                 dim::Dim,
                                 arch::DistributedArchitecture,
                                 grid::StructuredGrid,
-                                batch::ExchangeBatch; kwargs...)
-    exchange_halo!(side, dim, arch, grid, batch.fields...; kwargs...)
+                                batch::ExchangeBatch)
+    exchange_halo!(side, dim, arch, grid, batch.fields...)
     return
-end
\ No newline at end of file
+end
diff --git a/src/Distributed/stack_allocator.jl b/src/Distributed/stack_allocator.jl
index 1ea3567e..475eb28e 100644
--- a/src/Distributed/stack_allocator.jl
+++ b/src/Distributed/stack_allocator.jl
@@ -83,5 +83,7 @@ function allocate(sa::StackAllocator, T::DataType, dims, align::Integer=sizeof(T
     data_ptr = convert(pointertype(backend, T), pointer(sa.buffer) + aligned)
     sa.offset = aligned + nbytes
     sa.nallocs += 1
-    return unsafe_wrap(backend, data_ptr, dims)
+    buffer = unsafe_wrap(backend, data_ptr, dims)
+    disable_task_sync!(buffer)
+    return buffer
 end
diff --git a/src/Fields/field.jl b/src/Fields/field.jl
index 25d5fc80..a31817ee 100644
--- a/src/Fields/field.jl
+++ b/src/Fields/field.jl
@@ -128,7 +128,7 @@ end
     dst[I...] = fun(grid, loc, I..., args...)
 end
 
-function set!(f::Field{T,N}, grid::StructuredGrid{N}, fun::F; discrete=false, parameters=(), async=false) where {T,F,N}
+function set!(f::Field{T,N}, grid::StructuredGrid{N}, fun::F; discrete=false, parameters=()) where {T,F,N}
     loc = location(f)
     dst = interior(f)
     backend = KernelAbstractions.get_backend(dst)
@@ -137,7 +137,7 @@ function set!(f::Field{T,N}, grid::StructuredGrid{N}, fun::F; discrete=false, pa
     else
         _set_continuous!(backend, 256, size(dst))(dst, grid, loc, fun, parameters...)
     end
-    async || KernelAbstractions.synchronize(backend)
+    KernelAbstractions.synchronize(backend)
     return
 end
 
diff --git a/src/KernelLaunch.jl b/src/KernelLaunch.jl
index 3e456fd8..38f170a1 100644
--- a/src/KernelLaunch.jl
+++ b/src/KernelLaunch.jl
@@ -9,6 +9,8 @@ using Chmy.Grids
 using Chmy.BoundaryConditions
 using Chmy.Workers
 
+import Chmy.Architectures: enable_task_sync!, disable_task_sync!
+
 using KernelAbstractions
 
 """
@@ -32,7 +34,7 @@ Constructs a `Launcher` object configured based on the input parameters.
 
 !!! warning
 
-    worksize for the last dimension N takes into account only last outer width 
+    worksize for the last dimension N takes into account only last outer width
     W[N], N-1 uses W[N] and W[N-1], N-2 uses W[N], W[N-1], and W[N-2].
 """
 function Launcher(arch, grid; outer_width=nothing)
@@ -85,7 +87,7 @@ Base.@assume_effects :foldable function outer_offset(launcher::Launcher, ::Dim{D
 end
 
 """
-    (launcher::Launcher)(arch::Architecture, grid, kernel_and_args::Pair{F,Args}; bc=nothing, async=false) where {F,Args}
+    (launcher::Launcher)(arch::Architecture, grid, kernel_and_args::Pair{F,Args}; bc=nothing) where {F,Args}
 
 Launches a computational kernel using the specified `arch`, `grid`, `kernel_and_args`, and optional boundary conditions (`bc`).
 
@@ -94,14 +96,13 @@ Launches a computational kernel using the specified `arch`, `grid`, `kernel_and_
 - `grid`: The grid defining the computational domain.
 - `kernel_and_args::Pair{F,Args}`: A pair consisting of the computational kernel `F` and its arguments `Args`.
 - `bc=nothing`: Optional boundary conditions for the computation.
-- `async=false`: If `true`, launches the kernel asynchronously.
 
 !!! warning
     - `arch` should be compatible with the `Launcher`'s architecture.
     - If `bc` is `nothing`, the kernel is launched without boundary conditions.
-    - If `async` is `false` (default), the function waits for the computation to complete before returning.
+    - The function waits for the computation to complete before returning.
 """
-function (launcher::Launcher)(arch::Architecture, grid, kernel_and_args::Pair{F,Args}; bc=nothing, async=false) where {F,Args}
+function (launcher::Launcher)(arch::Architecture, grid, kernel_and_args::Pair{F,Args}; bc=nothing) where {F,Args}
     kernel, args = kernel_and_args
 
     backend = Architectures.get_backend(arch)
@@ -113,7 +114,7 @@ function (launcher::Launcher)(arch::Architecture, grid, kernel_and_args::Pair{F,
         launch_with_bc(arch, grid, launcher, offset, kernel, bc, args...)
     end
 
-    async || KernelAbstractions.synchronize(backend)
+    KernelAbstractions.synchronize(backend)
     return
 end
 
@@ -126,6 +127,28 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+# recursively apply a function to all fields of an object
+@generated function deepmap!(fn::F, x::T) where {F,T}
+    names = fieldnames(x)
+    N     = length(names)
+    quote
+        @inline
+        fn(x) # deepmap calls a function on the argument
+        Base.@nexprs $N i -> begin
+            args = getfield(x, $names[i])
+            deepmap!(fn, args) # deepmap calls a function on its fields
+        end
+    end
+end
+
+# helper function to temporarily disable task sync for arguments
+function with_no_task_sync!(fn::F, args::T) where {F,T}
+    deepmap!(disable_task_sync!, args)
+    fn()
+    deepmap!(enable_task_sync!, args)
+    return
+end
+
 @inline function launch_with_bc(arch, grid, launcher, offset, kernel, bc, args...)
     backend   = Architectures.get_backend(arch)
     groupsize = StaticSize(heuristic_groupsize(backend, Val(ndims(launcher))))
@@ -135,23 +158,25 @@ import KernelAbstractions.NDIteration.StaticSize
         fun(args..., offset)
         bc!(arch, grid, bc)
     else
-        inner_fun = kernel(backend, groupsize, StaticSize(inner_worksize(launcher)))
-        inner_fun(args..., offset + Offset(inner_offset(launcher)...))
-
-        N = ndims(grid)
-        ntuple(Val(N)) do J
-            Base.@_inline_meta
-            D = N - J + 1
-            outer_fun = kernel(backend, groupsize, StaticSize(outer_worksize(launcher, Dim(D))))
-            ntuple(Val(2)) do S
-                put!(launcher.workers[D][S]) do
-                    outer_fun(args..., offset + Offset(outer_offset(launcher, Dim(D), Side(S))...))
-                    bc!(Side(S), Dim(D), arch, grid, bc[D][S])
-                    KernelAbstractions.synchronize(backend)
+        with_no_task_sync!(args) do
+            inner_fun = kernel(backend, groupsize, StaticSize(inner_worksize(launcher)))
+            inner_fun(args..., offset + Offset(inner_offset(launcher)...))
+
+            N = ndims(grid)
+            ntuple(Val(N)) do J
+                Base.@_inline_meta
+                D = N - J + 1
+                outer_fun = kernel(backend, groupsize, StaticSize(outer_worksize(launcher, Dim(D))))
+                ntuple(Val(2)) do S
+                    put!(launcher.workers[D][S]) do
+                        outer_fun(args..., offset + Offset(outer_offset(launcher, Dim(D), Side(S))...))
+                        bc!(Side(S), Dim(D), arch, grid, bc[D][S])
+                        KernelAbstractions.synchronize(backend)
+                    end
                 end
+                wait(launcher.workers[D][1])
+                wait(launcher.workers[D][2])
             end
-            wait(launcher.workers[D][1])
-            wait(launcher.workers[D][2])
         end
     end
     return
diff --git a/test/test_kernel_launch.jl b/test/test_kernel_launch.jl
new file mode 100644
index 00000000..8fe11537
--- /dev/null
+++ b/test/test_kernel_launch.jl
@@ -0,0 +1,135 @@
+include("common.jl")
+
+import Chmy.Architectures: disable_task_sync!, enable_task_sync!
+import Chmy.KernelLaunch: deepmap!, with_no_task_sync!
+
+myfun!(::Any) = nothing
+
+const buff = DataType[]
+
+for backend in TEST_BACKENDS, T in TEST_TYPES
+    if !compatible(backend, T)
+        continue
+    end
+
+    @testset "$(basename(@__FILE__)) (backend: $backend, type: $T)" begin
+        if (backend isa CPU)
+            @testset "deepmap! on CPU backend" begin
+                calls = Symbol[]
+
+                struct CuArray end # Dummy struct to simulate CUDA array
+
+                Chmy.Architectures.disable_task_sync!(::Any) = nothing
+                Chmy.Architectures.enable_task_sync!(::Any) = nothing
+                Chmy.Architectures.disable_task_sync!(::CuArray) = push!(calls, :disable)
+                Chmy.Architectures.enable_task_sync!(::CuArray) = push!(calls, :enable)
+
+                empty!(calls)
+                foo1 = (1, 2, 3)
+                deepmap!(disable_task_sync!, foo1) # prints nothing
+                deepmap!(enable_task_sync!, foo1)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo2 = ()
+                deepmap!(disable_task_sync!, foo2) # prints nothing
+                deepmap!(enable_task_sync!, foo2)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo3 = (CuArray(), CuArray(), (1, ()))
+                deepmap!(disable_task_sync!, foo3) # prints "disable" twice
+                deepmap!(enable_task_sync!, foo3) # prints "enable" twice
+                @test calls == [:disable, :disable, :enable, :enable]
+
+                empty!(calls)
+                foo4 = (a=1, b=CuArray(), c=((), 1))
+                deepmap!(disable_task_sync!, foo4) # prints "disable" once
+                deepmap!(enable_task_sync!, foo4) # prints "enable" once
+                @test calls == [:disable, :enable]
+
+                empty!(calls)
+                struct Foo{T}
+                    a::Int
+                    b::CuArray
+                    c::T
+                end
+                foo5 = Foo(1, CuArray(), foo4)
+
+                with_no_task_sync!(foo5) do
+                    push!(calls, :in_out) # kernel execution for inner and outer regions goes here
+                end # will print "disable" twice, then "in_out" and "enable" twice
+                @test calls == [:disable, :disable, :in_out, :enable, :enable]
+            end
+        end
+
+        @testset "deepmap! on Field" begin
+            arch = Arch(backend)
+            grid = UniformGrid(arch; origin=(T(0.0), T(0.0)), extent=(T(1.0), T(1.0)), dims=(6, 4))
+
+            Δt = T(1.0)
+            C  = Field(backend, grid, Center())
+            q  = VectorField(backend, grid)
+
+            M = typeof(C.data)
+
+            @eval myfun!(x::$M) = push!(buff, typeof(x))
+
+            @testset "deepmap! on single Field" begin
+                empty!(buff)
+                deepmap!(myfun!, C)
+                @test length(buff) == 1
+                @test buff[1] == M
+            end
+            @testset "deepmap! on multiple args" begin
+                empty!(buff)
+                args = (q, Δt, grid)
+                deepmap!(myfun!, args)
+                @test length(buff) == 2
+                @test buff[1] == M
+                @test buff[2] == M
+            end
+
+            if haskey(ENV, "JULIA_CHMY_BACKEND_CUDA") && (backend isa CUDABackend)
+                @testset "CUDA implicit synchronization" begin
+                    @testset "Single Field" begin
+                        @testset "disable_task_sync!" begin
+                            @assert C.data.data[].synchronizing == true
+                            deepmap!(disable_task_sync!, C)
+                            @test C.data.data[].synchronizing == false
+                        end
+
+                        @testset "enable_task_sync!" begin
+                            @assert C.data.data[].synchronizing == false
+                            deepmap!(enable_task_sync!, C)
+                            @test C.data.data[].synchronizing == true
+                        end
+                    end
+
+                    @testset "Multiple args" begin
+                        args = (C, q, Δt, grid)
+                        @testset "disable_task_sync!" begin
+                            @assert C.data.data[].synchronizing == true
+                            @assert q.x.data.data[].synchronizing == true
+                            @assert q.y.data.data[].synchronizing == true
+                            deepmap!(disable_task_sync!, args)
+                            @test C.data.data[].synchronizing == false
+                            @test q.x.data.data[].synchronizing == false
+                            @test q.y.data.data[].synchronizing == false
+                        end
+
+                        @testset "enable_task_sync!" begin
+                            @assert C.data.data[].synchronizing == false
+                            @assert q.x.data.data[].synchronizing == false
+                            @assert q.y.data.data[].synchronizing == false
+                            deepmap!(enable_task_sync!, args)
+                            @test C.data.data[].synchronizing == true
+                            @test q.x.data.data[].synchronizing == true
+                            @test q.y.data.data[].synchronizing == true
+                        end
+                    end
+                end
+            end
+        end
+    end
+end


Reviews:
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:17:55Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:32:04Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:43:35Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:45:23Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:46:45Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:48:53Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:51:03Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T22:53:24Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T23:02:25Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-24T23:05:53Z
user: login: albert-de-montserrat
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T05:32:05Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:15:03Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:16:06Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:17:05Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:17:54Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:18:18Z
user: login: albert-de-montserrat
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T08:20:03Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T08:23:23Z
user: login: luraess
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T10:06:13Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T12:48:49Z
user: login: utkinis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-02T12:49:16Z

Review comments:
diff_hunk: @@ -135,6 +155,8 @@ import KernelAbstractions.NDIteration.StaticSize
         fun(args..., offset)
         bc!(arch, grid, bc)
     else
+        foreach(arg -> apply_disable_task_sync!(backend, arg, disable_task_sync!), args)
path: src/KernelLaunch.jl
user: login: luraess
body: What if we check on the fly for the args passed to the `launch_with_bc` function and disable task sync recursively there on any `Field` types, and enable task sync on any `Field` types before returning from the kernel launch function?
created_at: 2025-03-24T22:17:54Z
author_association: MEMBER
diff_hunk: @@ -135,6 +155,8 @@ import KernelAbstractions.NDIteration.StaticSize
         fun(args..., offset)
         bc!(arch, grid, bc)
     else
+        foreach(arg -> apply_disable_task_sync!(backend, arg, disable_task_sync!), args)
path: src/KernelLaunch.jl
user: login: luraess
body: This would still allow for using different task per BC kernel and avoid a heavy rework of workers launching workers and possibly paying the price of serialising all the compute first.
created_at: 2025-03-24T22:32:04Z
author_association: MEMBER
diff_hunk: @@ -2,12 +2,15 @@ module ChmyCUDAExt
 
 using CUDA, KernelAbstractions
 
-import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype
+import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype, disable_task_sync!, enable_task_sync!
 
 Base.unsafe_wrap(::CUDABackend, ptr::CuPtr, dims) = unsafe_wrap(CuArray, ptr, dims)
 
 pointertype(::CUDABackend, T::DataType) = CuPtr{T}
 
+disable_task_sync!(::CUDABackend, array) = CUDA.unsafe_disable_task_sync!(array)
path: ext/ChmyCUDAExt/ChmyCUDAExt.jl
user: login: utkinis
body: Why do we need to return the array here? This function doesn't create a new array, maybe it makes sense to highlight this in the API. Also then it becomes a no-op for backends that don't need this

Also, the dispatch could be performed directly on the array type, so no need to pass the backend I guess
created_at: 2025-03-24T22:43:35Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
+        return disable_task_sync!(backend, x.data)  # Apply disable_task_sync! to Field
+    elseif isa(x, NamedTuple)
+        return map(v -> apply_disable_task_sync!(backend, v, disable_task_sync!), x)  # Recurse for NamedTuple values
+    else
+        return x  # Leave other types unchanged
path: src/KernelLaunch.jl
user: login: utkinis
body: What about structs and simple tuples?
created_at: 2025-03-24T22:45:23Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
path: src/KernelLaunch.jl
user: login: utkinis
body: We don't want to only disable implicit sync for Fields, but for all GPU memory. What if a user passes just the GPU array with e.g. phase diagram?
created_at: 2025-03-24T22:46:44Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
+        return disable_task_sync!(backend, x.data)  # Apply disable_task_sync! to Field
+    elseif isa(x, NamedTuple)
+        return map(v -> apply_disable_task_sync!(backend, v, disable_task_sync!), x)  # Recurse for NamedTuple values
+    else
+        return x  # Leave other types unchanged
path: src/KernelLaunch.jl
user: login: luraess
body: Good point. I was hitting only this construct for now, but yes, maybe there would be a better and more abstract mechanism to recurse into more objects holding Fields?
created_at: 2025-03-24T22:48:53Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
path: src/KernelLaunch.jl
user: login: luraess
body: Would this be supported when using the `Launcher`? If not, we should then maybe check for GPU mem type, further recursing into Field (and other types)
created_at: 2025-03-24T22:51:03Z
author_association: MEMBER
diff_hunk: @@ -2,12 +2,15 @@ module ChmyCUDAExt
 
 using CUDA, KernelAbstractions
 
-import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype
+import Chmy.Architectures: heuristic_groupsize, set_device!, get_device, pointertype, disable_task_sync!, enable_task_sync!
 
 Base.unsafe_wrap(::CUDABackend, ptr::CuPtr, dims) = unsafe_wrap(CuArray, ptr, dims)
 
 pointertype(::CUDABackend, T::DataType) = CuPtr{T}
 
+disable_task_sync!(::CUDABackend, array) = CUDA.unsafe_disable_task_sync!(array)
path: ext/ChmyCUDAExt/ChmyCUDAExt.jl
user: login: luraess
body: Good point. Indeed, we could dispatch on the array type directly. And yes, for `CuArray`s, we actually don't return the array while my current proposition does so for other backends. This should indeed be fixed to a no-op for other backends.
created_at: 2025-03-24T22:53:24Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
+        return disable_task_sync!(backend, x.data)  # Apply disable_task_sync! to Field
+    elseif isa(x, NamedTuple)
+        return map(v -> apply_disable_task_sync!(backend, v, disable_task_sync!), x)  # Recurse for NamedTuple values
+    else
+        return x  # Leave other types unchanged
path: src/KernelLaunch.jl
user: login: luraess
body: Shall we handle "all" cases explicitly, or do you think there could be an abstract way of handling various types of noted objects holding GPU memory here?
created_at: 2025-03-24T23:02:24Z
author_association: MEMBER
diff_hunk: @@ -126,6 +126,26 @@ end
 
 import KernelAbstractions.NDIteration.StaticSize
 
+@inline function apply_disable_task_sync!(backend, x, disable_task_sync!)
+    if isa(x, Field)
path: src/KernelLaunch.jl
user: login: utkinis
body: `Launcher` doesn't check for the args types, so yes, the user could pass anything to the kernel, including the regular arrays. The recursion should stop at the level where the implicit sync happens, which for CUDA.jl is the array level
created_at: 2025-03-24T23:05:53Z
author_association: MEMBER
diff_hunk: @@ -25,8 +24,8 @@ import Chmy.KernelLaunch: modify_task_sync!
 
         @testset "Structs" begin
             struct Foo1
-                x::T
-                y::T
+                x::Int
+                y::Int
path: test/test_kernel_launch.jl
user: login: albert-de-montserrat
body: Could be simplified using one single `struct` such as
```julia
struct Foo{A,B}
   x::A
   y::B
end
```
created_at: 2025-04-01T05:31:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
+
+struct CuArray end # in reality this is replaced by "using CUDA"
+
+myfun!(::Any) = nothing
+
+for backend in TEST_BACKENDS, T in TEST_TYPES
+    if !compatible(backend, T)
+        continue
+    end
+
+    @testset "$(basename(@__FILE__)) (backend: $backend, type: $T)" begin
+        if (backend isa CPU)
+            @testset "deepmap! on CPU backend" begin
+                # struct CuArray end # in reality this is replaced by "using CUDA"
+
+                calls = Symbol[]
+
+                disable_task_sync!(::Any) = nothing
+                enable_task_sync!(::Any) = nothing
+                disable_task_sync!(::CuArray) = push!(calls, :disable) # actual unsafe_disable_task_sync! call
+                enable_task_sync!(::CuArray) = push!(calls, :enable) # actual unsafe_enable_task_sync! call
+
+                empty!(calls)
+                foo1 = (1, 2, 3)
+                deepmap!(disable_task_sync!, foo1) # prints nothing
+                deepmap!(enable_task_sync!, foo1)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo2 = ()
+                deepmap!(disable_task_sync!, foo2) # prints nothing
+                deepmap!(enable_task_sync!, foo2)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo3 = (CuArray(), CuArray(), (1, ()))
+                deepmap!(disable_task_sync!, foo3) # prints "disable" twice
+                deepmap!(enable_task_sync!, foo3) # prints "enable" twice
+                @test calls == [:disable, :disable, :enable, :enable]
+
+                empty!(calls)
+                foo4 = (a=1, b=CuArray(), c=((), 1))
+                deepmap!(disable_task_sync!, foo4) # prints "disable" once
+                deepmap!(enable_task_sync!, foo4) # prints "enable" once
+                @test calls == [:disable, :enable]
+
+                empty!(calls)
+                struct Foo{T}
+                    a::Int
+                    b::CuArray
+                    c::T
+                end
+                foo5 = Foo(1, CuArray(), foo4)
+
+                with_no_task_sync!(foo5) do
+                    push!(calls, :in_out) # kernel execution for inner and outer regions goes here
+                end # will print "disable" twice, then "in_out" and "enable" twice
+                @show calls
path: test/test_kernel_launch.jl
user: login: luraess
body: It fails here as it looks like `with_no_task_sync!` seems not to call the `deepmap` which would dispatch on `CuArray` but on `Any` which returns `nothing` instead of pushing the symbol.
created_at: 2025-04-02T08:15:02Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
+
+struct CuArray end # in reality this is replaced by "using CUDA"
+
+myfun!(::Any) = nothing
+
+for backend in TEST_BACKENDS, T in TEST_TYPES
+    if !compatible(backend, T)
+        continue
+    end
+
+    @testset "$(basename(@__FILE__)) (backend: $backend, type: $T)" begin
+        if (backend isa CPU)
+            @testset "deepmap! on CPU backend" begin
+                # struct CuArray end # in reality this is replaced by "using CUDA"
+
+                calls = Symbol[]
+
+                disable_task_sync!(::Any) = nothing
+                enable_task_sync!(::Any) = nothing
+                disable_task_sync!(::CuArray) = push!(calls, :disable) # actual unsafe_disable_task_sync! call
+                enable_task_sync!(::CuArray) = push!(calls, :enable) # actual unsafe_enable_task_sync! call
+
+                empty!(calls)
+                foo1 = (1, 2, 3)
+                deepmap!(disable_task_sync!, foo1) # prints nothing
+                deepmap!(enable_task_sync!, foo1)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo2 = ()
+                deepmap!(disable_task_sync!, foo2) # prints nothing
+                deepmap!(enable_task_sync!, foo2)  # prints nothing
+                @test isempty(calls)
+
+                empty!(calls)
+                foo3 = (CuArray(), CuArray(), (1, ()))
+                deepmap!(disable_task_sync!, foo3) # prints "disable" twice
+                deepmap!(enable_task_sync!, foo3) # prints "enable" twice
+                @test calls == [:disable, :disable, :enable, :enable]
+
+                empty!(calls)
+                foo4 = (a=1, b=CuArray(), c=((), 1))
+                deepmap!(disable_task_sync!, foo4) # prints "disable" once
+                deepmap!(enable_task_sync!, foo4) # prints "enable" once
+                @test calls == [:disable, :enable]
+
+                empty!(calls)
+                struct Foo{T}
+                    a::Int
+                    b::CuArray
+                    c::T
+                end
+                foo5 = Foo(1, CuArray(), foo4)
+
+                with_no_task_sync!(foo5) do
+                    push!(calls, :in_out) # kernel execution for inner and outer regions goes here
+                end # will print "disable" twice, then "in_out" and "enable" twice
+                @show calls
+                # @test calls == [:disable, :disable, :in_out, :enable, :enable]
+            end
+        end
+
+        @testset "deepmap! on Field" begin
+            count = DataType[]
+
+            arch = Arch(backend)
+            grid = UniformGrid(arch; origin=(T(0.0), T(0.0)), extent=(T(1.0), T(1.0)), dims=(6, 4))
+
+            Δt = T(1.0)
+            C  = Field(backend, grid, Center())
+            q  = VectorField(backend, grid)
+
+            M = typeof(C.data)
+            myfun!(x::M) where {M} = push!(count, typeof(x))
+            # myfun!(x::M) = push!(count, typeof(x))
+
+            @testset "deepmap! on single Field" begin
+                empty!(count)
+                deepmap!(myfun!, C)
+                @test length(count) == 1
path: test/test_kernel_launch.jl
user: login: luraess
body: Looks like `myfun!` does not dispatch on the correct type and thus count gets filled with the `myfun(::Any)` method
created_at: 2025-04-02T08:16:06Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
path: test/test_kernel_launch.jl
user: login: luraess
body: I don't get why besides `deepmap!`, I need to also import the other functions here as they should be exported
created_at: 2025-04-02T08:17:05Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
+
+struct CuArray end # in reality this is replaced by "using CUDA"
path: test/test_kernel_launch.jl
user: login: luraess
body: Setting this inside the test fails because of not being defined in `Main`
created_at: 2025-04-02T08:17:54Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
+
+struct CuArray end # in reality this is replaced by "using CUDA"
+
+myfun!(::Any) = nothing
path: test/test_kernel_launch.jl
user: login: luraess
body: Defining this in the `testset` also complains about scope
created_at: 2025-04-02T08:18:18Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
path: test/test_kernel_launch.jl
user: login: albert-de-montserrat
body: You are adding methods to `disable_task_sync!` and `enable_task_sync!`. To do so you need to import them, or do e.g. 
```julia
Chmy.Architectures.disable_task_sync!(::Any) = nothing
```
created_at: 2025-04-02T08:20:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,138 @@
+include("common.jl")
+
+import Chmy.Architectures: deepmap!, disable_task_sync!, enable_task_sync!
path: test/test_kernel_launch.jl
user: login: luraess
body: Ok, this one makes sense to me, yes.
created_at: 2025-04-02T08:23:23Z
author_association: MEMBER
diff_hunk: @@ -69,6 +69,10 @@ for backend in backends
     if backend != "CPU"
         Pkg.add(backend)
     end
+    # tmp fix to have the disable/enable task sync feature until merged in CUDA.jl
path: test/runtests.jl
user: login: luraess
body: Needs to be removed before merge
created_at: 2025-04-02T10:06:13Z
author_association: MEMBER
diff_hunk: @@ -1,7 +1,7 @@
 module Architectures
 
 export Architecture, SingleDeviceArchitecture
-export Arch, get_backend, get_device, activate!, set_device!, heuristic_groupsize, pointertype
+export Arch, get_backend, get_device, activate!, set_device!, heuristic_groupsize, pointertype, disable_task_sync!, enable_task_sync!, with_no_task_sync!
path: src/Architectures.jl
user: login: utkinis
body: I don't think we should export these functions as API, we only use it internally.
created_at: 2025-04-02T12:48:49Z
author_association: MEMBER
diff_hunk: @@ -3,13 +3,14 @@ module Chmy
 using MacroTools
 using KernelAbstractions
 
-export 
+export
     # utils
     Dim, Side, Left, Right, remove_dim, insert_dim, Offset,
 
     # Architectures
-    Architecture, SingleDeviceArchitecture, Arch, get_backend, get_device, activate!, set_device!, heuristic_groupsize, pointertype,
-    
+    Architecture, SingleDeviceArchitecture, Arch, get_backend, get_device, activate!, set_device!,
path: src/Chmy.jl
user: login: utkinis
body: Same as above, I suggest not exporting it.
created_at: 2025-04-02T12:49:16Z
author_association: MEMBER

PR comments:
user: login: luraess
created_at: 2025-02-20T21:39:24Z
body: According to our latest discussion, we may rework the communication / computation overlap in Chmy. By serialising all the compute (Outer -> BCs -> Inner) we could then spawn a new task only for doing IO, i.e., MPI comm preceded by copying from the array view into buffers and back into the array upon termination.

As I understand, this would only require opting out from implicit sync using e.g. `unsafe_disable_task_sync!` on the buffer arrays without needing further engineering.

It would still be nice to update CUDA.jl's proposed `unsafe_disable_task_sync!(arr)` to return the array instead of that status as to e.g. use the following construct:
```
a = CUDA.rand(2, 1) |> CUDA.unsafe_disable_task_sync!
```
author_association: MEMBER
user: login: luraess
created_at: 2025-03-24T22:25:23Z
body: @utkinis what if we use the above strategy to circumvent the implicit sync issue with CUDA.jl? This would actually make the `launch_with_bc` function "unsafe" wrt CUDA.jl, but would confine the unsafeness ti this function and not interfere with other parts of the code. Testing on ALPS, the approach seems to work and give good perf and ... overlap

![Screenshot 2025-03-24 at 23 12 36](https://github.com/user-attachments/assets/e60e877e-7573-4991-9eac-b6929c04ab41)

Happy to get your feedback and if positive to going this way, potential hints on improving my prototype attempt
author_association: MEMBER
user: login: utkinis
created_at: 2025-03-24T23:01:14Z
body: Thanks @luraess for digging further into this. Yes, checking the arguments at runtime could be a good way to solve the issue. One needs to recursively check all Julia composite types that could be passed to kernels, namely, structs and tuples (Not sure about arrays and refs, I guess only the bits types could be passed to GPU kernels, but might be wrong).

Also, the implicit sync happens at the array level, so the recursion must descend to the level of arrays and not just `Fields`, as the user might just pass a regular GPU array. I left the specific comments in the PR. Otherwise, I'm positive to make it work this way.
author_association: MEMBER
user: login: luraess
created_at: 2025-04-01T12:10:55Z
body: @utkinis could you have a look into this? We reworked the things with @albert-de-montserrat and it seems to do the job. Will test again on ALPS to see that all works as expected with the hide comm, but would be curious to get your feedback at this point 🙏 
author_association: MEMBER
user: login: luraess
created_at: 2025-04-01T13:15:44Z
body: Profiling on ALPS shows we are back on tracks
<img width="879" alt="Screenshot 2025-04-01 at 15 14 47" src="https://github.com/user-attachments/assets/a760933c-9981-4f8b-9f30-d2ac09b41dd6" />

author_association: MEMBER
user: login: utkinis
created_at: 2025-04-01T17:10:04Z
body: Thanks a lot @luraess and @albert-de-montserrat for the work. This looks like what we need. I'm thinking whether we could reduce the amount of code management a bit. So far all extensions must define bunch of no-ops, and defining both `enable/disable_task_sync!` and `modify_sync!` looks a bit repetitive. How about something like this:

```julia
##### this goes to Architectures.jl
disable_task_sync!(::Any) = nothing
enable_task_sync!(::Any)  = nothing

# function goes first according to Julia's conventions
# note that we don't need to define this for empty tuples
@generated function deepmap!(fn::F, x::T) where {F, T}
    names = fieldnames(x)
    N     = length(names)
    quote
        @inline
        fn(x) # deepmap calls a function on the argument...
        Base.@nexprs $N i -> begin
            args = getfield(x, $names[i])
            deepmap!(fn, args) # ...and on its fields
        end
    end
end

# a helper to temporarily disable task sync for arguments
function with_no_task_sync!(fn::F, args::T) where {F, T}
    deepmap!(disable_task_sync!, args)
    fn()
    deepmap!(enable_task_sync!, args)
    return
end

##### this goes only to CUDA extension, other extensions don't need anything special
struct CuArray end # in reality this is replaced by "using CUDA"

disable_task_sync!(::CuArray) = println("disable") # actual unsafe_disable_task_sync! call
enable_task_sync!(::CuArray) = println("enable") # actual unsafe_enable_task_sync! call

##### some random tests
foo1 = (1, 2, 3)
deepmap!(disable_task_sync!, foo1) # prints nothing
deepmap!(enable_task_sync!, foo1)  # prints nothing

foo2 = ()
deepmap!(disable_task_sync!, foo2) # prints nothing
deepmap!(enable_task_sync!, foo2)  # prints nothing

foo3 = (CuArray(), CuArray(), (1, ()))
deepmap!(disable_task_sync!, foo3) # prints "disable" twice
deepmap!(enable_task_sync!, foo3) # prints "enable" twice

foo4 = (a = 1, b = CuArray(), c = ((), 1))
deepmap!(disable_task_sync!, foo4) # prints "disable" once
deepmap!(enable_task_sync!, foo4) # prints "enable" once

struct Foo{T}
    a::Int
    b::CuArray
    c::T
end
foo5 = Foo(1, CuArray(), foo4)

with_no_task_sync!(foo5) do
    println("executing with no task sync") # kernel execution for inner and outer regions goes here
end # will print "disable" twice, then "executing with no task sync" and "enable" twice
```
author_association: MEMBER
user: login: albert-de-montserrat
created_at: 2025-04-01T17:37:23Z
body: Looks good to me!

author_association: CONTRIBUTOR
user: login: luraess
created_at: 2025-04-02T10:16:52Z
body: And the ALPS check that things overlap as they should
![Screenshot 2025-04-02 at 12 16 16](https://github.com/user-attachments/assets/6bf58f9d-5dd2-4ff6-97f8-06057340dbc8)

author_association: MEMBER
user: login: luraess
created_at: 2025-04-02T10:20:31Z
body: EDITED

From my side this could be ready to go. Before merging, one needs to:
- revert testing from a specific branch
- ~~bump minor (or patch?) version~~
- wait https://github.com/JuliaGPU/CUDA.jl/pull/2662 gets merged and release tagged
- update CUDA compat bounds
author_association: MEMBER
user: login: luraess
created_at: 2025-04-05T22:29:46Z
body: Upon final rework of https://github.com/JuliaGPU/CUDA.jl/pull/2662, the approach still works and successfully overlaps
![Screenshot 2025-04-06 at 00 24 09](https://github.com/user-attachments/assets/4cb517a5-635a-4607-899d-ada08af4a1f3)

author_association: MEMBER
user: login: luraess
created_at: 2025-04-07T08:10:35Z
body: This should be ready to go from my side.
author_association: MEMBER