Manually adding norecurse can significantly improve the performance on an AMD MI210 #1065

# MWE for an upstream issue (notes/upstream_norecurse_issue.md): on AMD GPUs, a kernel that merely contains a call to a
# `@noinline` device function is compiled with `.amdhsa_uses_dynamic_stack 1`, because GPUCompiler never infers `norecurse`
# for the callee. The call below is never executed, yet the kernel gets slower.
#
# Self-contained: needs AMDGPU.jl (which brings GPUCompiler.jl and LLVM.jl) and an AMD GPU. Usage:
#   julia --project=<env with AMDGPU, GPUCompiler, LLVM> mwe_amdgpu_norecurse.jl [none|infer|entry|force]
#     none  : stock compilation
#     infer : run LLVM's bottom-up `function-attrs` pass (PostOrderFunctionAttrsPass, part of LLVM's default -O pipelines) on
#             the optimized module, via the `GPUCompiler.optimize_module!` hook (not defined for GCN by GPUCompiler or AMDGPU).
#             LLVM < 18 infers `norecurse` only if every callee, intrinsics included, already has it (LLVM #76372 accepts
#             `nocallback` declarations from LLVM 18 on).
#     entry : mark the kernel entry point `norecurse` (cf. GPUCompiler #262), then run LLVM's top-down `rpo-function-attrs`
#             pass, which marks internal functions called only from `norecurse` functions
#     force : add `norecurse` to every function defined in the module (the workaround used in our benchmarks)
# Run each mode in a fresh process (kernels are cached per process).
using AMDGPU, GPUCompiler, LLVM
using Printf, Statistics

const MODE = get(ARGS, 1, "none")
MODE in ("none", "infer", "entry", "force") || error("mode: none | infer | entry | force")

function run_pipeline!(job, mod, pipeline)
    PB = isdefined(LLVM, :NewPMPassBuilder) ? LLVM.NewPMPassBuilder : LLVM.PassBuilder   # LLVM.jl 9 / 10
    @dispose tm=GPUCompiler.llvm_machine(job.config.target) pb=PB() begin
        LLVM.add!(pb, pipeline)
        LLVM.run!(pb, mod, tm)
    end
end

if MODE != "none"
    function GPUCompiler.optimize_module!(@nospecialize(job::GPUCompiler.CompilerJob{GPUCompiler.GCNCompilerTarget}),
                                          mod::LLVM.Module)
        if MODE == "infer"
            run_pipeline!(job, mod, "cgscc(function-attrs)")
        elseif MODE == "entry"
            # GPUCompiler's `finish_module!` has already given the entry point the kernel calling convention
            for f in LLVM.functions(mod)
                if LLVM.API.LLVMGetFunctionCallConv(f) == UInt32(LLVM.API.LLVMAMDGPUKERNELCallConv)
                    push!(LLVM.function_attributes(f), LLVM.EnumAttribute("norecurse", 0))
                end
            end
            run_pipeline!(job, mod, "rpo-function-attrs")
        else
            for f in LLVM.functions(mod)
                isempty(LLVM.blocks(f)) && continue      # skip declarations
                push!(LLVM.function_attributes(f), LLVM.EnumAttribute("norecurse", 0))
            end
        end
        return
    end
end

@inline function cold_body(x)
    y = x
    for _ in 1:8
        y = fma(y, y, x) / (abs(y) + one(x))
    end
    return y
end
@noinline cold(x) = cold_body(x)

@inline function hot(x)
    s = x
    for _ in 1:16
        s = fma(s, oftype(x, 0.999), oftype(x, 0.001))
    end
    return s
end

# identical kernels except for the rare branch: inlined vs. a call to the `@noinline` function
function kernel_inline!(y, x)
    i = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x
    if i <= length(x)
        @inbounds xi = x[i]
        s = hot(xi)
        if !(xi < floatmax(xi))      # never true for the finite inputs below
            s = cold_body(xi)
        end
        @inbounds y[i] = s
    end
    return
end

function kernel_call!(y, x)
    i = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x
    if i <= length(x)
        @inbounds xi = x[i]
        s = hot(xi)
        if !(xi < floatmax(xi))      # never true for the finite inputs below
            s = cold(xi)
        end
        @inbounds y[i] = s
    end
    return
end

function descriptor(kernel, tt)
    io = IOBuffer()
    # `dump_module = true` is essential: otherwise GPUCompiler empties all non-entry functions (`only_entry`), the callee
    # becomes a declaration, and LLVM reports a dynamic stack for any kernel with a call
    AMDGPU.code_gcn(io, kernel, tt; kernel = true, dump_module = true)
    asm = String(take!(io))
    keys = (".amdhsa_uses_dynamic_stack", ".amdhsa_private_segment_fixed_size", "SCRATCH_EN", "ScratchSize")
    return [strip(l) for l in split(asm, '\n') if any(k -> occursin(k, l), keys)]
end

function callee_attributes(kernel, tt)
    io = IOBuffer()
    AMDGPU.code_llvm(io, kernel, tt; kernel = true, dump_module = true)
    ll = String(take!(io))
    lines = split(ll, '\n')
    defs = [strip(l) for l in lines if startswith(l, "define") && occursin("julia_cold", l)]
    decls = [strip(l) for l in lines if startswith(l, "declare")]
    attrs = [strip(l) for l in lines if startswith(l, "attributes #")]
    return defs, decls, attrs
end

function main(; T = Float64, n = 2^24, groupsize = 256, reps = 200)
    println("AMDGPU.jl ", pkgversion(AMDGPU), ", GPUCompiler.jl ", pkgversion(GPUCompiler), ", LLVM.jl ", pkgversion(LLVM),
            ", Julia ", VERSION, ", libLLVM ", Base.libllvm_version)
    println("device: ", AMDGPU.device(), "; mode: ", MODE, "; T = ", T, ", n = ", n)
    x = ROCArray(T.(range(0.5, 2.0; length = n)))
    y_inline = similar(x)
    y_call = similar(x)
    tt = Tuple{typeof(AMDGPU.rocconvert(y_call)), typeof(AMDGPU.rocconvert(x))}
    gridsize = cld(n, groupsize)

    for (name, kernel) in (("kernel_inline!", kernel_inline!), ("kernel_call!", kernel_call!))
        println("\n", name, ": kernel descriptor")
        foreach(l -> println("  ", l), descriptor(kernel, tt))
    end
    defs, decls, attrs = callee_attributes(kernel_call!, tt)
    println("\nkernel_call!: definition of the callee, declarations and attribute groups (optimized IR)")
    foreach(l -> println("  ", l), defs)
    foreach(l -> println("  ", l), decls)
    foreach(l -> println("  ", l), attrs)

    k_inline = @roc launch=false kernel_inline!(y_inline, x)
    k_call = @roc launch=false kernel_call!(y_call, x)
    k_inline(y_inline, x; groupsize, gridsize)
    k_call(y_call, x; groupsize, gridsize)
    AMDGPU.synchronize()
    println("\nresults equal: ", Array(y_inline) == Array(y_call))

    t_inline = Float64[]
    t_call = Float64[]
    for _ in 1:reps      # interleaved
        push!(t_inline, AMDGPU.@elapsed k_inline(y_inline, x; groupsize, gridsize))
        push!(t_call, AMDGPU.@elapsed k_call(y_call, x; groupsize, gridsize))
    end
    @printf("median time (%d launches each): kernel_inline! %.4f ms, kernel_call! %.4f ms, ratio %.3f\n",
            reps, 1e3 * median(t_inline), 1e3 * median(t_call), median(t_call) / median(t_inline))
    return nothing
end

main()