applying map to each index results in Scalar Indexing is Disallowed Error

[ctessum]

Hello!

I have some code that uses a pattern like this:

using Reactant

input1 = Reactant.ConcreteRArray(rand(10))
input2 = Reactant.ConcreteRArray(rand(10))

function fff(i, x, y)
    c = max(x[i], y[i])
    x[i] + y[i] + c
end
function mctr(f, range, x, y)
    f2(i) = f(i, x, y)
    map(f2, range)
end

f = @compile mctr(fff, 1:10, input1, input2)

The code above gives the error "Scalar Indexing is Disallowed". However, this type of pattern works fine when using AcceleratedKernels.jl. Is this something that Reactant would at some point be expected to support? Apologies if this is a dumb question!

[avik-pal]

We should be able to support this. In principle we need to redirect map on abstractarrays via our broadcast implementation.

[ctessum]

Great!

FWIW broadcasting gives a different error:

using Reactant

input1 = Reactant.ConcreteRArray(rand(10))
input2 = Reactant.ConcreteRArray(rand(10))

function fff(i, x, y)
    c = max(x[i], y[i])
    x[i] + y[i] + c
end
function mctr(f, range, x, y)
    f2(i) = f(i, x, y)
    f2.(range)
end

mctr(fff, 1:10, rand(10), rand(10)) # This works.
f = @compile mctr(fff, 1:10, input1, input2) #doesn't work

Error:

ERROR: AssertionError: Could not recursively make tracer of object of type var"#f2#13"{typeof(fff), Reactant.TracedRArray{Float64, 1}, Reactant.TracedRArray{Float64, 1}} into var"#f2#13"{typeof(fff), Reactant.TracedRArray{Float64, 1}, Reactant.TracedRArray{Float64, 1}} at field 2 (named x), need object of type Reactant.TracedRArray{Float64, 1} found object of type Reactant.TracedRNumber{Float64} 
Stacktrace:
  [1] make_tracer_unknown(seen::Reactant.OrderedIdDict{…}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/Tracing.jl:1111
  [2] make_tracer_unknown
    @ ~/.julia/packages/Reactant/EsmaI/src/Tracing.jl:1009 [inlined]
  [3] #make_tracer#131
    @ ~/.julia/packages/Reactant/EsmaI/src/Tracing.jl:1146 [inlined]
  [4] make_tracer
    @ ~/.julia/packages/Reactant/EsmaI/src/Tracing.jl:1136 [inlined]
  [5] prepare_mlir_fn_args(args::Tuple{…}, name::String, concretein::Bool, toscalar::Bool, argprefix::Symbol, runtime::Nothing, optimize_then_pad::Bool, do_transpose::Bool, input_shardings::Nothing, verify_arg_names::Nothing)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:436
  [6] make_mlir_fn(f::typeof(Reactant.apply), 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::Nothing, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:308
  [7] make_mlir_fn
    @ ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:262 [inlined]
  [8] make_mlir_fn(f::var"#f2#13"{…}, 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::Nothing, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:284
  [9] elem_apply(f::Function, args::Reactant.TracedRArray{Int64, 1})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:1084
 [10] _copyto!
    @ ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:789 [inlined]
 [11] (::Nothing)(none::typeof(Reactant.TracedRArrayOverrides._copyto!), none::Vector{…}, none::Base.Broadcast.Broadcasted{…})
    @ Reactant ./<missing>:0
 [12] size
    @ ./array.jl:194 [inlined]
 [13] axes
    @ ./abstractarray.jl:98 [inlined]
 [14] _copyto!
    @ ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:782 [inlined]
 [15] call_with_reactant(::typeof(Reactant.TracedRArrayOverrides._copyto!), ::Vector{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [16] copyto!
    @ ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:727 [inlined]
 [17] copyto!
    @ ./broadcast.jl:925 [inlined]
 [18] copy
    @ ./broadcast.jl:897 [inlined]
 [19] materialize
    @ ./broadcast.jl:872 [inlined]
 [20] mctr
    @ .../test.jl:12 [inlined]
 [21] (::Nothing)(none::typeof(mctr), none::typeof(fff), none::UnitRange{…}, none::Reactant.TracedRArray{…}, none::Reactant.TracedRArray{…})
    @ Reactant ./<missing>:0
 [22] mctr
    @ ../test.jl:11 [inlined]
 [23] call_with_reactant(::typeof(mctr), ::typeof(fff), ::UnitRange{…}, ::Reactant.TracedRArray{…}, ::Reactant.TracedRArray{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [24] make_mlir_fn(f::typeof(mctr), args::Tuple{…}, kwargs::@NamedTuple{}, 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/EsmaI/src/TracedUtils.jl:332
 [25] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:1555
 [26] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:1522 [inlined]
 [27] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:3433
 [28] compile_xla
    @ ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:3406 [inlined]
 [29] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:3505
 [30] top-level scope
    @ ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:2586

[ctessum]

Also, to back up a bit, the original reason for setting up the function like that was to avoid allocating arrays for intermediate variables (c in the example above).

However, is it true that using Reactant is another way to avoid those allocations?

using Reactant

input1 = Reactant.ConcreteRArray(rand(1_000_000))
input2 = Reactant.ConcreteRArray(rand(1_000_000))
output = Reactant.ConcreteRArray(zeros(1_000_000))

function ff(output::AbstractArray, x::AbstractArray, y::AbstractArray)
    c = max.(x, y)
    output .= x .+ y .+ c
end
f = @compile ff(output, input1, input2)
@allocated ff(output, input1, input2) # = 16057936
@allocated f(output, input1, input2) # = 608

From the code above, it looks like the reactant version is non-allocating, but also I guess there maybe could be allocations happening somewhere that is not tracked by julia?

[wsmoses]

Reactant already optimizes away intermediate variables (x/ref https://c.wsmoses.com/presentations/reactant.pdf), but yes data allocated by reactant is not tracked by Julia

[avik-pal]

#1620 partially resolves this. We still need to figure out why the broadcast with closure fails

[wsmoses]

the broadcast handler is relatively simple and tries to downgrade the inputs to scalars, but it failed to do so for the closure

[wsmoses]

longer term, I think one potential solution to general broadcasting is to have broadcast emit a normal loop in IR, and then we run the loop-based batching we talked about earlier today.

[avik-pal]

With the loop to tensor passes #1799

function looped_elementwise(x, y, range)
    z = similar(x, length(range))
    @trace for i in range
        c = @allowscalar max(x[i], y[i])
        @allowscalar z[i] = x[i] + y[i] + c
    end
    return z
end

x = Reactant.to_rarray(rand(Float32, 10))
y = Reactant.to_rarray(rand(Float32, 10))
range = 1:10

@code_hlo looped_elementwise(x, y, range)

module @reactant_looped_... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<10xf32> {enzymexla.memory_effects = []}, %arg1: tensor<10xf32> {enzymexla.memory_effects = []}) -> tensor<10xf32> attributes {enzymexla.memory_effects = []} {
    %0 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
    %1 = stablehlo.maximum %arg0, %arg1 : tensor<10xf32>
    %2 = stablehlo.add %0, %1 : tensor<10xf32>
    return %2 : tensor<10xf32>
  }
}

We should be able to enable fallback lowering to a loop and raise it now