Extend threaded macro to use shared memory

Author: dennisYatunin

Project.toml

@@ -7,6 +7,7 @@ version = "0.6.7"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -17,9 +18,10 @@ ClimaCommsCUDAExt = "CUDA"
 ClimaCommsMPIExt = "MPI"
 
 [compat]
-CUDA = "3, 4, 5"
 Adapt = "3, 4"
+CUDA = "3, 4, 5"
 Logging = "1.9.4"
 LoggingExtras = "1.1.0"
 MPI = "0.20.18"
+StaticArrays = "1.9"
 julia = "1.9"

docs/Manifest.toml

@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.11.0"
+julia_version = "1.11.4"
 manifest_format = "2.0"
 project_hash = "d60839f726bd9115791d1a0807a21b61938765a9"
 
@@ -19,13 +19,11 @@ deps = ["LinearAlgebra", "Requires"]
 git-tree-sha1 = "50c3c56a52972d78e8be9fd135bfb91c9574c140"
 uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 version = "4.1.1"
+weakdeps = ["StaticArrays"]
 
     [deps.Adapt.extensions]
     AdaptStaticArraysExt = "StaticArrays"
 
-    [deps.Adapt.weakdeps]
-    StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
 version = "1.1.2"
@@ -39,10 +37,10 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 version = "1.11.0"
 
 [[deps.ClimaComms]]
-deps = ["Adapt", "Logging", "LoggingExtras"]
+deps = ["Adapt", "Logging", "LoggingExtras", "StaticArrays"]
 path = ".."
 uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
-version = "0.6.5"
+version = "0.6.7"
 
     [deps.ClimaComms.extensions]
     ClimaCommsCUDAExt = "CUDA"
@@ -361,6 +359,25 @@ version = "1.11.0"
 uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
 version = "1.11.0"
 
+[[deps.StaticArrays]]
+deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"]
+git-tree-sha1 = "0feb6b9031bd5c51f9072393eb5ab3efd31bf9e4"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "1.9.13"
+
+    [deps.StaticArrays.extensions]
+    StaticArraysChainRulesCoreExt = "ChainRulesCore"
+    StaticArraysStatisticsExt = "Statistics"
+
+    [deps.StaticArrays.weakdeps]
+    ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+    Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[deps.StaticArraysCore]]
+git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682"
+uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+version = "1.4.3"
+
 [[deps.StyledStrings]]
 uuid = "f489334b-da3d-4c2e-b8f0-e476e12c162b"
 version = "1.11.0"

docs/src/apis.md

@@ -28,13 +28,18 @@ ClimaComms.device
 ClimaComms.device_functional
 ClimaComms.array_type
 ClimaComms.allowscalar
-ClimaComms.@threaded
 ClimaComms.@time
 ClimaComms.@elapsed
 ClimaComms.@assert
 ClimaComms.@sync
 ClimaComms.@cuda_sync
 Adapt.adapt_structure(::Type{<:AbstractArray}, ::ClimaComms.AbstractDevice)
+ClimaComms.@threaded
+ClimaComms.@interdependent
+ClimaComms.InterdependentIteratorData
+ClimaComms.@sync_interdependent
+ClimaComms.synchronize_gpu_threads
+ClimaComms.static_shared_memory_array
 ```
 
 ## Contexts

ext/ClimaCommsCUDAExt.jl

@@ -5,6 +5,7 @@ import CUDA
 import Adapt
 import ClimaComms
 import ClimaComms: CUDADevice, threaded
+import ClimaComms: OneInterdependentItem, MultipleInterdependentItems
 
 function ClimaComms._assign_device(::CUDADevice, rank_number)
     CUDA.device!(rank_number % CUDA.ndevices())
@@ -50,17 +51,29 @@ ClimaComms.elapsed(f::F, ::CUDADevice, args...; kwargs...) where {F} =
 ClimaComms.assert(::CUDADevice, cond::C, text::T) where {C, T} =
     isnothing(text) ? (CUDA.@cuassert cond()) : (CUDA.@cuassert cond() text())
 
-# TODO: Generalize all of the following code to multi-dimensional thread blocks
-# and multiple iterators.
+ClimaComms.synchronize_gpu_threads(::CUDADevice) = CUDA.sync_threads()
 
-# The number of threads in the kernel being executed by the calling thread.
-threads_in_kernel() = CUDA.blockDim().x * CUDA.gridDim().x
+ClimaComms.static_shared_memory_array(
+    ::CUDADevice,
+    ::Type{T},
+    dims...,
+) where {T} = CUDA.CuStaticSharedArray(T, dims)
+
+# Number of blocks in kernel being executed and index of calling thread's block.
+blocks_in_kernel() = CUDA.gridDim().x
+block_idx_in_kernel() = CUDA.blockIdx().x
 
-# The index of the calling thread, which is between 1 and threads_in_kernel().
-thread_index() =
+# Number of threads in each block of kernel being executed and index of calling
+# thread within its block.
+threads_in_block() = CUDA.blockDim().x
+thread_idx_in_block() = CUDA.threadIdx().x
+
+# Total number of threads in kernel being executed and index of calling thread.
+threads_in_kernel() = CUDA.blockDim().x * CUDA.gridDim().x
+thread_idx_in_kernel() =
     (CUDA.blockIdx().x - 1) * CUDA.blockDim().x + CUDA.threadIdx().x
 
-# The maximum number of blocks that can fit on the GPU used for this kernel.
+# Maximum number of blocks that can fit on the GPU used for this kernel.
 grid_size_limit(kernel) = CUDA.attribute(
     CUDA.device(kernel.fun.mod.ctx),
     CUDA.DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
@@ -73,22 +86,22 @@ block_size_limit(max_threads_in_block::Int, _) = max_threads_in_block
 block_size_limit(::Val{:auto}, kernel) =
     CUDA.launch_configuration(kernel.fun).threads
 
-function threaded(f::F, ::CUDADevice, ::Val, itr; block_size) where {F}
+function threaded(f::F, device::CUDADevice, ::Val, itr; block_size) where {F}
     length(itr) > 0 || return nothing
     Base.require_one_based_indexing(itr)
 
-    function call_f_once_from_thread()
-        item_index = thread_index()
-        item_index <= length(itr) && @inbounds f(itr[item_index])
+    function thread_function()
+        itr_index = thread_idx_in_kernel()
+        itr_index <= length(itr) && @inbounds f(itr[itr_index])
         return nothing
     end
-    kernel = CUDA.@cuda launch=false call_f_once_from_thread()
+    kernel = CUDA.@cuda launch=false thread_function()
     max_blocks = grid_size_limit(kernel)
     max_threads_in_block = block_size_limit(block_size, kernel)
 
     # If there are too many items, coarsen by the smallest possible amount.
     length(itr) <= max_blocks * max_threads_in_block ||
-        return threaded(f, CUDADevice(), 1, itr)
+        return threaded(f, device, 1, itr; block_size)
 
     threads_in_block = min(max_threads_in_block, length(itr))
     blocks = cld(length(itr), threads_in_block)
@@ -102,17 +115,18 @@ function threaded(
     itr;
     block_size,
 ) where {F}
-    min_items_in_thread > 0 || throw(ArgumentError("`coarsen` is not positive"))
+    min_items_in_thread > 0 ||
+        throw(ArgumentError("integer `coarsen` value must be positive"))
     length(itr) > 0 || return nothing
     Base.require_one_based_indexing(itr)
 
     # Maximize memory coalescing with a "grid-stride loop"; for reference, see
     # https://developer.nvidia.com/blog/cuda-pro-tip-write-flexible-kernels-grid-stride-loops
-    call_f_multiple_times_from_thread() =
-        for item_index in thread_index():threads_in_kernel():length(itr)
-            @inbounds f(itr[item_index])
+    coarsened_thread_function() =
+        for itr_index in thread_idx_in_kernel():threads_in_kernel():length(itr)
+            @inbounds f(itr[itr_index])
         end
-    kernel = CUDA.@cuda launch=false call_f_multiple_times_from_thread()
+    kernel = CUDA.@cuda launch=false coarsened_thread_function()
     max_blocks = grid_size_limit(kernel)
     max_threads_in_block = block_size_limit(block_size, kernel)
 
@@ -129,4 +143,137 @@ function threaded(
     CUDA.@sync kernel(; blocks, threads = threads_in_block)
 end
 
+function threaded(
+    f::F,
+    device::CUDADevice,
+    ::Union{Val, NTuple{2, Val}},
+    independent_itr,
+    interdependent_itr;
+    block_size,
+) where {F}
+    length(independent_itr) > 0 || return nothing
+    length(interdependent_itr) > 0 || return nothing
+    Base.require_one_based_indexing(independent_itr)
+    Base.require_one_based_indexing(interdependent_itr)
+
+    function two_itr_thread_function()
+        block_index = block_idx_in_kernel()
+        thread_index = thread_idx_in_block()
+        (
+            block_index <= length(independent_itr) &&
+            thread_index <= length(interdependent_itr)
+        ) && @inbounds f(
+            independent_itr[block_index],
+            OneInterdependentItem(interdependent_itr[thread_index], device),
+        )
+        return nothing
+    end
+    kernel = CUDA.@cuda launch=false two_itr_thread_function()
+    max_blocks = grid_size_limit(kernel)
+    max_threads_in_block = block_size_limit(block_size, kernel)
+
+    # If there are too many items, coarsen by the smallest possible amount.
+    (
+        length(independent_itr) <= max_blocks &&
+        length(interdependent_itr) <= max_threads_in_block
+    ) || return threaded(
+        f,
+        device,
+        (1, 1),
+        independent_itr,
+        interdependent_itr;
+        block_size,
+    )
+
+    blocks = length(independent_itr)
+    threads_in_block = length(interdependent_itr)
+    CUDA.@sync kernel(; blocks, threads = threads_in_block)
+end
+
+# Use a default coarsen value of 1 for either iterator when a value is needed.
+threaded(
+    f::F,
+    device::CUDADevice,
+    min_independent_items_in_thread::Int,
+    independent_itr,
+    interdependent_itr;
+    block_size,
+) where {F} = threaded(
+    f,
+    device,
+    (min_independent_items_in_thread, 1),
+    independent_itr,
+    interdependent_itr;
+    block_size,
+)
+threaded(
+    f::F,
+    device::CUDADevice,
+    min_items_in_thread::Tuple{Val, Int},
+    independent_itr,
+    interdependent_itr;
+    block_size,
+) where {F} = threaded(
+    f,
+    device,
+    (1, min_items_in_thread[2]),
+    independent_itr,
+    interdependent_itr;
+    block_size,
+)
+
+function threaded(
+    f::F,
+    device::CUDADevice,
+    min_items_in_thread::NTuple{2, Int},
+    independent_itr,
+    interdependent_itr;
+    block_size,
+) where {F}
+    (min_items_in_thread[1] > 0 && min_items_in_thread[2] > 0) ||
+        throw(ArgumentError("all integer `coarsen` values must be positive"))
+    length(independent_itr) > 0 || return nothing
+    length(interdependent_itr) > 0 || return nothing
+    Base.require_one_based_indexing(independent_itr)
+    Base.require_one_based_indexing(interdependent_itr)
+
+    # Maximize memory coalescing with a "grid-stride loop" (reference is above).
+    function coarsened_two_itr_thread_function()
+        independent_itr_indices =
+            block_idx_in_kernel():blocks_in_kernel():length(independent_itr)
+        interdependent_itr_indices =
+            thread_idx_in_block():threads_in_block():length(interdependent_itr)
+        for independent_itr_index in independent_itr_indices
+            @inbounds f(
+                independent_itr[independent_itr_index],
+                MultipleInterdependentItems(
+                    interdependent_itr,
+                    interdependent_itr_indices,
+                    device,
+                ),
+            )
+        end
+    end
+    kernel = CUDA.@cuda launch=false coarsened_two_itr_thread_function()
+    max_blocks = grid_size_limit(kernel)
+    max_threads_in_block = block_size_limit(block_size, kernel)
+
+    # If there are too many items to use the specified coarsening, increase it
+    # by the smallest possible amount.
+    max_required_blocks = cld(length(independent_itr), min_items_in_thread[1])
+    max_required_threads_in_block =
+        cld(length(interdependent_itr), min_items_in_thread[2])
+    items_in_thread = (
+        max_required_blocks <= max_blocks ? min_items_in_thread[1] :
+        cld(length(independent_itr), max_blocks),
+        max_required_threads_in_block <= max_threads_in_block ?
+        min_items_in_thread[2] :
+        cld(length(interdependent_itr), max_threads_in_block),
+    )
+
+    blocks = cld(length(independent_itr), items_in_thread[1])
+    threads_in_block = cld(length(interdependent_itr), items_in_thread[2])
+    CUDA.@sync kernel(; blocks, threads = threads_in_block)
+end
+
 end