Add interface for stashing context in Logic (#2624)

Author: vchuravy

Project.toml

@@ -45,7 +45,7 @@ CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 DynamicPPL = "0.35, 0.36, 0.37"
 EnzymeCore = "0.8.14"
-Enzyme_jll = "0.0.201"
+Enzyme_jll = "0.0.202"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6.2"
 LLVM = "9.1"

src/Enzyme.jl

@@ -132,6 +132,9 @@ Base.convert(::Type{API.CDerivativeMode}, ::ForwardMode) = API.DEM_ForwardMode
 
 function guess_activity end
 
+mutable struct EnzymeContext
+end
+
 include("logic.jl")
 include("analyses/type.jl")
 include("typetree.jl")

src/api.jl

@@ -996,6 +996,14 @@ function FreeLogic(logic)
     ccall((:FreeEnzymeLogic, libEnzyme), Cvoid, (EnzymeLogicRef,), logic)
 end
 
+function LogicSetExternalContext(logic, ctx)
+    ccall((:EnzymeLogicSetExternalContext, libEnzyme), Cvoid, (EnzymeLogicRef, Ptr{Cvoid}), logic, ctx)
+end
+
+function LogicGetExternalContext(logic)
+    ccall((:EnzymeLogicGetExternalContext, libEnzyme), Ptr{Cvoid}, (EnzymeLogicRef,), logic)
+end
+
 function EnzymeExtractReturnInfo(ret, data, existed)
     @assert length(data) == length(existed)
     ccall(

src/compiler.jl

@@ -13,6 +13,7 @@ import Enzyme:
     guess_activity,
     eltype,
     API,
+    EnzymeContext,
     TypeTree,
     typetree,
     TypeTreeTable,
@@ -2497,7 +2498,10 @@ function enzyme!(
         convert(API.CDIFFE_TYPE, rt)
     end
 
-    logic = Logic()
+    enzyme_context = EnzymeContext()
+    GC.@preserve enzyme_context begin
+    LLVM.@dispose logic  = Logic(enzyme_context) begin
+
     TA = TypeAnalysis(logic)
 
     retTT = if !isa(actualRetType, Union) &&
@@ -2753,7 +2757,10 @@ function enzyme!(
     if DumpPostEnzyme[]
         API.EnzymeDumpModuleRef(mod.ref)
     end
+
     return adjointf, augmented_primalf, TapeType
+    end # @dispose logic
+    end # GC.preserve enzyme_context
 end
 
 function get_subprogram(f::LLVM.Function)

src/logic.jl

@@ -2,14 +2,28 @@ import LLVM: refcheck
 
 LLVM.@checked struct Logic
     ref::API.EnzymeLogicRef
-    function Logic()
+    ctx::EnzymeContext
+    function Logic(ctx::EnzymeContext)
         ref = API.CreateLogic()
-        new(ref)
+        GC.@preserve ctx begin
+            API.LogicSetExternalContext(ref, Base.pointer_from_objref(ctx))
+            return new(ref, ctx)
+        end
     end
 end
 Base.unsafe_convert(::Type{API.EnzymeLogicRef}, logic::Logic) = logic.ref
 LLVM.dispose(logic::Logic) = API.FreeLogic(logic)
 
+function enzyme_context(logic::Logic)
+    return logic.ctx::EnzymeContext
+end
+
+function enzyme_context(logic::API.EnzymeLogicRef)
+    ptr = API.LogicGetExternalContext(logic)
+    @assert ptr != C_NULL
+    return unsafe_pointer_to_objref(ptr)::EnzymeContext
+end
+
 # typedef bool (*CustomRuleType)(int /*direction*/, CTypeTree * /*return*/,
 #                                CTypeTree * /*args*/, size_t /*numArgs*/,
 #                                LLVMValueRef)=T