Broadcasting Operations: invalid MLIR generated

[glou-nes]

I don't know if that kind of invalid generation is interesting to fix.

using Reactant
Reactant.set_default_backend("cpu")
x = Reactant.ConcreteRArray(randn(Float64, 2, 2))
@code_hlo x .- x .+ x

module {
  func.func @main(%arg0: tensor<2x2xf64>, %arg1: tensor<2x2xf64>) -> tensor<2x2xf64> {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
    %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
    %2 = stablehlo.add %0, %1 : tensor<2x2xf64>
    %3 = stablehlo.transpose %2, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
    return %3 : tensor<2x2xf64>
  }
}

[avik-pal]

This is more of a parsing problem:

julia> @macroexpand @code_hlo x .- x .+ x
quote
    begin
        #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:473 =#
        var"##f#250" = begin
                #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:455 =#
                if isdefined(mod, :.+)
                    #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:456 =#
                    .+
                else
                    #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:458 =#
                    Base.Broadcast.BroadcastFunction(+)
                end
            end
        #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:474 =#
        var"##args#251" = (x .- x, x)
        #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:475 =#
        var"##compiled#252" = (Reactant.Compiler.compile_mlir)(var"##f#250", var"##args#251"; optimize = true)
    end
    #= /mnt/software/lux/Reactant.jl/src/Compiler.jl:404 =#
    (first)(var"##compiled#252")
end

julia> Meta.@dump x .- x .+ x
Expr
  head: Symbol call
  args: Array{Any}((3,))
    1: Symbol .+
    2: Expr
      head: Symbol call
      args: Array{Any}((3,))
        1: Symbol .-
        2: Symbol x
        3: Symbol x
    3: Symbol x

[glou-nes]

@avik-pal Thanks you! Should it be interesting to add a error case in the macros or insert closure to deal with complex expression such as this one?

[avik-pal]

Generating a warning might be a good start but I am unsure how one would distinguish between:

@compile .+(x_ra, <some operation>)
@compile x_ra .+ <some operation>

The former pattern (without broadcasting) is quite convenient to have IMO.

[mofeing]

they are both the same no?

julia> Meta.@dump .+(a,b)
Expr
  head: Symbol call
  args: Array{Any}((3,))
    1: Symbol .+
    2: Symbol a
    3: Symbol b

julia> Meta.@dump a .+ b
Expr
  head: Symbol call
  args: Array{Any}((3,))
    1: Symbol .+
    2: Symbol a
    3: Symbol b

[avik-pal]

Exactly that's why I think doing anything for the case in the top post is going to be hard without breaking current uses

[mofeing]

mmm it would be interesting to try the closure idea; i.e. take all the code, put it into a closure like () -> the_code, and take the closure parameters as the arguments.

[glou-nes]

Thank for the answers, my main problem was the discrepancy between @code_typed and @code_hlo for instance. I will try to get expressions which are not directly call with the closure approach!