Mock Enzyme plugin

Author: vchuravy

test/plugin_testsetup.jl

@@ -69,5 +69,89 @@ function inlining_handler(meta::InlineStateMeta, interp::GPUInterpreter, @nospec
     return nothing
 end
 
+struct MockEnzymeMeta end
 
+function autodiff end
+
+import GPUCompiler: DeferredCallInfo
+struct AutodiffCallInfo <: CallInfo
+    info::DeferredCallInfo
+end
+
+function abstract_call_known(meta::MockEnzymeMeta, interp::GPUInterpreter, @nospecialize(f),
+                             arginfo::ArgInfo, si::StmtInfo, sv::AbsIntState, max_methods::Int)
+    (; fargs, argtypes) = arginfo
+
+    if f === autodiff
+        if length(argtypes) >= 1
+            @static if VERSION < v"1.11.0-"
+                return CallMeta(Union{}, Effects(), NoCallInfo())
+            else
+                return CallMeta(Union{}, Union{}, Effects(), NoCallInfo())
+            end
+        end
+
+        other_fargs = fargs === nothing ? nothing : fargs[2:end]
+        other_arginfo = ArgInfo(other_fargs, argtypes[2:end])
+        call = Core.Compiler.abstract_call(interp, other_arginfo, si, sv, max_methods)
+        callinfo = DeferredCallInfo(MockEnzymeMeta(), call.rt, call.info)
+
+        # Real Enzyme must compute `rt` and `exct` according to enzyme semantics
+        # and likely perform a unwrapping of fargs...
+        rt = Nothing
+
+        # TODO: Edges? Effects?
+        @static if VERSION < v"1.11.0-"
+            return CallMeta(rt, call.effects, AutodiffCallInfo(callinfo))
+        else
+            return CallMeta(rt, call.exct, call.effects, AutodiffCallInfo(callinfo))
+        end
+    end
+
+    return nothing
+end
+
+const FlagType = VERSION >= v"1.11.0-" ? UInt32 : UInt8
+function CC.handle_call!(todo::Vector{Pair{Int,Any}}, ir::CC.IRCode, idx::CC.Int,
+                         stmt::Expr, info::AutodiffCallInfo, flag::FlagType,
+                         sig::CC.Signature, state::CC.InliningState)
+    # Goal:
+    # The IR we want to inline here is:
+    # unpack the args ...
+    # ptr = gpuc.deferred(MockEnzymeMeta(), f, primal_args...) 
+    # ret = ccall("extern __autodiff", llvmcall, RT, Tuple{Ptr{Cvoid, args...}}, ptr, adjoint_args...)
+    
+    push!(todo, idx=>(info::AutoDiffTodo))
+    
+    # # 1. Since Julia's inliner goes bottom up we need to pretend that we inlined the deferred call
+    # deferred_info = info.info
+    # # TODO: This is code duplication is unfortunate...
+    # minfo = deferred_info.info
+    # results = minfo.results
+    # if length(results.matches) != 1
+    #     return nothing
+    # end
+    # match = only(results.matches)
+
+    # # lookup the target mi with correct edge tracking
+    # case = CC.compileable_specialization(match, CC.Effects(), CC.InliningEdgeTracker(state),
+    #                                      info)
+    # @assert case isa CC.InvokeCase
+    # @assert stmt.head === :call
+
+    # stmt = Expr(:foreigncall, 
+    #     "extern gpuc.lookup",
+    #     Ptr{Cvoid},
+    #     Core.svec(Any, Any, Any, match.spec_types.parameters[2:end]...), # Must use Any for MethodInstance or ftype
+    #     0,
+    #     QuoteNode(:llvmcall),
+    #     info.meta,
+    #     case.invoke,
+    #     stmt.args[3:end]...
+    # )
+
+    # # 2. Form call to `__autodiff`
+    # # TODO!
+
+    return nothing
 end