allow @trace track_numbers=false ... (#1026)

* allow `@trace track_numbers=false ...`

* add kwarg

* version bumps

---------

Co-authored-by: William S. Moses <[email protected]>
Co-authored-by: Mosè Giordano <[email protected]>

Author: 3 people

Project.toml

@@ -85,7 +85,7 @@ Preferences = "1.4"
 PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.6"
+ReactantCore = "0.1.7"
 Reactant_jll = "0.0.100"
 Scratch = "1.2"
 Sockets = "1.10"

lib/ReactantCore/Project.toml

@@ -1,7 +1,7 @@
 name = "ReactantCore"
 uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>"]
-version = "0.1.6"
+version = "0.1.7"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -128,11 +128,25 @@ function fn(x)
 end
 ```
 """
-macro trace(expr)
+macro trace(args...)
+    track_numbers = true
+    expr = first(args)
+    if length(args) > 1 && Meta.isexpr(args[1], :(=))
+        tn_expr = args[1]
+        tn_expr.args[1] == :track_numbers ||
+            error("@trace supports setting track_numbers, but got $(tn_expr)")
+
+        track_numbers = tn_expr.args[2]
+        expr = only(args[2:end])
+    else
+        expr = only(args)
+    end
+    track_numbers = track_numbers ? Number : Union{}
     expr = macroexpand(__module__, expr)
+
     if Meta.isexpr(expr, :(=))
         if Meta.isexpr(expr.args[2], :if)
-            return esc(trace_if_with_returns(__module__, expr))
+            return esc(trace_if_with_returns(__module__, expr; track_numbers))
         end
     end
     Meta.isexpr(expr, :call) && return esc(trace_call(__module__, expr))
@@ -142,12 +156,12 @@ macro trace(expr)
         call = Expr(:call, fname, args...)
         return esc(trace_call(__module__, call))
     end
-    Meta.isexpr(expr, :if) && return esc(trace_if(__module__, expr))
-    Meta.isexpr(expr, :for) && return (esc(trace_for(__module__, expr)))
+    Meta.isexpr(expr, :if) && return esc(trace_if(__module__, expr; track_numbers))
+    Meta.isexpr(expr, :for) && return (esc(trace_for(__module__, expr; track_numbers)))
     return error("Only `if-elseif-else` blocks are currently supported by `@trace`")
 end
 
-function trace_for(mod, expr)
+function trace_for(mod, expr; track_numbers)
     Meta.isexpr(expr, :for, 2) || error("expected for expr")
     assign, body = expr.args
 
@@ -216,7 +230,9 @@ function trace_for(mod, expr)
                     ($counter + 1, results_...)
                 end
 
-            $(ReactantCore).traced_while(cond_fn, body_fn, args)
+            $(ReactantCore).traced_while(
+                cond_fn, body_fn, args; track_numbers=$(track_numbers)
+            )
         end
     end
 
@@ -232,9 +248,9 @@ function trace_for(mod, expr)
 end
 
 # ... = if ... style expressions
-function trace_if_with_returns(mod, expr)
+function trace_if_with_returns(mod, expr; track_numbers)
     new_expr, _, all_check_vars = trace_if(
-        mod, expr.args[2]; store_last_line=expr.args[1], depth=1
+        mod, expr.args[2]; store_last_line=expr.args[1], depth=1, track_numbers
     )
     cond_name = first(all_check_vars)
     original_cond = expr.args[2].args[1]
@@ -249,7 +265,7 @@ function trace_if_with_returns(mod, expr)
     end
 end
 
-function trace_if(mod, expr; store_last_line=nothing, depth=0)
+function trace_if(mod, expr; store_last_line=nothing, depth=0, track_numbers)
     discard_vars_from_expansion = []
     original_expr = expr
 
@@ -260,7 +276,9 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0)
         expr = MacroTools.prewalk(expr) do x
             counter += 1
             if x isa Expr && x.head == :if && counter > 1
-                ex_new, dv, _ = trace_if(mod, x; store_last_line, depth=depth + 1)
+                ex_new, dv, _ = trace_if(
+                    mod, x; store_last_line, depth=depth + 1, track_numbers
+                )
                 append!(discard_vars_from_expansion, dv)
                 return ex_new
             end
@@ -300,7 +318,7 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0)
         if !(expr.args[3] isa Expr) || expr.args[3].head != :elseif
             expr.args[3], [], nothing
         else
-            trace_if(mod, expr.args[3]; store_last_line, depth=depth + 1)
+            trace_if(mod, expr.args[3]; store_last_line, depth=depth + 1, track_numbers)
         end
     elseif length(expr.args) == 2
         tmp_expr = []
@@ -388,7 +406,8 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0)
             $(cond_name),
             $(true_branch_fn_name),
             $(false_branch_fn_name),
-            ($(all_input_vars...),),
+            ($(all_input_vars...),);
+            track_numbers=$(track_numbers),
         )
     end
 
@@ -458,7 +477,7 @@ function remove_shortcircuiting(expr)
 end
 
 # Generate this dummy function and later we remove it during tracing
-function traced_if(cond, true_fn, false_fn, args)
+function traced_if(cond, true_fn, false_fn, args; track_numbers)
     return cond ? true_fn(args) : false_fn(args)
 end
 

src/ControlFlow.jl

@@ -1,13 +1,15 @@
 function ReactantCore.traced_if(
-    cond::TracedRNumber{Bool}, true_fn::TFn, false_fn::FFn, args
+    cond::TracedRNumber{Bool}, true_fn::TFn, false_fn::FFn, args; track_numbers=Number
 ) where {TFn,FFn}
-    return Ops.if_condition(cond, true_fn, false_fn, args...)
+    return Ops.if_condition(cond, true_fn, false_fn, args...; track_numbers)
 end
 
 function ReactantCore.traced_call(f::Function, args...)
     return Ops.call(f, args...)
 end
 
-function ReactantCore.traced_while(cond_fn::CFn, body_fn::BFn, args) where {CFn,BFn}
-    return Ops.while_loop(cond_fn, body_fn, args...)
+function ReactantCore.traced_while(
+    cond_fn::CFn, body_fn::BFn, args; track_numbers=Number
+) where {CFn,BFn}
+    return Ops.while_loop(cond_fn, body_fn, args...; track_numbers)
 end

src/Ops.jl

@@ -1745,7 +1745,9 @@ use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
     )
 end
 
-@noinline function while_loop(cond_fn::CFn, body_fn::BFn, args...) where {CFn,BFn}
+@noinline function while_loop(
+    cond_fn::CFn, body_fn::BFn, args...; track_numbers
+) where {CFn,BFn}
     # TODO: detect and prevent mutation within the condition
 
     # Make all the args traced or concrete
@@ -1754,7 +1756,7 @@ end
     traced_args = Vector{Any}(undef, N)
     for i in 1:N
         @inbounds traced_args[i] = Reactant.make_tracer(
-            seen_args, args[i], (), Reactant.NoStopTracedTrack; track_numbers=Number
+            seen_args, args[i], (), Reactant.NoStopTracedTrack; track_numbers
         )
     end
 
@@ -1809,7 +1811,7 @@ end
 end
 
 @noinline function if_condition(
-    cond::TracedRNumber{Bool}, true_fn::TFn, false_fn::FFn, args...
+    cond::TracedRNumber{Bool}, true_fn::TFn, false_fn::FFn, args...; track_numbers
 ) where {TFn,FFn}
     true_fn_names = (gensym(:true_fn_args), gensym(:true_result), gensym(:true_fn_resargs))
     false_fn_names = (
@@ -1828,14 +1830,14 @@ end
             args[i],
             (true_fn_names[1], i),
             Reactant.TracedSetPath;
-            track_numbers=Number,
+            track_numbers,
         )
         @inbounds fb_traced_args[i] = Reactant.make_tracer(
             fb_seen_args,
             args[i],
             (false_fn_names[1], i),
             Reactant.TracedSetPath;
-            track_numbers=Number,
+            track_numbers,
         )
     end
 
@@ -1899,15 +1901,15 @@ end
         tb_result,
         (true_fn_names[2],),
         Reactant.NoStopTracedTrack;
-        track_numbers=Number,
+        track_numbers,
     )
     for i in eachindex(tb_linear_args)
         Reactant.make_tracer(
             seen_true_results,
             tb_linear_args[i],
             (true_fn_names[3], i),
             Reactant.NoStopTracedTrack;
-            track_numbers=Number,
+            track_numbers,
         )
     end
 
@@ -1964,15 +1966,15 @@ end
         fb_result,
         (false_fn_names[2],),
         Reactant.NoStopTracedTrack;
-        track_numbers=Number,
+        track_numbers,
     )
     for i in eachindex(fb_linear_args)
         Reactant.make_tracer(
             seen_false_results,
             fb_linear_args[i],
             (false_fn_names[3], i),
             Reactant.NoStopTracedTrack;
-            track_numbers=Number,
+            track_numbers,
         )
     end