Could not find unique name for Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar when compiling function with lots of matrix multiplications

[unit-of-inductance]

Consider the following MWE

using Reactant

n_points = 10000
dim = 10

points = Reactant.to_rarray(eachcol(rand(dim,n_points)))
matrix = Reactant.to_rarray(rand(1, dim))

function f(points)
  return sum(matrix * point for point in points)
end

f_compiled = @compile f(points)

This raises

Could not find unique name for Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar

Stracktrace

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] __lookup_unique_name_in_module(mod::Reactant.MLIR.IR.Module, name::String)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:998
  [3] (::Reactant.TracedUtils.var"#13#20"{…})()
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:842
  [4] block!(f::Reactant.TracedUtils.var"#13#20"{…}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR ~/.julia/packages/Reactant/EsmaI/src/mlir/IR/Block.jl:214
  [5] finalize_mlir_fn(result::Reactant.TracedRNumber{…}, traced_args::Vector{…}, linear_args::Vector{…}, skipped_args::Vector{…}, seen_args::Reactant.OrderedIdDict{…}, fnbody::Reactant.MLIR.IR.Block, func::Reactant.MLIR.IR.Operation, mod::Reactant.MLIR.IR.Module, name::String, in_tys::Vector{…}, do_transpose::Bool, optimize_then_pad::Bool, inv_map::IdDict{…}, args_in_result::Symbol, resprefix::Symbol, argprefix::Symbol, resargprefix::Symbol, verify_arg_names::Nothing, return_dialect::Symbol, traced_args_to_shardings::Reactant.OrderedIdDict{…}, output_shardings::Nothing, sym_visibility::Reactant.MLIR.IR.Attribute, num_replicas::Int64, runtime::Nothing, construct_function_without_args::Bool, args::Tuple{…}, N::Int64, concretein::Bool, toscalar::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:841
  [6] make_mlir_fn(f::Reactant.TracedUtils.TypeCast{…}, 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:352
  [7] elem_apply(f::Function, args::Reactant.TracedRArray{Float64, 2})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:1084
  [8] elem_apply(::Type{Float64}, x::Reactant.TracedRArray{Float64, 2})
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/EsmaI/src/TracedUtils.jl:1016
  [9] _copyto!(dest::Reactant.TracedRArray{Float64, 2}, bc::Base.Broadcast.Broadcasted{Nothing, Tuple{…}, Type{…}, Tuple{…}})
    @ Reactant.TracedRArrayOverrides ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:773
 [10] copyto!(dest::Reactant.TracedRArray{Float64, 2}, bc::Base.Broadcast.Broadcasted{Nothing, Tuple{…}, Type{…}, Tuple{…}})
    @ Reactant.TracedRArrayOverrides ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:727
 [11] copyto!
    @ ./broadcast.jl:925 [inlined]
 [12] copy(bc::Base.Broadcast.Broadcasted{…})
    @ Reactant.TracedRArrayOverrides ~/.julia/packages/Reactant/EsmaI/src/TracedRArray.jl:718
 [13] materialize(bc::Base.Broadcast.Broadcasted{…})
    @ Base.Broadcast ./broadcast.jl:872
 [14] overloaded_mul!(C::Reactant.TracedRArray{T, 2} where T, A::AbstractMatrix, B::AbstractMatrix, α::Bool, β::Bool)
    @ Reactant.TracedLinearAlgebra ~/.julia/packages/Reactant/EsmaI/src/stdlibs/LinearAlgebra.jl:258
 [15] overloaded_mul!(C::Reactant.TracedRArray{Float64, 1}, A::AbstractMatrix, B::AbstractVector, α::Bool, β::Bool)
    @ Reactant.TracedLinearAlgebra ~/.julia/packages/Reactant/EsmaI/src/stdlibs/LinearAlgebra.jl:220
 [16] #mul!
    @ ~/.julia/packages/Reactant/EsmaI/src/Overlay.jl:111 [inlined]
 [17] (::Nothing)(none::typeof(LinearAlgebra.mul!), none::Reactant.TracedRArray{…}, none::ConcretePJRTArray{…}, none::SubArray{…}, none::Bool, none::Bool)
    @ Reactant ./<missing>:0
 [18] call_with_reactant(::typeof(LinearAlgebra.mul!), ::Reactant.TracedRArray{…}, ::ConcretePJRTArray{…}, ::SubArray{…}, ::Bool, ::Bool)
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:501
 [19] mul!(C::Reactant.TracedRArray{…}, A::ConcretePJRTArray{…}, B::SubArray{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/Overlay.jl:126
 [20] *
    @ /nix/store/0f8r9l0fn26mjj96j08p92ak8gmbgwas-julia-bin-1.11.6/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:60 [inlined]
 [21] (::Nothing)(none::typeof(*), none::ConcretePJRTArray{…}, none::SubArray{…})
    @ Reactant ./<missing>:0
 [22] getproperty
    @ ~/.julia/packages/Reactant/EsmaI/src/Types.jl:11 [inlined]
 [23] size
    @ ~/.julia/packages/Reactant/EsmaI/src/ConcreteRArray.jl:92 [inlined]
 [24] axes
    @ ./abstractarray.jl:98 [inlined]
 [25] axes
    @ ./abstractarray.jl:77 [inlined]
 [26] *
    @ /nix/store/0f8r9l0fn26mjj96j08p92ak8gmbgwas-julia-bin-1.11.6/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:60 [inlined]
 [27] call_with_reactant(::typeof(*), ::ConcretePJRTArray{…}, ::SubArray{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [28] (::Nothing)(none::var"#28#29", none::SubArray{Reactant.TracedRNumber{…}, 1, Reactant.TracedRArray{…}, Tuple{…}, true})
    @ Reactant ./<missing>:0
 [29] #28
    @ ./Base.jl:0 [inlined]
 [30] call_with_reactant(::var"#28#29", ::SubArray{Reactant.TracedRNumber{…}, 1, Reactant.TracedRArray{…}, Tuple{…}, true})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [31] MappingRF
    @ ./reduce.jl:100 [inlined]
 [32] _foldl_impl
    @ ./reduce.jl:62 [inlined]
 [33] (::Nothing)(none::typeof(Base._foldl_impl), none::Base.MappingRF{…}, none::Base._InitialValue, none::ColumnSlices{…})
    @ Reactant ./<missing>:0
 [34] getproperty
    @ ./Base.jl:49 [inlined]
 [35] axes
    @ ./slicearray.jl:229 [inlined]
 [36] axes1
    @ ./abstractarray.jl:137 [inlined]
 [37] eachindex
    @ ./abstractarray.jl:321 [inlined]
 [38] iterate
    @ ./abstractarray.jl:1207 [inlined]
 [39] _foldl_impl
    @ ./reduce.jl:56 [inlined]
 [40] call_with_reactant(::typeof(Base._foldl_impl), ::Base.MappingRF{…}, ::Base._InitialValue, ::ColumnSlices{…})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [41] foldl_impl
    @ ./reduce.jl:48 [inlined]
 [42] mapfoldl_impl
    @ ./reduce.jl:44 [inlined]
 [43] mapfoldl
    @ ./reduce.jl:175 [inlined]
 [44] mapreduce
    @ ./reduce.jl:307 [inlined]
 [45] sum
    @ ./reduce.jl:532 [inlined]
 [46] sum
    @ ./reduce.jl:561 [inlined]
 [47] f
    @ ./REPL[11]:2 [inlined]
 [48] (::Nothing)(none::typeof(f), none::ColumnSlices{Reactant.TracedRArray{…}, Tuple{…}, SubArray{…}})
    @ Reactant ./<missing>:0
 [49] foldl_impl
    @ ./reduce.jl:48 [inlined]
 [50] mapfoldl_impl
    @ ./reduce.jl:44 [inlined]
 [51] mapfoldl
    @ ./reduce.jl:175 [inlined]
 [52] mapreduce
    @ ./reduce.jl:307 [inlined]
 [53] sum
    @ ./reduce.jl:532 [inlined]
 [54] sum
    @ ./reduce.jl:561 [inlined]
 [55] f
    @ ./REPL[11]:2 [inlined]
 [56] call_with_reactant(::typeof(f), ::ColumnSlices{Reactant.TracedRArray{…}, Tuple{…}, SubArray{…}})
    @ Reactant ~/.julia/packages/Reactant/EsmaI/src/utils.jl:0
 [57] make_mlir_fn(f::typeof(f), 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
 [58] 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
 [59] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:1522 [inlined]
 [60] 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
 [61] compile_xla
    @ ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:3406 [inlined]
 [62] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/EsmaI/src/Compiler.jl:3505
Some type information was truncated. Use `show(err)` to see complete types.

It also dumps some other info:

...
  "func.func"() <{function_type = (tensor<f64>, tensor<f64>) -> (tensor<f64>, tensor<f64>, tensor<f64>), sym_name = "+_broadcast_scalar_4997", sym_visibility = "private"}> ({
  ^bb0(%arg6: tensor<f64>, %arg7: tensor<f64>):
    %1 = "stablehlo.add"(%arg6, %arg7) : (tensor<f64>, tensor<f64>) -> tensor<f64>
    "func.return"(%1, %arg6, %arg7) : (tensor<f64>, tensor<f64>, tensor<f64>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>) -> tensor<f64>, sym_name = "Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar_9998", sym_visibility = "private"}> ({
  ^bb0(%arg5: tensor<f64>):
    "func.return"(%arg5) : (tensor<f64>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>) -> tensor<f64>, sym_name = "Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar_9999", sym_visibility = "private"}> ({
  ^bb0(%arg4: tensor<f64>):
    "func.return"(%arg4) : (tensor<f64>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>, tensor<f64>) -> (tensor<f64>, tensor<f64>, tensor<f64>), sym_name = "+_broadcast_scalar_4998", sym_visibility = "private"}> ({
  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
    %0 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
    "func.return"(%0, %arg2, %arg3) : (tensor<f64>, tensor<f64>, tensor<f64>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>) -> tensor<f64>, sym_name = "Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar_10000", sym_visibility = "private"}> ({
  ^bb0(%arg1: tensor<f64>):
    "func.return"(%arg1) : (tensor<f64>) -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<f64>) -> (), sym_name = "Reactant.TracedUtils.TypeCast{Float64}()_broadcast_scalar_tmp"}> ({
  ^bb0(%arg0: tensor<f64>):
    "func.return"(%arg0) : (tensor<f64>) -> ()
  }) : () -> ()
}) : () -> ()

Replicated in fresh environment with the following Pkg.status():

Status `/tmp/jl_qVYfKn/Project.toml`
  [3c362404] Reactant v0.2.160

[unit-of-inductance]

This is relevant for me when compiling loss functions for Lux models, where I'm using FiniteDifferences to compute the diagonal of the hessian up to relatively high order. Even for a single point this results in this error when applied for dim=10. I'd really prefer to be able to use Reactant here, since because Lux model applications are allocating, optimizing the loss function requires either multiple GB of memory or frequent manual GC invocations (the recommended way of avoiding allocations with Lux seems to be to compile with Reactant).

[avik-pal]

Reactant.jl/src/TracedUtils.jl

Line 993 in 71b744e

for i in 0:10000

did not expect anyone to exceed this limit 😓

[wsmoses]

I think that in part is because we really want @jumerckx function call PR

[avik-pal]

An aside issue but we might want to optimize this somehow

%1015 = stablehlo.slice %arg1 [9:10, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1016 = stablehlo.slice %arg1 [8:9, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1017 = stablehlo.slice %arg1 [7:8, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1018 = stablehlo.slice %arg1 [6:7, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1019 = stablehlo.slice %arg1 [5:6, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1020 = stablehlo.slice %arg1 [4:5, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1021 = stablehlo.slice %arg1 [3:4, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1022 = stablehlo.slice %arg1 [2:3, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1023 = stablehlo.slice %arg1 [1:2, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1024 = stablehlo.slice %arg1 [0:1, 0:10] : (tensor<1024x10xf64>) -> tensor<1x10xf64>
    %1025 = stablehlo.dot_general %0, %1024, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %1026 = stablehlo.dot_general %0, %1023, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %1027 = stablehlo.add %1025, %1026 : tensor<1x1xf64>
    %1028 = stablehlo.dot_general %0, %1022, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %1029 = stablehlo.add %1027, %1028 : tensor<1x1xf64>
    %1030 = stablehlo.dot_general %0, %1021, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %1031 = stablehlo.add %1029, %1030 : tensor<1x1xf64>

...

    %3038 = stablehlo.dot_general %0, %17, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3039 = stablehlo.add %3037, %3038 : tensor<1x1xf64>
    %3040 = stablehlo.dot_general %0, %16, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3041 = stablehlo.add %3039, %3040 : tensor<1x1xf64>
    %3042 = stablehlo.dot_general %0, %15, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3043 = stablehlo.add %3041, %3042 : tensor<1x1xf64>
    %3044 = stablehlo.dot_general %0, %14, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3045 = stablehlo.add %3043, %3044 : tensor<1x1xf64>
    %3046 = stablehlo.dot_general %0, %13, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3047 = stablehlo.add %3045, %3046 : tensor<1x1xf64>
    %3048 = stablehlo.dot_general %0, %12, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3049 = stablehlo.add %3047, %3048 : tensor<1x1xf64>
    %3050 = stablehlo.dot_general %0, %11, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3051 = stablehlo.add %3049, %3050 : tensor<1x1xf64>
    %3052 = stablehlo.dot_general %0, %10, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3053 = stablehlo.add %3051, %3052 : tensor<1x1xf64>
    %3054 = stablehlo.dot_general %0, %9, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3055 = stablehlo.add %3053, %3054 : tensor<1x1xf64>
    %3056 = stablehlo.dot_general %0, %8, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3057 = stablehlo.add %3055, %3056 : tensor<1x1xf64>
    %3058 = stablehlo.dot_general %0, %7, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3059 = stablehlo.add %3057, %3058 : tensor<1x1xf64>
    %3060 = stablehlo.dot_general %0, %6, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3061 = stablehlo.add %3059, %3060 : tensor<1x1xf64>
    %3062 = stablehlo.dot_general %0, %5, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3063 = stablehlo.add %3061, %3062 : tensor<1x1xf64>
    %3064 = stablehlo.dot_general %0, %4, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3065 = stablehlo.add %3063, %3064 : tensor<1x1xf64>
    %3066 = stablehlo.dot_general %0, %3, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3067 = stablehlo.add %3065, %3066 : tensor<1x1xf64>
    %3068 = stablehlo.dot_general %0, %2, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3069 = stablehlo.add %3067, %3068 : tensor<1x1xf64>
    %3070 = stablehlo.dot_general %0, %1, contracting_dims = [1] x [1], precision = [DEFAULT, DEFAULT] : (tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<1x1xf64>
    %3071 = stablehlo.add %3069, %3070 : tensor<1x1xf64>

[avik-pal]

This is actually the same issue where the equivalent ops don't end with a concat so we don't auto batch that.

Essentially we can do

1. tensor -> slice -> many dot_generals => tensor -> big dot_general -> slice
2. slice -> add => reduce

[jumerckx]

I think that in part is because we really want @jumerckx function call PR

I should look more closely but the most recent function call pr does not cache function definitions yet so not entirely sure.

[unit-of-inductance]

I saw that the seemingly related #1665 was merged, but rerunning my tests on 0.2.162 this issue here doesn't seem to have been fixed, though I'll try to retest on the most current version next week. Is there work going on somewhere else on this? Is there maybe a workaround, or is the best I can do to just use a patch that raises the 10k function definitions limit?