Fix global hostcall detection (#787)

Author: pxl-th

docs/src/tutorials/perf.md

@@ -44,96 +44,3 @@ julia> GPUArrays.unsafe_free!(cache)
 For a more sophisticated real-world example, see how
 [GaussianSplatting.jl](https://github.com/JuliaNeuralGraphics/GaussianSplatting.jl/blob/e4ef1324c187371e336bef875b053023afe7fb2c/src/training.jl#L183)
 handles it.
-
-## Avoid triggering Hostcalls
-
-Some functions in the kernel may cause an exception,
-capturing the original value of the variable that caused it.
-These are usually related to float-to-integer conversion, so functions like
-`Int(1.0), ceil(Int, 1.0), floor(Int, 1.0)` will cause it.
-
-This will perform dynamic memory allocation and launch a `Hostcall` for that,
-which will sit in the background thread until kernel finishes execution and the user synchronizes the `stream`.
-Having a hostcall unnecessarily slows execution down and you can avoid that by using
-"GPU-friendly" version of the function.
-
-!!! info "Hostcalls"
-    Hostcalls should be used mostly for debugging. When performance matters, they should be avoided.
-
-For example, let's see how we may deal with `ceil(Int, x)` and convert it to GPU-friendly version.
-
-Starting with the bad example:
-
-```jldoctest hostcall
-julia> function bad_kernel!(y, x)
-           @inbounds y[1] = ceil(Int, x[1])
-           return
-       end
-bad_kernel! (generic function with 1 method)
-
-julia> x = ROCArray(Float32[1.1f0]);
-
-julia> y = ROCArray(zeros(Int, 1));
-
-julia> @roc bad_kernel!(y, x);
-┌ Info: Global hostcalls detected!
-│ - Source: MethodInstance for bad_kernel!(::AMDGPU.Device.ROCDeviceVector{Int64, 1}, ::AMDGPU.Device.ROCDeviceVector{Float32, 1})
-│ - Hostcalls: [:malloc_hostcall]
-│
-│ Use `AMDGPU.synchronize(; stop_hostcalls=true)` to synchronize and stop them.
-└ Otherwise, performance might degrade if they keep running in the background.
-
-julia> y
-1-element ROCArray{Int64, 1, AMDGPU.Runtime.Mem.HIPBuffer}:
- 2
-
-julia> AMDGPU.synchronize(; stop_hostcalls=true)
-[ Info: Stopped global hostcall: `malloc_hostcall`.
-```
-
-Here we can see that using "un-optimized" version of `ceil(Int, x[1])`
-causes a `malloc_hostcall` to be launched.
-Which we then have to stop by passing `stop_hostcalls=true` to the synchronization functions.
-
-We can avoid this by using "unsafe" version that avoids checking for errors under-the-hood.
-
-```jldoctest hostcall
-julia> function good_kernel!(y, x)
-           @inbounds y[1] = unsafe_trunc(Int, ceil(x[1]))
-           return
-       end
-good_kernel! (generic function with 1 method)
-
-julia> fill!(y, 0);
-
-julia> @roc good_kernel!(y, x);
-
-julia> AMDGPU.synchronize(; stop_hostcalls=true) # Nothing is printed, so no hostcall was launched & stopped.
-
-julia> y
-1-element ROCArray{Int64, 1, AMDGPU.Runtime.Mem.HIPBuffer}:
- 2
-```
-
-By doing `ceil(x[1])` first, then "unsafely" converting `Float32` to `Int`
-we can avoid error-checking & hostcalls.
-
-We can compare LLVM IR of the function that converts `Float32` to `Int` to see how they differ:
-
-::: tabs
-
-== unsafe_trunc(Int, 1.0)
-
-```@example good-conversion
-using InteractiveUtils
-InteractiveUtils.@code_llvm unsafe_trunc(Int, 1.0)
-```
-
-== Int(1.0)
-
-```@example bad-conversion
-using InteractiveUtils
-InteractiveUtils.@code_llvm Int(1.0)
-```
-
-:::

src/compiler/codegen.jl

@@ -16,6 +16,8 @@ const _hip_compiler_cache = Dict{HIP.HIPDevice, Dict{Any, HIP.HIPFunction}}()
 
 # hash(fun, hash(f, hash(tt))) => HIPKernel
 const _kernel_instances = Dict{UInt, Runtime.HIPKernel}()
+# UInt (hash(job)) => Vector{Symbol} (global hostcall names)
+const _global_hostcalls = Dict{UInt, Vector{Symbol}}()
 
 function compiler_cache(dev::HIP.HIPDevice)
     get!(() -> Dict{UInt, Any}(), _hip_compiler_cache, dev)
@@ -34,6 +36,10 @@ function GPUCompiler.link_libraries!(
     invoke(GPUCompiler.link_libraries!,
         Tuple{CompilerJob{GCNCompilerTarget}, typeof(mod), typeof(undefined_fns)},
         job, mod, undefined_fns)
+
+    # Detect global hostcalls here, before optimizations & cleanup occur.
+    _global_hostcalls[hash(job)] = find_global_hostcalls(mod)
+
     # Link only if there are undefined functions.
     # Everything else was loaded in `finish_module!` stage.
     link_device_libs!(
@@ -189,21 +195,24 @@ function create_executable(obj)
     return bin
 end
 
-function hipcompile(@nospecialize(job::CompilerJob))
-    obj, meta = JuliaContext() do ctx
-        GPUCompiler.compile(:obj, job)
-    end
-
-    entry = LLVM.name(meta.entry)
-    globals = filter(isextinit, collect(LLVM.globals(meta.ir))) .|> LLVM.name
-
+function find_global_hostcalls(mod::LLVM.Module)
     global_hostcall_names = (
         :malloc_hostcall, :free_hostcall, :print_hostcall, :printf_hostcall)
+
     global_hostcalls = Symbol[]
-    for gbl in LLVM.globals(meta.ir), gbl_name in global_hostcall_names
+    for gbl in LLVM.globals(mod), gbl_name in global_hostcall_names
         occursin("__$gbl_name", LLVM.name(gbl)) || continue
         push!(global_hostcalls, gbl_name)
     end
+    return global_hostcalls
+end
+
+function hipcompile(@nospecialize(job::CompilerJob))
+    obj, meta = JuliaContext() do ctx
+        GPUCompiler.compile(:obj, job)
+    end
+
+    global_hostcalls = pop!(_global_hostcalls, hash(job))
     if !isempty(global_hostcalls)
         @info """Global hostcalls detected!
         - Source: $(job.source)
@@ -214,11 +223,13 @@ function hipcompile(@nospecialize(job::CompilerJob))
         """
     end
 
-    if !isempty(globals)
+    entry = LLVM.name(meta.entry)
+    extinit_globals = filter(isextinit, collect(LLVM.globals(meta.ir))) .|> LLVM.name
+    if !isempty(extinit_globals)
         @warn """
         HIP backend does not support setting extinit globals.
         But kernel `$entry` has following:
-        $globals
+        $extinit_globals
 
         Compilation will likely fail.
         """

src/hip/module.jl

@@ -4,7 +4,6 @@ mutable struct HIPModule
     function HIPModule(data)
         device_synchronize()
 
-        # TODO use alloc_retry?
         mod_ref = Ref{hipModule_t}()
         hipModuleLoadData(mod_ref, data)
         mod = new(mod_ref[])
@@ -21,9 +20,7 @@ struct HIPFunction
     mod::HIPModule
     global_hostcalls::Vector{Symbol}
 
-    function HIPFunction(
-        mod::HIPModule, name::String, global_hostcalls::Vector{Symbol},
-    )
+    function HIPFunction(mod::HIPModule, name::String, global_hostcalls::Vector{Symbol})
         fun_ref = Ref{hipFunction_t}()
         hipModuleGetFunction(fun_ref, mod, name)
         new(fun_ref[], mod, global_hostcalls)
@@ -32,9 +29,7 @@ end
 
 Base.unsafe_convert(::Type{hipFunction_t}, fun::HIPFunction) = fun.handle
 
-function launch_configuration(
-    fun::HIPFunction; shmem::Integer = 0, max_block_size::Integer = 0,
-)
+function launch_configuration(fun::HIPFunction; shmem::Integer = 0, max_block_size::Integer = 0)
     grid_size_ref, block_size_ref = Ref{Cint}(), Ref{Cint}()
     hipModuleOccupancyMaxPotentialBlockSize(
         grid_size_ref, block_size_ref, fun, shmem, max_block_size)