Skip to content

Resolve conflicting spectral broadcast styles for instantiated arguments#2551

Draft
haakon-e wants to merge 3 commits into
mainfrom
he/spectral-style-conflict
Draft

Resolve conflicting spectral broadcast styles for instantiated arguments#2551
haakon-e wants to merge 3 commits into
mainfrom
he/spectral-style-conflict

Conversation

@haakon-e

@haakon-e haakon-e commented Jul 17, 2026

Copy link
Copy Markdown
Member

Summary

Composing a spectral-element operator over an already-instantiated spectral broadcast throws conflicting broadcast rules defined (SpectralStyle vs CUDASpectralStyle on GPU, SpectralStyle vs SlabBlockSpectralStyle on 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 with LazyBroadcast.lazy, whose materialize is instantiate.

Mechanism

Base.Broadcast.broadcasted(op::SpectralElementOperator, args...) tags construction with the generic SpectralStyle() and merges it with the argument styles through result_style. instantiate re-tags spectral trees with the device-resolved style. The only rule ordering spectral styles is the symmetric BroadcastStyle(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 ambiguity result_style rejects.

Fix

Two changes in src/Operators/spectralelement.jl, mirroring the finite-difference design:

  • A rule making the device-resolved style take precedence over the generic construction-time style: 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.
  • The construction method broadcasted(::SpectralStyle, op::SpectralElementOperator, args...) is generalized to ::AbstractSpectralStyle, returning SpectralBroadcasted{typeof(style)}; instantiate re-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) and broadcasted is generic over AbstractStencilStyle (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 style SlabBlockSpectralStyle is 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

  • New unit test 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.
  • All style pairs were enumerated before and after: only the previously erroring (generic, device-resolved) pairs change behavior; Test.detect_ambiguities reports no new ambiguities; no other caller dispatches on the ::SpectralStyle construction method.
  • test/Operators/spectralelement/rectilinear.jl passes on the patched source.

Pairs of two different device-resolved styles (e.g. SlabBlockSpectralStyle vs CUDASpectralStyle) 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.

haakon-e added 3 commits July 16, 2026 19:09
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant