Fix attribute copy (#2407)

* Fix attribute copy

* mark roots

* fix returnroots

Author: wsmoses

src/compiler.jl

@@ -4232,6 +4232,22 @@ end
 
         wrapper_f = LLVM.Function(mod, safe_name(LLVM.name(llvmfn) * "mustwrap"), FT)
 
+        for idx in 1:length(collect(parameters(llvmfn)))
+            for attr in collect(parameter_attributes(llvmfn, idx))
+                push!(parameter_attributes(wrapper_f, idx), attr)
+            end
+        end
+
+        for attr in collect(function_attributes(llvmfn))
+            push!(function_attributes(wrapper_f), attr)
+        end
+
+        for attr in collect(return_attributes(llvmfn))
+            push!(return_attributes(wrapper_f), attr)
+        end
+
+        mi, rt = enzyme_custom_extract_mi(primalf)
+
         let builder = IRBuilder()
             entry = BasicBlock(wrapper_f, "entry")
             position!(builder, entry)
@@ -4248,7 +4264,7 @@ end
             else
                 EnumAttribute("sret")
             end)
-            for idx in length(collect(parameters(llvmfn)))
+            for idx in 1:length(collect(parameters(llvmfn)))
                 for attr in collect(parameter_attributes(llvmfn, idx))
                     if kind(attr) == sretkind
                         LLVM.API.LLVMAddCallSiteAttribute(
@@ -4260,6 +4276,14 @@ end
                 end
             end
 
+            _, _, returnRoots = get_return_info(rt)
+            returnRoots = returnRoots !== nothing
+            if returnRoots
+                attr = StringAttribute("enzymejl_returnRoots", "")
+                push!(parameter_attributes(wrapper_f, 2), attr)
+                LLVM.API.LLVMAddCallSiteAttribute(res, LLVM.API.LLVMAttributeIndex(2), attr)
+            end
+
             if LLVM.return_type(FT) == LLVM.VoidType()
                 ret!(builder)
             else
@@ -4270,7 +4294,6 @@ end
         end
         attributes = function_attributes(wrapper_f)
         push!(attributes, StringAttribute("enzymejl_world", string(job.world)))
-        mi, rt = enzyme_custom_extract_mi(primalf)
         push!(
             attributes,
             StringAttribute("enzymejl_mi", string(convert(UInt, pointer_from_objref(mi)))),

src/compiler/utils.jl

@@ -400,6 +400,18 @@ Base.@assume_effects :removable :foldable :nothrow function has_fn_attr(fn::LLVM
     return false
 end
 
+Base.@assume_effects :removable :foldable :nothrow function has_arg_attr(fn::LLVM.Function, i::Int, attr::LLVM.StringAttribute)::Bool
+    ekind = LLVM.kind(attr)
+    for attr in collect(parameter_attributes(fn, i))
+        if attr isa LLVM.StringAttribute
+            if kind(attr) == ekind
+                return true
+            end
+        end
+    end
+    return false
+end
+
 function eraseInst(bb::LLVM.BasicBlock, @nospecialize(inst::LLVM.Instruction))
     @static if isdefined(LLVM, Symbol("erase!"))
         LLVM.erase!(inst)

src/rules/customrules.jl

@@ -1528,7 +1528,7 @@ end
     fop = called_operand(orig)::LLVM.Function
     for (i, v) in enumerate(operands(orig)[1:end-1])
         if v == val
-            if !has_fn_attr(fop, StringAttribute("enzymejl_returnRoots"))
+            if !has_arg_attr(fop, i, StringAttribute("enzymejl_returnRoots"))
                 non_rooting_use = true
                 break
             end