Metal: Change approach to adding ptr arg ASs. (#340)

Turns out that using module cloning isn't working here. If we change
the address space of the arguments, and just remap those arguments
using a value mapper, those changes will get lost after, e.g., a
bitcast from the argument pointer to a different type (that destination
type won't have an address space, resulting in invalid operations).

We were working around this limitation by using a type remapper,
rewriting all pointer types to contain an address space, but that's
obviously invalid: Stack space (i.e. `alloca`s) should still reside
in address space 0.

So instead we're taking the approach we're using with byval lowering,
storing the arguments to a temporary slot. We're still changing the
function's pointer arguments to include an address space, but after
loading the argument and storing it in a stack slot we can dereference
it and get a valid pointer without an address space. We can then
re-use the existing IR with that pointer, without having to rewrite it.

This simplifies the pass a lot, because we don't have to worry about
rewriting intrinsics like memcpy.

Author: maleadt

src/metal.jl

@@ -109,7 +109,7 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
     entry_fn = LLVM.name(entry)
 
     if job.source.kernel
-        add_address_spaces!(mod, entry)
+        add_address_spaces!(job, mod, entry)
     end
 
     return functions(mod)[entry_fn]
@@ -147,103 +147,95 @@ end
 
 # generic pointer removal
 #
-# every pointer argument (e.g. byref objs) to a kernel needs an address space attached.
-function add_address_spaces!(mod::LLVM.Module, f::LLVM.Function)
+# every pointer argument (i.e. byref objs) to a kernel needs an address space attached.
+# this pass rewrites pointers to reference arguments to be located in address space 1.
+#
+# NOTE: this pass only rewrites byref objs, not plain pointers being passed; the user is
+# responsible for making sure these pointers have an address space attached (using LLVMPtr).
+#
+# NOTE: this pass also only rewrites pointers _without_ address spaces, which requires it to
+# be executed after optimization (where Julia's address spaces are stripped). If we ever
+# want to execute it earlier, adapt remapType to rewrite all pointer types.
+function add_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
     ctx = context(mod)
     ft = eltype(llvmtype(f))
 
+    # find the byref parameters
+    byref = BitVector(undef, length(parameters(ft)))
+    let args = classify_arguments(job, ft)
+        filter!(args) do arg
+            arg.cc != GHOST
+        end
+        for arg in args
+            byref[arg.codegen.i] = (arg.cc == BITS_REF)
+        end
+    end
+
     function remapType(src)
-        # TODO: recurse in structs
+        # TODO: cache?
+        # TODO: recurse in structs?
+        # TODO: when wrapping non-AS1 pointers, shouldn't the parent object use the same AS?
         dst = if src isa LLVM.PointerType && addrspace(src) == 0
             LLVM.PointerType(remapType(eltype(src)), #=device=# 1)
         else
             src
         end
-        # TODO: cache
         return dst
     end
 
     # generate the new function type & definition
-    new_types = LLVMType[remapType(typ) for typ in parameters(ft)]
+    new_types = LLVMType[]
+    for (i, param) in enumerate(parameters(ft))
+        if byref[i]
+            push!(new_types, remapType(param::LLVM.PointerType))
+        else
+            push!(new_types, param)
+        end
+    end
     new_ft = LLVM.FunctionType(LLVM.return_type(ft), new_types)
     new_f = LLVM.Function(mod, "", new_ft)
     linkage!(new_f, linkage(f))
     for (arg, new_arg) in zip(parameters(f), parameters(new_f))
         LLVM.name!(new_arg, LLVM.name(arg))
     end
 
-    # map the parameters
-    value_map = Dict{LLVM.Value, LLVM.Value}(
-        param => new_param for (param, new_param) in zip(parameters(f), parameters(new_f))
-    )
-    value_map[f] = new_f
-
-    # before D96531 (part of LLVM 13), clone_into! wants to duplicate debug metadata
-    # when the functions are part of the same module. that is invalid, because it
-    # results in desynchronized debug intrinsics (GPUCompiler#284), so remove those.
-    if LLVM.version() < v"13"
-        removals = LLVM.Instruction[]
-        for bb in blocks(f), inst in instructions(bb)
-            if inst isa LLVM.CallInst && LLVM.name(called_value(inst)) == "llvm.dbg.declare"
-                push!(removals, inst)
+    # we cannot simply remap the function arguments, because that will not propagate the
+    # address space changes across, e.g, bitcasts (the dest would still be in AS 0).
+    # using a type remapper on the other hand changes too much, including unrelated insts.
+    # so instead, we load the arguments in stack slots and dereference them so that we can
+    # keep on using the original IR that assumed pointers without address spaces
+    new_args = LLVM.Value[]
+    @dispose builder=Builder(ctx) begin
+        entry = BasicBlock(new_f, "conversion"; ctx)
+        position!(builder, entry)
+
+        # perform argument conversions
+        for (i, param) in enumerate(parameters(ft))
+            if byref[i]
+                # load the argument in a stack slot
+                val = load!(builder, parameters(new_f)[i])
+                ptr = alloca!(builder, eltype(param))
+                store!(builder, val, ptr)
+                push!(new_args, ptr)
+            else
+                push!(new_args, parameters(new_f)[i])
+            end
+            for attr in collect(parameter_attributes(f, i))
+                push!(parameter_attributes(new_f, i), attr)
             end
         end
-        for inst in removals
-            @assert isempty(uses(inst))
-            unsafe_delete!(LLVM.parent(inst), inst)
-        end
-        changes = LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges
-    else
-        changes = LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly
-    end
 
-    function type_mapper(typ)
-        remapType(typ)
-    end
-
-    clone_into!(new_f, f; value_map, changes, type_mapper)
-
-    # update calls to overloaded intrinsic, re-mangling their names
-    # XXX: shouldn't clone_into! do this?
-    LLVM.@dispose builder=Builder(ctx) begin
-        for bb in blocks(new_f), inst in instructions(bb)
-            if inst isa LLVM.CallBase
-                callee_f = called_value(inst)
-                LLVM.isintrinsic(callee_f) || continue
-                intr = Intrinsic(callee_f)
-                isoverloaded(intr) || continue
-
-                # get an appropriately-overloaded intrinsic instantiation
-                # NOTE: the overload types differs from the argument types
-                intr_f = if intr == Intrinsic("llvm.memcpy")
-                    LLVM.Function(mod, intr, llvmtype.(arguments(inst)[1:end-1]))
-                elseif intr == Intrinsic("llvm.lifetime.start") ||
-                       intr == Intrinsic("llvm.lifetime.end")
-                    LLVM.Function(mod, intr, [llvmtype(arguments(inst)[end])])
-                else
-                    # TODO: use matchIntrinsicSignature to do this generically
-                    error("""Unsupported intrinsic call:
-                                  $inst.
-                             Please file an issue with at https://github.com/JuliaGPU/GPUCompiler.jl""")
-                end
+        # map the arguments
+        value_map = Dict{LLVM.Value, LLVM.Value}(
+            param => new_args[i] for (i,param) in enumerate(parameters(f))
+        )
 
-                # create a call to the new intrinsic
-                # TODO: wrap setCalledFunction instead of using an IRBuilder
-                position!(builder, inst)
-                new_inst = if inst isa LLVM.CallInst
-                    call!(builder, intr_f, arguments(inst), operand_bundles(inst))
-                else
-                    # TODO: invoke and callbr
-                    error("Rewrite of $(typeof(inst))-based calls is not implemented: $inst")
-                end
-                callconv!(new_inst, callconv(inst))
+        value_map[f] = new_f
+        clone_into!(new_f, f; value_map,
+                    changes=LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges)
 
-                # replace the old call
-                replace_uses!(inst, new_inst)
-                @assert isempty(uses(inst))
-                unsafe_delete!(LLVM.parent(inst), inst)
-            end
-        end
+        # fall through
+        br!(builder, blocks(new_f)[2])
     end
 
     # remove the old function
@@ -253,6 +245,16 @@ function add_address_spaces!(mod::LLVM.Module, f::LLVM.Function)
     unsafe_delete!(mod, f)
     LLVM.name!(new_f, fn)
 
+    # clean-up after this pass (which runs after optimization)
+    @dispose pm=ModulePassManager() begin
+        cfgsimplification!(pm)
+        scalar_repl_aggregates!(pm)
+        early_cse!(pm)
+        instruction_combining!(pm)
+
+        run!(pm, mod)
+    end
+
     return new_f
 end
 

test/metal.jl

@@ -31,26 +31,6 @@ end
     @test occursin(r"@.*julia.*kernel.*\(({ i64 }|\[1 x i64\]) addrspace\(1\)\*", ir)
 end
 
-@testset "byref aggregates with memcpy" begin
-    ir = """
-        declare void @llvm.memcpy.p0i8.p0i8.i64(i8*, i8*, i64, i1 immarg)
-
-        define void @kernel(i8* %0, i8* %1) {
-        entry:
-            call void @llvm.memcpy.p0i8.p0i8.i64(i8* %0, i8* %1, i64 0, i1 false)
-            ret void
-        }
-    """
-    Context() do ctx
-        mod = parse(LLVM.Module, ir; ctx)
-        f = functions(mod)["kernel"]
-        GPUCompiler.add_address_spaces!(mod, f)
-        LLVM.verify(mod)
-    end
-    return
-
-end
-
 @testset "byref primitives" begin
     kernel(x) = return