Skip to content

Commit 68ae496

Browse files
dennisYatuninclaude
andcommitted
Make flat_mode type-level to fix GPU compilation [perf]
F1's `flat_mode(data::DataLayout)` decided the leaf-access strategy in the value domain, using a local `F = f_dim(data)` to slice `inferred_size(data)[F:end]`. This constant-folds under CPU inference but could fail to fold under the weaker device const-prop, leaving `_getidx_flat(flat_mode(data), ...)` a dynamic dispatch that compiled the Val(2) property-view-root path (parentindices / parent-of-parent) for unsupported layout types and emitted a runtime-checked `getfield` (ijl_get_nth_field_checked) -> GPU InvalidIRError. This regressed the Integrals (single_column_reduce!) and matrix-multiplication-recursion (eager_copyto_stencil_kernel!) GPU jobs on the previous commit. Redecide flat_mode entirely from the layout TYPE, mirroring `IndexStyle(::Type{D})` (which is GPU-safe), and thread F as a `Val` type parameter so the `inferred_size(D)[F:end]` slice is a compile-time literal. flat_mode is now a guaranteed compile-time constant on every backend, so exactly one _getidx_flat / _setidx_flat! method is compiled per concrete layout. Behavior is unchanged (same Val(0/1/2) selection); CPU getindex/setindex remain checked-getfield-free and the FD-stencil :kernel compile is unchanged (device-free checker). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent bd53967 commit 68ae496

1 file changed

Lines changed: 21 additions & 12 deletions

File tree

src/DataLayouts/indexing.jl

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -98,19 +98,28 @@ is_trivial_point(data, index) = isone(length(data)) && index == CartesianIndex()
9898
# Base reindex / `div` / `rem` is instantiated.
9999
# Val(0) everything else (dynamic `Nh`, or a true box `VIJFH` field view with
100100
# `Nh > 1`): the unchanged Cartesian path, which stays correct.
101-
@inline function flat_mode(data::DataLayout)
102-
has_inferred_size(data) || return Val(0)
103-
IndexStyle(data) isa IndexLinear && return Val(1)
104-
F = f_dim(data)
105-
F isa Integer || return Val(0)
106-
# The root path omits every logical coordinate at or beyond the F axis, so it
107-
# is only correct when those extents are all 1 (the column case). Unlike the
108-
# layout IndexStyle, a `ncomponents <= 1` escape is NOT valid here: a scalar
109-
# field view of a box (Nh > 1) still needs its h-term, which we drop.
110-
all_ones(inferred_size(data)[F:end]...) &&
111-
reducible_property_view(parent_type(data), Val(F)) && return Val(2)
112-
return Val(0)
101+
# Decided entirely from the layout TYPE, mirroring `Base.IndexStyle(::Type{D})`,
102+
# so the decision is a compile-time constant on every backend (the value-domain
103+
# form let `f_dim(data)` reach a runtime tuple slice that failed to fold under
104+
# GPU inference, turning `_getidx_flat(flat_mode(data), ...)` into a dynamic
105+
# dispatch that compiled the `Val(2)` root path for unsupported types and emitted
106+
# a checked `getfield`). `F` is threaded as a `Val` type parameter so
107+
# `inferred_size(D)[F:end]` is a literal slice.
108+
@inline flat_mode(data::DataLayout) = flat_mode(typeof(data))
109+
@inline function flat_mode(::Type{D}) where {D <: DataLayout}
110+
has_inferred_size(D) || return Val(0)
111+
IndexStyle(D) isa IndexLinear && return Val(1)
112+
return _flat_mode_cartesian(D, Val(f_dim(D)))
113113
end
114+
# The root path omits every logical coordinate at or beyond the F axis, so it is
115+
# only correct when those extents are all 1 (the column case). Unlike the layout
116+
# IndexStyle, a `ncomponents <= 1` escape is NOT valid here: a scalar field view
117+
# of a box (Nh > 1) still needs its h-term, which we drop.
118+
@inline _flat_mode_cartesian(::Type{D}, ::Val{nothing}) where {D} = Val(0)
119+
@inline _flat_mode_cartesian(::Type{D}, ::Val{F}) where {D, F} =
120+
F isa Integer &&
121+
all_ones(inferred_size(D)[F:end]...) &&
122+
reducible_property_view(parent_type(D), Val(F)) ? Val(2) : Val(0)
114123

115124
# A property-view SubArray (as built by struct_field_view) of an IndexLinear
116125
# array, whose indices are full `Slice`s except for a `UnitRange` at position F.

0 commit comments

Comments
 (0)