Skip to content

Commit e4a15ba

Browse files
dennisYatuninclaude
andcommitted
Fix GPU InvalidIR: compute _col_major_offset via unrolled_reduce [perf]
The real cause of the persisting Integrals/matrix GPU InvalidIRError (ijl_get_nth_field_checked in single_column_reduce!/bycolumn_kernel! and eager_copyto_stencil_kernel!) was F1's `_col_major_offset` helper, written as a tuple TAIL-RECURSION. It constant-folds on CPU, but under the deep inlining of a stencil / column-reduce kernel, device inference exhausts its recursion/const-prop budget, widens the peeled tuple argument to an abstract `Tuple`, and lowers `idx[1]` to a runtime bounds-checked `getfield` (ijl_get_nth_field_checked) that cannot compile on device. (The earlier type-level flat_mode change was not the culprit; it is kept as harmless hardening.) Replace the recursion with `unrolled_reduce` (UnrolledUtilities, already a DataLayouts dependency), keeping all unrolling isolated to that package rather than an ad-hoc @generated: it emits a non-self-recursive unrolled body that folds regardless of the surrounding inference budget. Carries (offset, stride) in a left fold over the (size, coord) pairs. Verified locally (device-free GPUCompiler IR check, manual CompilerJob on the real device kernel types — not the Adapt path): bycolumn_kernel! (extruded + column) and eager_copyto_stencil_kernel! now compile clean; DataLayouts unit_loops / unit_struct / unit_fill_and_copyto pass; CPU getindex/setindex values unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 68ae496 commit e4a15ba

1 file changed

Lines changed: 17 additions & 5 deletions

File tree

src/DataLayouts/indexing.jl

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -142,11 +142,23 @@ end
142142
end
143143

144144
# Column-major linear offset (0-based) of coordinate tuple `idx` in an array of
145-
# static size `sz`. Tail recursion over tuples inlines to a folded Horner form;
146-
# `sz` entries are type-domain literals, only the coordinates are runtime values.
147-
@inline _col_major_offset(::Tuple{}, ::Tuple{}) = 0
148-
@inline _col_major_offset(sz::Tuple, idx::Tuple) =
149-
(idx[1] - 1) + sz[1] * _col_major_offset(Base.tail(sz), Base.tail(idx))
145+
# static size `sz`, computed with `unrolled_reduce` so all unrolling lives in
146+
# UnrolledUtilities. A tuple tail-recursion folds on CPU but, under the deep
147+
# inlining of a stencil/column-reduce kernel, GPU inference hits its
148+
# recursion/const-prop budget, widens the peeled tuple to an abstract `Tuple`,
149+
# and lowers `idx[1]` to a runtime bounds-checked `getfield`
150+
# (`ijl_get_nth_field_checked`) that cannot compile on device. `unrolled_reduce`
151+
# emits a non-self-recursive unrolled body that folds regardless of the
152+
# surrounding inference budget. It carries `(offset, stride)`: at axis `k` the
153+
# offset gains `(idx[k] - 1) * stride` and the stride advances by `sz[k]`.
154+
@inline _col_major_offset(sz::Tuple, idx::Tuple) = first(
155+
unrolled_reduce(
156+
unrolled_map(tuple, sz, idx);
157+
init = (0, 1),
158+
) do (offset, stride), (s, i)
159+
(offset + (i - 1) * stride, stride * s)
160+
end,
161+
)
150162

151163
@inline _first_n(t::Tuple, ::Val{n}) where {n} = ntuple(k -> t[k], Val(n))
152164

0 commit comments

Comments
 (0)