Variables are not traced in while loops

[dominic-chang]

Here's some code

using Reactant

function ellip_k(m::T) where {T}
    oneT = one(T)
    twoT = T(2)

    a = oneT
    b = sqrt(oneT - m)
    pow2 = oneT

    @trace while (pow2 < (pow2+10.0))
        a = (a + b) / twoT
        pow2 = twoT * pow2
    end

    return a
end

ms = Float64[0.01, 0.1, 0.5, 0.9]
m_r = Reactant.to_rarray(ms)

Reactant.@jit ellip_k.(m_r)

That results in the following error

ERROR: Types do not match between function arguments and results.
The following arguments should be traced but were not: a (path=(2, 1), type=Reactant.TracedRNumber{Float64})
The following arguments should be returned but were not: 
argprefix = ##loop_bodyarg#511
resprefix = ##loop_bodyres#512
verify_arg_names = (:pow2, :a, :b, :twoT)
argtys = DataType[Base.RefValue{Reactant.TracedRNumber{Float64}}, Base.RefValue{Reactant.TracedRNumber{Float64}}, Base.RefValue{Reactant.TracedRNumber{Float64}}, Base.RefValue{Reactant.TracedRNumber{Float64}}]
Traced Arg Paths: 
(Ptr{Nothing} @0x000000031b191ef0, ((Symbol("##loop_bodyarg#511"), 1, 1),))
(Ptr{Nothing} @0x000000031b192510, ((Symbol("##loop_bodyarg#511"), 3, 1), (Symbol("##loop_bodyresarg#513"), 3, 1)))
(Ptr{Nothing} @0x000000031b1926f0, ((Symbol("##loop_bodyarg#511"), 4, 1), (Symbol("##loop_bodyresarg#513"), 4, 1)))

Traced Res Paths: 
(Ptr{Nothing} @0x000000031b0d4170, ((Symbol("##loop_bodyresarg#513"), 1, 1),))
(Ptr{Nothing} @0x000000031b11fbf0, ((Symbol("##loop_bodyresarg#513"), 2, 1),))
(Ptr{Nothing} @0x000000031b192510, ((Symbol("##loop_bodyarg#511"), 3, 1), (Symbol("##loop_bodyresarg#513"), 3, 1)))
(Ptr{Nothing} @0x000000031b1926f0, ((Symbol("##loop_bodyarg#511"), 4, 1), (Symbol("##loop_bodyresarg#513"), 4, 1)))

Traced Arg NI Paths: 
(Ptr{Nothing} @0x000000031b191ef0, ((Symbol("##loop_bodyarg#511"), 1, 1),))

Traced Res NI Paths: 
(Ptr{Nothing} @0x000000031b0d4170, ((Symbol("##loop_bodyresarg#513"), 1, 1),))
(Ptr{Nothing} @0x000000031b11fbf0, ((Symbol("##loop_bodyresarg#513"), 2, 1),))

traced_result : Nothing

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] finalize_mlir_fn(result::Nothing, 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::NTuple{…}, 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::Vector{…}, N::Int64, concretein::Bool, toscalar::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:831
  [3] make_mlir_fn(f::var"#92#94", args::Vector{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, within_autodiff::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Nothing, verify_arg_names::NTuple{…}, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:399
  [4] make_mlir_fn
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:299 [inlined]
  [5] while_loop(::var"#91#93", ::var"#92#94", ::Base.RefValue{…}, ::Vararg{…}; track_numbers::Type, verify_arg_names::NTuple{…}, checkpointing::Bool, mincut::Bool, location::Reactant.MLIR.IR.Location)
    @ Reactant.Ops ~/.julia/packages/Reactant/9BAHk/src/Ops.jl:2388
  [6] #traced_while#162
    @ ~/.julia/packages/Reactant/9BAHk/src/ControlFlow.jl:20 [inlined]
  [7] traced_while
    @ ~/.julia/packages/Reactant/9BAHk/src/ControlFlow.jl:11 [inlined]
  [8] macro expansion
    @ ~/.julia/packages/ReactantCore/WerTZ/src/ReactantCore.jl:412 [inlined]
  [9] ellip_k
    @ ~/Desktop/JacobiElliptic.jl/ext/bug.jl:11
 [10] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(ellip_k), ::Reactant.TracedRNumber{…})
    @ Reactant ~/.julia/packages/Reactant/9BAHk/src/utils.jl:1195
 [11] _copy
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedRArray.jl:320
 [12] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(Reactant.TracedRArrayOverrides._copy), ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/9BAHk/src/utils.jl:1195
 [13] copy
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedRArray.jl:330
 [14] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(copy), ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/.julia/packages/Reactant/9BAHk/src/utils.jl:1195
 [15] materialize
    @ ./broadcast.jl:903 [inlined]
 [16] BroadcastFunction
    @ ./broadcast.jl:1381 [inlined]
 [17] #apply#102
    @ ~/.julia/packages/Reactant/9BAHk/src/utils.jl:119
 [18] call_with_reactant(::Reactant.EnsureReturnType{…}, ::Reactant.var"##apply#102", ::@Kwargs{}, ::typeof(Reactant.apply), ::Base.Broadcast.BroadcastFunction{…}, ::Reactant.TracedRArray{…})
    @ Reactant ~/.julia/packages/Reactant/9BAHk/src/utils.jl:1195
 [19] apply
    @ ~/.julia/packages/Reactant/9BAHk/src/utils.jl:118
 [20] call_with_reactant(::typeof(Reactant.apply), ::Base.Broadcast.BroadcastFunction{…}, ::Reactant.TracedRArray{…})
    @ Reactant ~/.julia/packages/Reactant/9BAHk/src/utils.jl:1195
 [21] make_mlir_fn(f::typeof(Reactant.apply), 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, within_autodiff::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/9BAHk/src/TracedUtils.jl:376
 [22] make_mlir_fn
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:299 [inlined]
 [23] #make_mlir_fn#7
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:322 [inlined]
 [24] make_mlir_fn
    @ ~/.julia/packages/Reactant/9BAHk/src/TracedUtils.jl:299 [inlined]
 [25] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Base.Broadcast.BroadcastFunction{…}, args::Tuple{…}, compile_options::CompileOptions, debugcache::Vector{…}, callcache::Dict{…}, sdycache::Dict{…}, sdygroupidcache::Tuple{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, client::Reactant.XLA.PJRT.Client, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/9BAHk/src/compiler/Compiler.jl:279
 [26] compile_mlir!
    @ ~/.julia/packages/Reactant/9BAHk/src/compiler/Compiler.jl:238 [inlined]
 [27] compile_xla(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/9BAHk/src/compiler/Compiler.jl:1331
 [28] compile_xla
    @ ~/.julia/packages/Reactant/9BAHk/src/compiler/Compiler.jl:1306 [inlined]
 [29] compile(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/9BAHk/src/compiler/Compiler.jl:1423
 [30] macro expansion
    @ ~/.julia/packages/Reactant/9BAHk/src/compiler/Macros.jl:299 [inlined]
 [31] macro expansion
    @ ~/.julia/packages/LLVM/aXQtR/src/base.jl:113 [inlined]
 [32] macro expansion
    @ ~/.julia/packages/Reactant/9BAHk/src/compiler/Macros.jl:298 [inlined]

[ptiede]

On Reactant main I get

ERROR: Reactant.@trace: the aliasing pattern of loop-carried variables changed during the loop body.
`pow2` and `a` referred to the same object before the loop body but not after.
Fix: break the aliasing before the `@trace` loop, e.g. `a = copy(a)`.
Stacktrace:
  [1] call_with_reactant(::Reactant.EnsureReturnType{Union{}}, ::typeof(error), ::String)
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
  [2] _check_aliasing_unchanged
    @ ~/Research/Reactant.jl/lib/ReactantCore/src/ReactantCore.jl:354
  [3] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(ReactantCore._check_aliasing_unchanged), ::NTuple{…}, ::NTuple{…}, ::NTuple{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
  [4] #2
    @ ~/Research/Reactant.jl/lib/ReactantCore/src/ReactantCore.jl:456
  [5] call_with_reactant(::Reactant.EnsureReturnType{…}, ::var"#2#4"{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
  [6] call_with_reactant(::Reactant.EnsureReturnType{…}, ::Reactant.var"##apply#102", ::@Kwargs{}, ::typeof(Reactant.apply), ::var"#2#4"{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
  [7] call_with_reactant(::typeof(Reactant.apply), ::var"#2#4"{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…}, ::Base.RefValue{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
  [8] 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, within_autodiff::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Nothing, verify_arg_names::NTuple{…}, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/Research/Reactant.jl/src/TracedUtils.jl:376
  [9] make_mlir_fn(f::var"#2#4"{…}, args::Vector{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, within_autodiff::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Nothing, verify_arg_names::NTuple{…}, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/Research/Reactant.jl/src/TracedUtils.jl:322
 [10] make_mlir_fn
    @ ~/Research/Reactant.jl/src/TracedUtils.jl:299 [inlined]
 [11] while_loop(::var"#1#3", ::var"#2#4"{…}, ::Base.RefValue{…}, ::Vararg{…}; track_numbers::Type, verify_arg_names::NTuple{…}, checkpointing::Bool, mincut::Bool, location::Reactant.MLIR.IR.Location)
    @ Reactant.Ops ~/Research/Reactant.jl/src/Ops.jl:2388
 [12] #traced_while#162
    @ ~/Research/Reactant.jl/src/ControlFlow.jl:20 [inlined]
 [13] traced_while
    @ ~/Research/Reactant.jl/src/ControlFlow.jl:11 [inlined]
 [14] macro expansion
    @ ~/Research/Reactant.jl/lib/ReactantCore/src/ReactantCore.jl:468 [inlined]
 [15] ellip_k
    @ ./REPL[3]:9
 [16] call_with_reactant(::typeof(ellip_k), ::Reactant.TracedRNumber{Float64})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [17] make_mlir_fn(f::typeof(ellip_k), 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, within_autodiff::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 ~/Research/Reactant.jl/src/TracedUtils.jl:376
 [18] elem_apply(f::Function, args::Reactant.TracedRArray{Float64, 1})
    @ Reactant.TracedUtils ~/Research/Reactant.jl/src/TracedUtils.jl:1224
 [19] _copyto!
    @ ~/Research/Reactant.jl/src/TracedRArray.jl:401
 [20] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(Reactant.TracedRArrayOverrides._copyto!), ::Reactant.TracedRArray{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [21] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(copyto!), ::Reactant.TracedRArray{…}, ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [22] copyto!
    @ ./broadcast.jl:925 [inlined]
 [23] _copy
    @ ~/Research/Reactant.jl/src/TracedRArray.jl:326
 [24] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(Reactant.TracedRArrayOverrides._copy), ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [25] call_with_reactant(::Reactant.EnsureReturnType{…}, ::typeof(copy), ::Base.Broadcast.Broadcasted{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [26] materialize
    @ ./broadcast.jl:872 [inlined]
 [27] BroadcastFunction
    @ ./broadcast.jl:1365 [inlined]
 [28] #apply#102
    @ ~/Research/Reactant.jl/src/utils.jl:119
 [29] call_with_reactant(::Reactant.EnsureReturnType{…}, ::Reactant.var"##apply#102", ::@Kwargs{}, ::typeof(Reactant.apply), ::Base.Broadcast.BroadcastFunction{…}, ::Reactant.TracedRArray{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [30] call_with_reactant(::typeof(Reactant.apply), ::Base.Broadcast.BroadcastFunction{…}, ::Reactant.TracedRArray{…})
    @ Reactant ~/Research/Reactant.jl/src/utils.jl:1195
 [31] make_mlir_fn(f::typeof(Reactant.apply), 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, within_autodiff::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 ~/Research/Reactant.jl/src/TracedUtils.jl:376
 [32] make_mlir_fn
    @ ~/Research/Reactant.jl/src/TracedUtils.jl:299 [inlined]
 [33] #make_mlir_fn#7
    @ ~/Research/Reactant.jl/src/TracedUtils.jl:322 [inlined]
 [34] make_mlir_fn
    @ ~/Research/Reactant.jl/src/TracedUtils.jl:299 [inlined]
 [35] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Base.Broadcast.BroadcastFunction{…}, args::Tuple{…}, compile_options::CompileOptions, debugcache::Vector{…}, callcache::Dict{…}, sdycache::Dict{…}, sdygroupidcache::Tuple{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, client::Reactant.XLA.PJRT.Client, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/Research/Reactant.jl/src/compiler/Compiler.jl:279
 [36] compile_mlir!
    @ ~/Research/Reactant.jl/src/compiler/Compiler.jl:238 [inlined]
 [37] compile_xla(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/Research/Reactant.jl/src/compiler/Compiler.jl:1331
 [38] compile_xla
    @ ~/Research/Reactant.jl/src/compiler/Compiler.jl:1306 [inlined]
 [39] compile(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/Research/Reactant.jl/src/compiler/Compiler.jl:1423
 [40] macro expansion
    @ ~/Research/Reactant.jl/src/compiler/Macros.jl:299 [inlined]
 [41] macro expansion
    @ ~/.julia/packages/LLVM/aXQtR/src/base.jl:113 [inlined]
 [42] macro expansion
    @ ~/Research/Reactant.jl/src/compiler/Macros.jl:298 [inlined]
 [43] top-level scope
    @ REPL[6]:1
Some type information was truncated. Use `show(err)` to see complete types.

which I think was the resolution of #2853 (comment)

[ptiede]

Adding copy does fix this for me

function ellip_k(m::T) where {T}
           oneT = one(T)
           twoT = T(2)

           a = copy(oneT)
           b = sqrt(oneT - m)
           pow2 = copy(oneT)

           @trace while (pow2 < (pow2+10.0))
               a = (a + b) / twoT
               pow2 = twoT * pow2
           end

           return a
       end

Or just enforcing that a and pow2 don't refer to the same variable

function ellip_k(m::T) where {T}
           oneT = one(T)
           twoT = T(2)

           a = one(T)
           b = sqrt(one(T) - m)
           pow2 = one(T)

           @trace while (pow2 < (pow2+10.0))
               a = (a + b) / twoT
               pow2 = twoT * pow2
           end

           return a
       end