Cleanup solve rules and bigfloat (#1641)

* Cleanup solve rules and bigfloat

* fix

* Fix aug fwd msg

* fix cache ty

Author: wsmoses

src/compiler.jl

@@ -1999,7 +1999,7 @@ function julia_error(cstr::Cstring, val::LLVM.API.LLVMValueRef, errtype::API.Err
     end
 
     if errtype == API.ET_NoDerivative
-        if occursin("No create nofree of empty function", msg) || occursin("No forward mode derivative found for", msg) || occursin("No augmented forward mode derivative found for", msg) || occursin("No reverse pass found", msg)
+        if occursin("No create nofree of empty function", msg) || occursin("No forward mode derivative found for", msg) || occursin("No augmented forward pass", msg) || occursin("No reverse pass found", msg)
             ir = nothing
         end
         exc = NoDerivativeException(msg, ir, bt)

src/internal_rules.jl

@@ -390,11 +390,20 @@ else
     }
 end
 
-    cache = NamedTuple{(Symbol("1"),Symbol("2"), Symbol("3"), Symbol("4")), Tuple{typeof(res), typeof(dres), UT, typeof(cache_b)}}(
+    cache = NamedTuple{(Symbol("1"),Symbol("2"), Symbol("3"), Symbol("4")), Tuple{
+        eltype(RT),
+        EnzymeRules.needs_shadow(config) ? (EnzymeRules.width(config) == 1 ? eltype(RT) : NTuple{EnzymeRules.width(config), eltype(RT)}) : Nothing,
+        UT,
+        typeof(cache_b)
+        }}(
         (cache_res, dres, cache_A, cache_b)
     )
 
-    return EnzymeRules.AugmentedReturn{typeof(retres), typeof(dres), typeof(cache)}(retres, dres, cache)
+    return EnzymeRules.AugmentedReturn{
+        EnzymeRules.needs_primal(config) ? eltype(RT) : Nothing,
+        EnzymeRules.needs_shadow(config) ? (EnzymeRules.width(config) == 1 ? eltype(RT) : NTuple{EnzymeRules.width(config), eltype(RT)}) : Nothing,
+        typeof(cache)
+    }(retres, dres, cache)
 end
 
 function EnzymeRules.reverse(config, func::Const{typeof(\)}, ::Type{RT}, cache, A::Annotation{<:Array}, b::Annotation{<:Array}) where RT

src/typetree.jl

@@ -111,6 +111,10 @@ function typetree_inner(::Type{Float64}, ctx, dl, seen::TypeTreeTable)
     return TypeTree(API.DT_Double, -1, ctx)
 end
 
+function typetree_inner(::Type{BigFloat}, ctx, dl, seen::TypeTreeTable)
+    return TypeTree()
+end
+
 function typetree_inner(::Type{T}, ctx, dl, seen::TypeTreeTable) where {T<:AbstractFloat}
     GPUCompiler.@safe_warn "Unknown floating point type" T
     return TypeTree()