feat: @trace on function definition (#1685)

* feat: `@trace` on function definition

* feat: support functors

* feat: handle _ and no symbol functor

* feat: support typed kwargs

* test: add testcases

Co-authored-by: jumerckx <[email protected]>

* fix: single line fn

---------

Co-authored-by: jumerckx <[email protected]>

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.163"
+version = "0.2.164"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -102,7 +102,7 @@ Preferences = "1.4.3"
 PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.15"
+ReactantCore = "0.1.16"
 Reactant_jll = "0.0.240"
 ScopedValues = "1.3.0"
 Scratch = "1.2"

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.15"
+version = "0.1.16"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -1,7 +1,7 @@
 module ReactantCore
 
 using ExpressionExplorer: ExpressionExplorer
-using MacroTools: MacroTools
+using MacroTools: MacroTools, @capture
 
 export @trace, within_compile, MissingTracedValue, promote_to_traced
 
@@ -174,26 +174,123 @@ macro trace(args...)
     track_numbers = track_numbers ? Number : Union{}
     expr = macroexpand(__module__, expr)
 
+    #! format: off
+    if @capture(
+        expr,
+        (
+            fnname_(call_args__) where {Typs__} = fnbody_
+        ) | (
+            fnname_(call_args__) = fnbody_
+        ) | (
+            function fnname_(call_args__)
+                fnbody_
+            end
+        ) | (
+            function fnname_(call_args__) where {Typs__}
+                fnbody_
+            end
+        )
+    )
+        return esc(trace_function_definition(__module__, expr))
+    end
+    #! format: on
+
     if Meta.isexpr(expr, :(=))
         if Meta.isexpr(expr.args[2], :if)
             return esc(trace_if_with_returns(expr; track_numbers))
         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(expr; track_numbers))
+
     Meta.isexpr(expr, :for) &&
         return (esc(trace_for(expr; track_numbers, checkpointing, mincut)))
+
     Meta.isexpr(expr, :while) &&
         return (esc(trace_while(expr; track_numbers, checkpointing, mincut)))
+
     return error(
-        "Only `if-elseif-else` blocks, `for` and `while` loops are currently supported by `@trace`",
+        "Only `if-elseif-else` blocks, function definitions, `function calls`, `for` and \
+        `while` loops are currently supported by `@trace`"
+    )
+end
+
+function get_argname(expr)
+    if Meta.isexpr(expr, :(::))
+        length(expr.args) == 2 && return expr.args[1], expr
+        @assert length(expr.args) == 1
+        var = gensym(:_)
+        return var, Expr(:(::), var, expr.args[1])
+    end
+    Meta.isexpr(expr, :kw) && return get_argname(expr.args[1])[1], expr
+    Meta.isexpr(expr, :(...)) && return expr, expr
+    @assert expr isa Symbol
+    if expr == :_
+        var = gensym(:_)
+        return var, var
+    end
+    return expr, expr
+end
+
+function trace_function_definition(mod, expr)
+    internal_fn = MacroTools.splitdef(expr)
+    orig_fname = internal_fn[:name]
+
+    isfunctor = Meta.isexpr(orig_fname, :(::))
+    fname = gensym(Symbol(orig_fname, :internal))
+    internal_fn[:name] = fname
+
+    if isfunctor
+        if length(orig_fname.args) == 1
+            sym_name = gensym("functor")
+            orig_fname = Expr(:(::), sym_name, orig_fname.args[1])
+        end
+        @assert length(orig_fname.args) == 2
+        insert!(internal_fn[:args], 1, :($orig_fname))
+    end
+
+    new_fn = MacroTools.splitdef(expr)
+
+    standardized_argnames = get_argname.(new_fn[:args])
+    argnames = first.(standardized_argnames)
+    new_fn[:args] = last.(standardized_argnames)
+
+    if isfunctor
+        insert!(argnames, 1, orig_fname.args[1])
+    end
+
+    if isempty(new_fn[:kwargs])
+        traced_call_expr = :($(traced_call)($(fname), $(argnames...)))
+        untraced_call_expr = :($(fname)($(argnames...)))
+    else
+        kws = first.(get_argname.(new_fn[:kwargs]))
+        traced_call_expr =
+            :($(traced_call)(Core.kwcall, (; $(kws...)), $(fname), $(argnames...)))
+        untraced_call_expr = :(Core.kwcall((; $(kws...)), $(fname), $(argnames...)))
+    end
+
+    new_fn[:name] = orig_fname
+    new_fn[:body] = :(
+        if $(within_compile)() && $(any)($(is_traced), ($(argnames...),))
+            return $(traced_call_expr)
+        else
+            return $(untraced_call_expr)
+        end
     )
+
+    return quote
+        $(MacroTools.combinedef(new_fn))
+        $(MacroTools.combinedef(internal_fn))
+    end
 end
 
 function trace_while(expr; track_numbers, mincut, checkpointing, first_arg=nothing)

test/control_flow.jl

@@ -820,6 +820,115 @@ end
     @test a_ra == a
 end
 
+@testset "trace function definitions" begin
+    # Basic traced function definition
+    @trace function traced_add(x, y)
+        return x .+ y
+    end
+
+    @testset "basic traced function" begin
+        a = rand(2, 3)
+        b = rand(2, 3)
+        a_ra = Reactant.to_rarray(a)
+        b_ra = Reactant.to_rarray(b)
+
+        # Should work the same as untraced when JIT compiled
+        @test @jit(traced_add(a_ra, b_ra)) ≈ traced_add(a, b)
+        ir = @code_hlo optimize = false traced_add(a_ra, b_ra)
+        func_names = [
+            String(Reactant.MLIR.IR.attr(op, "sym_name")) for
+            op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))
+        ]
+        @test any(contains("traced_add"), func_names)
+
+        # Should also work with regular arrays (outside compile context)
+        @test traced_add(a, b) ≈ a .+ b
+    end
+
+    # Traced function with typed arguments
+    @trace function traced_multiply(x::AbstractArray, y::AbstractArray)
+        return x .* y
+    end
+
+    @testset "traced function with type annotations" begin
+        a = rand(3, 3)
+        b = rand(3, 3)
+        a_ra = Reactant.to_rarray(a)
+        b_ra = Reactant.to_rarray(b)
+
+        @test @jit(traced_multiply(a_ra, b_ra)) ≈ traced_multiply(a, b)
+        ir = @code_hlo optimize = false traced_multiply(a_ra, b_ra)
+        func_names = [
+            String(Reactant.MLIR.IR.attr(op, "sym_name")) for
+            op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))
+        ]
+        @test any(contains("traced_multiply"), func_names)
+
+        @test traced_multiply(a, b) ≈ a .* b
+    end
+
+    @trace singleline(x) = x .+ 1
+
+    @testset "single line function" begin
+        a = rand(2, 3)
+        a_ra = Reactant.to_rarray(a)
+
+        @test @jit(singleline(a_ra)) ≈ singleline(a)
+        ir = @code_hlo optimize = false singleline(a_ra)
+        func_names = [
+            String(Reactant.MLIR.IR.attr(op, "sym_name")) for
+            op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))
+        ]
+        @test any(contains("singleline"), func_names)
+
+        @test singleline(a) ≈ a .+ 1
+    end
+
+    struct FunctorTest1{X}
+        x::X
+    end
+
+    @trace function (f::FunctorTest1)(y)
+        return f.x .+ y
+    end
+
+    @testset "function with functor" begin
+        a = rand(2, 3)
+        a_ra = Reactant.to_rarray(a)
+
+        fn1 = FunctorTest1(2.0f0)
+
+        @test @jit(fn1(a_ra)) ≈ fn1(a)
+        ir = @code_hlo optimize = false fn1(a_ra)
+        func_names = [
+            String(Reactant.MLIR.IR.attr(op, "sym_name")) for
+            op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))
+        ]
+        @test any(contains("FunctorTest1"), func_names)
+
+        @test fn1(a) ≈ a .+ 2.0f0
+    end
+
+    @trace function func_with_kwargs(x; y=1)
+        return x .+ y
+    end
+
+    @testset "function with kwargs" begin
+        a = rand(2, 3)
+        a_ra = Reactant.to_rarray(a)
+
+        @test @jit(func_with_kwargs(a_ra; y=2.0f0)) ≈ func_with_kwargs(a; y=2.0f0)
+        ir = @code_hlo optimize = false func_with_kwargs(a_ra; y=2.0f0)
+        func_names = [
+            String(Reactant.MLIR.IR.attr(op, "sym_name")) for
+            op in Reactant.MLIR.IR.OperationIterator(Reactant.MLIR.IR.body(ir))
+        ]
+        @test any(contains("kwcall"), func_names)
+
+        @test func_with_kwargs(a; y=2.0f0) ≈ a .+ 2.0f0
+    end
+end
+
 mutable struct TestClock{I}
     iteration::I
 end