Skip to content

Use view instead of indexing in cpu_face_constructor_r#5376

Merged
glwagner merged 2 commits intoglw/reactant-correctnessfrom
glw/view-vertical-discretization
Mar 6, 2026
Merged

Use view instead of indexing in cpu_face_constructor_r#5376
glwagner merged 2 commits intoglw/reactant-correctnessfrom
glw/view-vertical-discretization

Conversation

@glwagner
Copy link
Copy Markdown
Member

@glwagner glwagner commented Mar 6, 2026

Summary

  • Replace cpu_nodes[1:Nz+1] with view(cpu_nodes, 1:Nz+1) in cpu_face_constructor_r
  • This avoids scalar indexing errors when grid data is stored on non-CPU architectures (e.g. Reactant's ConcreteIFRTArray)
  • Bump version to 0.105.4

Context

When using SplitExplicitFreeSurface with a TripolarGrid on Reactant, the materialize_free_surfacemaybe_extend_haloswith_halo call chain invokes cpu_face_constructor_r, which indexes into a ConcreteIFRTArray with [1:Nz+1]. This triggers Reactant's "Scalar indexing is disallowed" error. Using view avoids materializing the array.

🤖 Generated with Claude Code

Avoids scalar indexing errors when grid data is stored on
non-CPU architectures (e.g. Reactant ConcreteIFRTArray).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
glwagner added a commit to PRONTOLab/GB-25 that referenced this pull request Mar 6, 2026
Uses [sources] to pull from CliMA/Oceananigans.jl#5376 which replaces
scalar indexing with view() to avoid Reactant errors during with_halo.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
…ray})

The generic fallback on_architecture(arch, a) = a doesn't convert
OffsetArrays wrapping Reactant arrays. This adds a proper method
that converts the underlying Reactant array to Array while preserving
the OffsetArray wrapper.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
@glwagner glwagner changed the base branch from main to glw/reactant-correctness March 6, 2026 04:33
@glwagner glwagner merged commit 7b3a362 into glw/reactant-correctness Mar 6, 2026
69 of 77 checks passed
@glwagner glwagner deleted the glw/view-vertical-discretization branch March 6, 2026 04:33
glwagner added a commit that referenced this pull request Mar 8, 2026
* Add basic reactant correctness tests

* comment out other tests

* add more grids

* implement take Gu tendency

* test with different raise options

* test raise=true and raise=false

* raise=true

* add fake halo filling kernels

* filter triply periodic tests

* add reactant to weakdeps

* add more tests

* updates

* clean up

* update report

* fix

* fix

* update report

* Update test_reactant_correctness.jl

* Refactor new_data function to remove ReactantState

* Change total_size function parameter type to ShardedGrid

* try to restrict offset_array for sharding a bit more

* test raise_first too

* fix imports

* Comment out failing Reactant correctness tests and remove raise=true mode

Disable raise=true mode across all tests due to non-deterministic segfaults
in Reactant's CanonicalizeLoopsPass (MLIR bug). Comment out TripolarGrid
halo tests (Zipper BC compilation crashes) and HydrostaticFreeSurfaceModel
time-stepping tests (MLIR codegen error: incorrect operand count). The
remaining 102 tests (halo filling + Gu tendencies) pass reliably.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fix CI to actually run Reactant correctness tests

TEST_GROUP was set to "reactant" which matches no group in runtests.jl,
so CI was running 0 Reactant tests and reporting success.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Attempt fix

* Fix fix

* Whitespace

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Delete test/fake_halo_fill_kernels.jl

* Delete test/test_fake_halo_fill.jl

* Delete reactant_raise_true_report.md

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Update runtests.jl

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Update test_reactant_correctness.jl

* Apply suggestion from @glwagner

* Use halo=(3,3,3) in Reactant correctness tests and skip initialization_update_state\! for Reactant models

- Increase halo size from (1,1,1) to (3,3,3) in fill_halo_regions\! tests
- Re-enable Bounded LLG topology test
- Add no-op initialization_update_state\! for ReactantHFSM to defer
  kernel execution to first_time_step\! (runs inside @compile context)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Revert fill_halo tests to halo=(1,1,1) to avoid MLIR pass bug on Linux x64

halo=(3,3,3) with periodic boundaries triggers a Reactant MLIR
optimization pass failure on Linux x64. Keep halo=(1,1,1) for
fill_halo tests (which pass) while compute_simple_Gu\! retains
halo=(3,3,3) as required by WENO.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add MWE for Reactant MLIR pass failure on periodic halo kernels

Standalone reproducer (no Oceananigans dependency) for the MLIR
optimization pass bug on Linux x64 with halo size H=3.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Simplify MWE to single kernel

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Update MWE to reproduce ka_with_reactant error without CUDA

Remove `using CUDA` so the MWE reproduces the same MethodError
as Oceananigans: ka_with_reactant has no method without ReactantCUDAExt.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Comment out RectilinearGrid Gu tests with Periodic topologies

Reactant's MLIR raise pass fails on periodic halo-filling kernels with
halo >= 2 ("cannot raise if yet" error). Add standalone MWE reproducer.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Delete test/reactant_raise_periodic_halo_mwe.jl

* Update MWE to reproduce compute_simple_Gu\! CI test failure

Minimal script using Oceananigans that reproduces the fill_halo_regions\!
raise=true failure on periodic topologies with halo=(3,3,3).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fix MWE imports and usage instructions

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Use view instead of indexing in cpu_face_constructor_r (#5376)

* Use view instead of indexing in cpu_face_constructor_r

Avoids scalar indexing errors when grid data is stored on
non-CPU architectures (e.g. Reactant ConcreteIFRTArray).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add on_architecture(::CPU, ::OffsetArray{..., <:AnyConcreteReactantArray})

The generic fallback on_architecture(arch, a) = a doesn't convert
OffsetArrays wrapping Reactant arrays. This adds a proper method
that converts the underlying Reactant array to Array while preserving
the OffsetArray wrapper.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>

* Delete test/reactant_raise_periodic_halo_mwe.jl

* Add maybe_initialize_state\! no-op for Reactant models

The iteration == 0 check in maybe_initialize_state\! evaluates at trace
time (ConcreteRNumber(0) == 0 is true), causing a redundant update_state\!
to be compiled into every time_step\!. Make it a no-op for Reactant models
since first_time_step\! handles initialization explicitly.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Apply suggestion from @glwagner

* Replace bcs... splatting with explicit indexing in fill_halo_event\!

Extends the fix from 0807efe to polar_boundary_condition.jl and
fill_halo_regions_open.jl. Reactant cannot reconcile different call
arity from splatting variable-length tuples into the same kernel.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Pass clock and fields to only_local_halos fill_halo_regions\! calls

Avoids empty args tuple in fill_halo_event\!, which causes Reactant
MLIR compilation failures (empty tuples get eliminated at call sites
but not in function definitions, causing operand count mismatch).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add nothing to tendency kernels for Reactant compatibility

KA kernels with implicit Float64 return values (from array assignment
as last expression) cause Reactant to generate a Returns{Float64}
parameter in the LLVM function definition that doesn't match the call
site. Adding `nothing` as the last expression prevents this.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add kernel_clock workaround for Reactant MLIR/LLVM lowering bug

Reactant/Enzyme generates a phantom pointer parameter when mutable structs
containing ConcreteIFRTNumber scalar fields are placed in tuples passed to
KA kernels. This causes 'llvm.call incorrect number of operands' errors.

The workaround converts Clock to a NamedTuple (via kernel_clock) before
placing it in kernel argument tuples. The NamedTuple has identical field
access semantics and doesn't trigger the LLVM bug.

- kernel_clock(clock) defaults to identity (no-op on CPU/GPU)
- Reactant extension overrides to return NamedTuple
- Applied to all HFSM call sites: tendency kernels, fill_halo_regions\!,
  implicit_step\!, flux BCs, split-explicit substepping

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* revert

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Test more things

* is it about the halo size?

* run Reactant tests on 1.11.9

* put halo back to 3

* simplify SubArray convertion to arch

* emit something silly for prettytime of tracednumber

* try to run reactant tests with checkbounds=auto

* fix dispatch

* resolve multi region ambiguity

* Apply suggestion from @glwagner

* multiregion fill halo fix

---------

Co-authored-by: William Moses <[email protected]>
Co-authored-by: dkytezab <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
loisbaker pushed a commit to loisbaker/Oceananigans.jl that referenced this pull request Mar 10, 2026
…A#5093)

* Add basic reactant correctness tests

* comment out other tests

* add more grids

* implement take Gu tendency

* test with different raise options

* test raise=true and raise=false

* raise=true

* add fake halo filling kernels

* filter triply periodic tests

* add reactant to weakdeps

* add more tests

* updates

* clean up

* update report

* fix

* fix

* update report

* Update test_reactant_correctness.jl

* Refactor new_data function to remove ReactantState

* Change total_size function parameter type to ShardedGrid

* try to restrict offset_array for sharding a bit more

* test raise_first too

* fix imports

* Comment out failing Reactant correctness tests and remove raise=true mode

Disable raise=true mode across all tests due to non-deterministic segfaults
in Reactant's CanonicalizeLoopsPass (MLIR bug). Comment out TripolarGrid
halo tests (Zipper BC compilation crashes) and HydrostaticFreeSurfaceModel
time-stepping tests (MLIR codegen error: incorrect operand count). The
remaining 102 tests (halo filling + Gu tendencies) pass reliably.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fix CI to actually run Reactant correctness tests

TEST_GROUP was set to "reactant" which matches no group in runtests.jl,
so CI was running 0 Reactant tests and reporting success.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Attempt fix

* Fix fix

* Whitespace

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Delete test/fake_halo_fill_kernels.jl

* Delete test/test_fake_halo_fill.jl

* Delete reactant_raise_true_report.md

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Update runtests.jl

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Update test_reactant_correctness.jl

* Apply suggestion from @glwagner

* Use halo=(3,3,3) in Reactant correctness tests and skip initialization_update_state\! for Reactant models

- Increase halo size from (1,1,1) to (3,3,3) in fill_halo_regions\! tests
- Re-enable Bounded LLG topology test
- Add no-op initialization_update_state\! for ReactantHFSM to defer
  kernel execution to first_time_step\! (runs inside @compile context)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Revert fill_halo tests to halo=(1,1,1) to avoid MLIR pass bug on Linux x64

halo=(3,3,3) with periodic boundaries triggers a Reactant MLIR
optimization pass failure on Linux x64. Keep halo=(1,1,1) for
fill_halo tests (which pass) while compute_simple_Gu\! retains
halo=(3,3,3) as required by WENO.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add MWE for Reactant MLIR pass failure on periodic halo kernels

Standalone reproducer (no Oceananigans dependency) for the MLIR
optimization pass bug on Linux x64 with halo size H=3.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Simplify MWE to single kernel

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Update MWE to reproduce ka_with_reactant error without CUDA

Remove `using CUDA` so the MWE reproduces the same MethodError
as Oceananigans: ka_with_reactant has no method without ReactantCUDAExt.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Comment out RectilinearGrid Gu tests with Periodic topologies

Reactant's MLIR raise pass fails on periodic halo-filling kernels with
halo >= 2 ("cannot raise if yet" error). Add standalone MWE reproducer.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Delete test/reactant_raise_periodic_halo_mwe.jl

* Update MWE to reproduce compute_simple_Gu\! CI test failure

Minimal script using Oceananigans that reproduces the fill_halo_regions\!
raise=true failure on periodic topologies with halo=(3,3,3).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Fix MWE imports and usage instructions

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Use view instead of indexing in cpu_face_constructor_r (CliMA#5376)

* Use view instead of indexing in cpu_face_constructor_r

Avoids scalar indexing errors when grid data is stored on
non-CPU architectures (e.g. Reactant ConcreteIFRTArray).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add on_architecture(::CPU, ::OffsetArray{..., <:AnyConcreteReactantArray})

The generic fallback on_architecture(arch, a) = a doesn't convert
OffsetArrays wrapping Reactant arrays. This adds a proper method
that converts the underlying Reactant array to Array while preserving
the OffsetArray wrapper.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>

* Delete test/reactant_raise_periodic_halo_mwe.jl

* Add maybe_initialize_state\! no-op for Reactant models

The iteration == 0 check in maybe_initialize_state\! evaluates at trace
time (ConcreteRNumber(0) == 0 is true), causing a redundant update_state\!
to be compiled into every time_step\!. Make it a no-op for Reactant models
since first_time_step\! handles initialization explicitly.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Apply suggestion from @glwagner

* Replace bcs... splatting with explicit indexing in fill_halo_event\!

Extends the fix from 0807efe to polar_boundary_condition.jl and
fill_halo_regions_open.jl. Reactant cannot reconcile different call
arity from splatting variable-length tuples into the same kernel.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Pass clock and fields to only_local_halos fill_halo_regions\! calls

Avoids empty args tuple in fill_halo_event\!, which causes Reactant
MLIR compilation failures (empty tuples get eliminated at call sites
but not in function definitions, causing operand count mismatch).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add nothing to tendency kernels for Reactant compatibility

KA kernels with implicit Float64 return values (from array assignment
as last expression) cause Reactant to generate a Returns{Float64}
parameter in the LLVM function definition that doesn't match the call
site. Adding `nothing` as the last expression prevents this.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* Add kernel_clock workaround for Reactant MLIR/LLVM lowering bug

Reactant/Enzyme generates a phantom pointer parameter when mutable structs
containing ConcreteIFRTNumber scalar fields are placed in tuples passed to
KA kernels. This causes 'llvm.call incorrect number of operands' errors.

The workaround converts Clock to a NamedTuple (via kernel_clock) before
placing it in kernel argument tuples. The NamedTuple has identical field
access semantics and doesn't trigger the LLVM bug.

- kernel_clock(clock) defaults to identity (no-op on CPU/GPU)
- Reactant extension overrides to return NamedTuple
- Applied to all HFSM call sites: tendency kernels, fill_halo_regions\!,
  implicit_step\!, flux BCs, split-explicit substepping

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* revert

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Apply suggestion from @glwagner

* Test more things

* is it about the halo size?

* run Reactant tests on 1.11.9

* put halo back to 3

* simplify SubArray convertion to arch

* emit something silly for prettytime of tracednumber

* try to run reactant tests with checkbounds=auto

* fix dispatch

* resolve multi region ambiguity

* Apply suggestion from @glwagner

* multiregion fill halo fix

---------

Co-authored-by: William Moses <[email protected]>
Co-authored-by: dkytezab <[email protected]>
Co-authored-by: Claude Opus 4.6 <[email protected]>
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