Scalar indexing triggered when requesting a view of a traced array

[ptiede]

It looks like Reactant is missing a method when I index an array with an array of indices. @avik-pal

using Reactant
fr = Reactant.to_rarray(rand(4, 4))
idxr= Reactant.to_rarray([1,2])
test(fr, idxr) = fr[idxr]
@compile test(fr, idxr)


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
    @ ./error.jl:35 [inlined]
  [2] error(s::String)
    @ Reactant ./<missing>:0
  [3] ErrorException
    @ ./boot.jl:282 [inlined]
  [4] error
    @ ./error.jl:35 [inlined]
  [5] call_with_reactant(::typeof(error), ::String)
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
  [6] errorscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151 [inlined]
  [7] errorscalar(op::String)
    @ Reactant ./<missing>:0
  [8] scalardesc
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:134 [inlined]
  [9] errorscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:150 [inlined]
 [10] call_with_reactant(::typeof(GPUArraysCore.errorscalar), ::String)
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
 [11] _assertscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124 [inlined]
 [12] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ Reactant ./<missing>:0
 [13] _assertscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:123 [inlined]
 [14] call_with_reactant(::typeof(GPUArraysCore._assertscalar), ::String, ::GPUArraysCore.ScalarIndexing)
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
 [15] assertscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112 [inlined]
 [16] assertscalar(op::String)
    @ Reactant ./<missing>:0
 [17] current_task
    @ ./task.jl:150 [inlined]
 [18] task_local_storage
    @ ./task.jl:269 [inlined]
 [19] assertscalar
    @ ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:97 [inlined]
 [20] call_with_reactant(::typeof(GPUArraysCore.assertscalar), ::String)
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
 [21] getindex
    @ ~/.julia/dev/Reactant/src/TracedRArray.jl:41 [inlined]
 [22] getindex(a::Reactant.TracedRArray{Int64, 1}, index::Tuple{Int64})
    @ Reactant ./<missing>:0
 [23] getindex
    @ ~/.julia/dev/Reactant/src/TracedRArray.jl:41 [inlined]
 [24] call_with_reactant(::typeof(getindex), ::Reactant.TracedRArray{Int64, 1}, ::Int64)
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
 [25] iterate
    @ ./abstractarray.jl:1216 [inlined]
 [26] iterate
    @ ./abstractarray.jl:1214 [inlined]
 [27] checkindex
    @ ./abstractarray.jl:774 [inlined]
 [28] checkbounds
    @ ./abstractarray.jl:687 [inlined]
 [29] checkbounds
    @ ./abstractarray.jl:702 [inlined]
 [30] _getindex
    @ ./multidimensional.jl:888 [inlined]
 [31] getindex
    @ ./abstractarray.jl:1290 [inlined]
 [32] test
    @ ~/Research/ReactantVLBI/imager.jl:39 [inlined]
 [33] test(fr::Reactant.TracedRArray{Float64, 2}, idxr::Reactant.TracedRArray{Int64, 1})
    @ Reactant ./<missing>:0
 [34] _getindex
    @ ./multidimensional.jl:888 [inlined]
 [35] getindex
    @ ./abstractarray.jl:1290 [inlined]
 [36] test
    @ ~/Research/ReactantVLBI/imager.jl:39 [inlined]
 [37] call_with_reactant(::typeof(test), ::Reactant.TracedRArray{Float64, 2}, ::Reactant.TracedRArray{Int64, 1})
    @ Reactant ~/.julia/dev/Reactant/src/utils.jl:0
 [38] (::Reactant.TracedUtils.var"#8#18"{Bool, Bool, typeof(test), Tuple{…}, Vector{…}, Tuple{…}})()
    @ Reactant.TracedUtils ~/.julia/dev/Reactant/src/TracedUtils.jl:178
 [39] block!(f::Reactant.TracedUtils.var"#8#18"{…}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR ~/.julia/dev/Reactant/src/mlir/IR/Block.jl:201
 [40] make_mlir_fn(f::Function, args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, no_args_in_result::Bool, construct_function_without_args::Bool, do_transpose::Bool)
    @ Reactant.TracedUtils ~/.julia/dev/Reactant/src/TracedUtils.jl:165
 [41] make_mlir_fn
    @ ~/.julia/dev/Reactant/src/TracedUtils.jl:85 [inlined]
 [42] #10
    @ ~/.julia/dev/Reactant/src/Compiler.jl:298 [inlined]
 [43] block!(f::Reactant.Compiler.var"#10#15"{typeof(test), Tuple{…}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR ~/.julia/dev/Reactant/src/mlir/IR/Block.jl:201
 [44] #9
    @ ~/.julia/dev/Reactant/src/Compiler.jl:297 [inlined]
 [45] mmodule!(f::Reactant.Compiler.var"#9#14"{Reactant.MLIR.IR.Module, typeof(test), Tuple{…}}, blk::Reactant.MLIR.IR.Module)
    @ Reactant.MLIR.IR ~/.julia/dev/Reactant/src/mlir/IR/Module.jl:92
 [46] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}; optimize::Bool)
    @ Reactant.Compiler ~/.julia/dev/Reactant/src/Compiler.jl:294
 [47] compile_mlir!
    @ ~/.julia/dev/Reactant/src/Compiler.jl:293 [inlined]
 [48] (::Reactant.Compiler.var"#34#36"{Bool, typeof(test), Tuple{ConcreteRArray{Float64, 2}, ConcreteRArray{Int64, 1}}})()
    @ Reactant.Compiler ~/.julia/dev/Reactant/src/Compiler.jl:701
 [49] context!(f::Reactant.Compiler.var"#34#36"{Bool, typeof(test), Tuple{…}}, ctx::Reactant.MLIR.IR.Context)
    @ Reactant.MLIR.IR ~/.julia/dev/Reactant/src/mlir/IR/Context.jl:76
 [50] compile_xla(f::Function, args::Tuple{ConcreteRArray{…}, ConcreteRArray{…}}; client::Nothing, optimize::Bool)
    @ Reactant.Compiler ~/.julia/dev/Reactant/src/Compiler.jl:698
 [51] compile_xla
    @ ~/.julia/dev/Reactant/src/Compiler.jl:693 [inlined]
 [52] compile(f::Function, args::Tuple{ConcreteRArray{…}, ConcreteRArray{…}}; client::Nothing, optimize::Bool, sync::Bool)
    @ Reactant.Compiler ~/.julia/dev/Reactant/src/Compiler.jl:725
Some type information was truncated. Use `show(err)` to see complete types.

[avik-pal]

Incidentally I was looking at a similar issue the other day. While we should definitely fix this, if you are in need of an immediate workaround doing an allowscalar here would be fine. For the tests I ran, our opt passes were able to generate a single dynamic slice