Report IndexLinear for contiguous-order DataLayouts [perf]

A DataLayout's IndexStyle now reports IndexLinear directly whenever a point's
linear index equals its parent array's linear index -- when there is at most
one dimension, a single component, or the F axis and every axis after it are
singleton -- instead of deferring to the parent array's IndexStyle in the
latter two cases. This is determined by the logical column-major shape, so it
holds for any parent array type. A contiguous property-view SubArray reports
IndexCartesian even though its linear order matches the layout's, so deferring
to it routed every linear point access through a needless CartesianIndices
conversion; reporting IndexLinear directly removes that conversion, which
shrinks the typed IR of each access and lowers compilation memory. The change
does not affect operator inlining, so it introduces no runtime regression, and
it never reports IndexLinear where a point's order differs from its parent's
(e.g. a multi-component layout with a non-singleton axis after F), which is
verified by a new correctness test.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: dennisYatunin

src/DataLayouts/indexing.jl

@@ -3,11 +3,18 @@ const PointIndex = Union{Integer, CartesianIndex}
 
 all_ones(params...) = params isa Tuple{Vararg{Integer}} && unrolled_all(isone, params)
 
-# Allow linear indexing if parent(data)[1:length(data)] has one value per point.
+# Use linear indexing whenever a point's linear index equals its parent's linear
+# index: when there is at most one dimension, when there is a single component
+# (one parent entry per point), or when the F axis and every axis after it are
+# singleton (so points are stored contiguously). This is determined by the
+# logical column-major shape, so it holds for any parent array type, and we
+# report IndexLinear directly rather than deferring to the parent's IndexStyle.
+# A contiguous property-view SubArray reports IndexCartesian even though its
+# linear order matches the layout's, and deferring to it would route every
+# linear access through a needless CartesianIndices conversion.
 Base.IndexStyle(::Type{D}) where {D <: DataLayout} =
-    ndims(D) <= 1 ? IndexLinear() :
-    ncomponents(D) <= 1 || all_ones(inferred_size(D)[f_dim(D):end]...) ?
-    IndexStyle(parent_type(D)) : IndexCartesian()
+    ndims(D) <= 1 || ncomponents(D) <= 1 ||
+    all_ones(inferred_size(D)[f_dim(D):end]...) ? IndexLinear() : IndexCartesian()
 
 Base.IndexStyle(bc::MaybeFusedDataLayoutBroadcast) = IndexStyle(layout_args(bc)...)
 

test/DataLayouts/unit_loops.jl

@@ -175,3 +175,36 @@ end
     end
     @test data.unit != test_data(device, T, 1, 11).unit
 end
+
+@testset "linear indexing style and correctness" begin
+    device = ClimaComms.device()
+    arr(dims...) = device_array(device, Float64.(rand(1:(2^20), dims...)))
+    S = Tuple{Float64, Float64}
+
+    # A point's linear index equals its parent's linear index (so IndexLinear is
+    # valid) whenever there is a single component, or the F axis and everything
+    # after it are singleton. Multi-component layouts with a non-singleton axis
+    # after F must use Cartesian indexing.
+    linear = [
+        DataLayouts.VIJFH{Float64, 10, 4, 4, 5}(arr(10, 4, 4, 1, 5)),  # 1 component
+        DataLayouts.VIJHF{S, 10, 4, 4, 5}(arr(10, 4, 4, 5, 2)),         # F axis last
+        DataLayouts.VIJFH{S, 10, 4, 4, 1}(arr(10, 4, 4, 2, 1)),         # Nh = 1
+        DataLayouts.VIJFH{S, 10, 4, 4, 5}(arr(10, 4, 4, 2, 5)).:1,      # property view
+    ]
+    cartesian = [DataLayouts.VIJFH{S, 10, 4, 4, 5}(arr(10, 4, 4, 2, 5))] # F axis interior
+
+    for data in linear
+        @test IndexStyle(data) == IndexLinear()
+    end
+    for data in cartesian
+        @test IndexStyle(data) == IndexCartesian()
+    end
+    # Linear and Cartesian accesses must return identical values everywhere, so
+    # that choosing IndexLinear never silently reads the wrong element.
+    for data in vcat(linear, cartesian)
+        @test all(
+            data[i] == data[I] for
+            (i, I) in enumerate(CartesianIndices(size(data)))
+        )
+    end
+end