add_input_arguments! for other backends

Allows other backends to pass additional hidden arguments that can be accessed through intrinsics. Required for OpenCL device-side RNG support, where additional shared memory must be passed as arguments to the kernel.

Replaces #717

Author: simeonschaub

src/irgen.jl

@@ -790,44 +790,47 @@ function cleanup_kernel_state!(mod::LLVM.Module)
 end
 CleanupKernelStatePass() = NewPMModulePass("CleanupKernelStatePass", cleanup_kernel_state!)
 
-function kernel_state_intr(mod::LLVM.Module, T_state)
-    state_intr = if haskey(functions(mod), "julia.gpu.state_getter")
-        functions(mod)["julia.gpu.state_getter"]
+function custom_intr(mod::LLVM.Module, T::LLVMType, name::String)
+    custom_intr = if haskey(functions(mod), name)
+        functions(mod)[name]
     else
-        LLVM.Function(mod, "julia.gpu.state_getter", LLVM.FunctionType(T_state))
+        LLVM.Function(mod, name, LLVM.FunctionType(T))
     end
-    push!(function_attributes(state_intr), EnumAttribute("readnone", 0))
+    push!(function_attributes(custom_intr), EnumAttribute("readnone", 0))
 
-    return state_intr
+    return custom_intr
 end
 
 # run-time equivalent
-function kernel_state_value(state)
+function call_custom_intrinsic(T::Type, name::String, call_name::String)
     @dispose ctx=Context() begin
-        T_state = convert(LLVMType, state)
+        T_llvm = convert(LLVMType, T)
 
         # create function
-        llvm_f, _ = create_function(T_state)
+        llvm_f, _ = create_function(T_llvm)
         mod = LLVM.parent(llvm_f)
 
         # get intrinsic
-        state_intr = kernel_state_intr(mod, T_state)
-        state_intr_ft = function_type(state_intr)
+        _custom_intr = custom_intr(mod, T_llvm, name)
+        custom_intr_ft = function_type(_custom_intr)
 
         # generate IR
         @dispose builder=IRBuilder() begin
             entry = BasicBlock(llvm_f, "entry")
             position!(builder, entry)
 
-            val = call!(builder, state_intr_ft, state_intr, Value[], "state")
+            val = call!(builder, custom_intr_ft, _custom_intr, Value[], call_name)
 
             ret!(builder, val)
         end
 
-        call_function(llvm_f, state)
+        call_function(llvm_f, T)
     end
 end
 
+kernel_state_intr(mod::LLVM.Module, T::LLVMType) = custom_intr(mod, T, "julia.gpu.state_getter")
+kernel_state_value(T::Type) = call_custom_intrinsic(T, "julia.gpu.state_getter", "state")
+
 # convert kernel state argument from pass-by-value to pass-by-reference
 #
 # the kernel state argument is always passed by value to avoid codegen issues with byval.
@@ -923,3 +926,160 @@ function kernel_state_to_reference!(@nospecialize(job::CompilerJob), mod::LLVM.M
         return new_f
     end
 end
+
+function add_input_arguments!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
+                              entry::LLVM.Function, kernel_intrinsics::Dict)
+    entry_fn = LLVM.name(entry)
+
+    # figure out which intrinsics are used and need to be added as arguments
+    used_intrinsics = filter(keys(kernel_intrinsics)) do intr_fn
+        haskey(functions(mod), intr_fn)
+    end |> collect
+    nargs = length(used_intrinsics)
+
+    # determine which functions need these arguments
+    worklist = Set{LLVM.Function}([entry])
+    for intr_fn in used_intrinsics
+        push!(worklist, functions(mod)[intr_fn])
+    end
+    worklist_length = 0
+    while worklist_length != length(worklist)
+        # iteratively discover functions that use an intrinsic or any function calling it
+        worklist_length = length(worklist)
+        additions = LLVM.Function[]
+        for f in worklist, use in uses(f)
+            inst = user(use)::Instruction
+            bb = LLVM.parent(inst)
+            new_f = LLVM.parent(bb)
+            in(new_f, worklist) || push!(additions, new_f)
+        end
+        for f in additions
+            push!(worklist, f)
+        end
+    end
+    for intr_fn in used_intrinsics
+        delete!(worklist, functions(mod)[intr_fn])
+    end
+
+    # add the arguments
+    # NOTE: we don't need to be fine-grained here, as unused args will be removed during opt
+    workmap = Dict{LLVM.Function, LLVM.Function}()
+    for f in worklist
+        fn = LLVM.name(f)
+        ft = function_type(f)
+        LLVM.name!(f, fn * ".orig")
+        # create a new function
+        new_param_types = LLVMType[parameters(ft)...]
+
+        for intr_fn in used_intrinsics
+            llvm_typ = convert(LLVMType, kernel_intrinsics[intr_fn].typ)
+            push!(new_param_types, llvm_typ)
+        end
+        new_ft = LLVM.FunctionType(return_type(ft), new_param_types)
+        new_f = LLVM.Function(mod, fn, 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
+        for (intr_fn, new_arg) in zip(used_intrinsics, parameters(new_f)[end-nargs+1:end])
+            LLVM.name!(new_arg, kernel_intrinsics[intr_fn].name)
+        end
+
+        workmap[f] = new_f
+    end
+
+    # clone and rewrite the function bodies.
+    # we don't need to rewrite much as the arguments are added last.
+    for (f, new_f) in workmap
+        # map the arguments
+        value_map = Dict{LLVM.Value, LLVM.Value}()
+        for (param, new_param) in zip(parameters(f), parameters(new_f))
+            LLVM.name!(new_param, LLVM.name(param))
+            value_map[param] = new_param
+        end
+
+        value_map[f] = new_f
+        clone_into!(new_f, f; value_map,
+                    changes=LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly)
+
+        # we can't remove this function yet, as we might still need to rewrite any called,
+        # but remove the IR already
+        empty!(f)
+    end
+
+    # drop unused constants that may be referring to the old functions
+    # XXX: can we do this differently?
+    for f in worklist
+        prune_constexpr_uses!(f)
+    end
+
+    # update other uses of the old function, modifying call sites to pass the arguments
+    function rewrite_uses!(f, new_f)
+        # update uses
+        @dispose builder=IRBuilder() begin
+            for use in uses(f)
+                val = user(use)
+                if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
+                    callee_f = LLVM.parent(LLVM.parent(val))
+                    # forward the arguments
+                    position!(builder, val)
+                    new_val = if val isa LLVM.CallInst
+                        call!(builder, function_type(new_f), new_f,
+                              [arguments(val)..., parameters(callee_f)[end-nargs+1:end]...],
+                              operand_bundles(val))
+                    else
+                        # TODO: invoke and callbr
+                        error("Rewrite of $(typeof(val))-based calls is not implemented: $val")
+                    end
+                    callconv!(new_val, callconv(val))
+
+                    replace_uses!(val, new_val)
+                    @assert isempty(uses(val))
+                    erase!(val)
+                elseif val isa LLVM.ConstantExpr && opcode(val) == LLVM.API.LLVMBitCast
+                    # XXX: why isn't this caught by the value materializer above?
+                    target = operands(val)[1]
+                    @assert target == f
+                    new_val = LLVM.const_bitcast(new_f, value_type(val))
+                    rewrite_uses!(val, new_val)
+                    # we can't simply replace this constant expression, as it may be used
+                    # as a call, taking arguments (so we need to rewrite it to pass the input arguments)
+
+                    # drop the old constant if it is unused
+                    # XXX: can we do this differently?
+                    if isempty(uses(val))
+                        LLVM.unsafe_destroy!(val)
+                    end
+                else
+                    error("Cannot rewrite unknown use of function: $val")
+                end
+            end
+        end
+    end
+    for (f, new_f) in workmap
+        rewrite_uses!(f, new_f)
+        @assert isempty(uses(f))
+        erase!(f)
+    end
+
+    # replace uses of the intrinsics with references to the input arguments
+    for (i, intr_fn) in enumerate(used_intrinsics)
+        intr = functions(mod)[intr_fn]
+        for use in uses(intr)
+            val = user(use)
+            callee_f = LLVM.parent(LLVM.parent(val))
+            if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
+                replace_uses!(val, parameters(callee_f)[end-nargs+i])
+            else
+                error("Cannot rewrite unknown use of function: $val")
+            end
+
+            @assert isempty(uses(val))
+            erase!(val)
+        end
+        @assert isempty(uses(intr))
+        erase!(intr)
+    end
+
+    return functions(mod)[entry_fn]
+end

src/metal.jl

@@ -53,9 +53,7 @@ function finish_module!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mo
     # update calling conventions
     if job.config.kernel
         entry = pass_by_reference!(job, mod, entry)
-
-        add_input_arguments!(job, mod, entry)
-        entry = LLVM.functions(mod)[entry_fn]
+        entry = add_input_arguments!(job, mod, entry, kernel_intrinsics)
     end
 
     # emit the AIR and Metal version numbers as constants in the module. this makes it
@@ -553,164 +551,6 @@ function argument_type_name(typ)
     end
 end
 
-function add_input_arguments!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
-                              entry::LLVM.Function)
-    entry_fn = LLVM.name(entry)
-
-    # figure out which intrinsics are used and need to be added as arguments
-    used_intrinsics = filter(keys(kernel_intrinsics)) do intr_fn
-        haskey(functions(mod), intr_fn)
-    end |> collect
-    nargs = length(used_intrinsics)
-
-    # determine which functions need these arguments
-    worklist = Set{LLVM.Function}([entry])
-    for intr_fn in used_intrinsics
-        push!(worklist, functions(mod)[intr_fn])
-    end
-    worklist_length = 0
-    while worklist_length != length(worklist)
-        # iteratively discover functions that use an intrinsic or any function calling it
-        worklist_length = length(worklist)
-        additions = LLVM.Function[]
-        for f in worklist, use in uses(f)
-            inst = user(use)::Instruction
-            bb = LLVM.parent(inst)
-            new_f = LLVM.parent(bb)
-            in(new_f, worklist) || push!(additions, new_f)
-        end
-        for f in additions
-            push!(worklist, f)
-        end
-    end
-    for intr_fn in used_intrinsics
-        delete!(worklist, functions(mod)[intr_fn])
-    end
-
-    # add the arguments
-    # NOTE: we don't need to be fine-grained here, as unused args will be removed during opt
-    workmap = Dict{LLVM.Function, LLVM.Function}()
-    for f in worklist
-        fn = LLVM.name(f)
-        ft = function_type(f)
-        LLVM.name!(f, fn * ".orig")
-        # create a new function
-        new_param_types = LLVMType[parameters(ft)...]
-
-        for intr_fn in used_intrinsics
-            llvm_typ = convert(LLVMType, kernel_intrinsics[intr_fn].typ)
-            push!(new_param_types, llvm_typ)
-        end
-        new_ft = LLVM.FunctionType(return_type(ft), new_param_types)
-        new_f = LLVM.Function(mod, fn, 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
-        for (intr_fn, new_arg) in zip(used_intrinsics, parameters(new_f)[end-nargs+1:end])
-            LLVM.name!(new_arg, kernel_intrinsics[intr_fn].name)
-        end
-
-        workmap[f] = new_f
-    end
-
-    # clone and rewrite the function bodies.
-    # we don't need to rewrite much as the arguments are added last.
-    for (f, new_f) in workmap
-        # map the arguments
-        value_map = Dict{LLVM.Value, LLVM.Value}()
-        for (param, new_param) in zip(parameters(f), parameters(new_f))
-            LLVM.name!(new_param, LLVM.name(param))
-            value_map[param] = new_param
-        end
-
-        value_map[f] = new_f
-        clone_into!(new_f, f; value_map,
-                    changes=LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly)
-
-        # we can't remove this function yet, as we might still need to rewrite any called,
-        # but remove the IR already
-        empty!(f)
-    end
-
-    # drop unused constants that may be referring to the old functions
-    # XXX: can we do this differently?
-    for f in worklist
-        prune_constexpr_uses!(f)
-    end
-
-    # update other uses of the old function, modifying call sites to pass the arguments
-    function rewrite_uses!(f, new_f)
-        # update uses
-        @dispose builder=IRBuilder() begin
-            for use in uses(f)
-                val = user(use)
-                if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
-                    callee_f = LLVM.parent(LLVM.parent(val))
-                    # forward the arguments
-                    position!(builder, val)
-                    new_val = if val isa LLVM.CallInst
-                        call!(builder, function_type(new_f), new_f,
-                              [arguments(val)..., parameters(callee_f)[end-nargs+1:end]...],
-                              operand_bundles(val))
-                    else
-                        # TODO: invoke and callbr
-                        error("Rewrite of $(typeof(val))-based calls is not implemented: $val")
-                    end
-                    callconv!(new_val, callconv(val))
-
-                    replace_uses!(val, new_val)
-                    @assert isempty(uses(val))
-                    erase!(val)
-                elseif val isa LLVM.ConstantExpr && opcode(val) == LLVM.API.LLVMBitCast
-                    # XXX: why isn't this caught by the value materializer above?
-                    target = operands(val)[1]
-                    @assert target == f
-                    new_val = LLVM.const_bitcast(new_f, value_type(val))
-                    rewrite_uses!(val, new_val)
-                    # we can't simply replace this constant expression, as it may be used
-                    # as a call, taking arguments (so we need to rewrite it to pass the input arguments)
-
-                    # drop the old constant if it is unused
-                    # XXX: can we do this differently?
-                    if isempty(uses(val))
-                        LLVM.unsafe_destroy!(val)
-                    end
-                else
-                    error("Cannot rewrite unknown use of function: $val")
-                end
-            end
-        end
-    end
-    for (f, new_f) in workmap
-        rewrite_uses!(f, new_f)
-        @assert isempty(uses(f))
-        erase!(f)
-    end
-
-    # replace uses of the intrinsics with references to the input arguments
-    for (i, intr_fn) in enumerate(used_intrinsics)
-        intr = functions(mod)[intr_fn]
-        for use in uses(intr)
-            val = user(use)
-            callee_f = LLVM.parent(LLVM.parent(val))
-            if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
-                replace_uses!(val, parameters(callee_f)[end-nargs+i])
-            else
-                error("Cannot rewrite unknown use of function: $val")
-            end
-
-            @assert isempty(uses(val))
-            erase!(val)
-        end
-        @assert isempty(uses(intr))
-        erase!(intr)
-    end
-
-    return
-end
-
-
 # argument metadata generation
 #
 # module metadata is used to identify buffers that are passed as kernel arguments.