EnzymeAD / EnzymeAD/Reactant.jl
Scalar indexing during reduction
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
This code errors with Scalar indexing error:
```julia
using Oceananigans, Reactant
grid = RectilinearGrid(Oceananigans.Architectures.ReactantState(), size=(2, 2, 2), extent=(1, 1, 1))
c = CenterField(grid)
maximum(c)
```
but should work with `GPU()` instead of `ReactantState()` as the architecture.
From the stacktrace is appears the code does enter Reactant, (eg calling `mymapreducedim!`), but nevertheless ends up on an invalid codepath.
Note that Oceananigans makes extensive use of patterns that looks something like
```julia
sum!(interior(r), operation)
```
where `interior(r)` is a `SubArray`, ie something like
```julia
julia> typeof(interior(c))
SubArray{Float64, 3, ConcreteIFRTArray{Float64, 3, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}, Tuple{Int64, Int64, Int64}}, Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}, false}
```
The basic design is that we expect `SubArray{T, N, A}` to dispatch correctly for array type `A`. Then we can build our `Field` / array wrapper, which often include "halo points" that we don't want to touch during a computation, but call into external functionality using the the SubArray trick.
We can change Oceananigans design too but this is the current state of affairs. It's not of utmost priority but does impact user experience,
because users typically want to compute reductions while a simulation is running in order to monitor the status of things. Error below:
```
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.
If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:35
[2] errorscalar(op::String)
@ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
[3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
@ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
[4] assertscalar(op::String)
@ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
[5] getindex(::Reactant.TracedRArray{Float64, 3}, ::Int64, ::Int64, ::Int64)
@ Reactant.TracedRArrayOverrides ~/.julia/packages/Reactant/UYyVU/src/TracedRArray.jl:62
[6] getindex
@ ~/.julia/packages/Reactant/UYyVU/src/TracedRArray.jl:213 [inlined]
[7] getindex
@ ~/.julia/packages/Reactant/UYyVU/ext/ReactantOffsetArraysExt.jl:52 [inlined]
[8] _getindex
@ ./abstractarray.jl:1358 [inlined]
[9] getindex
@ ./abstractarray.jl:1312 [inlined]
[10] getindex
@ ~/Projects/Oceananigans.jl/src/Fields/field.jl:405 [inlined]
[11] iterate
@ ./abstractarray.jl:1209 [inlined]
[12] iterate
@ ./abstractarray.jl:1207 [inlined]
[13] _foldl_impl(op::Base.BottomRF{…}, init::Base._InitialValue, itr::Field{…})
@ Base ./reduce.jl:56
[14] foldl_impl
@ ./reduce.jl:48 [inlined]
[15] mapfoldl_impl(f::typeof(identity), op::typeof(max), nt::Base._InitialValue, itr::Field{…})
@ Base ./reduce.jl:44
[16] mapfoldl(f::Function, op::Function, itr::Field{…}; init::Base._InitialValue)
@ Base ./reduce.jl:175
[17] mapfoldl
@ ./reduce.jl:175 [inlined]
[18] _mapreduce
@ ./reduce.jl:450 [inlined]
[19] _mapreduce_dim
@ ./reducedim.jl:337 [inlined]
[20] mapreduce(f::Function, op::Function, A::Field{…})
@ Base ./reducedim.jl:329
[21] reducedim_init(f::Function, op::typeof(min), A::Field{…}, region::Vector{…})
@ Base ./reducedim.jl:106
[22] _mapreduce_dim(f::Function, op::Function, ::Base._InitialValue, A::Field{…}, dims::Vector{…})
@ Base ./reducedim.jl:343
[23] #mapreduce#926
@ ./reducedim.jl:329 [inlined]
[24] mapreduce(f::Function, op::Function, A::Field{…}; kwargs::@Kwargs{…})
@ Reactant ~/.julia/packages/Reactant/UYyVU/src/Overlay.jl:170
[25] #kwcall
@ ~/.julia/packages/Reactant/UYyVU/src/Overlay.jl:164 [inlined]
[26] kwcall(none::@NamedTuple{…}, none::typeof(mapreduce), none::typeof(identity), none::typeof(min), none::Field{…})
@ Reactant ./:0
[27] Pairs
@ ./essentials.jl:483 [inlined]
[28] pairs
@ ./iterators.jl:279 [inlined]
[29] #kwcall
@ ~/.julia/packages/Reactant/UYyVU/src/Overlay.jl:164 [inlined]
[30] call_with_reactant(::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(mapreduce), ::typeof(identity), ::typeof(min), ::Field{…})
@ Reactant ~/.julia/packages/Reactant/UYyVU/src/utils.jl:0
[31] mapreducedim!
@ ~/.julia/packages/Reactant/UYyVU/src/TracedRArray.jl:570 [inlined]
[32] mapreducedim!(none::Any, none::Any, none::SubArray{…}, none::Field{…})
@ Reactant ./:0
[33] call_with_reactant(::typeof(Base.mapreducedim!), ::Any, ::Any, ::SubArray{…}, ::Field{…})
@ Reactant ~/.julia/packages/Reactant/UYyVU/src/utils.jl:790
[34] mymapreducedim!
@ ~/.julia/packages/Reactant/UYyVU/src/ConcreteRArray.jl:575 [inlined]
[35] mymapreducedim!(none::typeof(identity), none::typeof(min), none::SubArray{…}, none::Field{…})
@ Reactant ./:0
[36] mymapreducedim!
@ ~/.julia/packages/Reactant/UYyVU/src/ConcreteRArray.jl:575 [inlined]
[37] call_with_reactant(::typeof(Reactant.mymapreducedim!), ::typeof(identity), ::typeof(min), ::SubArray{…}, ::Field{…})
@ Reactant ~/.julia/packages/Reactant/UYyVU/src/utils.jl:0
[38] make_mlir_fn(f::typeof(Reactant.mymapreducedim!), args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
@ Reactant.TracedUtils ~/.julia/packages/Reactant/UYyVU/src/TracedUtils.jl:370
[39] make_mlir_fn
@ ~/.julia/packages/Reactant/UYyVU/src/TracedUtils.jl:198 [inlined]
[40] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(Reactant.mymapreducedim!), args::Tuple{…}, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::Tuple{}, optimize::Bool, shardy_passes::Symbol, no_nan::Bool, transpose_propagate::Symbol, reshape_propagate::Symbol, assert_nonallocating::Bool, backend::String, raise::Bool, donated_args::Symbol, optimize_then_pad::Bool, runtime::Val{…}, kwargs::@Kwargs{})
@ Reactant.Compiler ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:995
[41] compile_mlir!
@ ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:952 [inlined]
[42] compile_xla(f::Function, args::Tuple{…}; client::Nothing, serializable::Bool, kwargs::@Kwargs{})
@ Reactant.Compiler ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:2610
[43] compile_xla
@ ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:2592 [inlined]
[44] compile(f::Function, args::Tuple{typeof(identity), typeof(min), SubArray{…}, Field{…}}; sync::Bool, kwargs::@Kwargs{})
@ Reactant.Compiler ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:2663
[45] compile(f::Function, args::Tuple{typeof(identity), typeof(min), SubArray{…}, Field{…}})
@ Reactant.Compiler ~/.julia/packages/Reactant/UYyVU/src/Compiler.jl:2662
[46] mapreducedim!(f::Function, op::Function, R::SubArray{…}, A::Field{…})
@ Reactant ~/.julia/packages/Reactant/UYyVU/src/ConcreteRArray.jl:585
[47] #minimum!#969
@ ./reducedim.jl:1006 [inlined]
[48] minimum!(f::Function, r::Field{…}, a::Field{…}; condition::Nothing, mask::Float64, kwargs::@Kwargs{…})
@ Oceananigans.Fields ~/Projects/Oceananigans.jl/src/Fields/field.jl:684
[49] minimum(f::Function, c::Field{…}; condition::Nothing, mask::Float64, dims::Function)
@ Oceananigans.Fields ~/Projects/Oceananigans.jl/src/Fields/field.jl:714
[50] minimum
@ ~/Projects/Oceananigans.jl/src/Fields/field.jl:703 [inlined]
[51] minimum(c::Field{…})
@ Oceananigans.Fields ~/Projects/Oceananigans.jl/src/Fields/field.jl:723
Some type information was truncated. Use `show(err)` to see complete types.
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.