@trace: Promote loop variables to traced numbers (#1381)

* `@trace`: Promote loop variables to traced numbers

* fmt

* rm

* Update ReactantCore.jl

Author: Pangoraw

docs/src/.vitepress/config.mts

@@ -155,6 +155,7 @@ export default defineConfig({
           { text: "Profiling", link: "/tutorials/profiling" },
           { text: "Distributed", link: "/tutorials/multihost" },
           { text: "Local build", link: "/tutorials/local-build" },
+          { text: "Control Flow", link: "/tutorials/control-flow" },
         ],
       }
     ],

lib/ReactantCore/src/ReactantCore.jl

@@ -231,8 +231,8 @@ function trace_while(expr; track_numbers, first_arg=nothing)
                 $(cond_fn_sym),
                 $(body_fn_sym),
                 $(args_sym);
-                track_numbers=$(track_numbers),
-                verify_arg_names=$(verify_arg_names_sym),
+                track_numbers=($(track_numbers)),
+                verify_arg_names=($(verify_arg_names_sym)),
             )
         end
     end
@@ -289,10 +289,26 @@ function trace_for(expr; track_numbers)
         end
     end
 
-    quote
+    return quote
         local $start_sym, $limit_sym, $step_sym
         $bounds_defs
-        local $counter = 0
+
+        if $(within_compile)()
+            $start_sym = Reactant.TracedUtils.promote_to(
+                Reactant.TracedRNumber{Reactant.unwrapped_eltype(typeof($start_sym))},
+                $start_sym,
+            )
+            $limit_sym = Reactant.TracedUtils.promote_to(
+                Reactant.TracedRNumber{Reactant.unwrapped_eltype(typeof($limit_sym))},
+                $limit_sym,
+            )
+            $step_sym = Reactant.TracedUtils.promote_to(
+                Reactant.TracedRNumber{Reactant.unwrapped_eltype(typeof($step_sym))},
+                $step_sym,
+            )
+        end
+
+        local $counter = zero($start_sym)
 
         $(trace_while(
             Expr(
@@ -471,7 +487,7 @@ function trace_if(expr; store_last_line=nothing, depth=0, track_numbers)
             $(true_branch_fn_name),
             $(false_branch_fn_name),
             ($(all_input_vars...),);
-            track_numbers=$(track_numbers),
+            track_numbers=($(track_numbers)),
         )
     end
 

test/control_flow.jl

@@ -461,7 +461,7 @@ function condition_with_structure(x)
     @trace if sum(y) > 0
         z = (; a=y, b=(y .- 1, y))
     else
-        z = (; a=-y, b=(y, y .+ 1))
+        z = (; a=(-y), b=(y, y .+ 1))
     end
     return z
 end
@@ -676,6 +676,27 @@ function while_convergence(x, y)
     return diff
 end
 
+function for_no_track_numbers(x, n)
+    @trace track_numbers = false for i in n:16
+        x = x .+ 1
+    end
+    return x
+end
+
+@testset "for: track_numbers=false" begin
+    x = [1, 2, 3]
+    x_ra = Reactant.to_rarray(x)
+
+    n = 12
+    n_ra = Reactant.ConcreteRNumber(n)
+
+    # set optimize to only do enzyme-batch to prevent crash in opt
+    for_no_track_numbers_ra = @compile optimize="enzyme-batch" for_no_track_numbers(
+        x_ra, n_ra
+    )
+    for_no_track_numbers_ra(x_ra, n_ra) == for_no_track_numbers(x, n)
+end
+
 @testset "while: convergence" begin
     x = [1.0, 10.0, 20.0]
     y = [0.0, -2.0, -3.0]