Merge pull request #327 from JuliaGPU/tb/late_kernel_state

Perform all kernel-state transformations late

Author: maleadt

Manifest.toml

@@ -39,15 +39,15 @@ version = "1.4.1"
 
 [[LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "c9b86064be5ae0f63e50816a5a90b08c474507ae"
+git-tree-sha1 = "dd58421009014ff1ffacaa0db2a9a392114d75ee"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "4.9.1"
+version = "4.11.0"
 
 [[LLVMExtra_jll]]
-deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"]
-git-tree-sha1 = "5558ad3c8972d602451efe9d81c78ec14ef4f5ef"
+deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
+git-tree-sha1 = "771bfe376249626d3ca12bcd58ba243d3f961576"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.14+2"
+version = "0.0.16+0"
 
 [[LazyArtifacts]]
 deps = ["Artifacts", "Pkg"]

Project.toml

@@ -14,6 +14,6 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
 ExprTools = "0.1"
-LLVM = "4.8"
+LLVM = "4.11"
 TimerOutputs = "0.5"
 julia = "1.6"

src/driver.jl

@@ -257,11 +257,10 @@ const __llvm_initialized = Ref(false)
     # deferred code generation
     do_deferred_codegen = !only_entry && deferred_codegen &&
                           haskey(functions(ir), "deferred_codegen")
+    deferred_jobs = Dict{CompilerJob, String}(job => entry_fn)
     if do_deferred_codegen
         dyn_marker = functions(ir)["deferred_codegen"]
 
-        cache = Dict{CompilerJob, String}(job => entry_fn)
-
         # iterative compilation (non-recursive)
         changed = true
         while changed
@@ -286,7 +285,7 @@ const __llvm_initialized = Ref(false)
             # compile and link
             for dyn_job in keys(worklist)
                 # cached compilation
-                dyn_entry_fn = get!(cache, dyn_job) do
+                dyn_entry_fn = get!(deferred_jobs, dyn_job) do
                     dyn_ir, dyn_meta = codegen(:llvm, dyn_job; optimize=false,
                                                deferred_codegen=false, parent_job=job, ctx)
                     dyn_entry_fn = LLVM.name(dyn_meta.entry)
@@ -317,6 +316,14 @@ const __llvm_initialized = Ref(false)
     end
 
     @timeit_debug to "IR post-processing" begin
+        # mark the kernel entry-point functions (optimization may need it)
+        if job.source.kernel
+            push!(metadata(ir)["julia.kernel"], MDNode([entry]; ctx))
+
+            # IDEA: save all jobs, not only kernels, and save other attributes
+            #       so that we can reconstruct the CompileJob instead of setting it globally
+        end
+
         if optimize
             @timeit_debug to "optimization" begin
                 optimize!(job, ir)
@@ -361,7 +368,18 @@ const __llvm_initialized = Ref(false)
             end
         end
 
-        entry = finish_ir!(job, ir, entry)
+        # finish the module
+        #
+        # we want to finish the module after optimization, so we cannot do so during
+        # deferred code generation. instead, process the deferred jobs here.
+        if deferred_codegen
+            entry = finish_ir!(job, ir, entry)
+
+            for (deferred_job, deferred_fn) in deferred_jobs
+                deferred_job == job && continue
+                finish_ir!(deferred_job, ir, functions(ir)[deferred_fn])
+            end
+        end
 
         # replace non-entry function definitions with a declaration
         # NOTE: we can't do this before optimization, because the definitions of called

src/interface.jl

@@ -233,15 +233,7 @@ optimize_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module) = return
 
 # finalization of the module, before deferred codegen and optimization
 function finish_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry::LLVM.Function)
-    ctx = context(mod)
-    entry_fn = LLVM.name(entry)
-
-    # add the kernel state, and lower calls to the `julia.gpu.state_getter` intrinsic.
-    if job.source.kernel
-        add_kernel_state!(job, mod, entry)
-    end
-
-    return functions(mod)[entry_fn]
+    return entry
 end
 
 # final processing of the IR, right before validation and machine-code generation

src/irgen.jl

@@ -404,21 +404,14 @@ function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.
         # XXX: byval is not round-trippable on LLVM < 12 (see maleadt/LLVM.jl#186)
         #      so we need to re-classify the Julia arguments.
         #      remove this once we only support 1.7.
-        has_kernel_state = kernel_state_type(job) !== Nothing
-        orig_ft = if has_kernel_state
-            # the kernel state has been added here already, so strip the first parameter
-            LLVM.FunctionType(LLVM.return_type(ft), parameters(ft)[2:end]; vararg=isvararg(ft))
-        else
-            ft
-        end
-        args = classify_arguments(job, orig_ft)
+        args = classify_arguments(job, ft)
         filter!(args) do arg
             arg.cc != GHOST
         end
         for arg in args
             if arg.cc == BITS_REF
                 # NOTE: +1 since this pass runs after introducing the kernel state
-                byval[arg.codegen.i+has_kernel_state] = true
+                byval[arg.codegen.i] = true
             end
         end
     end
@@ -510,6 +503,7 @@ function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.
     # NOTE: if we ever have legitimate uses of the old function, create a shim instead
     fn = LLVM.name(f)
     @assert isempty(uses(f))
+    replace_metadata_uses!(f, new_f)
     unsafe_delete!(mod, f)
     LLVM.name!(new_f, fn)
 
@@ -535,10 +529,9 @@ end
 # so that the julia.gpu.state_getter` can be simplified to return an opaque pointer.
 
 # add a state argument to every function in the module, starting from the kernel entry point
-function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
-                           entry::LLVM.Function)
+function add_kernel_state!(mod::LLVM.Module)
+    job = current_job::CompilerJob
     ctx = context(mod)
-    entry_fn = LLVM.name(entry)
 
     # check if we even need a kernel state argument
     state = kernel_state_type(job)
@@ -552,12 +545,18 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
     # this is both for extern uses, and to make this transformation a two-step process.
     state_intr = kernel_state_intr(mod, T_state)
 
+    kernels = []
+    kernels_md = metadata(mod)["julia.kernel"]
+    for kernel_md in operands(kernels_md)
+        push!(kernels, Value(operands(kernel_md)[1]; ctx))
+    end
+
     # determine which functions need a kernel state argument
     #
     # previously, we add the argument to every function and relied on unused arg elim to
     # clean-up the IR. however, some libraries do Funny Stuff, e.g., libdevice bitcasting
     # function pointers. such IR is hard to rewrite, so instead be more conservative.
-    worklist = Set{LLVM.Function}([entry, state_intr])
+    worklist = Set{LLVM.Function}([state_intr, kernels...])
     worklist_length = 0
     while worklist_length != length(worklist)
         # iteratively discover functions that use the intrinsic or any function calling it
@@ -669,6 +668,7 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                 error("old function still has uses")
             end
         end
+        replace_metadata_uses!(f, workmap[f])
         unsafe_delete!(mod, f)
     end
 
@@ -707,10 +707,12 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                 elseif val isa LLVM.CallBase
                     # the function is being passed as an argument, which we'll just permit,
                     # because we expect to have rewritten the call down the line separately.
+                elseif val isa LLVM.StoreInst
+                    # the function is being stored, which again we'll permit like before.
                 elseif val isa ConstantExpr
                     rewrite_uses!(val)
                 else
-                    error("Cannot rewrite unknown use of function: $val")
+                    error("Cannot rewrite $(typeof(val)) use of function: $val")
                 end
             end
         end

src/optim.jl

@@ -192,6 +192,8 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
         if job.source.kernel
             # GC lowering is the last pass that may introduce calls to the runtime library,
             # and thus additional uses of the kernel state intrinsic.
+            # TODO: now that all kernel state-related passes are being run here, merge some?
+            add!(pm, ModulePass("AddKernelState", add_kernel_state!))
             add!(pm, FunctionPass("LowerKernelState", lower_kernel_state!))
             add!(pm, ModulePass("CleanupKernelState", cleanup_kernel_state!))
         end
@@ -253,8 +255,7 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
     ModulePassManager() do pm
         addTargetPasses!(pm, tm, triple)
 
-        # - remove unused kernel state arguments
-        # - simplify function calls that don't use the returned value
+        # simplify function calls that don't use the returned value
         dead_arg_elimination!(pm)
 
         run!(pm, mod)
@@ -354,16 +355,8 @@ function lower_gc_frame!(fun::LLVM.Function)
 
             # replace with PTX alloc_obj
             Builder(ctx) do builder
-                # NOTE: this happens late during the pipeline, where we may have to
-                #       pass a kernel state arguments to the runtime function.
-                state = if job.source.kernel
-                    kernel_state_type(job)
-                else
-                    Nothing
-                end
-
                 position!(builder, call)
-                ptr = call!(builder, Runtime.get(:gc_pool_alloc), [sz]; state)
+                ptr = call!(builder, Runtime.get(:gc_pool_alloc), [sz])
                 replace_uses!(call, ptr)
             end
 

src/ptx.jl

@@ -185,8 +185,17 @@ function finish_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
     if job.source.kernel
         # work around bad byval codegen (JuliaGPU/GPUCompiler.jl#92)
         entry = lower_byval(job, mod, entry)
-        # TODO: optimization passes to clean-up byval
+    end
+
+    return entry
+end
 
+function finish_ir!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
+                        mod::LLVM.Module, entry::LLVM.Function)
+    ctx = context(mod)
+    entry = invoke(finish_ir!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)
+
+    if job.source.kernel
         # add metadata annotations for the assembler to the module
 
         # property annotations

src/rtlib.jl

@@ -28,8 +28,7 @@ end
 
 ## higher-level functionality to work with runtime functions
 
-function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[];
-                    state::Type=Nothing)
+function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[])
     bb = position(builder)
     f = LLVM.parent(bb)
     mod = LLVM.parent(f)
@@ -40,21 +39,10 @@ function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[
         f = functions(mod)[rt.llvm_name]
         ft = eltype(llvmtype(f))
     else
-        ft = convert(LLVM.FunctionType, rt; ctx, state)
+        ft = convert(LLVM.FunctionType, rt; ctx)
         f = LLVM.Function(mod, rt.llvm_name, ft)
     end
 
-    # we may be calling this function after kernel state lowering,
-    # in which case we need to manually get and pass the state.
-    args = Value[args...]
-    if state !== Nothing
-        T_state = convert(LLVMType, state; ctx)
-
-        state_intr = kernel_state_intr(mod, T_state)
-        state_val = call!(builder, state_intr, Value[], "state")
-        pushfirst!(args, state_val)
-    end
-
     # runtime functions are written in Julia, while we're calling from LLVM,
     # this often results in argument type mismatches. try to fix some here.
     for (i,arg) in enumerate(args)

src/runtime.jl

@@ -34,19 +34,13 @@ struct RuntimeMethodInstance
 end
 
 function Base.convert(::Type{LLVM.FunctionType}, rt::RuntimeMethodInstance;
-                      ctx::LLVM.Context, state::Type=Nothing)
+                      ctx::LLVM.Context)
     types = if rt.llvm_types === nothing
         LLVMType[convert(LLVMType, typ; ctx, allow_boxed=true) for typ in rt.types]
     else
         rt.llvm_types(ctx)
     end
 
-    # if we're running post-optimization, prepend the kernel state to the argument list
-    if state !== Nothing
-        T_state = convert(LLVMType, state; ctx)
-        pushfirst!(types, T_state)
-    end
-
     return_type = if rt.llvm_return_type === nothing
         convert(LLVMType, rt.return_type; ctx, allow_boxed=true)
     else

src/spirv.jl

@@ -232,26 +232,15 @@ function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.F
         end
     else
         # XXX: byval is not round-trippable on LLVM < 12 (see maleadt/LLVM.jl#186)
-        has_kernel_state = kernel_state_type(job) !== Nothing
-        orig_ft = if has_kernel_state
-            # the kernel state has been added here already, so strip the first parameter
-            LLVM.FunctionType(LLVM.return_type(ft), parameters(ft)[2:end]; vararg=isvararg(ft))
-        else
-            ft
-        end
-        args = classify_arguments(job, orig_ft)
+        args = classify_arguments(job, ft)
         filter!(args) do arg
             arg.cc != GHOST
         end
         for arg in args
             if arg.cc == BITS_REF
-                # NOTE: +1 since this pass runs after introducing the kernel state
-                byval[arg.codegen.i+has_kernel_state] = true
+                byval[arg.codegen.i] = true
             end
         end
-        if has_kernel_state
-            byval[1] = true
-        end
     end
 
     # generate the wrapper function type & definition
@@ -317,6 +306,7 @@ function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.F
     # NOTE: if we ever have legitimate uses of the old function, create a shim instead
     fn = LLVM.name(f)
     @assert isempty(uses(f))
+    replace_metadata_uses!(f, new_f)
     unsafe_delete!(mod, f)
     LLVM.name!(new_f, fn)