Rewrite byval passes using clone utils.

Also have them look for the byval attribute instead of
processing arguments from scratch again.

Author: maleadt

Manifest.toml

@@ -39,15 +39,17 @@ version = "1.3.0"
 
 [[LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "d6041ad706cf458b2c9f3e501152488a26451e9c"
+git-tree-sha1 = "effe3552dba16b1e9c3cc9beac454a3566f09637"
+repo-rev = "b4dfdfcf86dde4563c87e90d130c78dbbe8550f8"
+repo-url = "https://github.com/maleadt/LLVM.jl.git"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
 version = "4.2.0"
 
 [[LLVMExtra_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "a9b1130c4728b0e462a1c28772954650039eb847"
+git-tree-sha1 = "873e7962f14f6bdd8a0e10552d964ec0a7c69f3b"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.7+0"
+version = "0.0.9+0"
 
 [[LibCURL]]
 deps = ["LibCURL_jll", "MozillaCACerts_jll"]

src/gcn.jl

@@ -44,6 +44,8 @@ function process_module!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module)
 end
 
 function process_entry!(job::CompilerJob{GCNCompilerTarget}, mod::LLVM.Module, entry::LLVM.Function)
+    invoke(process_entry!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
+
     if job.source.kernel
         entry = lower_byval(job, mod, entry)
 

src/irgen.jl

@@ -382,114 +382,78 @@ end
 # some back-ends don't support byval, or support it badly
 # https://reviews.llvm.org/D79744
 
-# generate a kernel wrapper to fix & improve argument passing
-function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry_f::LLVM.Function)
+# modify the kernel function to fix & improve argument passing
+function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
     ctx = context(mod)
-    entry_ft = eltype(llvmtype(entry_f)::LLVM.PointerType)::LLVM.FunctionType
-    @compiler_assert return_type(entry_ft) == LLVM.VoidType(ctx) job
-
-    args = classify_arguments(job, entry_f)
-    filter!(args) do arg
-        arg.cc != GHOST
+    ft = eltype(llvmtype(f)::LLVM.PointerType)::LLVM.FunctionType
+    @compiler_assert return_type(ft) == LLVM.VoidType(ctx) job
+
+    # find the byval parameters
+    byval = BitVector(undef, length(parameters(ft)))
+    for i in 1:length(byval)
+        attrs = collect(parameter_attributes(f, i))
+        byval[i] = any(attrs) do attr
+            kind(attr) == kind(EnumAttribute("byval", 0; ctx))
+        end
     end
 
-    # generate the wrapper function type & definition
-    wrapper_types = LLVM.LLVMType[]
-    for arg in args
-        typ = if arg.cc == BITS_REF
-            eltype(arg.codegen.typ)
+    # generate the new function type & definition
+    new_types = LLVM.LLVMType[]
+    for (i, param) in enumerate(parameters(ft))
+        if byval[i]
+            push!(new_types, eltype(param::LLVM.PointerType))
         else
-            convert(LLVMType, arg.typ; ctx)
+            push!(new_types, param)
         end
-        push!(wrapper_types, typ)
     end
-    wrapper_fn = LLVM.name(entry_f)
-    LLVM.name!(entry_f, wrapper_fn * ".inner")
-    wrapper_ft = LLVM.FunctionType(LLVM.VoidType(ctx), wrapper_types)
-    wrapper_f = LLVM.Function(mod, wrapper_fn, wrapper_ft)
+    new_ft = LLVM.FunctionType(return_type(ft), new_types)
+    new_f = LLVM.Function(mod, "", new_ft)
+    linkage!(new_f, linkage(f))
 
     # emit IR performing the "conversions"
-    let builder = Builder(ctx)
-        entry = BasicBlock(wrapper_f, "entry"; ctx)
+    new_args = LLVM.Value[]
+    Builder(ctx) do builder
+        entry = BasicBlock(new_f, "entry"; ctx)
         position!(builder, entry)
 
-        wrapper_args = Vector{LLVM.Value}()
-
         # perform argument conversions
-        for arg in args
-            if arg.cc == BITS_REF
+        for (i, param) in enumerate(parameters(ft))
+            if byval[i]
                 # copy the argument value to a stack slot, and reference it.
-                ptr = alloca!(builder, eltype(arg.codegen.typ))
-                if LLVM.addrspace(arg.codegen.typ) != 0
-                    ptr = addrspacecast!(builder, ptr, arg.codegen.typ)
+                ptr = alloca!(builder, eltype(param))
+                if LLVM.addrspace(param) != 0
+                    ptr = addrspacecast!(builder, ptr, param)
                 end
-                store!(builder, parameters(wrapper_f)[arg.codegen.i], ptr)
-                push!(wrapper_args, ptr)
+                store!(builder, parameters(new_f)[i], ptr)
+                push!(new_args, ptr)
             else
-                push!(wrapper_args, parameters(wrapper_f)[arg.codegen.i])
-                for attr in collect(parameter_attributes(entry_f, arg.codegen.i))
-                    push!(parameter_attributes(wrapper_f, arg.codegen.i), attr)
+                push!(new_args, parameters(new_f)[i])
+                for attr in collect(parameter_attributes(f, i))
+                    push!(parameter_attributes(new_f, i), attr)
                 end
             end
         end
 
-        call!(builder, entry_f, wrapper_args)
-
-        ret!(builder)
-
-        dispose(builder)
-    end
-
-    # early-inline the original entry function into the wrapper
-    push!(function_attributes(entry_f), EnumAttribute("alwaysinline", 0; ctx))
-    linkage!(entry_f, LLVM.API.LLVMInternalLinkage)
+        # inline the old IR
+        value_map = Dict{LLVM.Value, LLVM.Value}(
+            param => new_args[i] for (i,param) in enumerate(parameters(f))
+        )
+        clone_into!(new_f, f; value_map,
+                    changes=LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges)
+        # NOTE: we need global changes because LLVM 12 wants to clone debug metadata
 
-    # copy debug info
-    sp = LLVM.get_subprogram(entry_f)
-    if sp !== nothing
-        LLVM.set_subprogram!(wrapper_f, sp)
+        # fall through
+        br!(builder, collect(blocks(new_f))[2])
     end
 
-    fixup_metadata!(entry_f)
-    ModulePassManager() do pm
-        always_inliner!(pm)
-        run!(pm, mod)
-    end
+    # remove the old function
+    # NOTE: if we ever have legitimate uses of the old function, create a shim instead
+    fn = LLVM.name(f)
+    @assert isempty(uses(f))
+    # XXX: there may still be metadata using this function. RAUW updates those,
+    #      but asserts on a debug build due to the updated function type.
+    unsafe_delete!(mod, f)
+    LLVM.name!(new_f, fn)
 
-    return wrapper_f
-end
-
-# HACK: get rid of invariant.load and const TBAA metadata on loads from pointer args,
-#       since storing to a stack slot violates the semantics of those attributes.
-# TODO: can we emit a wrapper that doesn't violate Julia's metadata?
-function fixup_metadata!(f::LLVM.Function)
-    for param in parameters(f)
-        if isa(llvmtype(param), LLVM.PointerType)
-            # collect all uses of the pointer
-            worklist = Vector{LLVM.Instruction}(user.(collect(uses(param))))
-            while !isempty(worklist)
-                value = popfirst!(worklist)
-
-                # remove the invariant.load attribute
-                md = metadata(value)
-                if haskey(md, LLVM.MD_invariant_load)
-                    delete!(md, LLVM.MD_invariant_load)
-                end
-                if haskey(md, LLVM.MD_tbaa)
-                    delete!(md, LLVM.MD_tbaa)
-                end
-
-                # recurse on the output of some instructions
-                if isa(value, LLVM.BitCastInst) ||
-                   isa(value, LLVM.GetElementPtrInst) ||
-                   isa(value, LLVM.AddrSpaceCastInst)
-                    append!(worklist, user.(collect(uses(value))))
-                end
-
-                # IMPORTANT NOTE: if we ever want to inline functions at the LLVM level,
-                # we need to recurse into call instructions here, and strip metadata from
-                # called functions (see CUDAnative.jl#238).
-            end
-        end
-    end
+    return new_f
 end

src/ptx.jl

@@ -114,8 +114,9 @@ end
 
 function process_entry!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                         mod::LLVM.Module, entry::LLVM.Function)
-    ctx = context(mod)
+    invoke(process_entry!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
 
+    ctx = context(mod)
     if job.source.kernel
         # work around bad byval codegen (JuliaGPU/GPUCompiler.jl#92)
         entry = lower_byval(job, mod, entry)

src/spirv.jl

@@ -40,6 +40,8 @@ function process_module!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module
 end
 
 function process_entry!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module, entry::LLVM.Function)
+    invoke(process_entry!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
+
     if job.source.kernel
         # HACK: Intel's compute runtime doesn't properly support SPIR-V's byval attribute.
         #       they do support struct byval, for OpenCL, so wrap byval parameters in a struct.
@@ -193,75 +195,83 @@ function rm_freeze!(mod::LLVM.Module)
 end
 
 # wrap byval pointers in a single-value struct
-function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry_f::LLVM.Function)
+function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
     ctx = context(mod)
-    entry_ft = eltype(llvmtype(entry_f)::LLVM.PointerType)::LLVM.FunctionType
-    @compiler_assert return_type(entry_ft) == LLVM.VoidType(ctx) job
-
-    args = classify_arguments(job, entry_f)
-    filter!(args) do arg
-        arg.cc != GHOST
+    ft = eltype(llvmtype(f)::LLVM.PointerType)::LLVM.FunctionType
+    @compiler_assert return_type(ft) == LLVM.VoidType(ctx) job
+
+    # find the byval parameters
+    byval = BitVector(undef, length(parameters(ft)))
+    for i in 1:length(byval)
+        attrs = collect(parameter_attributes(f, i))
+        byval[i] = any(attrs) do attr
+            kind(attr) == kind(EnumAttribute("byval", 0; ctx))
+        end
     end
 
     # generate the wrapper function type & definition
-    wrapper_types = LLVM.LLVMType[]
-    for arg in args
-        typ = if arg.cc == BITS_REF
-            st = LLVM.StructType([eltype(arg.codegen.typ)]; ctx)
-            LLVM.PointerType(st, addrspace(arg.codegen.typ))
+    new_types = LLVM.LLVMType[]
+    for (i, param) in enumerate(parameters(ft))
+        typ = if byval[i]
+            st = LLVM.StructType([eltype(param)]; ctx)
+            LLVM.PointerType(st, addrspace(param))
         else
-            convert(LLVMType, arg.typ; ctx)
+            param
         end
-        push!(wrapper_types, typ)
+        push!(new_types, typ)
     end
-    wrapper_fn = LLVM.name(entry_f)
-    LLVM.name!(entry_f, wrapper_fn * ".inner")
-    wrapper_ft = LLVM.FunctionType(LLVM.VoidType(ctx), wrapper_types)
-    wrapper_f = LLVM.Function(mod, wrapper_fn, wrapper_ft)
+    new_ft = LLVM.FunctionType(LLVM.VoidType(ctx), new_types)
+    new_f = LLVM.Function(mod, "", new_ft)
+    linkage!(new_f, linkage(f))
 
     # emit IR performing the "conversions"
-    let builder = Builder(ctx)
-        entry = BasicBlock(wrapper_f, "entry"; ctx)
+    new_args = Vector{LLVM.Value}()
+    Builder(ctx) do builder
+        entry = BasicBlock(new_f, "entry"; ctx)
         position!(builder, entry)
 
-        wrapper_args = Vector{LLVM.Value}()
-
         # perform argument conversions
-        for arg in args
-            param = parameters(wrapper_f)[arg.codegen.i]
-            attrs = parameter_attributes(wrapper_f, arg.codegen.i)
-            if arg.cc == BITS_REF
+        for (i, param) in enumerate(parameters(new_f))
+            if byval[i]
+                ptr = struct_gep!(builder, param, 0)
+                push!(new_args, ptr)
+            else
+                push!(new_args, param)
+            end
+        end
+
+        # inline the old IR
+        value_map = Dict{LLVM.Value, LLVM.Value}(
+            param => new_args[i] for (i,param) in enumerate(parameters(f))
+        )
+        clone_into!(new_f, f; value_map,
+                    changes=LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges)
+        # NOTE: we need global changes because LLVM 12 wants to clone debug metadata
+
+        # apply byval attributes again (`clone_into!` didn't due to the type mismatch)
+        for i in 1:length(byval)
+            attrs = parameter_attributes(new_f, i)
+            if byval[i]
                 if LLVM.version() >= v"12"
-                    push!(attrs, TypeAttribute("byval", eltype(wrapper_types[arg.codegen.i]); ctx))
+                    push!(attrs, TypeAttribute("byval", eltype(new_types[i]); ctx))
                 else
                     push!(attrs, EnumAttribute("byval", 0; ctx))
                 end
-                ptr = struct_gep!(builder, param, 0)
-                push!(wrapper_args, ptr)
-            else
-                push!(wrapper_args, param)
-                for attr in collect(attrs)
-                    push!(parameter_attributes(wrapper_f, arg.codegen.i), attr)
-                end
             end
         end
 
-        call!(builder, entry_f, wrapper_args)
-
-        ret!(builder)
-
-        dispose(builder)
+        # fall through
+        br!(builder, collect(blocks(new_f))[2])
     end
 
-    # early-inline the original entry function into the wrapper
-    delete!(function_attributes(entry_f), EnumAttribute("noinline", 0; ctx))
-    push!(function_attributes(entry_f), EnumAttribute("alwaysinline", 0; ctx))
-    linkage!(entry_f, LLVM.API.LLVMInternalLinkage)
-
-    ModulePassManager() do pm
-        always_inliner!(pm)
-        run!(pm, mod)
-    end
+    # remove the old function
+    # NOTE: if we ever have legitimate uses of the old function, create a shim instead
+    fn = LLVM.name(f)
+    @assert isempty(uses(f))
+    # XXX: there may still be metadata using this function. RAUW updates those,
+    #      but asserts on a debug build due to the updated function type.
+    unsafe_delete!(mod, f)
+    LLVM.name!(new_f, fn)
 
-    return wrapper_f
+    return new_f
 end