Fix GPU reductions, kernel-side slice lengths, and mask device placement

Fixes for the CI failures on build 7043, all in code that can only run with
a GPU:

- Compute slice lengths from shape parameters instead of inferring the
  types of slices; the return_type of a slice operator cannot be obtained
  from within a GPU kernel, and the previous approach also failed on the
  host for slices of broadcasted expressions and for single-point views of
  expressions with mixtures of parent array types
- Generate the butterfly shuffle offsets in warp reductions from a static
  thread count, since tuples with runtime lengths cannot be constructed in
  GPU kernels
- Launch at most one reduction thread per point, so that every thread has
  at least one value to reduce; threads without values would need
  placeholders for warp shuffles, but reductions without neutral elements
  (like minimum) cannot generate placeholders
- Use the active mask instead of the full mask for warp shuffles, so that
  partial warps can also execute them, and exclude the lanes of partial
  warps from reduction results
- Construct each grid's mask from its device-side local geometry, so that
  masks of GPU spaces are not stored in CPU arrays
- Qualify the CUDA intrinsics and import slab in the GPU DataLayouts test

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: dennisYatunin

ext/cuda/loops.jl

@@ -31,7 +31,15 @@ function DataLayouts.reduce_points(::ThisHost, op::O, arg; kwargs...) where {O}
     end
     T = return_type(op, NTuple{2, eltype(arg)})
     empty_results = DataLayouts.scoped_array(ThisHost(), T, 0)
-    (; threads, blocks) = config_via_occupancy(kernel, length(arg), (empty_results, arg))
+    # Launch at most one thread per point, so that every thread has at least
+    # one value to reduce (indices are divided among threads through strided
+    # ranges, which are only guaranteed to be nonempty when the total thread
+    # count does not exceed the number of points). Threads without values would
+    # need placeholders for the warp shuffles, but reduction operations without
+    # neutral elements (like the min in minimum) cannot generate placeholders.
+    max_threads = threads_via_occupancy(kernel, (empty_results, arg))
+    threads = min(length(arg), max_threads)
+    blocks = max(fld(length(arg), threads), 1)
     num_results = min(max_resident_blocks(threads), blocks)
     results = similar(empty_results, num_results)
     @cuda always_inline = true threads = threads blocks = num_results kernel(results, arg)
@@ -69,29 +77,46 @@ function DataLayouts.reduce_points(::ThisBlock, op::O, arg; kwargs...) where {O}
     return @inbounds results[1]
 end
 
-DataLayouts.reduce_points(scope::ThisSubBlock, op::O, arg; kwargs...) where {O} =
-    DataLayouts.num_threads(scope) <= THREADS_PER_WARP ?
-    shuffle_reduce(scope, op, DataLayouts.reduce_points(ThisThread(), op, arg; kwargs...)) :
+DataLayouts.reduce_points(scope::ThisSubBlock{N}, op::O, arg; kwargs...) where {N, O} =
+    N <= THREADS_PER_WARP ?
+    shuffle_reduce(
+        scope,
+        op,
+        DataLayouts.reduce_points(ThisThread(), op, arg; kwargs...),
+        # Partial warps at the ends of partial blocks have fewer than N threads.
+        min(N, count_ones(CUDA.active_mask())),
+    ) :
     DataLayouts.reduce_points(
         ThisWarp(),
         op,
         DataLayouts.reassign(arg, ThisWarp());
         kwargs...,
     )
 
+# Generate the tuple of butterfly shuffle offsets (n ÷ 2, n ÷ 4, ..., 1) from a
+# static thread count, since a tuple whose length depends on a runtime value
+# cannot be constructed in a GPU kernel.
+@generated shuffle_offsets(::Val{n}) where {n} =
+    ntuple(Base.Fix1(>>, n), 8 * sizeof(n) - leading_zeros(n) - 1)
+
 # Use warp shuffles to perform binary tree reductions over the first
 # num_threads values in a scope. Every thread in the scope must execute the
 # shuffles, but values from threads with ranks above num_threads are ignored.
-function shuffle_reduce(scope, op::O, value, num_threads = THREADS_PER_WARP) where {O}
-    DataLayouts.is_subscope(scope, ThisWarp()) ||
-        throw(ArgumentError(DataLayouts.invalid_subscope_string(scope, ThisWarp())))
-    num_threads <= THREADS_PER_WARP ||
+# The active mask is used instead of the full mask so that partial warps can
+# also execute the shuffles (values from lanes outside the mask are undefined,
+# but such lanes always have ranks above num_threads).
+function shuffle_reduce(
+    scope::ThisSubBlock{n},
+    op::O,
+    value,
+    num_threads = THREADS_PER_WARP,
+) where {n, O}
+    n <= THREADS_PER_WARP ||
         throw(ArgumentError("Number of threads is too large for warp shuffle"))
-    n = DataLayouts.num_threads(scope)
+    mask = CUDA.active_mask()
     rank = DataLayouts.thread_rank(scope)
-    log2_n = 8 * sizeof(n) - Base.ctlz_int(n) - 1
-    for offset in ntuple(Base.Fix1(>>, n), Val(log2_n)) # n ÷ 2, n ÷ 4, ..., 1
-        shuffled_value = CUDA.shfl_xor_sync(CUDA.FULL_MASK, value, offset)
+    for offset in shuffle_offsets(Val(n))
+        shuffled_value = CUDA.shfl_xor_sync(mask, value, offset)
         if rank <= num_threads && xor(rank - 1, offset) + 1 <= num_threads
             value = op(value, shuffled_value)
         end

src/DataLayouts/loops.jl

@@ -18,11 +18,19 @@ end
 @generated invalid_mask_string(::M, ::O) where {M, O} =
     "$M cannot be applied to $(O.instance) slices"
 
+# Compute slice lengths directly from shape parameters instead of inferring the
+# types of slices, since the return_type of a slice operator cannot be obtained
+# from within a GPU kernel.
+@inline slice_size_params(::typeof(view), arg) = ()
+@inline slice_size_params(::typeof(column), arg) = (vijh_params(arg).Nv,)
+@inline slice_size_params(::typeof(slab), arg) =
+    (vijh_params(arg).Ni, vijh_params(arg).Nj)
+@inline slice_size_params(::typeof(level), arg) =
+    (vijh_params(arg).Ni, vijh_params(arg).Nj, vijh_params(arg).Nh)
 @inline function inferred_slice_length(op::O, arg) where {O}
-    index = one(eltype(each_slice_index(op, arg)))
-    slice_type = return_type(op, typeof((arg, Tuple(index)...)))
-    has_inferred_size(slice_type) && return prod(inferred_size(slice_type))
-    throw(ArgumentError("Size of view from slice operator must be inferrable"))
+    params = slice_size_params(op, arg)
+    params isa Tuple{Vararg{Integer}} && return prod(params; init = 1)
+    throw(ArgumentError("Size of slice from slice operator must be inferrable"))
 end
 
 """

src/Grids/spectralelement.jl

@@ -249,7 +249,6 @@ function _SpectralElementGrid2D(
     LG = Geometry.LocalGeometryType(CoordType2D, FT, AIdx)
 
     local_geometry = VIJH{LG, 1, Nq, Nq, nothing}(Array{FT}, Nh)
-    mask = enable_mask ? DataLayouts.IJHMask(local_geometry) : DataLayouts.NoMask()
 
     _, quad_weights = Quadratures.quadrature_points(FT, quadrature_style)
     _, high_order_quad_weights =
@@ -424,6 +423,11 @@ function _SpectralElementGrid2D(
     end
 
     device_local_geometry = DataLayouts.rebuild(local_geometry, DA)
+    # Construct the mask from the device-side geometry, so that its data is
+    # stored on the same device as the rest of the grid.
+    mask =
+        enable_mask ? DataLayouts.IJHMask(device_local_geometry) :
+        DataLayouts.NoMask()
     return SpectralElementGrid2D(
         topology,
         quadrature_style,

test/DataLayouts/cuda.jl

@@ -1,12 +1,13 @@
 using Test
 import ClimaComms
+import ClimaCore: slab
 import ClimaCore.DataLayouts: VIJFH
 ClimaComms.@import_required_backends
 
 function knl_copy!(dest, src)
-    i = threadIdx().x
-    j = threadIdx().y
-    h = blockIdx().x
+    i = CUDA.threadIdx().x
+    j = CUDA.threadIdx().y
+    h = CUDA.blockIdx().x
     p_dest = slab(dest, h)
     p_src = slab(src, h)
     @inbounds p_dest[1, i, j, 1] = p_src[1, i, j, 1]