@trace function calls (#366)

Author: jumerckx

docs/src/api/api.md

@@ -13,6 +13,10 @@ Reactant.@jit
 
 ## ReactantCore API
 
+```@docs
+within_compile
+```
+
 ```@docs
 @trace
 ```

lib/ReactantCore/src/ReactantCore.jl

@@ -3,7 +3,7 @@ module ReactantCore
 using ExpressionExplorer: ExpressionExplorer
 using MacroTools: MacroTools
 
-export @trace, MissingTracedValue
+export @trace, within_compile, MissingTracedValue
 
 # Traits
 is_traced(x) = false
@@ -21,6 +21,13 @@ const SPECIAL_SYMBOLS = [
     :(:), :nothing, :missing, :Inf, :Inf16, :Inf32, :Inf64, :Base, :Core
 ]
 
+"""
+    within_compile()
+
+Returns true if this function is executed in a Reactant compilation context, otherwise false.
+"""
+@inline within_compile() = false # behavior is overwritten in Interpreter.jl
+
 # Code generation
 """
     @trace <expr>
@@ -117,6 +124,13 @@ macro trace(expr)
             return esc(trace_if_with_returns(__module__, expr))
         end
     end
+    Meta.isexpr(expr, :call) && return esc(trace_call(__module__, expr))
+    if Meta.isexpr(expr, :(.), 2) && Meta.isexpr(expr.args[2], :tuple)
+        fname = :($(Base.Broadcast.BroadcastFunction)($(expr.args[1])))
+        args = only(expr.args[2:end]).args
+        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)))
     return error("Only `if-elseif-else` blocks are currently supported by `@trace`")
@@ -196,7 +210,9 @@ function trace_for(mod, expr)
     end
 
     return quote
-        if any($(is_traced), $(Expr(:tuple, cond_val.(all_syms.args[(begin + 1):end])...)))
+        if $(within_compile)() && $(any)(
+            $(is_traced), $(Expr(:tuple, cond_val.(all_syms.args[(begin + 1):end])...))
+        )
             $(reactant_code_block)
         else
             $(expr)
@@ -210,7 +226,7 @@ function trace_if_with_returns(mod, expr)
         mod, expr.args[2]; store_last_line=expr.args[1], depth=1
     )
     return quote
-        if any($(is_traced), ($(all_check_vars...),))
+        if $(within_compile)() && $(any)($(is_traced), ($(all_check_vars...),))
             $(new_expr)
         else
             $(expr)
@@ -356,14 +372,41 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0)
     )
 
     return quote
-        if any($(is_traced), ($(all_check_vars...),))
+        if $(within_compile)() && $(any)($(is_traced), ($(all_check_vars...),))
             $(reactant_code_block)
         else
             $(original_expr)
         end
     end
 end
 
+function correct_maybe_bcast_call(fname)
+    startswith(string(fname), '.') || return false, fname, fname
+    return true, Symbol(string(fname)[2:end]), fname
+end
+
+function trace_call(mod, call)
+    bcast, fname, fname_full = correct_maybe_bcast_call(call.args[1])
+    f = if bcast
+        quote
+            if isdefined(mod, $(Meta.quot(fname_full)))
+                $(fname_full)
+            else
+                Base.Broadcast.BroadcastFunction($(fname))
+            end
+        end
+    else
+        :($(fname))
+    end
+    return quote
+        if $(within_compile)()
+            $(traced_call)($f, $(call.args[2:end]...))
+        else
+            $(call)
+        end
+    end
+end
+
 function remove_shortcircuiting(expr)
     return MacroTools.prewalk(expr) do x
         if MacroTools.@capture(x, a_ && b_)
@@ -382,6 +425,8 @@ end
 
 function traced_while end # defined inside Reactant.jl
 
+traced_call(f, args...; kwargs...) = f(args...; kwargs...)
+
 function cleanup_expr_to_avoid_boxing(expr, prepend::Symbol, all_vars)
     return MacroTools.postwalk(expr) do x
         if Meta.isexpr(x, :kw) # undo lhs rewriting

src/Compiler.jl

@@ -20,6 +20,8 @@ import ..Reactant:
     ancestor,
     TracedType
 
+import ..ReactantCore: correct_maybe_bcast_call
+
 @inline function traced_getfield(@nospecialize(obj), field)
     return Base.getfield(obj, field)
 end
@@ -440,18 +442,34 @@ const DEBUG_KERNEL = Ref{Bool}(false)
 const DUMP_LLVMIR = Ref{Bool}(false)
 
 function compile_mlir!(
-    mod, f, args; optimize::Union{Bool,Symbol}=true, no_nan::Bool=false, backend="gpu"
+    mod,
+    f,
+    args,
+    callcache=Dict{
+        Vector,
+        @NamedTuple{
+            f_name::String,
+            mlir_result_types::Vector{MLIR.IR.Type},
+            traced_result::Any,
+            mutated::Vector{Int},
+        }
+    }();
+    optimize::Union{Bool,Symbol}=true,
+    no_nan::Bool=false,
+    backend="gpu",
 )
     # Explicitly don't use block! to avoid creating a closure, which creates
     # both compile-time and relocatability issues
 
     MLIR.IR.activate!(mod)
     MLIR.IR.activate!(MLIR.IR.body(mod))
+    activate_callcache!(callcache)
     fnwrapped,
     func2, traced_result, result, seen_args, ret, linear_args, in_tys,
     linear_results = try
         Reactant.TracedUtils.make_mlir_fn(f, args, (), "main", true)
     finally
+        deactivate_callcache!(callcache)
         MLIR.IR.deactivate!(MLIR.IR.body(mod))
         MLIR.IR.deactivate!(mod)
     end
@@ -716,11 +734,6 @@ function compile_call_expr(mod, compiler, options, args...)
     (; compiled=compiled_symbol, args=args_symbol)
 end
 
-function correct_maybe_bcast_call(fname)
-    startswith(string(fname), '.') || return false, fname, fname
-    return true, Symbol(string(fname)[2:end]), fname
-end
-
 """
     codegen_flatten!
 
@@ -1167,4 +1180,40 @@ function register_thunk(
     return Thunk{Core.Typeof(f),tag,argtys,isclosure}(f)
 end
 
+function activate_callcache!(callcache)
+    stack = get!(task_local_storage(), :callcache) do
+        return []
+    end
+    push!(stack, callcache)
+    return nothing
+end
+
+function deactivate_callcache!(callcache)
+    callcache === last(task_local_storage(:callcache)) ||
+        error("Deactivating wrong callcache")
+    return pop!(task_local_storage(:callcache))
+end
+
+function _has_callcache()
+    return haskey(task_local_storage(), :callcache) &&
+           !Base.isempty(task_local_storage(:callcache))
+end
+
+function callcache(; throw_error::Bool=true)
+    if !_has_callcache()
+        throw_error && error("No callcache is active")
+        return nothing
+    end
+    return last(task_local_storage(:callcache))
+end
+
+function callcache!(f, callcache)
+    activate_callcache!(callcache)
+    try
+        return f()
+    finally
+        deactivate_callcache!(callcache)
+    end
+end
+
 end

src/ControlFlow.jl

@@ -4,6 +4,10 @@ function ReactantCore.traced_if(
     return Ops.if_condition(cond, true_fn, false_fn, args...)
 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...)
 end

src/Interpreter.jl

@@ -39,6 +39,25 @@ function set_reactant_abi(
 )
     (; fargs, argtypes) = arginfo
 
+    if f === ReactantCore.within_compile
+        if length(argtypes) != 1
+            @static if VERSION < v"1.11.0-"
+                return CallMeta(Union{}, Effects(), NoCallInfo())
+            else
+                return CallMeta(Union{}, Union{}, Effects(), NoCallInfo())
+            end
+        end
+        @static if VERSION < v"1.11.0-"
+            return CallMeta(
+                Core.Const(true), Core.Compiler.EFFECTS_TOTAL, MethodResultPure()
+            )
+        else
+            return CallMeta(
+                Core.Const(true), Union{}, Core.Compiler.EFFECTS_TOTAL, MethodResultPure()
+            )
+        end
+    end
+
     # Improve inference by considering call_with_reactant as having the same results as
     # the original call
     if f === Reactant.call_with_reactant
@@ -236,7 +255,7 @@ function overload_autodiff(
     primf = f.val
     primargs = ((v.val for v in args)...,)
 
-    fnwrap, func2, traced_result, result, seen_args, ret, linear_args, in_tys, linear_results = TracedUtils.make_mlir_fn(
+    fnwrap, func2, traced_result, result, seen_args, ret, linear_args, in_tys, linear_results, _ = TracedUtils.make_mlir_fn(
         primf, primargs, (), string(f) * "_autodiff", false
     )
 

src/Ops.jl

@@ -1070,7 +1070,7 @@ end
         (sample_inputs...,),
         (),
         "comparator";
-        no_args_in_result=true,
+        args_in_result=:none,
         return_dialect=:stablehlo,
     )[2]
     @assert MLIR.IR.nregions(func) == 1
@@ -1679,7 +1679,7 @@ end
         string(gensym("cond_fn")),
         false;
         return_dialect=:stablehlo,
-        no_args_in_result=true,
+        args_in_result=:none,
         do_transpose=false,
     )
 
@@ -1690,7 +1690,7 @@ end
         string(gensym("body_fn")),
         false;
         return_dialect=:stablehlo,
-        no_args_in_result=true,
+        args_in_result=:none,
         do_transpose=false,
     )
 
@@ -2060,4 +2060,75 @@ end
     return corrected_traced_results
 end
 
+@noinline function call(f, args...)
+    seen_cache = Reactant.OrderedIdDict()
+    Reactant.make_tracer(
+        seen_cache,
+        args,
+        (), # we have to insert something here, but we remove it immediately below.
+        Reactant.TracedTrack;
+        toscalar=false,
+    )
+    linear_args = []
+    mlir_caller_args = Reactant.MLIR.IR.Value[]
+    for (k, v) in seen_cache
+        v isa Reactant.TracedType || continue
+        push!(linear_args, v)
+        push!(mlir_caller_args, v.mlir_data)
+        # make tracer inserted `()` into the path, here we remove it:
+        v.paths = v.paths[1:(end - 1)]
+    end
+
+    seen = Dict()
+    cache_key = []
+    Reactant.make_tracer(seen, (f, args...), cache_key, Reactant.TracedToTypes)
+    cache = Reactant.Compiler.callcache()
+    if haskey(cache, cache_key)
+        # cache lookup:
+        (; f_name, mlir_result_types, traced_result, mutated) = cache[cache_key]
+    else
+        f_name = String(gensym(Symbol(f)))
+        temp = Reactant.TracedUtils.make_mlir_fn(
+            f, args, (), f_name, false; args_in_result=:mutated, do_transpose=false
+        )
+        traced_result, ret, mutated = temp[[3, 6, 10]]
+        mlir_result_types = [
+            MLIR.IR.type(MLIR.IR.operand(ret, i)) for i in 1:MLIR.IR.noperands(ret)
+        ]
+        cache[cache_key] = (; f_name, mlir_result_types, traced_result, mutated)
+    end
+
+    call_op = MLIR.Dialects.func.call(
+        mlir_caller_args;
+        result_0=mlir_result_types,
+        callee=MLIR.IR.FlatSymbolRefAttribute(f_name),
+    )
+
+    seen_results = Reactant.OrderedIdDict()
+    traced_result = Reactant.make_tracer(
+        seen_results,
+        traced_result,
+        (), # we have to insert something here, but we remove it immediately below.
+        Reactant.TracedSetPath;
+        toscalar=false,
+    )
+    i = 1
+    for (k, v) in seen_results
+        v isa Reactant.TracedType || continue
+        # this mutates `traced_result`, which is what we want:
+        v.mlir_data = MLIR.IR.result(call_op, i)
+        # make tracer inserted `()` into the path, here we remove it:
+        v.paths = v.paths[1:(end - 1)]
+        i += 1
+    end
+    nres = MLIR.IR.nresults(call_op)
+    # mutated args are included as the last ones in the call op results
+    for (result_i, arg_i) in zip((nres - length(mutated)):nres, mutated)
+        Reactant.TracedUtils.set_mlir_data!(
+            linear_args[arg_i], MLIR.IR.result(call_op, result_i + 1)
+        )
+    end
+    return traced_result
+end
+
 end # module Ops

src/Reactant.jl

@@ -1,6 +1,6 @@
 module Reactant
 
-using ReactantCore: ReactantCore, @trace, MissingTracedValue
+using ReactantCore: ReactantCore, @trace, within_compile, MissingTracedValue
 
 using LinearAlgebra: LinearAlgebra
 using Random: Random, AbstractRNG
@@ -231,7 +231,7 @@ function Enzyme.make_zero(
 end
 
 using .Compiler: @compile, @code_hlo, @jit, traced_getfield, create_result, compile
-export ConcreteRArray, ConcreteRNumber, @compile, @code_hlo, @jit, @trace
+export ConcreteRArray, ConcreteRNumber, @compile, @code_hlo, @jit, @trace, within_compile
 
 const registry = Ref{Union{Nothing,MLIR.IR.DialectRegistry}}()