Address review feedback: fix comments, add tests, Runic formatting

- Fix addrspace(4) comment: it's the "constant" address space, not "kernarg"
- Rewrite doc comment to accurately describe the Julia → AS(11) → AS(0) → AS(4) flow
- Add InferAddressSpaces + SROA + InstCombine + EarlyCSE after kernarg rewrite
- Add tests: byref attribute preservation, mixed byref/scalar params,
  programmatic IR inspection via compile(:llvm), zero scratch spills
- Apply Runic formatting to new code

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

Author: gbaraldi

src/gcn.jl

@@ -54,15 +54,17 @@ function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
     return entry
 end
 
-function finish_ir!(@nospecialize(job::CompilerJob{GCNCompilerTarget}), mod::LLVM.Module,
-                    entry::LLVM.Function)
+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)
 
         # optimize after address space rewriting: propagate addrspace(4) through
         # the addrspacecast chains, then clean up newly-exposed opportunities
         tm = llvm_machine(job.config.target)
-        @dispose pb=NewPMPassBuilder() tm begin
+        @dispose pb = NewPMPassBuilder() tm begin
             add!(pb, NewPMFunctionPassManager()) do fpm
                 add!(fpm, InferAddressSpacesPass())
                 add!(fpm, SROAPass())
@@ -76,23 +78,26 @@ function finish_ir!(@nospecialize(job::CompilerJob{GCNCompilerTarget}), mod::LLV
     return entry
 end
 
-# Rewrite byref kernel parameters from flat (addrspace 0) to kernarg (addrspace 4).
+# Rewrite byref kernel parameters from flat (addrspace 0) to constant (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.
+# On AMDGPU, kernel arguments reside in the constant address space (addrspace 4),
+# which is scalar-loadable via s_load. Julia initially emits byref parameters as
+# pointers in addrspace(11) (tracked/derived), but RemoveJuliaAddrspacesPass strips
+# all non-integral address spaces to flat (addrspace 0) during optimization. This pass
+# restores addrspace(4) on byref parameters so that the backend can emit s_load
+# instead of flat_load for struct field accesses.
 #
 # 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)
+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)
+    args = classify_arguments(job, ft; post_optimization = job.config.optimize)
     filter!(args) do arg
         arg.cc != GHOST
     end
@@ -114,7 +119,7 @@ function add_kernarg_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.
     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))
+            push!(new_types, LLVM.PointerType(#=constant=# 4))
         else
             push!(new_types, param)
         end
@@ -130,7 +135,7 @@ function add_kernarg_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.
     # (which expects flat pointers) continues to work. The AMDGPU backend's
     # AMDGPULowerKernelArguments traces these casts and produces s_load.
     new_args = LLVM.Value[]
-    @dispose builder=IRBuilder() begin
+    @dispose builder = IRBuilder() begin
         entry_bb = BasicBlock(new_f, "conversion")
         position!(builder, entry_bb)
 
@@ -148,8 +153,10 @@ function add_kernarg_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.
             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)
+        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])
@@ -174,7 +181,7 @@ function add_kernarg_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.
     LLVM.name!(new_f, fn)
 
     # clean up the extra conversion block
-    @dispose pb=NewPMPassBuilder() begin
+    @dispose pb = NewPMPassBuilder() begin
         add!(pb, NewPMFunctionPassManager()) do fpm
             add!(fpm, SimplifyCFGPass())
         end

test/gcn.jl

@@ -63,6 +63,96 @@ end
     end
 end
 
+@testset "byref attribute preserved on kernarg parameters" begin
+    mod = @eval module $(gensym())
+        struct LargeStruct
+            a::Float64
+            b::Float64
+            c::Float64
+            d::Float64
+        end
+
+        function kernel(s::LargeStruct, out::Ptr{Float64})
+            unsafe_store!(out, s.a + s.b + s.c + s.d)
+            return
+        end
+    end
+
+    # the byref attribute must survive the addrspace rewrite (clone_into! can drop it)
+    @test @filecheck begin
+        check"CHECK: byref"
+        check"CHECK: addrspace(4)"
+        GCN.code_llvm(mod.kernel, Tuple{mod.LargeStruct, Ptr{Float64}};
+                       dump_module=true, kernel=true)
+    end
+end
+
+@testset "mixed byref and scalar kernel parameters" begin
+    mod = @eval module $(gensym())
+        struct Params
+            x::Float64
+            y::Float64
+        end
+
+        function kernel(a::Float64, s::Params, out::Ptr{Float64})
+            unsafe_store!(out, a + s.x + s.y)
+            return
+        end
+    end
+
+    # scalar Float64 and Ptr should NOT be in addrspace(4),
+    # only the struct byref param should be
+    @test @filecheck begin
+        check"CHECK: define amdgpu_kernel void"
+        check"CHECK-SAME: double"
+        check"CHECK-SAME: ptr addrspace(4)"
+        check"CHECK-SAME: ptr"
+        GCN.code_llvm(mod.kernel, Tuple{Float64, mod.Params, Ptr{Float64}};
+                       dump_module=true, kernel=true)
+    end
+end
+
+@testset "add_kernarg_address_spaces! rewrites IR correctly" begin
+    mod = @eval module $(gensym())
+        struct KernelArgs
+            x::Float64
+            y::Float64
+            z::Float64
+        end
+
+        function kernel(s::KernelArgs, scale::Float64, out::Ptr{Float64})
+            unsafe_store!(out, (s.x + s.y + s.z) * scale)
+            return
+        end
+    end
+
+    job, _ = GCN.create_job(mod.kernel, Tuple{mod.KernelArgs, Float64, Ptr{Float64}};
+                             kernel=true)
+    JuliaContext() do ctx
+        ir, meta = GPUCompiler.compile(:llvm, job)
+
+        entry = meta.entry
+        ft = function_type(entry)
+        params = parameters(ft)
+
+        # the struct byref param should be ptr addrspace(4)
+        has_as4 = any(p -> p isa LLVM.PointerType && addrspace(p) == 4, params)
+        @test has_as4
+
+        # non-struct params (double, ptr) should NOT be in addrspace(4)
+        non_as4_ptrs = filter(params) do p
+            p isa LLVM.PointerType && addrspace(p) != 4
+        end
+        @test !isempty(non_as4_ptrs)  # the Ptr{Float64} out param
+
+        # byref attribute must be present
+        ir_str = string(ir)
+        @test occursin("byref", ir_str)
+
+        dispose(ir)
+    end
+end
+
 @testset "https://github.com/JuliaGPU/AMDGPU.jl/issues/846" begin
     ir, rt = GCN.code_typed((Tuple{Tuple{Val{4}}, Tuple{Float32}},); always_inline=true) do t
         t[1]
@@ -88,13 +178,35 @@ end
         end
     end
 
+    # struct field loads from kernarg should use s_load, not flat_load
     @test @filecheck begin
         check"CHECK: s_load_dwordx"
         check"CHECK-NOT: flat_load"
         GCN.code_native(mod.kernel, Tuple{mod.MyStruct, Ptr{Float64}}; kernel=true)
     end
 end
 
+@testset "no scratch spills for small struct kernarg" begin
+    mod = @eval module $(gensym())
+        struct SmallStruct
+            x::Float64
+            y::Float64
+        end
+
+        function kernel(s::SmallStruct, out::Ptr{Float64})
+            unsafe_store!(out, s.x + s.y)
+            return
+        end
+    end
+
+    # a small struct kernel should not need scratch memory
+    @test @filecheck begin
+        check"CHECK: .private_segment_fixed_size: 0"
+        GCN.code_native(mod.kernel, Tuple{mod.SmallStruct, Ptr{Float64}};
+                         dump_module=true, kernel=true)
+    end
+end
+
 @testset "skip scalar trap" begin
     mod = @eval module $(gensym())
         workitem_idx_x() = ccall("llvm.amdgcn.workitem.id.x", llvmcall, Int32, ())