Merge pull request #1837 from CliMA/ck/gpu_kernel_tweaks

Some minor improvement to cuda kernel utils

ext/cuda/cuda_utils.jl

@@ -97,38 +97,44 @@ function auto_launch!(
 end
 
 """
-    kernel_indexes(n)
+    thread_index()
+
+Return the threadindex:
+```
+(CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x + CUDA.threadIdx().x
+```
+"""
+@inline thread_index() =
+    (CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x + CUDA.threadIdx().x
+
+"""
+    kernel_indexes(tidx, n)
 Return a tuple of indexes from the kernel,
-where `n` is a tuple of max lengths along each
+where `tidx` is the cuda thread index and
+`n` is a tuple of max lengths along each
 dimension of the accessed data.
 """
-function kernel_indexes(n::Tuple)
-    tidx = (CUDA.blockIdx().x - 1) * CUDA.blockDim().x + CUDA.threadIdx().x
-    inds = if 1 ≤ tidx ≤ prod(n)
-        CartesianIndices(map(x -> Base.OneTo(x), n))[tidx].I
-    else
-        ntuple(x -> -1, length(n))
-    end
-    return inds
-end
+Base.@propagate_inbounds kernel_indexes(tidx, n::Tuple) =
+    CartesianIndices(map(x -> Base.OneTo(x), n))[tidx]
 
 """
-    valid_range(inds, n)
+    valid_range(tidx, n::Int)
+
 Returns a `Bool` indicating if the thread index
-is in the valid range, based on `inds` (the result
-of `kernel_indexes`) and `n`, a tuple of max lengths
-along each dimension of the accessed data.
+(`tidx`) is in the valid range, based on `n`, a
+tuple of max lengths along each dimension of the
+
+accessed data.
 ```julia
 function kernel!(data, n)
-    inds = kernel_indexes(n)
-    if valid_range(inds, n)
-        do_work!(data[inds...])
+    @inbounds begin
+        tidx = thread_index()
+        if valid_range(tidx, n)
+            I = kernel_indexes(tidx, n)
+            do_work!(data[I])
+        end
     end
 end
 ```
 """
-valid_range(inds::NTuple, n::Tuple) = all(i -> 1 ≤ inds[i] ≤ n[i], 1:length(n))
-function valid_range(n::Tuple)
-    inds = kernel_indexes(n)
-    return all(i -> 1 ≤ inds[i] ≤ n[i], 1:length(n))
-end
+@inline valid_range(tidx, n) = 1 ≤ tidx ≤ n

ext/cuda/fill.jl

@@ -1,11 +1,11 @@
-cartesian_index(::AbstractData, inds) = CartesianIndex(inds)
-
 function knl_fill_flat!(dest::AbstractData, val)
-    n = DataLayouts.universal_size(dest)
-    inds = kernel_indexes(n)
-    if valid_range(inds, n)
-        I = cartesian_index(dest, inds)
-        @inbounds dest[I] = val
+    @inbounds begin
+        tidx = thread_index()
+        n = DataLayouts.universal_size(dest)
+        if valid_range(tidx, prod(n))
+            I = kernel_indexes(tidx, n)
+            @inbounds dest[I] = val
+        end
     end
     return nothing
 end

ext/cuda/limiters.jl

@@ -54,9 +54,9 @@ function compute_element_bounds_kernel!(
     ::Val{Nj},
 ) where {Ni, Nj}
     n = (Nv, Nh)
-    if valid_range(n)
-        inds = kernel_indexes(n)
-        (v, h) = inds
+    tidx = thread_index()
+    @inbounds if valid_range(tidx, prod(n))
+        (v, h) = kernel_indexes(tidx, n).I
         (; q_bounds) = limiter
         local q_min, q_max
         slab_ρq = slab(ρq, v, h)
@@ -118,9 +118,9 @@ function compute_neighbor_bounds_local_kernel!(
 ) where {Ni, Nj}
 
     n = (Nv, Nh)
-    if valid_range(n)
-        inds = kernel_indexes(n)
-        (v, h) = inds
+    tidx = thread_index()
+    @inbounds if valid_range(tidx, prod(n))
+        (v, h) = kernel_indexes(tidx, n).I
         (; q_bounds, q_bounds_nbr, ghost_buffer, rtol) = limiter
         slab_q_bounds = slab(q_bounds, v, h)
         q_min = slab_q_bounds[1]
@@ -188,9 +188,9 @@ function apply_limiter_kernel!(
     (; q_bounds_nbr, rtol) = limiter
     converged = true
     n = (Nv, Nh)
-    if valid_range(n)
-        inds = kernel_indexes(n)
-        (v, h) = inds
+    tidx = thread_index()
+    @inbounds if valid_range(tidx, prod(n))
+        (v, h) = kernel_indexes(tidx, n).I
 
         slab_ρ = slab(ρ_data, v, h)
         slab_ρq = slab(ρq_data, v, h)

ext/cuda/matrix_fields_multiple_field_solve.jl

@@ -85,10 +85,10 @@ function multiple_field_solve_kernel!(
 ) where {Nnames}
     @inbounds begin
         Ni, Nj, _, _, Nh = size(Fields.field_values(x1))
-        tidx = (CUDA.blockIdx().x - 1) * CUDA.blockDim().x + CUDA.threadIdx().x
-        if 1 ≤ tidx ≤ prod((Ni, Nj, Nh, Nnames))
-            (i, j, h, iname) =
-                CartesianIndices((1:Ni, 1:Nj, 1:Nh, 1:Nnames))[tidx].I
+        tidx = thread_index()
+        n = (Ni, Nj, Nh, Nnames)
+        if valid_range(tidx, prod(n))
+            (i, j, h, iname) = kernel_indexes(tidx, n).I
             generated_single_field_solve!(
                 device,
                 caches,

src/DataLayouts/DataLayouts.jl

@@ -46,6 +46,21 @@ abstract type AbstractData{S} end
 
 Base.size(data::AbstractData, i::Integer) = size(data)[i]
 
+"""
+    (Ni, Nj, Nf, Nv, Nh) = universal_size(data::AbstractData)
+
+Returns dimensions in a universal
+format for all data layouts:
+ - `Ni` number of spectral element nodal degrees of freedom in first horizontal direction
+ - `Nj` number of spectral element nodal degrees of freedom in second horizontal direction
+ - `Nf` number of field components
+ - `Nv` number of vertical degrees of freedom
+ - `Nh` number of horizontal elements
+
+Note: this is similar to `Base.size`, except
+      that `universal_size` does not return 1
+      for the number of field components.
+"""
 function universal_size(data::AbstractData)
     s = size(data)
     return (s[1], s[2], ncomponents(data), s[4], s[5])

test/Spaces/distributed_cuda/ddss3.jl

@@ -147,5 +147,5 @@ partition numbers
     end
 #! format: on
     p = @allocated Spaces.weighted_dss!(y0, dss_buffer)
-    iamroot && @test p ≤ 10816
+    iamroot && @test p ≤ 11072
 end