Merge pull request #335 from lcw/lcw/always_inline

Add `always_inline` target property to ptx backend

Author: maleadt

src/interface.jl

@@ -192,7 +192,7 @@ isintrinsic(@nospecialize(job::CompilerJob), fn::String) = false
 
 # provide a specific interpreter to use.
 get_interpreter(@nospecialize(job::CompilerJob)) =
-    GPUInterpreter(ci_cache(job), method_table(job), job.source.world)
+    GPUInterpreter(ci_cache(job), method_table(job), job.source.world, inference_params(job), optimization_params(job))
 
 # does this target support throwing Julia exceptions with jl_throw?
 # if not, calls to throw will be replaced with calls to the GPU runtime
@@ -265,11 +265,34 @@ link_libraries!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
 valid_function_pointer(@nospecialize(job::CompilerJob), ptr::Ptr{Cvoid}) = false
 
 # the codeinfo cache to use
-ci_cache(@nospecialize(job::CompilerJob)) = GLOBAL_CI_CACHE
+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
+            CodeCache()
+        end
+        return cache
+    end
+end
 
 # the method table to use
 method_table(@nospecialize(job::CompilerJob)) = GLOBAL_METHOD_TABLE
 
+# the inference parameters to use when constructing the GPUInterpreter
+function inference_params(@nospecialize(job::CompilerJob))
+    return InferenceParams(;unoptimize_throw_blocks=false)
+end
+
+# the optimization parameters to use when constructing the GPUInterpreter
+function optimization_params(@nospecialize(job::CompilerJob))
+    kwargs = NamedTuple()
+
+    if VERSION < v"1.8.0-DEV.486"
+        kwargs = (kwargs..., unoptimize_throw_blocks=false)
+    end
+
+    return OptimizationParams(;kwargs...)
+end
+
 # how much debuginfo to emit
 function llvm_debug_info(@nospecialize(job::CompilerJob))
     if Base.JLOptions().debug_level == 0

src/jlgen.jl

@@ -2,7 +2,7 @@
 
 ## cache
 
-using Core.Compiler: CodeInstance, MethodInstance
+using Core.Compiler: CodeInstance, MethodInstance, InferenceParams, OptimizationParams
 
 struct CodeCache
     dict::Dict{MethodInstance,Vector{CodeInstance}}
@@ -39,7 +39,8 @@ end
 
 Base.empty!(cc::CodeCache) = empty!(cc.dict)
 
-const GLOBAL_CI_CACHE = CodeCache()
+const GLOBAL_CI_CACHES = Dict{Tuple{DataType, InferenceParams, OptimizationParams}, CodeCache}()
+const GLOBAL_CI_CACHES_LOCK = ReentrantLock()
 
 
 ## method invalidations
@@ -182,7 +183,8 @@ struct GPUInterpreter <: AbstractInterpreter
     inf_params::InferenceParams
     opt_params::OptimizationParams
 
-    function GPUInterpreter(cache::CodeCache, mt::Union{Nothing,Core.MethodTable}, world::UInt)
+
+    function GPUInterpreter(cache::CodeCache, mt::Union{Nothing,Core.MethodTable}, world::UInt, ip::InferenceParams, op::OptimizationParams)
         @assert world <= Base.get_world_counter()
 
         return new(
@@ -196,9 +198,8 @@ struct GPUInterpreter <: AbstractInterpreter
             world,
 
             # parameters for inference and optimization
-            InferenceParams(unoptimize_throw_blocks=false),
-            VERSION >= v"1.8.0-DEV.486" ? OptimizationParams() :
-                                          OptimizationParams(unoptimize_throw_blocks=false),
+            ip,
+            op
         )
     end
 end

src/ptx.jl

@@ -21,6 +21,7 @@ Base.@kwdef struct PTXCompilerTarget <: AbstractCompilerTarget
     maxthreads::Union{Nothing,Int,NTuple{<:Any,Int}} = nothing
     blocks_per_sm::Union{Nothing,Int} = nothing
     maxregs::Union{Nothing,Int} = nothing
+    always_inline::Bool = false
 end
 
 function Base.hash(target::PTXCompilerTarget, h::UInt)
@@ -35,6 +36,7 @@ function Base.hash(target::PTXCompilerTarget, h::UInt)
     h = hash(target.maxthreads, h)
     h = hash(target.blocks_per_sm, h)
     h = hash(target.maxregs, h)
+    h = hash(target.always_inline, h)
 
     h
 end
@@ -74,6 +76,7 @@ function Base.show(io::IO, @nospecialize(job::CompilerJob{PTXCompilerTarget}))
     job.target.maxthreads !== nothing && print(io, ", maxthreads=$(job.target.maxthreads)")
     job.target.blocks_per_sm !== nothing && print(io, ", blocks_per_sm=$(job.target.blocks_per_sm)")
     job.target.maxregs !== nothing && print(io, ", maxregs=$(job.target.maxregs)")
+    job.target.always_inline !== nothing && print(io, ", always_inline=$(job.target.always_inline)")
 end
 
 const ptx_intrinsics = ("vprintf", "__assertfail", "malloc", "free")
@@ -86,6 +89,20 @@ runtime_slug(@nospecialize(job::CompilerJob{PTXCompilerTarget})) =
        "-debuginfo=$(Int(llvm_debug_info(job)))" *
        "-exitable=$(job.target.exitable)"
 
+function optimization_params(@nospecialize(job::CompilerJob{PTXCompilerTarget}))
+    kwargs = NamedTuple()
+
+    if VERSION < v"1.8.0-DEV.486"
+        kwargs = (kwargs..., unoptimize_throw_blocks=false)
+    end
+
+    if job.target.always_inline
+        kwargs = (kwargs..., inline_cost_threshold=typemax(Int))
+    end
+
+    return OptimizationParams(;kwargs...)
+end
+
 function process_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}), mod::LLVM.Module)
     ctx = context(mod)
 

src/validation.jl

@@ -32,7 +32,9 @@ function check_method(@nospecialize(job::CompilerJob))
     if job.source.kernel
         cache = ci_cache(job)
         mt = method_table(job)
-        interp = GPUInterpreter(cache, mt, world)
+        ip = inference_params(job)
+        op = optimization_params(job)
+        interp = GPUInterpreter(cache, mt, world, ip, op)
         rt = return_type(only(ms); interp)
 
         if rt != Nothing

test/definitions/ptx.jl

@@ -39,11 +39,12 @@ GPUCompiler.runtime_module(::PTXCompilerJob) = PTXTestRuntime
 
 function ptx_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false,
                  minthreads=nothing, maxthreads=nothing, blocks_per_sm=nothing,
-                 maxregs=nothing, kwargs...)
+                 maxregs=nothing, always_inline=false, kwargs...)
     source = FunctionSpec(func, Base.to_tuple_type(types), kernel)
     target = PTXCompilerTarget(cap=v"7.0",
                                minthreads=minthreads, maxthreads=maxthreads,
-                               blocks_per_sm=blocks_per_sm, maxregs=maxregs)
+                               blocks_per_sm=blocks_per_sm, maxregs=maxregs,
+                               always_inline=always_inline)
     params = TestCompilerParams()
     CompilerJob(target, source, params), kwargs
 end

test/ptx.jl

@@ -174,6 +174,32 @@ end
 end
 end
 
+@testset "always_inline" begin
+    @eval f_expensive(x) = $(foldl((e, _) -> :(sink($e) + sink(x)), 1:100; init=:x))
+    function g(x)
+        f_expensive(x)
+        return
+    end
+    function h(x)
+        f_expensive(x)
+        return
+    end
+
+    asm = sprint(io->ptx_code_native(io, g, Tuple{Int64}; kernel=true))
+    @test occursin(r"\.func .*julia_f_expensive", asm)
+
+    asm = sprint(io->ptx_code_native(io, g, Tuple{Int64};
+                                     kernel=true, always_inline=true))
+    @test !occursin(r"\.func .*julia_f_expensive", asm)
+
+    asm = sprint(io->ptx_code_native(io, h, Tuple{Int64};
+                                     kernel=true, always_inline=true))
+    @test !occursin(r"\.func .*julia_f_expensive", asm)
+
+    asm = sprint(io->ptx_code_native(io, h, Tuple{Int64}; kernel=true))
+    @test occursin(r"\.func .*julia_f_expensive", asm)
+end
+
 @testset "child function reuse" begin
     # bug: depending on a child function from multiple parents resulted in
     #      the child only being present once