fix: handle traced array returns inside objects

Author: avik-pal

src/Compiler.jl

@@ -315,7 +315,7 @@ function compile_mlir!(mod, f, args; optimize::Union{Bool,Symbol}=true)
 
     toolkit = ""
     if isdefined(Reactant_jll, :ptxas_path)
-	 toolkit = Reactant_jll.ptxas_path[1:end-length("/bin/ptxas")]
+        toolkit = Reactant_jll.ptxas_path[1:(end - length("/bin/ptxas"))]
     end
     kern = "lower-kernel{toolkitPath=$toolkit cuLaunchKernelPtr=$(cuLaunch[]) cuModuleLoadDataPtr=$(cuModule[]) cuModuleGetFunctionPtr=$(cuFunc[])}"
     if optimize === :all
@@ -329,7 +329,7 @@ function compile_mlir!(mod, f, args; optimize::Union{Bool,Symbol}=true)
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
                     opt_passes,
-		    kern
+                    kern,
                 ],
                 ',',
             ),
@@ -385,7 +385,7 @@ function compile_mlir!(mod, f, args; optimize::Union{Bool,Symbol}=true)
                     "remove-unnecessary-enzyme-ops",
                     "enzyme-simplify-math",
                     opt_passes,
-		    kern
+                    kern,
                 ],
                 ',',
             ),
@@ -394,7 +394,7 @@ function compile_mlir!(mod, f, args; optimize::Union{Bool,Symbol}=true)
         run_pass_pipeline!(mod, join([opt_passes, "enzyme-batch", opt_passes], ","))
         run_pass_pipeline!(mod, "enzyme,arith-raise{stablehlo=true}"; enable_verifier=false)
         run_pass_pipeline!(
-            mod, "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math,"*kern
+            mod, "canonicalize,remove-unnecessary-enzyme-ops,enzyme-simplify-math," * kern
         )
     elseif optimize !== :none
         error("Invalid optimize option: $(Meta.quot(optimize))")
@@ -638,10 +638,30 @@ function codegen_unflatten!(
             end
 
             # unroll path tree
-            for p in path
+            for p in path[1:(end - 1)]
                 unflatcode = :(traced_getfield($unflatcode, $(Meta.quot(p))))
             end
-            unflatcode = :($unflatcode.data = $concrete_res_name)
+            final_val = gensym("final_val")
+            unflatcode = quote
+                $final_val = traced_getfield($unflatcode, $(Meta.quot(path[end])))
+                if $final_val isa TracedRArray
+                    setfield!(
+                        $unflatcode,
+                        $(Meta.quot(path[end])),
+                        ConcreteRArray{eltype($final_val),ndims($final_val)}(
+                            $concrete_res_name, size($final_val)
+                        ),
+                    )
+                elseif $final_val isa TracedRNumber
+                    setfield!(
+                        $unflatcode,
+                        $(Meta.quot(path[end])),
+                        ConcreteRNumber{eltype($final_val)}($concrete_res_name),
+                    )
+                else
+                    setfield!($final_val, :data, $concrete_res_name)
+                end
+            end
 
             push!(unflatten_code, unflatcode)
         end

src/Interpreter.jl

@@ -235,7 +235,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 = Reactant.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
     )
 
@@ -302,7 +302,7 @@ function overload_autodiff(
                 cst = MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1)
                 push!(ad_inputs, cst)
             end
-        else
+        elseif TracedUtils.has_argidx(a)
             idx, path = TracedUtils.get_argidx(a)
             if idx == 1 && fnwrap
                 act = act_from_type(f, reverse, true)
@@ -322,6 +322,12 @@ function overload_autodiff(
                 end
                 TracedUtils.push_val!(ad_inputs, args[idx].dval, path[3:end])
             end
+        else
+            act = act_from_type(Enzyme.Const, reverse, true)
+            push!(ret_activity, act)
+            if act != enzyme_out && act != enzyme_outnoneed
+                continue
+            end
         end
     end
 
@@ -385,7 +391,7 @@ function overload_autodiff(
                 end
                 residx += 1
             end
-        else
+        elseif TracedUtils.has_argidx(a)
             idx, path = TracedUtils.get_argidx(a)
             if idx == 1 && fnwrap
                 TracedUtils.set!(
@@ -405,6 +411,11 @@ function overload_autodiff(
                 )
                 residx += 1
             end
+        else
+            TracedUtils.set!(
+                a, (), TracedUtils.transpose_val(MLIR.IR.result(res, residx))
+            )
+            residx += 1
         end
     end
 

src/TracedUtils.jl

@@ -341,6 +341,18 @@ function get_argidx(x)
     throw(AssertionError("No path found for $x"))
 end
 
+function has_argidx(x)
+    for path in x.paths
+        if length(path) == 0
+            continue
+        end
+        if path[1] == :args
+            return true
+        end
+    end
+    return false
+end
+
 function set!(x, path, tostore; emptypath=false)
     for p in path
         x = Reactant.Compiler.traced_getfield(x, p)

src/utils.jl

@@ -229,9 +229,9 @@ function make_oc(
     sig::Type, rt::Type, src::Core.CodeInfo, nargs::Int, isva::Bool, f::Any
 )::Core.OpaqueClosure
     key = (sig, rt, src, nargs, isva, f)
-    if haskey(oc_captures, key)
-        return oc_captures[key]
-    else
+    # if haskey(oc_captures, key)
+    #     return oc_captures[key]
+    # else
         ores = ccall(
             :jl_new_opaque_closure_from_code_info,
             Any,
@@ -250,7 +250,7 @@ function make_oc(
         )::Core.OpaqueClosure
         oc_captures[key] = ores
         return ores
-    end
+    # end
 end
 
 function safe_print(name, x)