Add differential use handler

src/api.jl

@@ -230,6 +230,8 @@
 EnzymeRegisterCallHandler(name, fwdhandle, revhandle) = ccall((:EnzymeRegisterCallHandler, libEnzyme), Cvoid, (Cstring, CustomAugmentedForwardPass, CustomReversePass), name, fwdhandle, revhandle)
 EnzymeRegisterFwdCallHandler(name, fwdhandle) = ccall((:EnzymeRegisterFwdCallHandler, libEnzyme), Cvoid, (Cstring, CustomForwardPass), name, fwdhandle)
 
+const CustomDiffUse = Ptr{Cvoid}
+EnzymeRegisterDiffUseCallHandler(name, handle) = ccall((:EnzymeRegisterDiffUseCallHandler, libEnzyme), Cvoid, (Cstring, CustomDiffUse), name, handle)
 EnzymeSetCalledFunction(ci::LLVM.CallInst, fn::LLVM.Function, toremove) = ccall((:EnzymeSetCalledFunction, libEnzyme), Cvoid, (LLVMValueRef, LLVMValueRef, Ptr{Int64}, Int64), ci, fn, toremove, length(toremove))
 EnzymeCloneFunctionWithoutReturnOrArgs(fn::LLVM.Function, keepret, args) = ccall((:EnzymeCloneFunctionWithoutReturnOrArgs, libEnzyme), LLVMValueRef, (LLVMValueRef,UInt8,Ptr{Int64}, Int64), fn, keepret, args, length(args))
 EnzymeGetShadowType(width, T) = ccall((:EnzymeGetShadowType, libEnzyme), LLVMTypeRef, (UInt64,LLVMTypeRef), width, T)
@@ -258,7 +260,7 @@
 
 EnzymeGradientUtilsAllocAndGetTypeTree(gutils, val) = ccall((:EnzymeGradientUtilsAllocAndGetTypeTree, libEnzyme), CTypeTreeRef, (EnzymeGradientUtilsRef,LLVMValueRef), gutils, val)
 
-EnzymeGradientUtilsGetUncacheableArgs(gutils, orig, uncacheable, size) = ccall((:EnzymeGradientUtilsGetUncacheableArgs, libEnzyme), Cvoid, (EnzymeGradientUtilsRef,LLVMValueRef, Ptr{UInt8}, UInt64), gutils, orig, uncacheable, size)
+EnzymeGradientUtilsGetUncacheableArgs(gutils, orig, uncacheable, size) = ccall((:EnzymeGradientUtilsGetUncacheableArgs, libEnzyme), UInt8, (EnzymeGradientUtilsRef,LLVMValueRef, Ptr{UInt8}, UInt64), gutils, orig, uncacheable, size)
 
 EnzymeGradientUtilsGetDiffeType(gutils, op, isforeign) = ccall((:EnzymeGradientUtilsGetDiffeType, libEnzyme), CDIFFE_TYPE, (EnzymeGradientUtilsRef,LLVMValueRef, UInt8), gutils, op, isforeign)
 

src/gradientutils.jl

@@ -24,7 +24,9 @@ function get_shadow_type(gutils::GradientUtils, T::LLVM.Type)
 end
 function get_uncacheable(gutils::GradientUtils, orig::LLVM.CallInst)    
     uncacheable = Vector{UInt8}(undef, length(collect(LLVM.operands(orig)))-1)
-    API.EnzymeGradientUtilsGetUncacheableArgs(gutils, orig, uncacheable, length(uncacheable))
+    if API.EnzymeGradientUtilsGetUncacheableArgs(gutils, orig, uncacheable, length(uncacheable)) != 1
+        uncacheable .= 1
+    end
     return uncacheable
 end
 

src/rules/customrules.jl

@@ -1,5 +1,5 @@
 
-function enzyme_custom_setup_args(B, orig, gutils, mi, RT, reverse, isKWCall)
+function enzyme_custom_setup_args(B, orig::LLVM.CallInst, gutils::GradientUtils, mi, @nospecialize(RT), reverse::Bool, isKWCall::Bool)
     ops = collect(operands(orig))
     called = ops[end]
     ops = ops[1:end-1]
@@ -207,7 +207,7 @@ function enzyme_custom_setup_args(B, orig, gutils, mi, RT, reverse, isKWCall)
     return args, activity, (overwritten...,), actives, kwtup
 end
 
-function enzyme_custom_setup_ret(gutils, orig, mi, RealRt)
+function enzyme_custom_setup_ret(gutils::GradientUtils, orig::LLVM.CallInst, mi, @nospecialize(RealRt))
     width = get_width(gutils)
     mode = get_mode(gutils)
 
@@ -216,7 +216,23 @@ function enzyme_custom_setup_ret(gutils, orig, mi, RealRt)
     needsShadowP = Ref{UInt8}(0)
     needsPrimalP = Ref{UInt8}(0)
 
-    activep = API.EnzymeGradientUtilsGetReturnDiffeType(gutils, orig, needsPrimalP, needsShadowP, mode)
+    # Conditionally use the get return. This is done because EnzymeGradientUtilsGetReturnDiffeType
+    # calls differential use analysis to determine needsprimal/shadow. However, since now this function
+    # is used as part of differential use analysis, we need to avoid an ininite recursion. Thus use
+    # the version without differential use if actual unreachable results are not available anyways.
+    uncacheable = Vector{UInt8}(undef, length(collect(LLVM.operands(orig)))-1)
+    activep = if mode == API.DEM_ForwardMode || API.EnzymeGradientUtilsGetUncacheableArgs(gutils, orig, uncacheable, length(uncacheable)) == 1
+        API.EnzymeGradientUtilsGetReturnDiffeType(gutils, orig, needsPrimalP, needsShadowP, mode)
+    else
+        actv = API.EnzymeGradientUtilsGetDiffeType(gutils, orig, false)
+        if !isghostty(RealRt)
+            needsPrimalP[] = 1
+            if actv == API.DFT_DUP_ARG || actv == API.DFT_DUP_NONEED
+                needsShadowP[] = 1
+            end
+        end
+        actv
+    end
     needsPrimal = needsPrimalP[] != 0
     origNeedsPrimal = needsPrimal
     _, sret, _ = get_return_info(RealRt)
@@ -479,7 +495,7 @@ function enzyme_custom_fwd(B, orig, gutils, normalR, shadowR)
     return false
 end
 
-@inline function aug_fwd_mi(orig, gutils)
+@inline function aug_fwd_mi(orig::LLVM.CallInst, gutils::GradientUtils)
     width = get_width(gutils)
 
     # 1) extract out the MI from attributes
@@ -568,7 +584,7 @@ function enzyme_custom_common_rev(forward::Bool, B, orig::LLVM.CallInst, gutils,
     end
 
     C = EnzymeRules.Config{Bool(needsPrimal), Bool(needsShadowJL), Int(width), overwritten}
-    
+
     alloctx = LLVM.IRBuilder()
     position!(alloctx, LLVM.BasicBlock(API.EnzymeGradientUtilsAllocationBlock(gutils)))
 
@@ -949,3 +965,12 @@ function enzyme_custom_rev(B, orig, gutils, tape)
     enzyme_custom_common_rev(#=forward=#false, B, orig, gutils, #=normalR=#C_NULL, #=shadowR=#C_NULL, #=tape=#tape)
     return nothing
 end
+
+function enzyme_custom_diffuse(orig, gutils, val, isshadow, mode)
+    # use default
+    if is_constant_value(gutils, orig) && is_constant_inst(gutils, orig) && !has_aug_fwd_rule(orig, gutils)
+        return (false, true)
+    end
+    # don't use default and always require the arg
+    return (true, false)
+end

src/rules/llvmrules.jl

@@ -1149,13 +1149,24 @@
 end
 
 macro fwdfunc(f)
    :(@cfunction((B, OrigCI, gutils, normalR, shadowR) -> begin
      UInt8($f(LLVM.IRBuilder(B), LLVM.CallInst(OrigCI), GradientUtils(gutils), normalR, shadowR)::Bool)
     end, UInt8, (LLVM.API.LLVMBuilderRef, LLVM.API.LLVMValueRef, API.EnzymeGradientUtilsRef, Ptr{LLVM.API.LLVMValueRef}, Ptr{LLVM.API.LLVMValueRef})
     ))
 end
 
+
+macro diffusefunc(f)
+   :(@cfunction((OrigCI, gutils, val, shadow, mode, useDefault) -> begin
+     res = $f(LLVM.CallInst(OrigCI), GradientUtils(gutils), LLVM.Value(val), shadow != 0, mode)::Tuple{Bool, Bool}
+     unsafe_store!(useDefault, UInt8(res[2]))
+     UInt8(res[1])
+    end, UInt8, (LLVM.API.LLVMValueRef, API.EnzymeGradientUtilsRef, LLVM.API.LLVMValueRef, UInt8, API.CDerivativeMode, Ptr{UInt8})
+    ))
+end
+
 @noinline function register_llvm_rules()
+    API.EnzymeRegisterDiffUseCallHandler("enzyme_custom", @diffusefunc(enzyme_custom_diffuse))
     register_handler!(
         ("julia.call",),
         @augfunc(jlcall_augfwd),