Add always_inline parameter to ptx backend

This add the parameter `always_inline` to force the inlining of
functions in the ptx backend.

Author: lcw

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)
 

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