Fix kw rrule closure argument index (#1543)

Author: wsmoses

src/rules/customrules.jl

@@ -734,7 +734,7 @@ function enzyme_custom_common_rev(forward::Bool, B, orig::LLVM.CallInst, gutils,
         end
     end
 
-    # push!(function_attributes(llvmf), EnumAttribute("alwaysinline", 0))
+    push!(function_attributes(llvmf), EnumAttribute("alwaysinline", 0))
 
     needsTape = !isghostty(TapeT) && !Core.Compiler.isconstType(TapeT)
 
@@ -765,11 +765,11 @@ function enzyme_custom_common_rev(forward::Bool, B, orig::LLVM.CallInst, gutils,
         funcTy = rev_TT.parameters[isKWCall ? 4 : 2] 
         if needsTape
             @assert tape != C_NULL
-            tape_idx = 1+(kwtup!==nothing && !isghostty(kwtup))+(isKWCall && !isghostty(rev_TT.parameters[4])) + !isghostty(funcTy)
-            trueidx = tape_idx+(sret !== nothing)+(returnRoots !== nothing)+swiftself+(RT <: Active)
+            tape_idx = 1+(kwtup!==nothing && !isghostty(kwtup)) + !isghostty(funcTy) 
+            trueidx = tape_idx+(sret !== nothing)+(returnRoots !== nothing)+swiftself + (RT <: Active)
             innerTy = value_type(parameters(llvmf)[trueidx])
             if innerTy != value_type(tape)
-                if isabstracttype(TapeT) || TapeT == Tuple || TapeT.layout == C_NULL
+                if isabstracttype(TapeT) || TapeT == Tuple || TapeT.layout == C_NULL || TapeT == Array
                     msg = sprint() do io
                         println(io, "Enzyme : mismatch between innerTy $innerTy and tape type $(value_type(tape))")
                         println(io, "tape_idx=", tape_idx)
@@ -831,7 +831,7 @@ function enzyme_custom_common_rev(forward::Bool, B, orig::LLVM.CallInst, gutils,
             if any_jltypes(llty)
                 emit_writebarrier!(B, get_julia_inner_types(B, al0, val))
             end
-            insert!(args, 1+(!isghostty(funcTy))+(kwtup!==nothing && !isghostty(kwtup))+(isKWCall && !isghostty(rev_TT.parameters[4])), al)
+            insert!(args, 1+(!isghostty(funcTy))+(kwtup!==nothing && !isghostty(kwtup)), al)
         end
     end
 

test/kwrrules.jl

@@ -111,5 +111,47 @@ g4(x, y) = f_kw4(x; y)
 @test autodiff(Reverse, g4, Active(2.0), Const(42.0))[1][1] ≈ 42004.0
 @test_throws Enzyme.Compiler.EnzymeRuntimeException autodiff(Reverse, g4, Active(2.0), Active(42.0))[1]
 
+struct Closure2
+    v::Vector{Float64}
+    str::String
+end
+
+function (cl::Closure2)(x; width=7)
+    val = cl.v[1] * x * width
+    cl.v[1] = 0.0
+    return val
+end
+
+function wrapclos(cl, x)
+    cl(x; width=9)
+end
+
+function EnzymeRules.augmented_primal(config::ConfigWidth{1}, func::Const{Closure2},
+    ::Type{<:Active}, args::Vararg{Active,N}; width=7) where {N}
+    vec = copy(func.val.v)
+    pval = func.val(args[1].val)
+    primal = if EnzymeRules.needs_primal(config)
+        pval
+    else
+        nothing
+    end
+    return AugmentedReturn(primal, nothing, vec)
+end
+
+function EnzymeRules.reverse(config::ConfigWidth{1}, func::Const{Closure2},
+    dret::Active, tape, args::Vararg{Active,N}; width=7) where {N}
+    dargs = ntuple(Val(N)) do i
+        7 * args[1].val * dret.val + tape[1] * 1000 + width * 100000
+    end
+    return dargs
+end
+
+@testset "KWClosure rule" begin
+    cl = Closure2([3.14], "3.14")
+    res = autodiff(Reverse, wrapclos, Active, Const(cl), Active(2.7))[1][2]
+    @test res ≈ 7 * 2.7 + 3.14 * 1000 + 9 * 100000
+    @test cl.v[1] ≈ 0.0
+end
+
 end # KWReverseRules