EnzymeAD / EnzymeAD/Reactant.jl

Scalar indexing during reduction

Open
#1,146 20 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.