Revert "PTX: only lower byval arguments the back-end would copy." (#838)

This reverts commit 189c66e.

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "1.19.0"
+version = "1.19.1"
 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 for kernel arguments?
+# Should the target use byref instead of byval+lower_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,53 +660,24 @@ end
 
 # byval lowering
 #
-# 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
-
+# some back-ends don't support byval, or support it badly, so lower it eagerly ourselves
+# https://reviews.llvm.org/D79744
 function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
     ft = function_type(f)
     @tracepoint "lower byval" begin
 
-    # find the byval parameters that need lowering
+    # find the byval parameters
     byval = BitVector(undef, length(parameters(ft)))
     types = Vector{LLVMType}(undef, length(parameters(ft)))
-    for (i, param) in enumerate(parameters(f))
+    for i in 1:length(byval)
         byval[i] = false
         for attr in collect(parameter_attributes(f, i))
             if kind(attr) == kind(TypeAttribute("byval", LLVM.VoidType()))
-                byval[i] = !loads_parameter_directly(param)
+                byval[i] = true
                 types[i] = value(attr)
             end
         end
     end
-    any(byval) || return f
 
     # fixup metadata
     #
@@ -807,7 +778,6 @@ 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.,
@@ -1213,7 +1183,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.
+# serves to convert that argument (and is conceptually the inverse of `lower_byval`).
 function kernel_state_to_reference!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                                     f::LLVM.Function)
     ft = function_type(f)

src/ptx.jl

@@ -190,7 +190,7 @@ function finish_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
     end
 
     if job.config.kernel
-        # materialize byval arguments the back-end would otherwise copy to local memory
+        # work around bad byval codegen (JuliaGPU/GPUCompiler.jl#92)
         entry = lower_byval(job, mod, entry)
 
         # emit kernel property annotations into the module. these have to be in

test/ptx.jl

@@ -28,9 +28,9 @@ end
     end
 
     @test @filecheck begin
-        # aggregates are passed `byval`; the back-end loads them from parameter space
         @check_label "define ptx_kernel void @_Z6kernel9Aggregate"
-        @check_same "byval({{({ i64 }|\\[1 x i64\\])}})"
+        @check_not cond=typed_ptrs "*"
+        @check_not cond=opaque_ptrs "ptr"
         PTX.code_llvm(mod.kernel, Tuple{mod.Aggregate}; kernel=true)
     end
 end