Shorten code comments and tighten reduction utilities

Long explanatory comments introduced during the DataLayouts rewrite are
condensed (142 lines to 83) while keeping every load-bearing rationale.
safe_mapreduce now asserts that collections are non-empty when init is
not provided, PAIRWISE_BLOCKSIZE cites Base.pairwise_blocksize (which
uses the same 1024 for Base's pairwise mapreduce), and
running_in_external_threaded_loop reads the task storage once.

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

Author: dennisYatunin

ext/ClimaCoreCUDAExt.jl

@@ -32,11 +32,10 @@ include(joinpath("cuda", "matrix_fields_single_field_solve.jl"))
 include(joinpath("cuda", "matrix_fields_multiple_field_solve.jl"))
 include(joinpath("cuda", "operators_spectral_element.jl"))
 
-# Disable the recursion limit for the device reduce_points, which recursively
-# reduces over warps and sub-warps while forwarding its keyword arguments. This
-# resembles unbounded recursion to the compiler, which otherwise widens and boxes
-# the arguments (requiring dynamic dispatch). The limit must also be lifted on the
-# keyword-argument body functions, since that is where the recursion occurs.
+# Lift the recursion limit for the device reduce_points, whose recursion over warps
+# and sub-warps forwards kwargs and looks unbounded to the compiler, which would widen
+# and box the arguments (requiring dynamic dispatch). The limit must also be lifted on
+# the keyword-argument body functions, since that is where the recursion occurs.
 @static if hasfield(Method, :recursion_relation)
     for method in methods(ClimaCore.DataLayouts.reduce_points)
         method.module === (@__MODULE__) || continue

ext/cuda/loops.jl

@@ -31,12 +31,9 @@ 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)
-    # 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.
+    # Launch at most one thread per point, so that every thread's strided range of
+    # indices is nonempty. Threads without values would need warp-shuffle placeholders,
+    # which reductions without neutral elements (like min) cannot generate.
     max_threads = threads_via_occupancy(kernel, (empty_results, arg))
     threads = min(length(arg), max_threads)
     blocks = max(fld(length(arg), threads), 1)
@@ -50,10 +47,8 @@ function DataLayouts.reduce_points(::ThisHost, op::O, arg; kwargs...) where {O}
     return CUDA.@allowscalar @inbounds results[1]
 end
 
-# If given a warp or a partition of a warp, reduce each threads's values, then
-# reduce the results using warp shuffles. Otherwise, reduce each warp's values,
-# then reduce the results in the first warp. Limit warp shuffles to the active
-# threads in each warp, since the results from inactive threads are undefined.
+# Reduce a warp or sub-warp with warp shuffles, limited to active threads (inactive
+# results are undefined); otherwise reduce each warp, then combine in the first warp.
 DataLayouts.reduce_points(scope::ThisCooperativeGroup, op::O, arg; kwargs...) where {O} =
     if scope != ThisBlock() && DataLayouts.num_threads(scope) <= THREADS_PER_WARP
         thread_result =
@@ -87,10 +82,8 @@ DataLayouts.reduce_points(scope::ThisCooperativeGroup, op::O, arg; kwargs...) wh
 @generated num_reductions(::ThisSubBlock{N}) where {N} =
     8 * sizeof(N) - leading_zeros(N) - 1
 
-# Use warp shuffles to perform binary tree reductions over the first num_values
-# threads in a warp or sub-warp. All active threads in the warp must execute the
-# shuffles, but values from threads with ranks above num_values are ignored. The
-# mask identifies active lanes of the warp, which correspond to its lowest bits.
+# Binary-tree warp-shuffle reduction over the first num_values threads: every active
+# lane (the mask's lowest bits) must shuffle, but ranks above num_values are ignored.
 function shuffle_reduce(scope, op::O, value, num_values) where {O}
     num_offsets = num_reductions(scope)
     num_inactive = THREADS_PER_WARP - num_active_threads(ThisWarp())
@@ -104,10 +97,8 @@ function shuffle_reduce(scope, op::O, value, num_values) where {O}
     return value
 end
 
-# CUDA's warp shuffle intrinsics only support scalar values, so register how they
-# should shuffle the AutoBroadcaster-wrapped values produced by multi-component
-# reductions: shuffle each of the underlying values, recursing through any nested
-# AutoBroadcasters. shfl_recurse is CUDA's documented extension point for value
-# types that the shuffle intrinsics do not natively support.
+# CUDA's warp shuffle intrinsics only support scalar values; shfl_recurse is their
+# documented extension point, used here to recursively shuffle each value wrapped in
+# the AutoBroadcasters produced by multi-component reductions.
 CUDA.shfl_recurse(op::O, x::Utilities.AutoBroadcaster) where {O} =
     Utilities.AutoBroadcaster(UnrolledUtilities.unrolled_map(op, Utilities.unwrap(x)))

ext/cuda/scopes.jl

@@ -115,13 +115,10 @@ function DataLayouts.scoped_static_array(scope::ThisSubBlock, ::Type{T}, dims) w
     return @inbounds view(array, ntuple(Returns(:), Val(length(dims)))..., subblock_index)
 end
 
-# The last sub-block in a block may be only partially filled, so its active
-# thread count is computed from the total number of threads in the block. This
-# cannot be determined from CUDA.active_mask, which is not guaranteed to be
-# consistent across different lanes of the same warp (since the Volta
-# architecture introduced independent thread scheduling, the lanes of a warp are
-# no longer guaranteed to be synchronized when they poll their mask; see
-# https://stackoverflow.com/questions/54055195 for more details).
+# The last sub-block in a block may be only partially filled, so its active thread
+# count is computed from the block's total. CUDA.active_mask is unusable here: since
+# Volta's independent thread scheduling, it is not guaranteed to be consistent across
+# lanes of a warp (see https://stackoverflow.com/questions/54055195).
 num_active_threads(scope::ThisBlock) = DataLayouts.num_threads(ThisBlock())
 function num_active_threads(scope::ThisSubBlock)
     num_threads = DataLayouts.num_threads(scope)
@@ -130,11 +127,9 @@ function num_active_threads(scope::ThisSubBlock)
     return clamp(DataLayouts.num_threads(ThisBlock()) - subblock_offset, 0, num_threads)
 end
 
-# Iterate over a reshape-free generator on GPUs, since the default view of a
-# strided range of multidimensional CartesianIndices builds a ReshapedArray
-# whose bounds-checking machinery cannot be compiled in a GPU kernel. One-
-# dimensional indices (from getindex slices) are left to the default view, which
-# is already a reshape-free range that supports the indexing used by reductions.
+# A strided view of multidimensional CartesianIndices is a ReshapedArray whose bounds
+# checking cannot be compiled in a GPU kernel, so iterate a reshape-free generator here.
+# One-dimensional indices (from getindex slices) keep the default view, an indexable range.
 Base.@propagate_inbounds DataLayouts.subscope_index_view(
     ::Union{ThisKernel, ThisCooperativeGroup},
     indices::CartesianIndices,

ext/cuda/topologies_dss.jl

@@ -127,11 +127,9 @@ function Topologies.dss_local!(
             (v, vertex_index) = CartesianIndices((Nv, nlocalvertices))[gidx].I
             first_offset = local_vertex_offset[vertex_index]
             last_offset = local_vertex_offset[vertex_index + 1] - 1
-            # Accumulate in a loop instead of calling sum with a closure; the
-            # closure would box the variable v (which is also assigned in the
-            # other branch of this conditional), and the empty-collection error
-            # path of sum cannot be compiled in a GPU kernel. Every vertex has
-            # at least one entry in local_vertices, so the loop is not empty.
+            # Accumulate in a loop: sum with a closure would box v (also assigned in the
+            # other branch), and sum's empty-collection error path cannot be compiled in
+            # a GPU kernel. Every vertex has a local_vertices entry, so the loop is nonempty.
             (h, vert) = local_vertices[first_offset]
             p = Topologies.perimeter_vertex_node_index(vert)
             sum_data = perimeter_data[v, p, 1, h]

src/DataLayouts/DataLayouts.jl

@@ -283,16 +283,12 @@ end
 
 @inline Base.propertynames(data::DataLayout) = fieldnames(eltype(data))
 
-# Wrap the field index in a Val as soon as it is available, so that all of the
-# lookups required to generate a field view are resolved through specialization
-# instead of constant propagation. The index or name of a field is a constant
-# at every getproperty call site, and converting it to a Val only requires one
-# level of constant propagation, which the compiler performs by default; in
-# contrast, propagating the constant through every fieldtype lookup can exhaust
-# the compiler's budget in expressions with several getproperty calls, which
-# results in runtime allocations from dynamic types. Use Base's fieldindex to
-# convert the field name into an index, since it is guaranteed to constant-fold
-# whenever the name can be statically inferred.
+# Wrap the field index in a Val as soon as it is available, so field-view lookups are
+# resolved through specialization: making the Val needs only one level of constant
+# propagation (done by default), whereas propagating the index through every fieldtype
+# lookup can exhaust the compiler's budget when an expression has several getproperty
+# calls, causing runtime allocations from dynamic types. Base's fieldindex is used
+# because it is guaranteed to constant-fold whenever the name is statically inferred.
 @inline function property_view(data, ::Val{i}) where {i}
     1 <= i <= fieldcount(eltype(data)) || throw(BoundsError(data, i))
     array =
@@ -315,10 +311,8 @@ end
     size(array) == array_size ? array :
     throw(DimensionMismatch("Array size is not consistent with layout type"))
 
-# A dynamic Nh is marked with nothing instead of missing because missing cannot
-# appear in the type parameters of a GPU kernel argument; GPUCompiler mangles
-# kernel names by comparing each type parameter with ==, and the three-valued
-# logic of missing turns those comparisons into non-boolean conditions.
+# A dynamic Nh is marked with nothing rather than missing: GPUCompiler mangles kernel
+# names by comparing type parameters with ==, and missing's three-valued == is non-boolean.
 @inline check_Nh_dynamic(Nh_dynamic) =
     !isnothing(Nh_dynamic) ||
     throw(ArgumentError("Nh_dynamic must be specified to construct layout type"))
@@ -542,13 +536,10 @@ include("masks.jl")
 include("loops.jl")
 include("deprecated.jl")
 
-# Disable the recursion limit for every method of Core.kwcall in this module,
-# and also for the recursive slice_subscope function. This is needed for keyword
-# argument functions like fill! and column_reduce! to be composed with each
-# other an arbitrary number of times, and for slice_subscope to recursively
-# partition a scope multiple times before terminating. The default limit
-# prevents the compiler from fully analyzing such nested calls, instead making
-# it use widened argument types that are boxed and require dynamic dispatch.
+# Lift the recursion limit for this module's Core.kwcall methods and the recursive
+# scope functions, so that keyword-argument functions like fill! and column_reduce!
+# can compose arbitrarily and slice_subscope can repeatedly partition a scope. The
+# default limit makes the compiler widen and box arguments, requiring dynamic dispatch.
 @static if hasfield(Method, :recursion_relation)
     for func in (Core.kwcall, slice_subscope, is_subscope), method in methods(func)
         method.module === (@__MODULE__) || continue

src/DataLayouts/broadcast.jl

@@ -9,10 +9,9 @@ DataStyle(::Type{D}) where {D} = DataStyle{ndims(D), layout_type(D)}()
 
 Broadcast.BroadcastStyle(::Type{D}) where {D <: DataLayout} = DataStyle(D)
 
-# When two styles share a typename but have different dimensionalities, Base's
-# fallback for AbstractArrayStyle evaluates typeof(style)(Val(N)) with the
-# larger dimensionality N. DataStyle cannot implement such a Val constructor
-# because it needs the layout type D, so it bypasses this fallback.
+# For styles with equal typenames but different dimensionalities, Base's fallback for
+# AbstractArrayStyle calls typeof(style)(Val(N)); DataStyle needs the layout type D, so
+# it cannot define that constructor and bypasses the fallback instead.
 Broadcast.BroadcastStyle(::DataStyle{<:Any, D1}, ::DataStyle{<:Any, D2}) where {D1, D2} =
     D1 == D2 || iszero(ndims(D2)) ? DataStyle(D1) :
     iszero(ndims(D1)) ? DataStyle(D2) : Broadcast.Unknown()
@@ -37,12 +36,9 @@ A [`DataStyle`](@ref) broadcast expression whose [`layout_type`](@ref) is `D`.
 """
 const LazyDataLayout{D} = Broadcast.Broadcasted{<:DataStyle{<:Any, D}}
 
-# Optimize axes(::LazyDataLayout) to use statically inferrable size information.
-# When the size is dynamic, combine the axes of the arguments without Base's
-# combine_axes, whose dimension-mismatch check has an error path that cannot be
-# compiled in GPU kernels; the arguments are assumed to be broadcast-compatible,
-# as they are in any valid broadcast, so each dimension takes the non-singleton
-# axis of the two being combined.
+# Optimize axes(::LazyDataLayout) with statically inferrable sizes. Dynamic sizes avoid
+# Base's combine_axes, whose dimension-mismatch error path cannot compile in GPU kernels:
+# assuming broadcast-compatible arguments, each dimension takes the non-singleton axis.
 @inline Broadcast._axes(bc::LazyDataLayout, ::Nothing) =
     has_inferred_size(bc) ? unrolled_map(Base.OneTo, inferred_size(bc)) :
     unrolled_reduce(broadcast_axes, unrolled_map(axes, bc.args))

src/DataLayouts/indexing.jl

@@ -9,14 +9,11 @@ Base.IndexStyle(bc::FusedMultiBroadcast) = IndexStyle(unrolled_map(first, bc.pai
 
 const IndexableData = Union{DataLayout, LazyDataLayout, FusedMultiBroadcast}
 
-# Preserve linear indices when they are used to access broadcast arguments.
-# Base's fallback method for newindex reinterprets an integer index as a
-# CartesianIndex along the first dimension, which would cause a linear index to
-# silently read the wrong element from any argument with multiple dimensions.
-# Linear indices are only used when IndexStyle confirms that every layout with
-# a nonzero number of dimensions in a broadcast expression has the same shape,
-# so the only required index conversion is for 0-dimensional data, whose single
-# point is accessed by every index.
+# Preserve linear indices into broadcast arguments: Base's newindex fallback
+# reinterprets an integer as a CartesianIndex along the first dimension, silently
+# reading the wrong element from multidimensional arguments. IndexStyle only permits
+# linear indices when all nonzero-dimensional layouts share a shape, so only
+# 0-dimensional data (its single point read by every index) needs conversion.
 @inline Broadcast.newindex(arg::IndexableData, index::Integer) =
     iszero(ndims(arg)) ? CartesianIndex() : index
 
@@ -92,13 +89,10 @@ end
     indices::Integer...,
 ) = view(arg, CartesianIndex(index1, index2, indices...))
 
-# Return a DataF instead of a layout with singleton dimensions, since a view of
-# a single point is 0-dimensional. This also avoids reshaping the 1-dimensional
-# view of the point's entries, which would require the construction of division
-# helpers that cannot be optimized into SIMD instructions in pointwise loops.
-# The inner DataF constructor is used because the view of the entries is
-# already shaped correctly; the checking constructor's runtime size comparison
-# would make the parent array type a Union of the checked and reshaped types.
+# A single-point view is 0-dimensional, so return a DataF instead of a layout with
+# singleton dimensions, avoiding a reshape of the 1-D entry view whose division helpers
+# block SIMD in pointwise loops. The inner DataF constructor is used since the checking
+# one's runtime size comparison would union the checked and reshaped parent array types.
 @propagate_inbounds function Base.view(data::DataLayout, index::PointIndex)
     is_invalid_linear(data, index) && return view(data, CartesianIndices(data)[index])
     array = view_struct(parent(data), eltype(data), index, Val(f_dim(data)))