Merge pull request #249 from JuliaGPU/tb/pass_state_ccall

Fixes for late optimization of kernel state arguments

Author: maleadt

src/driver.jl

@@ -179,6 +179,25 @@ const __llvm_initialized = Ref(false)
         end
     end
 
+    # mark everything internal except for the entry and any exported global variables.
+    # this makes sure that the optimizer can, e.g., touch function signatures.
+    ModulePassManager() do pm
+        # NOTE: this needs to happen after linking libraries to remove unused functions,
+        #       but before deferred codegen so that all kernels remain available.
+        exports = String[entry_fn]
+        for gvar in globals(ir)
+            if linkage(gvar) == LLVM.API.LLVMExternalLinkage
+                push!(exports, LLVM.name(gvar))
+            end
+        end
+        internalize!(pm, exports)
+        run!(pm, ir)
+    end
+
+    # finalize the current module. this needs to happen before linking deferred modules,
+    # since those modules have been finalized themselves, and we don't want to re-finalize.
+    entry = finish_module!(job, ir, entry)
+
     # deferred code generation
     do_deferred_codegen = !only_entry && deferred_codegen &&
                           haskey(functions(ir), "deferred_codegen")
@@ -242,39 +261,9 @@ const __llvm_initialized = Ref(false)
     end
 
     @timeit_debug to "IR post-processing" begin
-        entry = finish_module!(job, ir, entry)
-
-        if optimize
-            @timeit_debug to "optimization" optimize!(job, ir)
-
-            # optimization may have replaced functions, so look the entry point up again
-            entry = functions(ir)[entry_fn]
-        end
-
-        if ccall(:jl_is_debugbuild, Cint, ()) == 1
-            @timeit_debug to "verification" verify(ir)
-        end
-
+        # some early clean-up to reduce the amount of code to optimize
         @timeit_debug to "clean-up" begin
-            # replace non-entry function definitions with a declaration
-            if only_entry
-                for f in functions(ir)
-                    f == entry && continue
-                    isdeclaration(f) && continue
-                    LLVM.isintrinsic(f) && continue
-                    empty!(f)
-                end
-            end
-
-            # remove everything except for the entry and any exported global variables
-            exports = String[entry_fn]
-            for gvar in globals(ir)
-                push!(exports, LLVM.name(gvar))
-            end
-
             ModulePassManager() do pm
-                internalize!(pm, exports)
-
                 # eliminate all unused internal functions
                 global_optimizer!(pm)
                 global_dce!(pm)
@@ -283,9 +272,20 @@ const __llvm_initialized = Ref(false)
                 # merge constants (such as exception messages)
                 constant_merge!(pm)
 
-                if do_deferred_codegen
-                    # inline and optimize the call to the deferred code. in particular we want to
-                    # remove unnecessary alloca's that are created by pass-by-ref semantics.
+                run!(pm, ir)
+            end
+        end
+
+        if optimize
+            @timeit_debug to "optimization" begin
+                optimize!(job, ir)
+
+                # deferred codegen has some special optimization requirements,
+                # which also need to happen _after_ regular optimization.
+                # XXX: make these part of the optimizer pipeline?
+                do_deferred_codegen && ModulePassManager() do pm
+                    # inline and optimize the call to e deferred code. in particular we want
+                    # to remove unnecessary alloca's created by pass-by-ref semantics.
                     instruction_combining!(pm)
                     always_inliner!(pm)
                     scalar_repl_aggregates_ssa!(pm)
@@ -295,11 +295,30 @@ const __llvm_initialized = Ref(false)
                     # merge duplicate functions, since each compilation invocation emits everything
                     # XXX: ideally we want to avoid emitting these in the first place
                     merge_functions!(pm)
+
+                    run!(pm, ir)
                 end
+            end
 
-                run!(pm, ir)
+            # optimization may have replaced functions, so look the entry point up again
+            entry = functions(ir)[entry_fn]
+        end
+
+        # replace non-entry function definitions with a declaration
+        # NOTE: we can't do this before optimization, because the definitions of called
+        #       functions may affect optimization.
+        if only_entry
+            for f in functions(ir)
+                f == entry && continue
+                isdeclaration(f) && continue
+                LLVM.isintrinsic(f) && continue
+                empty!(f)
             end
         end
+
+        if ccall(:jl_is_debugbuild, Cint, ()) == 1
+            @timeit_debug to "verification" verify(ir)
+        end
     end
 
     return ir, (; entry, compiled)

src/interface.jl

@@ -205,20 +205,8 @@ function finish_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry
     entry_fn = LLVM.name(entry)
 
     # add the kernel state, and lower calls to the `julia.gpu.state_getter` intrinsic.
-    # we do this _after_ optimization, because the runtime is linked after optimization too.
     if job.source.kernel
-        state = kernel_state_type(job)
-        if state !== Nothing
-            T_state = convert(LLVMType, state; ctx)
-            add_kernel_state!(job, mod, entry, T_state)
-        end
-
-        # don't pass the state when unnecessary
-        # XXX: only apply in add_kernel_state! when needed?
-        ModulePassManager() do pm
-            dead_arg_elimination!(pm)
-            run!(pm, mod)
-        end
+        add_kernel_state!(job, mod, entry)
     end
 
     return functions(mod)[entry_fn]