Resolve conflicting spectral broadcast styles for instantiated arguments#2551
Draft
haakon-e wants to merge 3 commits into
Draft
Resolve conflicting spectral broadcast styles for instantiated arguments#2551haakon-e wants to merge 3 commits into
haakon-e wants to merge 3 commits into
Conversation
The device-resolved style assigned by instantiate takes precedence over the generic construction-time SpectralStyle, and operator construction accepts any AbstractSpectralStyle, so an instantiated spectral broadcast (e.g. one captured with LazyBroadcast.lazy) can be used as an argument to a new operator broadcast.
Adds CPU and CUDA Buildkite steps for the new test, an assertion using LazyBroadcast.lazy alongside the explicit instantiate, and trims the style-rule comment.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Composing a spectral-element operator over an already-instantiated spectral broadcast throws
conflicting broadcast rules defined(SpectralStylevsCUDASpectralStyleon GPU,SpectralStylevsSlabBlockSpectralStyleon CPU). This arises whenever a broadcast containing spectral operators is instantiated first and then used as an argument to a new operator broadcast, for example a tree captured withLazyBroadcast.lazy, whosematerializeisinstantiate.Mechanism
Base.Broadcast.broadcasted(op::SpectralElementOperator, args...)tags construction with the genericSpectralStyle()and merges it with the argument styles throughresult_style.instantiatere-tags spectral trees with the device-resolved style. The only rule ordering spectral styles is the symmetricBroadcastStyle(style::AbstractSpectralStyle, ::Fields.AbstractFieldStyle) = style; since every concrete spectral style is a subtype of both argument types, the pair (generic, device-resolved) matches in both argument orders with different winners, which is the ambiguityresult_stylerejects.Fix
Two changes in
src/Operators/spectralelement.jl, mirroring the finite-difference design:BroadcastStyle(::SpectralStyle, style::AbstractSpectralStyle) = style. Both call orders then agree: no explicit reverse-order method is needed because the pair (device style,SpectralStyle) already resolves through the pre-existing symmetric rule, which returns its first argument; the new rule is strictly more specific than that rule in both arguments, so there is no ambiguity.broadcasted(::SpectralStyle, op::SpectralElementOperator, args...)is generalized to::AbstractSpectralStyle, returningSpectralBroadcasted{typeof(style)};instantiatere-resolves the style parameter from the device as before, so materialization is unchanged.The finite-difference hierarchy already handles the analogous pairing: the CUDA extension defines
BroadcastStyle(::ColumnStencilStyle, ::CUDAColumnStencilStyle)andbroadcastedis generic overAbstractStencilStyle(src/Operators/finitedifference.jl,ext/cuda/operators_finite_difference.jl). The spectral rule lives in the main package rather than the extension because, unlike finite differences, the spectral construction style is distinct from both device styles and the CPU device styleSlabBlockSpectralStyleis a main-package type, so the CPU half of the conflict can only be resolved in the main package; the single rule covers both device styles and needs no extension change.Validation
test/Operators/spectralelement/lazy_composition.jl: the style-combination lattice, and a weak divergence applied to a pre-instantiated gradient sum, checked against the inline composition. Passes on CPU and GPU.Test.detect_ambiguitiesreports no new ambiguities; no other caller dispatches on the::SpectralStyleconstruction method.test/Operators/spectralelement/rectilinear.jlpasses on the patched source.Pairs of two different device-resolved styles (e.g.
SlabBlockSpectralStylevsCUDASpectralStyle) continue to error; that pairing indicates a genuine device mismatch. A more descriptive message could be added separately.This fix resolves the style ambiguity only. Compiling a single large fused spectral kernel remains subject to pre-existing GPU kernel-size limitations, so the fix does not by itself make arbitrarily large lazy spectral compositions viable on GPU.