Properly rewrite values and constant expressions.

Doing more with the CloneFunction APIs, and making sure we rewrite
uses of constant expressions.

Author: maleadt

src/irgen.jl

@@ -605,7 +605,39 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
         workmap[f] = new_f
     end
 
-    # clone and rewrite the function bodies
+    # clone and rewrite the function bodies, replacing uses of the old stateless function
+    # with the newly created definition that includes the state argument.
+    #
+    # most uses are rewritten by LLVM by putting the functions in the value map.
+    # a separate value materializer is used to recreate constant expressions.
+    #
+    # note that this only _replaces_ the uses of these functions, we'll still need to
+    # _correct_ the uses (i.e. actually add the state argument) afterwards.
+    function materializer(val)
+        if val isa ConstantExpr
+            if opcode(val) == LLVM.API.LLVMBitCast
+                target = operands(val)[1]
+                if target isa LLVM.Function && haskey(workmap, target)
+                    # the function is being bitcasted to a different function type.
+                    # we need to mutate that function type to include the state argument,
+                    # or we'd be invoking the original function in an invalid way.
+                    #
+                    # XXX: ptrtoint/inttoptr pairs can also lose the state argument...
+                    #      is all this even sound?
+                    typ = llvmtype(val)::LLVM.PointerType
+                    ft = eltype(typ)::LLVM.FunctionType
+                    new_ft = LLVM.FunctionType(return_type(ft), [T_state, parameters(ft)...])
+                    return const_bitcast(workmap[target], LLVM.PointerType(new_ft, addrspace(typ)))
+                end
+            elseif opcode(val) == LLVM.API.LLVMPtrToInt
+                target = operands(val)[1]
+                if target isa LLVM.Function && haskey(workmap, target)
+                    return const_ptrtoint(workmap[target], llvmtype(val))
+                end
+            end
+        end
+        return val
+    end
     for (f, new_f) in workmap
         # use a value mapper for rewriting function arguments
         value_map = Dict{LLVM.Value, LLVM.Value}()
@@ -614,33 +646,43 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
             value_map[param] = new_param
         end
 
-        value_map[f] = new_f
-        # XXX: do we want this? we're adding a new arg, after all
-        clone_into!(new_f, f; value_map,
+        # rewrite references to the old function
+        merge!(value_map, workmap)
+
+        clone_into!(new_f, f; value_map, materializer,
                     changes=LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges)
 
-        # we can't remove this function yet, as we might still need to rewrite any called,
-        # but remove the IR already
+        # remove the function IR so that we won't have any uses left after this pass.
         empty!(f)
     end
 
-    # update other uses of the old function, modifying call sites to pass the state argument
-    # TODO: why isn't this covered by the value mapper above? because we need to add an arg!
-    # XXX: do this with a value mapper, and a materialize (?) to rewrite calls.
-    function rewrite_uses!(f, new_f)
+    # ensure the old (stateless) functions don't have uses anymore, and remove them
+    for f in keys(workmap)
+        for use in uses(f)
+            val = user(use)
+            if val isa ConstantExpr
+                # XXX: shouldn't clone_into! remove unused CEs?
+                isempty(uses(val)) || error("old function still has uses (via a constant expr)")
+                LLVM.unsafe_destroy!(val)
+            else
+                error("old function still has uses")
+            end
+        end
+        unsafe_delete!(mod, f)
+    end
+
+    # update uses of the new function, modifying call sites to include the kernel state
+    function rewrite_uses!(f)
         # update uses
         Builder(ctx) do builder
             for use in uses(f)
                 val = user(use)
-                if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
-                    # NOTE: we unconditionally add the state argument, even if there's no uses,
-                    #       assuming we'll perform dead arg elimination during optimization.
-
+                if val isa LLVM.CallBase && called_value(val) == f
                     # forward the state argument
                     position!(builder, val)
                     state = call!(builder, state_intr, Value[], "state")
                     new_val = if val isa LLVM.CallInst
-                        call!(builder, new_f, [state, arguments(val)...], operand_bundles(val))
+                        call!(builder, f, [state, arguments(val)...], operand_bundles(val))
                     else
                         # TODO: invoke and callbr
                         error("Rewrite of $(typeof(val))-based calls is not implemented: $val")
@@ -650,26 +692,19 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                     replace_uses!(val, new_val)
                     @assert isempty(uses(val))
                     unsafe_delete!(LLVM.parent(val), val)
-                elseif val isa ConstantExpr && opcode(val) == LLVM.API.LLVMPtrToInt
-                    # XXX: are these safe? we're not adding an arg
-                    # XXX: in addition, won't this RAUW assert in debug mode?
-                    replace_uses!(val, const_ptrtoint(new_f, llvmtype(val)))
-                    LLVM.unsafe_destroy!(val)
-                elseif val isa ConstantExpr && opcode(val) == LLVM.API.LLVMBitCast
-                    # XXX: are these safe? we're not adding an arg
-                    # XXX: in addition, won't this RAUW assert in debug mode?
-                    replace_uses!(val, const_bitcast(new_f, llvmtype(val)))
-                    LLVM.unsafe_destroy!(val)
+                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 ConstantExpr
+                    rewrite_uses!(val)
                 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))
-        unsafe_delete!(mod, f)
+    for f in values(workmap)
+        rewrite_uses!(f)
     end
 
     return true