GCN: fold null checks of stack objects

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "2.13.0"
+version = "2.13.1"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

src/gcn.jl

@@ -85,6 +85,7 @@ have_fma(@nospecialize(target::GCNCompilerTarget), T::Type) = true
 function finish_module!(@nospecialize(job::CompilerJob{GCNCompilerTarget}),
                         mod::LLVM.Module, entry::LLVM.Function)
     lower_throw_extra!(job, mod)
+    fold_alloca_null_checks!(mod)
 
     if job.config.kernel
         # calling convention
@@ -108,6 +109,9 @@ function finish_ir!(
         @nospecialize(job::CompilerJob{GCNCompilerTarget}), mod::LLVM.Module,
         entry::LLVM.Function
     )
+    # (optimization can reduce a phi of an alloca and null to the alloca, exposing more)
+    fold_alloca_null_checks!(mod)
+
     if job.config.kernel
         entry = add_kernarg_address_spaces!(job, mod, entry)
 
@@ -133,6 +137,37 @@ function finish_ir!(
     return entry
 end
 
+# Fold comparisons of a stack object's address against null. Julia 1.10 represents phis of
+# immutables kept in memory as phis of pointers, with null for undefined incoming values, and
+# guards their loads with `ptr == null ? undef : load(ptr)`. Elsewhere, LLVM folds these checks
+# once the phi reduces to an alloca, but AMDGPU allocates in the private address space, where
+# null is a valid address: the check survives, and the stack object at offset 0 reads as null,
+# i.e. as `undef`. No stack object Julia allocates is at null, so fold them all.
+function fold_alloca_null_checks!(mod::LLVM.Module)
+    worklist = LLVM.ICmpInst[]
+    for f in mod.functions, bb in f.blocks, inst in bb.instructions
+        inst isa LLVM.ICmpInst || continue
+        inst.predicate in (LLVM.IntPredicate.EQ, LLVM.IntPredicate.NE) || continue
+        lhs, rhs = inst.operands
+        ptr = isnull(rhs) ? lhs : isnull(lhs) ? rhs : continue
+        ptr.value_type isa LLVM.PointerType || continue
+        # (in-bounds offsets from a stack object do not reach null either)
+        while true
+            ptr = strip_pointer_casts(ptr)
+            ptr isa LLVM.GetElementPtrInst && ptr.inbounds || break
+            ptr = ptr.operands[1]
+        end
+        ptr isa LLVM.AllocaInst && push!(worklist, inst)
+    end
+
+    for inst in worklist
+        is_ne = inst.predicate == LLVM.IntPredicate.NE
+        replace_uses!(inst, ConstantInt(LLVM.Int1Type(), Int(is_ne)))
+        erase!(inst)
+    end
+    return !isempty(worklist)
+end
+
 # Rewrite byref kernel parameters from flat (addrspace 0) to constant (addrspace 4).
 #
 # On AMDGPU, kernel arguments reside in the constant address space (addrspace 4),

test/gcn.jl

@@ -56,6 +56,39 @@ end
     end
 end
 
+@testset "stack objects are not null" begin
+    # Julia 1.10 keeps `acc` in a stack slot and guards its loads with a null check, which
+    # holds for the slot at private address 0
+    mod = @eval module $(gensym())
+        function kernel(dst::Ptr{ComplexF64}, src::Ptr{ComplexF64}, init::ComplexF64,
+                        len::Int, i::Int)
+            acc = init
+            j = 0
+            while j + 128 <= len
+                for s in 0:3
+                    acc += unsafe_load(src, j + 32s + i + 1)
+                end
+                j += 128
+            end
+            if j < len
+                for s in 0:3
+                    k = j + 32s + i
+                    k < len && (acc += unsafe_load(src, k + 1))
+                end
+            end
+            unsafe_store!(dst, acc)
+            return
+        end
+    end
+
+    @test @filecheck begin
+        @check_label "define amdgpu_kernel void @_Z6kernel"
+        @check_not "icmp {{(eq|ne)}} {{.*}}addrspace(5){{.*}}, null"
+        GCN.code_llvm(mod.kernel, Tuple{Ptr{ComplexF64}, Ptr{ComplexF64}, ComplexF64, Int, Int};
+                      kernel=true)
+    end
+end
+
 @testset "synchronization scopes" begin
     # front-ends spell scopes like LLVM's SPIR-V back-end; they're renamed for the target
     mod = @eval module $(gensym())
@@ -90,8 +123,14 @@ end
         source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter())
         target = GCNCompilerTarget(; dev_isa="gfx90a", backend)
         job = CompilerJob(source, CompilerConfig(target, GCN.CompilerParams(); kernel=true))
-        datalayout = @dispose dl=GPUCompiler.llvm_datalayout(target) begin
-            string(dl)
+        # (a string for the external back-end)
+        dl = GPUCompiler.llvm_datalayout(target)
+        datalayout = if dl isa LLVM.DataLayout
+            @dispose dl=dl begin
+                string(dl)
+            end
+        else
+            dl
         end
         Context(; opaque_pointers=true) do ctx
             mod = parse(LLVM.Module, """