Optimize device lib linking

.buildkite/pipeline.yml

@@ -22,26 +22,7 @@ steps:
     if: build.message !~ /\[skip docs\]/
     timeout_in_minutes: 10
 
-  # - label: "Julia 1.9 - GPUArrays 8"
-  #   plugins:
-  #     - JuliaCI/julia#v1:
-  #         version: 1.9-nightly
-  #     - JuliaCI/julia-test#v1:
-  #     - JuliaCI/julia-coverage#v1:
-  #         codecov: true
-  #   agents:
-  #     queue: "juliagpu"
-  #     julia_1.9: "*"
-  #     rocm: "*"
-  #     rocmgpu: "*"
-  #   if: build.message !~ /\[skip tests\]/
-  #   command: "julia --project -e 'using Pkg; Pkg.update()'"
-  #   timeout_in_minutes: 180
-  #   env:
-  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-
-  - label: "Julia 1.9 - GPUArrays 8 - No Artifacts"
+  - label: "Julia 1.9 - No Artifacts"
     plugins:
       - JuliaCI/julia#v1:
           version: 1.9-nightly
@@ -50,7 +31,6 @@ steps:
           codecov: true
     agents:
       queue: "juliagpu"
-      julia_1.9: "*"
       rocm: "*"
       rocmgpu: "gfx1031"
     if: build.message !~ /\[skip tests\]/

src/compiler/codegen.jl

@@ -27,15 +27,16 @@ function GPUCompiler.link_libraries!(
     invoke(GPUCompiler.link_libraries!,
         Tuple{CompilerJob{GCNCompilerTarget}, typeof(mod), typeof(undefined_fns)},
         job, mod, undefined_fns)
-    link_device_libs!(job.config.target, mod)
+    link_device_libs!(job.config.target, mod, undefined_fns)
 end
 
 # FIXME
 function GPUCompiler.finish_ir!(
     @nospecialize(job::HIPCompilerJob), mod::LLVM.Module, entry::LLVM.Function,
 )
-    isempty(collect(GPUCompiler.decls(mod))) && return entry
-    link_device_libs!(job.config.target, mod)
+    undefined_fns = GPUCompiler.decls(mod)
+    isempty(undefined_fns) && return entry
+    link_device_libs!(job.config.target, mod, LLVM.name.(undefined_fns))
     return entry
 end
 

src/compiler/device_libs.jl

@@ -1,37 +1,9 @@
-## ROCm device library
-
 import AMDGPU: libdevice_libs
 
-function load_and_link!(mod, path)
-    lib = parse(LLVM.Module, read(path))
-
-    for f in LLVM.functions(lib)
-        # FIXME: We should be able to inline this, that we can't means
-        #        we are inserting calls to it late.
-        startswith(LLVM.name(f), "__ockl_hsa_signal") && continue
-
-        attrs = function_attributes(f)
-        inline = true
-        noinline_attr = EnumAttribute("noinline")
-        for attr in collect(attrs)
-            if kind(attr) == kind(noinline_attr)
-                inline = false
-                break
-            end
-        end
-        inline && push!(attrs, EnumAttribute("alwaysinline"))
-    end
-
-    # override triple and datalayout to avoid warnings
-    triple!(lib, triple(mod))
-    datalayout!(lib, datalayout(mod))
-    LLVM.link!(mod, lib)
-end
-
 function locate_lib(file)
-    file_path = joinpath(libdevice_libs, file*".bc")
+    file_path = joinpath(libdevice_libs, file * ".bc")
     if !ispath(file_path)
-        file_path = joinpath(libdevice_libs, file*".amdgcn.bc")
+        file_path = joinpath(libdevice_libs, file * ".amdgcn.bc")
         if !ispath(file_path)
             # failed to find matching bitcode file
             return nothing
@@ -40,61 +12,112 @@ function locate_lib(file)
     return file_path
 end
 
-function link_device_libs!(target, mod::LLVM.Module)
-    # TODO: only link if used
-    # TODO: make these globally/locally configurable
+mutable struct DevLib
+    name::String
+    path::String
+    data::Vector{UInt8}
+    fn_names::Set{String}
 
+    DevLib(name::String, path::String) = new(name, path, read(path), Set{String}())
+    DevLib(name::String, ::Nothing) = new(name, "", UInt8[], Set{String}())
+end
+
+const DEVICE_LIBS::Dict{String, DevLib} = Dict{String, DevLib}()
+
+function link_device_libs!(
+    target::GCNCompilerTarget, mod::LLVM.Module, undefined_fns::Vector{String},
+)
     isnothing(libdevice_libs) && return
+    isempty(undefined_fns) && return
 
-    # https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/9420f6380990b09851edc2a5f9cbfaa88742b449/doc/OCML.md#controls
-    # Note: It seems we need to load in reverse order, to avoid LLVM deleting the globals from the module, before we use them.
-
-    # 1. Load other libraries
-    libs = ("hc", "hip", "irif", "ockl", "opencl", "ocml")
-    for lib in libs
-        lib_path = locate_lib(lib)
-        lib_path === nothing && continue
-        try
-            load_and_link!(mod, lib_path)
-        catch err
-            @warn "Failed to load/link $lib" err=err
+    # 1. Load other libraries.
+    lib_names = ("hc", "hip", "irif", "ockl", "opencl", "ocml")
+    for lib_name in lib_names
+        devlib = get!(DEVICE_LIBS, lib_name) do
+            DevLib(lib_name, locate_lib(lib_name))
         end
+        load_and_link!(devlib, mod, undefined_fns)
     end
 
-    # 2. Load OCLC library
-    isa_short = replace(target.dev_isa, "gfx"=>"")
-    lib = locate_lib("oclc_isa_version_$isa_short")
-    @assert lib !== nothing
-    try
-        load_and_link!(mod, lib)
-    catch err
-        @warn "Failed to load/link OCLC core library `$lib` for ISA $(target.dev_isa)." err=err
+    # 2. Load OCLC library.
+    devlib = get!(DEVICE_LIBS, "oclc") do
+        isa_short = replace(target.dev_isa, "gfx"=>"")
+        DevLib("oclc", locate_lib("oclc_isa_version_$isa_short"))
     end
+    load_and_link!(devlib, mod)
+
+    # 3. Load OCLC ABI library.
+    devlib = get!(DEVICE_LIBS, "oclc_abi") do
+        DevLib("oclc_abi", locate_lib("oclc_abi_version_500"))
+    end
+    load_and_link!(devlib, mod)
+
+    # 4. Load options libraries.
+    options = (
+        (:finite_only, false),
+        (:unsafe_math, false),
+        (:correctly_rounded_sqrt, true),
+        (:daz_opt, false),
+        (:wavefrontsize64, true))
 
-    # 3. Load OCLC ABI library (required for printing).
-    lib = locate_lib("oclc_abi_version_500")
-    @assert lib !== nothing
-    load_and_link!(mod, lib)
-
-    # 4. Load options libraries
-    options = Dict(
-        :finite_only => false,
-        :unsafe_math => false,
-        :correctly_rounded_sqrt => true,
-        :daz_opt => false,
-        :wavefrontsize64 => true
-    )
     for (option, value) in options
         toggle = value ? "on" : "off"
-        lib = locate_lib("oclc_$(option)_$(toggle)")
-        if lib === nothing
-            @warn "Could not find OCLC library for option $option=$value"
-            continue
+        name = "oclc_$(option)_$(toggle)"
+        devlib = get!(DEVICE_LIBS, name) do
+            DevLib(name, locate_lib(name))
         end
-        try
-            load_and_link!(mod, lib)
-        catch err
-            @warn "Failed to load/link OCLC library for option $option=$value" err=err
+        load_and_link!(devlib, mod)
+    end
+end
+
+function load_and_link!(
+    devlib::DevLib, mod::LLVM.Module, undefined_fns::Vector{String} = String[],
+)
+    isempty(devlib.path) && return
+
+    fill_fn_names = isempty(devlib.fn_names)
+    do_linking = false
+
+    if !fill_fn_names && !isempty(undefined_fns)
+        for undef_fn in undefined_fns
+            undef_fn ∈ devlib.fn_names && (do_linking = true; break)
+        end
+        do_linking || return
+    end
+
+    lib = parse(LLVM.Module, devlib.data)
+    inline_attr = EnumAttribute("alwaysinline")
+    noinline_attr = EnumAttribute("noinline")
+
+    for f in LLVM.functions(lib)
+        fn_name = LLVM.name(f)
+        fill_fn_names && push!(devlib.fn_names, fn_name)
+
+        # FIXME: We should be able to inline this, that we can't means
+        #        we are inserting calls to it late.
+        startswith(fn_name, "__ockl_hsa_signal") && continue
+
+        attrs = function_attributes(f)
+        inline = true
+        for attr in collect(attrs)
+            if kind(attr) == kind(noinline_attr)
+                inline = false
+                break
+            end
+        end
+        inline && push!(attrs, inline_attr)
+    end
+
+    if !do_linking && !isempty(undefined_fns)
+        for undef_fn in undefined_fns
+            undef_fn ∈ devlib.fn_names && (do_linking = true; break)
         end
+        do_linking || return
     end
+
+    # override triple and datalayout to avoid warnings
+    triple!(lib, triple(mod))
+    datalayout!(lib, datalayout(mod))
+    LLVM.link!(mod, lib)
+    return
 end

test/rocarray/base.jl

@@ -77,6 +77,7 @@ end
 
         # GPU pointer works.
         RA .+= 1.0
+        AMDGPU.synchronize()
 
         # Host pointer is updated.
         @test A ≈ A_orig .+ 1.0