should fold and similar operations be unrolled?

[avik-pal]

(needs LuxDL/Lux.jl#1026 for the dump)

using Lux, Reactant, Random

model = Recurrence(RNNCell(4 => 4))
ps, st = Lux.setup(Xoshiro(123), model) |> Reactant.to_rarray
x = rand(Float32, 4, 2, 12) |> Reactant.ConcreteRArray

@code_hlo model(x, ps, st)

This currently gets unrolled into the following monstrosity: https://pastebin.com/QE00SCqb

[avik-pal]

I will create a MWE later but opening this so that I don't forget. Recurrence internally just calls foldl, which we should be able to @reactant_override?

[avik-pal]

using Reactant

function custom_op(x::AbstractArray)
    function inner_op(xᵢ::AbstractVector, yᵢ::AbstractVector)
        return xᵢ .+ yᵢ
    end
    return foldl(inner_op, eachcol(x))
end

x_ra = Reactant.to_rarray(rand(2, 4))

@code_hlo custom_op(x_ra)

Module:
module attributes {transform.with_named_sequence} {
  func.func @main(%arg0: tensor<4x2xf64>) -> tensor<2xf64> {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x2xf64>) -> tensor<2x4xf64>
    %1 = stablehlo.slice %0 [0:2, 0:1] : (tensor<2x4xf64>) -> tensor<2x1xf64>
    %2 = stablehlo.reshape %1 : (tensor<2x1xf64>) -> tensor<2xf64>
    %3 = stablehlo.slice %0 [0:2, 1:2] : (tensor<2x4xf64>) -> tensor<2x1xf64>
    %4 = stablehlo.reshape %3 : (tensor<2x1xf64>) -> tensor<2xf64>
    %5 = stablehlo.add %2, %4 : tensor<2xf64>
    %6 = stablehlo.slice %0 [0:2, 2:3] : (tensor<2x4xf64>) -> tensor<2x1xf64>
    %7 = stablehlo.reshape %6 : (tensor<2x1xf64>) -> tensor<2xf64>
    %8 = stablehlo.add %5, %7 : tensor<2xf64>
    %9 = stablehlo.slice %0 [0:2, 3:4] : (tensor<2x4xf64>) -> tensor<2x1xf64>
    %10 = stablehlo.reshape %9 : (tensor<2x1xf64>) -> tensor<2xf64>
    %11 = stablehlo.add %8, %10 : tensor<2xf64>
    return %11 : tensor<2xf64>
  }
}

[wsmoses]

Yeah we should probably lower this directly as a reduction of sorts