If the lowered return of active really can't be inferred, just assume… (#2390)

* If the lowered return of active really can't be inferred, just assume float64 and emit a cast error

* set art

* fix

* fixup

* fix

Author: wsmoses

src/Enzyme.jl

@@ -344,11 +344,11 @@ Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))
 @inline function autodiff(
     mode::ReverseMode{ReturnPrimal,RuntimeActivity,RABI,Holomorphic,ErrIfFuncWritten},
     f::FA,
-    ::Type{A},
+    ::Type{A0},
     args::Vararg{Annotation,Nargs},
 ) where {
     FA<:Annotation,
-    A<:Annotation,
+    A0<:Annotation,
     ReturnPrimal,
     RuntimeActivity,
     RABI<:ABI,
@@ -369,13 +369,14 @@ Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))
 
     FTy = Core.Typeof(f.val)
 
-    rt = if A isa UnionAll
-        Compiler.primal_return_type(Reverse, FTy, tt)
+    rt, A = if A0 isa UnionAll
+        rt0 = Compiler.primal_return_type(Reverse, FTy, tt)
+        rt0, A0{rt0}
     else
-        eltype(A)
+        eltype(A0), A0
     end
 
-    if A <: Active
+    if A0 <: Active
         if (!allocatedinline(rt) || rt isa Union) && rt != Union{}
             forward, adjoint = autodiff_thunk(
                 ReverseModeSplit{
@@ -401,11 +402,11 @@ Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))
                 return adjoint(f, args..., tape)
             end
         end
-    elseif A <: Duplicated ||
-           A <: DuplicatedNoNeed ||
-           A <: BatchDuplicated ||
-           A <: BatchDuplicatedNoNeed ||
-           A <: BatchDuplicatedFunc
+    elseif A0 <: Duplicated ||
+           A0 <: DuplicatedNoNeed ||
+           A0 <: BatchDuplicated ||
+           A0 <: BatchDuplicatedNoNeed ||
+           A0 <: BatchDuplicatedFunc
         throw(ErrorException("Duplicated Returns not yet handled"))
     end
 
@@ -415,7 +416,7 @@ Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))
         Val(0)
     end
 
-    if (A <: Active && rt <: Complex) && rt != Union{}
+    if (A0 <: Active && rt <: Complex) && rt != Union{}
         if Holomorphic
             seen = IdDict()
             seen2 = IdDict()
@@ -497,7 +498,7 @@ Enzyme.autodiff(ReverseWithPrimal, x->x*x, Active(3.0))
         Val(RuntimeActivity),
     ) #=ShadowInit=#
 
-    if A <: Active
+    if A0 <: Active
         args = (args..., Compiler.default_adjoint(rt))
     end
     thunk(f, args...)

src/compiler.jl

@@ -2714,10 +2714,15 @@ function lower_convention(
     returnRoots = returnRoots !== nothing
 
     loweredReturn = RetActivity <: Active && (actualRetType === Any)
+    expected_RT = Nothing
     if loweredReturn
         @assert !sret
         @assert !returnRoots
-        RT = convert(LLVMType, eltype(RetActivity))
+        expected_RT = eltype(RetActivity)
+        if expected_RT === Any
+            expected_RT = Float64
+        end
+        RT = convert(LLVMType, expected_RT)
     end
 
     # TODO removed implications
@@ -3170,7 +3175,7 @@ function lower_convention(
                     return_attributes(wrapper_f),
                     StringAttribute(
                         "enzymejl_parmtype",
-                        string(convert(UInt, unsafe_to_pointer(eltype(RetActivity)))),
+                        string(convert(UInt, unsafe_to_pointer(expected_RT))),
                     ),
                 )
                 push!(
@@ -3181,7 +3186,7 @@ function lower_convention(
                     ),
                 )
                 ty = emit_jltypeof!(builder, res)
-                cmp = icmp!(builder, LLVM.API.LLVMIntEQ, ty, unsafe_to_llvm(builder, eltype(RetActivity)))
+                cmp = icmp!(builder, LLVM.API.LLVMIntEQ, ty, unsafe_to_llvm(builder, expected_RT))
                 cmpret = BasicBlock(wrapper_f, "ret")
                 failure = BasicBlock(wrapper_f, "fail")
                 br!(builder, cmp, cmpret, failure)
@@ -3194,7 +3199,7 @@ function lower_convention(
 
                 position!(builder, failure)
 
-                emit_error(builder, nothing, "Expected return type of primal to be "*string(eltype(RetActivity))*" but did not find a value of that type")
+                emit_error(builder, nothing, "Expected return type of primal to be "*string(expected_RT)*" but did not find a value of that type")
                 unreachable!(builder)
             else
                 push!(
@@ -3428,7 +3433,7 @@ function lower_convention(
         end
         throw(LLVM.LLVMException(msg))
     end
-    return wrapper_f, returnRoots, boxedArgs, loweredArgs, loweredReturn ? eltype(RetActivity) : actualRetType
+    return wrapper_f, returnRoots, boxedArgs, loweredArgs, loweredReturn ? expected_RT : actualRetType
 end
 
 using Random