Remove GCN alloca hacks when using Julia 1.9 (#342)

Co-authored-by: Julian P Samaroo <jpsamaroo@jpsamaroo.me>

Author: collinwarner

src/gcn.jl

@@ -63,12 +63,14 @@ function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
 end
 
 function optimize!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module)
-    # we have to fake our target early in the pipeline because Julia's
-    # optimization passes weren't designed for a non-0 stack addrspace, and the
-    # AMDGPU target is very strict about which addrspaces are permitted for
-    # various code patterns
-    triple!(mod, llvm_triple(NativeCompilerTarget()))
-    datalayout!(mod, julia_datalayout(NativeCompilerTarget()))
+    @static if VERSION < v"1.9.0-DEV.1018"
+        # we have to fake our target early in the pipeline because Julia's
+        # optimization passes weren't designed for a non-0 stack addrspace, and the
+        # AMDGPU target is very strict about which addrspaces are permitted for
+        # various code patterns
+        triple!(mod, llvm_triple(NativeCompilerTarget()))
+        datalayout!(mod, julia_datalayout(NativeCompilerTarget()))
+    end
 
     invoke(optimize!, Tuple{CompilerJob, LLVM.Module}, job, mod)
 end
@@ -80,18 +82,20 @@ end
 # 2. We don't want any chance of messing with Julia's optimizations, since they
 #    eliminate target-unsafe IR patterns
 function optimize_module!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module)
-    # revert back to the AMDGPU target
-    triple!(mod, llvm_triple(job.target))
-    datalayout!(mod, julia_datalayout(job.target))
+    @static if VERSION < v"1.9.0-DEV.1018"
+        # revert back to the AMDGPU target
+        triple!(mod, llvm_triple(job.target))
+        datalayout!(mod, julia_datalayout(job.target))
 
-    tm = llvm_machine(job.target)
-    @dispose pm=ModulePassManager() begin
-        add_library_info!(pm, triple(mod))
-        add_transform_info!(pm, tm)
+        tm = llvm_machine(job.target)
+        @dispose pm=ModulePassManager() begin
+            add_library_info!(pm, triple(mod))
+            add_transform_info!(pm, tm)
 
-        add!(pm, FunctionPass("FixAllocaAddrspace", fix_alloca_addrspace!))
+            add!(pm, FunctionPass("FixAllocaAddrspace", fix_alloca_addrspace!))
 
-        run!(pm, mod)
+            run!(pm, mod)
+        end
     end
 end
 

test/gcn.jl

@@ -1,3 +1,9 @@
+if VERSION >= v"1.9.0-DEV.1018"
+@inline sink_gcn(i) = sink(i, Val(5))
+else
+@inline sink_gcn(i) = sink(i, Val(0))
+end
+
 @testset "GCN" begin
 
 @test GCNCompilerTarget(dev_isa="gfx900") == GCNCompilerTarget("gfx900")
@@ -19,16 +25,18 @@ include("definitions/gcn.jl")
     @test occursin("amdgpu_kernel", ir)
 end
 
+if VERSION < v"1.9.0-DEV.1018"
 @testset "alloca addrspace" begin
     function kernel(i)
-        sink(i) # sink provides an alloca in addrspace 0
+        sink(i, Val(0)) # sink provides an alloca in addrspace 0
         return
     end
 
     ir = sprint(io->gcn_code_llvm(io, kernel, Tuple{Int64}; dump_module=true))
     @test occursin(r"alloca i64, (align 8, )?addrspace\(5\)$"m, ir)
     @test !occursin(r"alloca i64(, align \d)?$"m, ir)
 end
+end
 
 end
 
@@ -57,7 +65,7 @@ end
 @testset "child functions" begin
     # we often test using @noinline child functions, so test whether these survive
     # (despite not having side-effects)
-    @noinline child(i) = sink(i)
+    @noinline child(i) = sink_gcn(i)
     function parent(i)
         child(i)
         return
@@ -69,7 +77,7 @@ end
 end
 
 @testset "kernel functions" begin
-    @noinline nonentry(i) = sink(i)
+    @noinline nonentry(i) = sink_gcn(i)
     function entry(i)
         nonentry(i)
         return
@@ -85,7 +93,7 @@ end
     # bug: depending on a child function from multiple parents resulted in
     #      the child only being present once
 
-    @noinline child(i) = sink(i)
+    @noinline child(i) = sink_gcn(i)
     function parent1(i)
         child(i)
         return
@@ -106,8 +114,8 @@ end
 @testset "child function reuse bis" begin
     # bug: similar, but slightly different issue as above
     #      in the case of two child functions
-    @noinline child1(i) = sink(i)
-    @noinline child2(i) = sink(i+1)
+    @noinline child1(i) = sink_gcn(i)
+    @noinline child2(i) = sink_gcn(i+1)
     function parent1(i)
         child1(i) + child2(i)
         return

test/util.jl

@@ -23,16 +23,22 @@ macro test_throws_message(f, typ, ex...)
 end
 
 # helper function for sinking a value to prevent the callee from getting optimized away
-@inline sink(i::T) where T <: Union{Int32,UInt32} =
-    Base.llvmcall("""%slot = alloca i32
-                     store volatile i32 %0, i32* %slot
-                     %value = load volatile i32, i32* %slot
-                     ret i32 %value""", T, Tuple{T}, i)
-@inline sink(i::T) where T <: Union{Int64,UInt64} =
-    Base.llvmcall("""%slot = alloca i64
-                     store volatile i64 %0, i64* %slot
-                     %value = load volatile i64, i64* %slot
-                     ret i64 %value""", T, Tuple{T}, i)
+@inline @generated function sink(i::T, ::Val{addrspace}=Val(0)) where {T <: Union{Int32,UInt32}, addrspace}
+    as_str = addrspace > 0 ? " addrspace($addrspace)" : ""
+    llvmcall_str = """%slot = alloca i32$(addrspace > 0 ? ", addrspace($addrspace)" : "")
+                     store volatile i32 %0, i32$(as_str)* %slot
+                     %value = load volatile i32, i32$(as_str)* %slot
+                     ret i32 %value"""
+    return :(Base.llvmcall($llvmcall_str, T, Tuple{T}, i))
+end
+@inline @generated function sink(i::T, ::Val{addrspace}=Val(0)) where {T <: Union{Int64,UInt64}, addrspace}
+    as_str = addrspace > 0 ? " addrspace($addrspace)" : ""
+    llvmcall_str = """%slot = alloca i64$(addrspace > 0 ? ", addrspace($addrspace)" : "")
+                     store volatile i64 %0, i64$(as_str)* %slot
+                     %value = load volatile i64, i64$(as_str)* %slot
+                     ret i64 %value"""
+    return :(Base.llvmcall($llvmcall_str, T, Tuple{T}, i))
+end
 
 
 # the GPU runtime library