Support Julia 1.11 (#1372)

* Test against v1.11

* WIP: adapt to 1.11 changes

* fix constructor

* Update interpreter.jl

* add cache_token

* fixup! add cache_token

* Apply suggestions from code review

* fixup! Apply suggestions from code review

---------

Co-authored-by: William Moses <[email protected]>

Author: vchuravy

.github/workflows/CI.yml

@@ -26,6 +26,7 @@ jobs:
           - '1.8'
           - '1.9'
           - '1.10'
+          - ~1.11.0-0
           - 'nightly'
         os:
           - ubuntu-20.04
@@ -86,6 +87,11 @@ jobs:
             libEnzyme: packaged
             version: '1.10'
             assertions: true
+          - os: ubuntu-20.04
+            arch: x64
+            libEnzyme: packaged
+            version: '1.11'
+            assertions: true
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
@@ -170,6 +176,7 @@ jobs:
           - '1.8'
           - '1.9'
           - '1.10'
+          - ~1.11.0-0
           - 'nightly'
         os:
           - ubuntu-latest

src/compiler.jl

@@ -2992,8 +2992,27 @@ GPUCompiler.runtime_module(::CompilerJob{<:Any,<:AbstractEnzymeCompilerParams})
 GPUCompiler.runtime_slug(job::CompilerJob{EnzymeTarget}) = "enzyme"
 
 # provide a specific interpreter to use.
+if VERSION >= v"1.11.0-DEV.1552"
+struct EnzymeCacheToken
+    target_type::Type
+    always_inline
+    method_table::Core.MethodTable
+    param_type::Type
+    mode::API.CDerivativeMode
+end
+
+GPUCompiler.ci_cache_token(job::CompilerJob{<:Any,<:AbstractEnzymeCompilerParams}) =
+    EnzymeCacheToken(
+        typeof(job.config.target), job.config.always_inline, GPUCompiler.method_table(job),
+        typeof(job.config.params), job.config.params.mode,
+    )
+
+GPUCompiler.get_interpreter(job::CompilerJob{<:Any,<:AbstractEnzymeCompilerParams}) =
+    Interpreter.EnzymeInterpreter(GPUCompiler.ci_cache_token(job), GPUCompiler.method_table(job), job.world, job.config.params.mode)
+else
 GPUCompiler.get_interpreter(job::CompilerJob{<:Any,<:AbstractEnzymeCompilerParams}) =
     Interpreter.EnzymeInterpreter(GPUCompiler.ci_cache(job), GPUCompiler.method_table(job), job.world, job.config.params.mode)
+end
 
 include("compiler/passes.jl")
 include("compiler/optimize.jl")

src/compiler/interpreter.jl

@@ -1,12 +1,28 @@
 module Interpreter
 import Enzyme: API
 using Core.Compiler: AbstractInterpreter, InferenceResult, InferenceParams, InferenceState, OptimizationParams, MethodInstance
-using GPUCompiler: CodeCache, WorldView, @safe_debug
+using GPUCompiler: @safe_debug
+if VERSION < v"1.11.0-DEV.1552"
+    using GPUCompiler: CodeCache, WorldView, @safe_debug
+end
+const HAS_INTEGRATED_CACHE = VERSION >= v"1.11.0-DEV.1552"
+
 import ..Enzyme
 import ..EnzymeRules
 
+@static if VERSION ≥ v"1.11.0-DEV.1498"
+    import Core.Compiler: get_inference_world
+    using Base: get_world_counter
+else
+    import Core.Compiler: get_world_counter, get_world_counter as get_inference_world
+end
+
 struct EnzymeInterpreter <: AbstractInterpreter
-    global_cache::CodeCache
+@static if HAS_INTEGRATED_CACHE
+    token::Any
+else
+    code_cache::CodeCache
+end
     method_table::Union{Nothing,Core.MethodTable}
 
     # Cache of inference results for this particular interpreter
@@ -19,34 +35,38 @@ struct EnzymeInterpreter <: AbstractInterpreter
     opt_params::OptimizationParams
 
     mode::API.CDerivativeMode
+end
 
-    function EnzymeInterpreter(cache::CodeCache, mt::Union{Nothing,Core.MethodTable}, world::UInt, mode::API.CDerivativeMode)
-        @assert world <= Base.get_world_counter()
+function EnzymeInterpreter(cache_or_token, mt::Union{Nothing,Core.MethodTable}, world::UInt, mode::API.CDerivativeMode)
+    @assert world <= Base.get_world_counter()
 
-        return new(
-            cache,
-            mt,
+    return EnzymeInterpreter(
+        cache_or_token,
+        mt,
 
-            # Initially empty cache
-            Vector{InferenceResult}(),
+        # Initially empty cache
+        Vector{InferenceResult}(),
 
-            # world age counter
-            world,
+        # world age counter
+        world,
 
-            # parameters for inference and optimization
-            InferenceParams(unoptimize_throw_blocks=false),
-            VERSION >= v"1.8.0-DEV.486" ? OptimizationParams() :
-                                          OptimizationParams(unoptimize_throw_blocks=false),
-            mode
-        )
-    end
+        # parameters for inference and optimization
+        InferenceParams(unoptimize_throw_blocks=false),
+        VERSION >= v"1.8.0-DEV.486" ? OptimizationParams() :
+                                        OptimizationParams(unoptimize_throw_blocks=false),
+        mode
+    )
 end
 
 Core.Compiler.InferenceParams(interp::EnzymeInterpreter) = interp.inf_params
 Core.Compiler.OptimizationParams(interp::EnzymeInterpreter) = interp.opt_params
-Core.Compiler.get_world_counter(interp::EnzymeInterpreter) = interp.world
+get_inference_world(interp::EnzymeInterpreter) = interp.world
 Core.Compiler.get_inference_cache(interp::EnzymeInterpreter) = interp.local_cache
-Core.Compiler.code_cache(interp::EnzymeInterpreter) = WorldView(interp.global_cache, interp.world)
+@static if HAS_INTEGRATED_CACHE
+    Core.Compiler.cache_owner(interp::EnzymeInterpreter) = interp.token
+else
+    Core.Compiler.code_cache(interp::EnzymeInterpreter) = WorldView(interp.code_cache, interp.world)
+end
 
 # No need to do any locking since we're not putting our results into the runtime cache
 Core.Compiler.lock_mi_inference(interp::EnzymeInterpreter, mi::MethodInstance) = nothing