GCN: rewrite byref kernel params to addrspace(4) for s_load

Add finish_ir! for GCN that rewrites byref kernel parameters from
flat (addrspace 0) to kernarg (addrspace 4) after optimization.

Clang emits byref params as ptr addrspace(4) from the frontend,
but Julia's RemoveJuliaAddrspacesPass strips them to flat. This
causes struct field loads to use flat_load instead of s_load.

The pass creates a new function with ptr addrspace(4) parameters,
inserts addrspacecasts back to flat for the cloned IR, then runs
InferAddressSpaces to propagate addrspace(4) through all GEPs and
loads. The result is that all kernel argument struct field accesses
become s_load (scalar, cached, one per wavefront) instead of
flat_load (per-lane, address disambiguation overhead).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: gbaraldi

src/gcn.jl

@@ -49,15 +49,126 @@ function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
     if job.config.kernel
         # calling convention
         callconv!(entry, LLVM.API.LLVMAMDGPUKERNELCallConv)
-
-        # with byref, the AMDGPU backend's AMDGPULowerKernelArguments pass
-        # will handle loading from the kernarg segment directly.
-        # no need for lower_byval or manual rewriting.
     end
 
     return entry
 end
 
+function finish_ir!(@nospecialize(job::CompilerJob{GCNCompilerTarget}), mod::LLVM.Module,
+                    entry::LLVM.Function)
+    if job.config.kernel
+        entry = add_kernarg_address_spaces!(job, mod, entry)
+    end
+    return entry
+end
+
+# Rewrite byref kernel parameters from flat (addrspace 0) to kernarg (addrspace 4).
+#
+# On AMDGPU, the kernarg segment is in address space 4 and is scalar-loadable via s_load.
+# Clang emits byref parameters as `ptr addrspace(4)` from the frontend, but Julia's
+# RemoveJuliaAddrspacesPass strips all address spaces to flat. This pass restores the
+# correct address space so that struct field loads from byref arguments become s_load
+# instead of flat_load.
+#
+# NOTE: must run after optimization, where RemoveJuliaAddrspacesPass has already
+# converted Julia's addrspace(11) to flat (addrspace 0) on these parameters.
+function add_kernarg_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
+                                     f::LLVM.Function)
+    ft = function_type(f)
+
+    # find the byref parameters
+    byref_mask = BitVector(undef, length(parameters(ft)))
+    args = classify_arguments(job, ft; post_optimization=job.config.optimize)
+    filter!(args) do arg
+        arg.cc != GHOST
+    end
+    for arg in args
+        byref_mask[arg.idx] = (arg.cc == BITS_REF || arg.cc == KERNEL_STATE)
+    end
+
+    # check if any flat pointer byref params need rewriting
+    needs_rewrite = false
+    for (i, param) in enumerate(parameters(ft))
+        if byref_mask[i] && param isa LLVM.PointerType && addrspace(param) == 0
+            needs_rewrite = true
+            break
+        end
+    end
+    needs_rewrite || return f
+
+    # generate the new function type with kernarg address space on byref params
+    new_types = LLVMType[]
+    for (i, param) in enumerate(parameters(ft))
+        if byref_mask[i] && param isa LLVM.PointerType && addrspace(param) == 0
+            push!(new_types, LLVM.PointerType(#=kernarg=# 4))
+        else
+            push!(new_types, param)
+        end
+    end
+    new_ft = LLVM.FunctionType(return_type(ft), new_types)
+    new_f = LLVM.Function(mod, "", new_ft)
+    linkage!(new_f, linkage(f))
+    for (arg, new_arg) in zip(parameters(f), parameters(new_f))
+        LLVM.name!(new_arg, LLVM.name(arg))
+    end
+
+    # insert addrspacecasts from kernarg (4) back to flat (0) so that the cloned IR
+    # (which expects flat pointers) continues to work. InferAddressSpaces will then
+    # propagate addrspace(4) through GEPs and loads, eliminating the casts.
+    new_args = LLVM.Value[]
+    @dispose builder=IRBuilder() begin
+        entry_bb = BasicBlock(new_f, "conversion")
+        position!(builder, entry_bb)
+
+        for (i, param) in enumerate(parameters(ft))
+            if byref_mask[i] && param isa LLVM.PointerType && addrspace(param) == 0
+                cast = addrspacecast!(builder, parameters(new_f)[i], LLVM.PointerType(0))
+                push!(new_args, cast)
+            else
+                push!(new_args, parameters(new_f)[i])
+            end
+            for attr in collect(parameter_attributes(f, i))
+                push!(parameter_attributes(new_f, i), attr)
+            end
+        end
+
+        # clone the original function body
+        value_map = Dict{LLVM.Value, LLVM.Value}(
+            param => new_args[i] for (i, param) in enumerate(parameters(f))
+        )
+        value_map[f] = new_f
+        clone_into!(new_f, f; value_map,
+                    changes=LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges)
+
+        # fall through from conversion block to cloned entry
+        br!(builder, blocks(new_f)[2])
+    end
+
+    # replace the old function
+    fn = LLVM.name(f)
+    prune_constexpr_uses!(f)
+    @assert isempty(uses(f))
+    replace_metadata_uses!(f, new_f)
+    erase!(f)
+    LLVM.name!(new_f, fn)
+
+    # propagate addrspace(4) through GEPs and loads, then clean up
+    @dispose pb=NewPMPassBuilder() begin
+        add!(pb, NewPMFunctionPassManager()) do fpm
+            add!(fpm, InferAddressSpacesPass())
+        end
+        add!(pb, NewPMFunctionPassManager()) do fpm
+            add!(fpm, SimplifyCFGPass())
+            add!(fpm, SROAPass())
+            add!(fpm, EarlyCSEPass())
+            add!(fpm, InstCombinePass())
+        end
+        run!(pb, mod)
+    end
+
+    return functions(mod)[fn]
+end
+
 
 ## LLVM passes