Split CompilerJob in dynamic and static part. (#395)

Author: maleadt

examples/kernel.jl

@@ -19,7 +19,8 @@ function main()
     source = FunctionSpec(typeof(kernel), Tuple{})
     target = NativeCompilerTarget()
     params = TestCompilerParams()
-    job = CompilerJob(source, target, params)
+    config = CompilerConfig(target, params)
+    job = CompilerJob(config, source)
 
     println(GPUCompiler.compile(:asm, job)[1])
 end

src/bpf.jl

@@ -32,4 +32,4 @@ const bpf_intrinsics = () # TODO
 isintrinsic(::CompilerJob{BPFCompilerTarget}, fn::String) = in(fn, bpf_intrinsics)
 
 valid_function_pointer(job::CompilerJob{BPFCompilerTarget}, ptr::Ptr{Cvoid}) =
-    reinterpret(UInt, ptr) in job.target.function_pointers
+    reinterpret(UInt, ptr) in job.config.target.function_pointers

src/cache.jl

@@ -130,7 +130,7 @@ end
 const cache_lock = ReentrantLock()
 
 """
-    cached_compilation(cache::Dict, job::CompilerJob, compiler, linker)
+    cached_compilation(cache::Dict{UInt}, job::CompilerJob, compiler, linker)
 
 Compile `job` using `compiler` and `linker`, and store the result in `cache`.
 
@@ -140,46 +140,57 @@ and return data that can be cached across sessions (e.g., LLVM IR). This data is
 forwarded, along with the `CompilerJob`, to the `linker` function which is allowed to create
 session-dependent objects (e.g., a `CuModule`).
 """
-function cached_compilation(cache::AbstractDict,
-                            @nospecialize(job::CompilerJob),
-                            compiler::Function, linker::Function)
-    # NOTE: it is OK to index the compilation cache directly with the compilation job, i.e.,
-    #       using a world age instead of intersecting world age ranges, because we expect
-    #       that the world age is aquired through calling `get_world` and thus will only
-    #       ever change when the kernel function is redefined.
-    #
-    #       if we ever want to be able to index the cache using a compilation job that
-    #       contains a more recent world age, yet still return an older cached object that
-    #       would still be valid, we'd need the cache to store world ranges instead and
-    #       use an invalidation callback to add upper bounds to entries.
-    key = hash(job)
-
-    force_compilation = compile_hook[] !== nothing
-
-    # NOTE: no use of lock(::Function)/@lock/get! to keep stack traces clean
+function cached_compilation(cache::AbstractDict{UInt,V},
+                            cfg::CompilerConfig,
+                            ft::Type, tt::Type,
+                            compiler::Function, linker::Function) where {V}
+    # NOTE: it is OK to index the compilation cache directly with the world age, instead of
+    #       intersecting world age ranges, because we the world age is aquired by calling
+    #       `get_world` and thus will only change when the kernel function is redefined.
+    world = get_world(ft, tt)
+    key = hash(ft)
+    key = hash(tt, key)
+    key = hash(world, key)
+    key = hash(cfg, key)
+
+    # NOTE: no use of lock(::Function)/@lock/get! to avoid try/catch and closure overhead
     lock(cache_lock)
-    try
-        obj = get(cache, key, nothing)
-        if obj === nothing || force_compilation
-            asm = nothing
+    obj = get(cache, key, nothing)
+    unlock(cache_lock)
 
-            # compile
-            if asm === nothing
-                if compile_hook[] !== nothing
-                    compile_hook[](job)
-                end
-
-                asm = compiler(job)
-            end
-
-            # link (but not if we got here because of forced compilation)
-            if obj === nothing
-                obj = linker(job, asm)
-                cache[key] = obj
-            end
+    LLVM.Interop.assume(isassigned(compile_hook))
+    if obj === nothing || compile_hook[] !== nothing
+        obj = actual_compilation(cache, key, cfg, ft, tt, world, compiler, linker)::V
+    end
+    return obj::V
+end
+
+@noinline function actual_compilation(cache::AbstractDict, key::UInt,
+                                      cfg::CompilerConfig,
+                                      ft::Type, tt::Type, world,
+                                      compiler::Function, linker::Function)
+    src = FunctionSpec(ft, tt, world)
+    job = CompilerJob(cfg, src)
+
+    asm = nothing
+    # TODO: consider loading the assembly from an on-disk cache here
+
+    # compile
+    if asm === nothing
+        if compile_hook[] !== nothing
+            compile_hook[](job)
         end
+
+        asm = compiler(job)
+    end
+
+    # link (but not if we got here because of forced compilation,
+    # in which case the cache will already be populated)
+    lock(cache_lock) do
+        haskey(cache, key) && return cache[key]
+
+        obj = linker(job, asm)
+        cache[key] = obj
         obj
-    finally
-        unlock(cache_lock)
     end
 end

src/driver.jl

@@ -227,7 +227,7 @@ const __llvm_initialized = Ref(false)
 
     @timeit_debug to "IR generation" begin
         ir, compiled = irgen(job, method_instance; ctx)
-        if job.entry_abi === :specfunc
+        if job.config.entry_abi === :specfunc
             entry_fn = compiled[method_instance].specfunc
         else
             entry_fn = compiled[method_instance].func
@@ -300,11 +300,11 @@ const __llvm_initialized = Ref(false)
 
                 # get a job in the appopriate world
                 dyn_job = if dyn_val isa CompilerJob
-                    dyn_spec = FunctionSpec(dyn_val.source; world=job.source.world)
-                    CompilerJob(dyn_val; source=dyn_spec)
+                    dyn_src = FunctionSpec(dyn_val.source; world=job.source.world)
+                    CompilerJob(dyn_val.config, dyn_src)
                 elseif dyn_val isa FunctionSpec
-                    dyn_spec = FunctionSpec(dyn_val; world=job.source.world)
-                    CompilerJob(job; source=dyn_spec)
+                    dyn_src = FunctionSpec(dyn_val; world=job.source.world)
+                    CompilerJob(job.config, dyn_src)
                 else
                     error("invalid deferred job type $(typeof(dyn_val))")
                 end
@@ -349,7 +349,7 @@ const __llvm_initialized = Ref(false)
 
     @timeit_debug to "IR post-processing" begin
         # mark the kernel entry-point functions (optimization may need it)
-        if job.source.kernel
+        if job.config.kernel
             push!(metadata(ir)["julia.kernel"], MDNode([entry]; ctx=unwrap_context(ctx)))
 
             # IDEA: save all jobs, not only kernels, and save other attributes

src/gcn.jl

@@ -30,15 +30,15 @@ end
 
 # TODO: encode debug build or not in the compiler job
 #       https://github.com/JuliaGPU/CUDAnative.jl/issues/368
-runtime_slug(job::CompilerJob{GCNCompilerTarget}) = "gcn-$(job.target.dev_isa)$(job.target.features)"
+runtime_slug(job::CompilerJob{GCNCompilerTarget}) = "gcn-$(job.config.target.dev_isa)$(job.config.target.features)"
 
 const gcn_intrinsics = () # TODO: ("vprintf", "__assertfail", "malloc", "free")
 isintrinsic(::CompilerJob{GCNCompilerTarget}, fn::String) = in(fn, gcn_intrinsics)
 
 function process_entry!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module, entry::LLVM.Function)
     entry = invoke(process_entry!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
 
-    if job.source.kernel
+    if job.config.kernel
         # calling convention
         callconv!(entry, LLVM.API.LLVMAMDGPUKERNELCallConv)
     end
@@ -54,7 +54,7 @@ function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
                         mod::LLVM.Module, entry::LLVM.Function)
     entry = invoke(finish_module!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
 
-    if job.source.kernel
+    if job.config.kernel
         # work around bad byval codegen (JuliaGPU/GPUCompiler.jl#92)
         entry = lower_byval(job, mod, entry)
     end
@@ -84,10 +84,10 @@ end
 function optimize_module!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module)
     @static if VERSION < v"1.9.0-DEV.1018"
         # revert back to the AMDGPU target
-        triple!(mod, llvm_triple(job.target))
-        datalayout!(mod, julia_datalayout(job.target))
+        triple!(mod, llvm_triple(job.config.target))
+        datalayout!(mod, julia_datalayout(job.config.target))
 
-        tm = llvm_machine(job.target)
+        tm = llvm_machine(job.config.target)
         @dispose pm=ModulePassManager() begin
             add_library_info!(pm, triple(mod))
             add_transform_info!(pm, tm)

src/interface.jl

@@ -66,34 +66,24 @@ struct FunctionSpec
     tt::Type
     world::UInt
 
-    kernel::Bool
-    name::Union{Nothing,String}
-
-    FunctionSpec(ft::Type, tt::Type, world::Integer=get_world(ft, tt);
-                 kernel=true, name=nothing) =
-        new(ft, tt, world, kernel, name)
+    FunctionSpec(ft::Type, tt::Type, world::Integer=get_world(ft, tt)) =
+        new(ft, tt, world)
 end
 
 # copy constructor
-FunctionSpec(spec::FunctionSpec; ft=spec.ft, tt=spec.tt, world=spec.world,
-                                 kernel=spec.kernel, name=spec.name) =
-    FunctionSpec(ft, tt, world; kernel, name)
+FunctionSpec(spec::FunctionSpec; ft=spec.ft, tt=spec.tt, world=spec.world) =
+    FunctionSpec(ft, tt, world)
 
 function Base.hash(spec::FunctionSpec, h::UInt)
     h = hash(spec.ft, h)
     h = hash(spec.tt, h)
     h = hash(spec.world, h)
 
-    h = hash(spec.kernel, h)
-    h = hash(spec.name, h)
-
     return h
 end
 
 function signature(@nospecialize(spec::FunctionSpec))
-    fn = if spec.name !== nothing
-        spec.name
-    elseif spec.ft.name.mt == Symbol.name.mt
+    fn = if spec.ft.name.mt == Symbol.name.mt
         # uses shared method table, so name is not unique to this function type
         nameof(spec.ft)
     else
@@ -104,22 +94,21 @@ function signature(@nospecialize(spec::FunctionSpec))
 end
 
 function Base.show(io::IO, @nospecialize(spec::FunctionSpec))
-    spec.kernel ? print(io, "kernel ") : print(io, "function ")
     print(io, signature(spec), " in world ", spec.world)
 end
 
 
-## job
+## config
 
-export CompilerJob
+export CompilerConfig
 
-# a specific invocation of the compiler, bundling everything needed to generate code
+# the configuration of the compiler
 
 """
-    CompilerJob(source, target, params; entry_abi=:specfunc, always_inline=false)
+    CompilerConfig(target, params; kernel=true, entry_abi=:specfunc, always_inline=false)
 
-Construct a `CompilerJob` for `source` that will be used to drive compilation for the given
-`target` and `params`.
+Construct a `CompilerConfig` that will be used to drive compilation for the given `target`
+and `params`.
 
 The `entry_abi` can be either `:specfunc` the default, or `:func`. `:specfunc` expects the
 arguments to be passed in registers, simple return values are returned in registers as well,
@@ -133,46 +122,62 @@ generally easier to invoke directly.
 `always_inline` specifies if the Julia front-end should inline all functions into one if
 possible.
 """
-struct CompilerJob{T,P}
+struct CompilerConfig{T,P}
     target::T
     params::P
-    source::FunctionSpec
 
+    kernel::Bool
     entry_abi::Symbol
     always_inline::Bool
 
-    function CompilerJob(source::FunctionSpec,
-                         target::AbstractCompilerTarget,
-                         params::AbstractCompilerParams;
-                         entry_abi::Symbol=:specfunc, always_inline=false)
+    function CompilerConfig(target::AbstractCompilerTarget,
+                            params::AbstractCompilerParams;
+                            kernel::Bool=true,
+                            entry_abi::Symbol=:specfunc,
+                            always_inline=false)
         if entry_abi ∉ (:specfunc, :func)
             error("Unknown entry_abi=$entry_abi")
         end
-        new{typeof(target), typeof(params)}(target, params, source, entry_abi, always_inline)
+        new{typeof(target), typeof(params)}(target, params, kernel, entry_abi, always_inline)
     end
 end
 
 # copy constructor
-CompilerJob(job::CompilerJob; source=job.source, target=job.target, params=job.params,
-            entry_abi=job.entry_abi, always_inline=job.always_inline) =
-    CompilerJob(source, target, params; entry_abi, always_inline)
+CompilerConfig(cfg::CompilerConfig; target=cfg.target, params=cfg.params,
+               kernel=cfg.kernel, entry_abi=cfg.entry_abi, always_inline=cfg.always_inline) =
+    CompilerConfig(target, params; kernel, entry_abi, always_inline)
 
-function Base.show(io::IO, @nospecialize(job::CompilerJob{T})) where {T}
-    print(io, "CompilerJob of ", job.source, " for ", T)
+function Base.show(io::IO, @nospecialize(cfg::CompilerConfig{T})) where {T}
+    print(io, "CompilerConfig for ", T)
 end
 
-function Base.hash(job::CompilerJob, h::UInt)
-    h = hash(job.source, h)
-    h = hash(job.target, h)
-    h = hash(job.params, h)
+function Base.hash(cfg::CompilerConfig, h::UInt)
+    h = hash(cfg.target, h)
+    h = hash(cfg.params, h)
 
-    h = hash(job.entry_abi, h)
-    h = hash(job.always_inline, h)
+    h = hash(cfg.kernel, h)
+    h = hash(cfg.entry_abi, h)
+    h = hash(cfg.always_inline, h)
 
     return h
 end
 
 
+## job
+
+export CompilerJob
+
+# a specific invocation of the compiler, bundling everything needed to generate code
+
+struct CompilerJob{T,P}
+    config::CompilerConfig{T,P}
+    source::FunctionSpec
+
+    CompilerJob(cfg::CompilerConfig{T,P}, src::FunctionSpec) where {T,P} =
+        new{T,P}(cfg, src)
+end
+
+
 ## contexts
 
 if VERSION >= v"1.9.0-DEV.516"
@@ -233,7 +238,7 @@ function process_entry!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                         entry::LLVM.Function)
     ctx = context(mod)
 
-    if job.source.kernel && needs_byval(job)
+    if job.config.kernel && needs_byval(job)
         # pass all bitstypes by value; by default Julia passes aggregates by reference
         # (this improves performance, and is mandated by certain back-ends like SPIR-V).
         args = classify_arguments(job, eltype(llvmtype(entry)))
@@ -279,7 +284,7 @@ valid_function_pointer(@nospecialize(job::CompilerJob), ptr::Ptr{Cvoid}) = false
 # the codeinfo cache to use
 function ci_cache(@nospecialize(job::CompilerJob))
     lock(GLOBAL_CI_CACHES_LOCK) do
-        cache = get!(GLOBAL_CI_CACHES, (typeof(job.target), inference_params(job), optimization_params(job))) do
+        cache = get!(GLOBAL_CI_CACHES, (typeof(job.config.target), inference_params(job), optimization_params(job))) do
             CodeCache()
         end
         return cache
@@ -302,7 +307,7 @@ function optimization_params(@nospecialize(job::CompilerJob))
         kwargs = (kwargs..., unoptimize_throw_blocks=false)
     end
 
-    if job.always_inline
+    if job.config.always_inline
         kwargs = (kwargs..., inline_cost_threshold=typemax(Int))
     end