Support Julia's internal cache (#545)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: vchuravy

.github/workflows/ci.yml

@@ -18,7 +18,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        version: ['1.8', '1.9', '1.10'] # 'nightly'
+        version: ['1.8', '1.9', '1.10', '^1.11.0-alpha1'] # 'nightly'
         os: [ubuntu-latest, macOS-latest, windows-latest]
         arch: [x64]
         llvm_args: ['']
@@ -36,6 +36,18 @@ jobs:
             os: 'windows-latest'
             arch: 'x64'
             llvm_args: '--opaque-pointers'
+          - version: '^1.11.0-alpha1'
+            os: 'ubuntu-latest'
+            arch: 'x64'
+            llvm_args: '--opaque-pointers'
+          - version: '^1.11.0-alpha1'
+            os: 'macOS-latest'
+            arch: 'x64'
+            llvm_args: '--opaque-pointers'
+          - version: '^1.11.0-alpha1'
+            os: 'windows-latest'
+            arch: 'x64'
+            llvm_args: '--opaque-pointers'
           #- version: 'nightly'
           #  os: 'ubuntu-latest'
           #  arch: 'x64'
@@ -96,9 +108,9 @@ jobs:
         include:
           - version: '1.11'
             pipeline: 'julia-release-1-dot-11'
-            build: 'x86_64-linux-gnu'
+            build: 'x86_64-linux-gnuassert'
             branch: 'release-1.11'
-          #- version: '1.12'
+          #- version: 'master'
           #  pipeline: 'julia-master'
           #  build: 'x86_64-linux-gnuassert'
           #  branch: 'master'

.gitignore

@@ -1,2 +1,3 @@
 test/Manifest.toml
 Manifest.toml
+Manifest-*.toml

src/interface.jl

@@ -176,10 +176,17 @@ runtime_module(@nospecialize(job::CompilerJob)) = error("Not implemented")
 isintrinsic(@nospecialize(job::CompilerJob), fn::String) = false
 
 # provide a specific interpreter to use.
+if VERSION >= v"1.11.0-DEV.1552"
+get_interpreter(@nospecialize(job::CompilerJob)) =
+    GPUInterpreter(job.world; method_table=method_table(job),
+                   token=ci_cache_token(job), inf_params=inference_params(job),
+                   opt_params=optimization_params(job))
+else
 get_interpreter(@nospecialize(job::CompilerJob)) =
     GPUInterpreter(job.world; method_table=method_table(job),
                    code_cache=ci_cache(job), inf_params=inference_params(job),
                    opt_params=optimization_params(job))
+end
 
 # does this target support throwing Julia exceptions with jl_throw?
 # if not, calls to throw will be replaced with calls to the GPU runtime
@@ -207,7 +214,26 @@ needs_byval(@nospecialize(job::CompilerJob)) = true
 # whether pointer is a valid call target
 valid_function_pointer(@nospecialize(job::CompilerJob), ptr::Ptr{Cvoid}) = false
 
-# the codeinfo cache to use
+# Care is required for anything that impacts:
+#   - method_table
+#   - inference_params
+#   - optimization_params
+# By default that is just always_inline
+# the cache token is compared with jl_egal
+struct GPUCompilerCacheToken
+    target_type::Type
+    always_inline::Bool
+    method_table::Core.MethodTable
+end
+
+ci_cache_token(@nospecialize(job::CompilerJob)) =
+    GPUCompilerCacheToken(typeof(job.config.target), job.config.always_inline, method_table(job))
+
+# the codeinfo cache to use -- should only be used for the constructor
+if VERSION >= v"1.11.0-DEV.1552"
+    # Soft deprecated user should use `CC.code_cache(get_interpreter(job))`
+    ci_cache(@nospecialize(job::CompilerJob)) = CC.code_cache(get_interpreter(job))
+else
 function ci_cache(@nospecialize(job::CompilerJob))
     lock(GLOBAL_CI_CACHES_LOCK) do
         cache = get!(GLOBAL_CI_CACHES, job.config) do
@@ -216,6 +242,7 @@ function ci_cache(@nospecialize(job::CompilerJob))
         return cache
     end
 end
+end
 
 # the method table to use
 method_table(@nospecialize(job::CompilerJob)) = GLOBAL_METHOD_TABLE

src/jlgen.jl

@@ -6,12 +6,15 @@
 # `tls_world_age` should be used to look up the current world age. in most cases, this is
 # what you should use to invoke the compiler with.
 
-tls_world_age() = ccall(:jl_get_tls_world_age, UInt, ())
-
+if isdefined(Base, :tls_world_age)
+    import Base: tls_world_age
+else
+    tls_world_age() = ccall(:jl_get_tls_world_age, UInt, ())
+end
 
 ## looking up method instances
 
-export methodinstance
+export methodinstance, generic_methodinstance
 
 @inline function signature_type_by_tt(ft::Type, tt::Type)
     u = Base.unwrap_unionall(tt)::DataType
@@ -49,21 +52,38 @@ end
 Look up the method instance that corresponds to invoking the function with type `ft` with
 argument typed `tt`. If the `world` argument is specified, the look-up is static and will
 always return the same result. If the `world` argument is not specified, the look-up is
-dynamic and the returned method instance will depende on the current world age.
+dynamic and the returned method instance will depende on the current world age. If no method
+is found, a `MethodError` is thrown.
+
+This function is highly optimized, and results do not need to be cached additionally.
 
-This call is highly optimized, and does not need to be cached additionally.
+Only use this function with concrete signatures, i.e., using the types of values you would
+pass at run time. For non-concrete signatures, use `generic_methodinstance` instead.
 
-If the method is not found, a `MethodError` is thrown.
 """
 methodinstance
 
+function generic_methodinstance(@nospecialize(ft::Type), @nospecialize(tt::Type),
+                                world::Integer=tls_world_age())
+    sig = signature_type_by_tt(ft, tt)
+
+    match, _ = CC._findsup(sig, nothing, world)
+    match === nothing && throw(MethodError(ft, tt, world))
+
+    mi = CC.specialize_method(match)
+
+    return mi::MethodInstance
+end
+
 # on 1.11 (JuliaLang/julia#52572, merged as part of JuliaLang/julia#52233) we can use
 # Julia's cached method lookup to simply look up method instances at run time.
 if VERSION >= v"1.11.0-DEV.1552"
 
 # XXX: version of Base.method_instance that uses a function type
-@inline function methodinstance(@nospecialize(ft::Type), @nospecialize(tt::Type), world::Integer=tls_world_age())
+@inline function methodinstance(@nospecialize(ft::Type), @nospecialize(tt::Type),
+                                world::Integer=tls_world_age())
     sig = signature_type_by_tt(ft, tt)
+    @assert Base.isdispatchtuple(sig)   # JuliaLang/julia#52233
 
     mi = ccall(:jl_method_lookup_by_tt, Any,
                (Any, Csize_t, Any),
@@ -79,19 +99,10 @@ if VERSION >= v"1.11.0-DEV.1552"
     return mi
 end
 
-# on older versions of Julia, the run-time lookup is much slower, so we'll need to cache it
+# on older versions of Julia, we always need to use the generic lookup
 else
 
-function methodinstance(ft::Type, tt::Type, world::Integer)
-    sig = signature_type_by_tt(ft, tt)
-
-    match, _ = CC._findsup(sig, nothing, world)
-    match === nothing && throw(MethodError(ft, tt, world))
-
-    mi = CC.specialize_method(match)
-
-    return mi::MethodInstance
-end
+const methodinstance = generic_methodinstance
 
 # on 1.10 (JuliaLang/julia#48611) generated functions know which world to generate code for.
 # we can use this to cache and automatically invalidate method instance look-ups.
@@ -153,18 +164,15 @@ end
     $(Expr(:meta, :generated, methodinstance_generator))
 end
 
-# on really old versions, we can't cache the run-time lookup
-else
-
-methodinstance(f, tt) = methodinstance(f, tt, tls_world_age())
-
 end
 
 end
 
 
 ## code instance cache
+const HAS_INTEGRATED_CACHE = VERSION >= v"1.11.0-DEV.1552"
 
+if !HAS_INTEGRATED_CACHE
 struct CodeCache
     dict::IdDict{MethodInstance,Vector{CodeInstance}}
 
@@ -292,6 +300,8 @@ function (callback::CodeCacheCallback)(replaced::MethodInstance, max_world::UInt
 end
 
 end
+end # !HAS_INTEGRATED_CACHE
+
 
 ## method overrides
 
@@ -323,13 +333,47 @@ struct GPUInterpreter <: CC.AbstractInterpreter
     world::UInt
     method_table::GPUMethodTableView
 
+@static if HAS_INTEGRATED_CACHE
+    token::Any
+else
     code_cache::CodeCache
+end
     inf_cache::Vector{CC.InferenceResult}
 
     inf_params::CC.InferenceParams
     opt_params::CC.OptimizationParams
 end
 
+@static if HAS_INTEGRATED_CACHE
+function GPUInterpreter(world::UInt=Base.get_world_counter();
+                        method_table::MTType,
+                        token::Any,
+                        inf_params::CC.InferenceParams,
+                        opt_params::CC.OptimizationParams)
+    @assert world <= Base.get_world_counter()
+
+    method_table = get_method_table_view(world, method_table)
+    inf_cache = Vector{CC.InferenceResult}()
+
+    return GPUInterpreter(world, method_table,
+                          token, inf_cache,
+                          inf_params, opt_params)
+end
+
+function GPUInterpreter(interp::GPUInterpreter;
+                        world::UInt=interp.world,
+                        method_table::GPUMethodTableView=interp.method_table,
+                        token::Any=interp.token,
+                        inf_cache::Vector{CC.InferenceResult}=interp.inf_cache,
+                        inf_params::CC.InferenceParams=interp.inf_params,
+                        opt_params::CC.OptimizationParams=interp.opt_params)
+    return GPUInterpreter(world, method_table,
+                          token, inf_cache,
+                          inf_params, opt_params)
+end
+
+else
+
 function GPUInterpreter(world::UInt=Base.get_world_counter();
                         method_table::MTType,
                         code_cache::CodeCache,
@@ -356,12 +400,17 @@ function GPUInterpreter(interp::GPUInterpreter;
                           code_cache, inf_cache,
                           inf_params, opt_params)
 end
+end # HAS_INTEGRATED_CACHE
 
 CC.InferenceParams(interp::GPUInterpreter) = interp.inf_params
 CC.OptimizationParams(interp::GPUInterpreter) = interp.opt_params
 #=CC.=#get_inference_world(interp::GPUInterpreter) = interp.world
 CC.get_inference_cache(interp::GPUInterpreter) = interp.inf_cache
-CC.code_cache(interp::GPUInterpreter) = WorldView(interp.code_cache, interp.world)
+if HAS_INTEGRATED_CACHE
+    CC.cache_owner(interp::GPUInterpreter) = interp.token
+else
+    CC.code_cache(interp::GPUInterpreter) = WorldView(interp.code_cache, interp.world)
+end
 
 # No need to do any locking since we're not putting our results into the runtime cache
 CC.lock_mi_inference(interp::GPUInterpreter, mi::MethodInstance) = nothing
@@ -413,9 +462,10 @@ end
 
 
 ## world view of the cache
-
 using Core.Compiler: WorldView
 
+if !HAS_INTEGRATED_CACHE
+
 function CC.haskey(wvc::WorldView{CodeCache}, mi::MethodInstance)
     CC.get(wvc, mi, nothing) !== nothing
 end
@@ -454,6 +504,7 @@ function CC.setindex!(wvc::WorldView{CodeCache}, ci::CodeInstance, mi::MethodIns
     CC.setindex!(wvc.cache, ci, mi)
 end
 
+end # HAS_INTEGRATED_CACHE
 
 ## codegen/inference integration
 
@@ -526,8 +577,8 @@ end
 
 function compile_method_instance(@nospecialize(job::CompilerJob))
     # populate the cache
-    cache = ci_cache(job)
     interp = get_interpreter(job)
+    cache = CC.code_cache(interp)
     if ci_cache_lookup(cache, job.source, job.world, job.world) === nothing
         ci_cache_populate(interp, cache, job.source, job.world, job.world)
     end

src/rtlib.jl

@@ -55,7 +55,7 @@ end
 
 function emit_function!(mod, config::CompilerConfig, f, method)
     tt = Base.to_tuple_type(method.types)
-    source = methodinstance(f, tt)
+    source = generic_methodinstance(f, tt)
     new_mod, meta = codegen(:llvm, CompilerJob(source, config);
                             optimize=false, libraries=false, validate=false)
     ft = function_type(meta.entry)

test/Project.toml

@@ -1,6 +1,5 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"