Merge pull request #836 from JuliaGPU/tb/ptx_llvm22

PTX: modernize compilation now that LLVM 22 is used

Author: maleadt

Project.toml

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

src/interface.jl

@@ -296,7 +296,7 @@ kernel_state_type(@nospecialize(job::CompilerJob)) = Nothing
 # Does the target need to pass kernel arguments by value?
 pass_by_value(@nospecialize(job::CompilerJob)) = true
 
-# Should the target use byref instead of byval+lower_byval for kernel arguments?
+# Should the target use byref instead of byval for kernel arguments?
 # When true, aggregate arguments are passed as pointers with the byref attribute,
 # allowing the backend to load fields directly from the argument memory (e.g. kernarg
 # segment on AMDGPU) instead of materializing the entire struct via first-class aggregates.

src/irgen.jl

@@ -660,24 +660,53 @@ end
 
 # byval lowering
 #
-# some back-ends don't support byval, or support it badly, so lower it eagerly ourselves
-# https://reviews.llvm.org/D79744
+# the NVPTX back-end accesses byval kernel arguments directly in parameter space when all
+# uses are simple loads (possibly via GEPs), but otherwise falls back to copying the
+# argument to local memory during machine-code generation, too late for the optimizer to
+# clean up. eagerly materialize such arguments ourselves instead, leaving the directly-
+# accessible ones to the back-end. (we historically lowered all byval arguments, working
+# around bad codegen in old back-ends; see #92 and https://reviews.llvm.org/D79744.)
+
+# can the back-end service all uses of this argument straight from parameter space?
+function loads_parameter_directly(param::LLVM.Argument)
+    worklist = LLVM.Value[param]
+    while !isempty(worklist)
+        value = popfirst!(worklist)
+        for use in uses(value)
+            inst = user(use)
+            if inst isa LLVM.LoadInst
+                # serviced by ld.param
+            elseif inst isa LLVM.GetElementPtrInst ||
+                   inst isa LLVM.BitCastInst ||
+                   inst isa LLVM.AddrSpaceCastInst
+                push!(worklist, inst)
+            else
+                # phis, selects, stores, calls (e.g. memcpy), etc. make the back-end
+                # fall back to a local copy
+                return false
+            end
+        end
+    end
+    return true
+end
+
 function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
     ft = function_type(f)
     @tracepoint "lower byval" begin
 
-    # find the byval parameters
+    # find the byval parameters that need lowering
     byval = BitVector(undef, length(parameters(ft)))
     types = Vector{LLVMType}(undef, length(parameters(ft)))
-    for i in 1:length(byval)
+    for (i, param) in enumerate(parameters(f))
         byval[i] = false
         for attr in collect(parameter_attributes(f, i))
             if kind(attr) == kind(TypeAttribute("byval", LLVM.VoidType()))
-                byval[i] = true
+                byval[i] = !loads_parameter_directly(param)
                 types[i] = value(attr)
             end
         end
     end
+    any(byval) || return f
 
     # fixup metadata
     #
@@ -778,6 +807,7 @@ function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.
 end
 
 
+
 # kernel state arguments
 #
 # to facilitate passing stateful information to kernels without having to recompile, e.g.,
@@ -1183,7 +1213,7 @@ end
 #
 # the kernel state argument is always passed by value to avoid codegen issues with byval.
 # some back-ends however do not support passing kernel arguments by value, so this pass
-# serves to convert that argument (and is conceptually the inverse of `lower_byval`).
+# serves to convert that argument.
 function kernel_state_to_reference!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                                     f::LLVM.Function)
     ft = function_type(f)