Merge pull request #168 from JuliaDiff/sf/globalref_types

Improve type inference of non-`const` GlobalRef's

Author: oxinabox

src/codegen/forward_demand.jl

@@ -256,11 +256,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             # TODO: Should we remember whether the callbacks wanted the arg?
             return transform!(ir, arg, order)
         elseif isa(arg, GlobalRef)
-            if !isconst(arg)
-                # Non-const GlabalRefs need to need to be accessed as seperate statements
-                arg = insert_node!(ir, ssa, NewInstruction(arg, Any))
-            end
-
+            @assert isconst(arg)
             return insert_node!(ir, ssa, NewInstruction(Expr(:call, ZeroBundle{order}, arg), Any))
         elseif isa(arg, QuoteNode)
             return ZeroBundle{order}(arg.value)
@@ -302,7 +298,15 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 # TODO: New PiNode that discriminates based on primal?
                 inst[:inst] = maparg(stmt.val, SSAValue(ssa), order)
                 inst[:type] = Any
-            elseif isa(stmt, GlobalRef) || isa(stmt, SSAValue) || isa(stmt, QuoteNode)
+            elseif isa(stmt, GlobalRef)
+                if !isconst(stmt)
+                    # Non-const GlobalRefs need to need to be accessed as seperate statements
+                    stmt = insert_node!(ir, ssa, NewInstruction(inst))
+                end
+
+                inst[:inst] = Expr(:call, ZeroBundle{order}, stmt)
+                inst[:type] = Any
+            elseif isa(stmt, SSAValue) || isa(stmt, QuoteNode)
                 inst[:inst] = maparg(stmt, SSAValue(ssa), order)
                 inst[:type] = Any
             elseif isa(stmt, Expr) || isa(stmt, PhiNode) || isa(stmt, PhiCNode) ||

test/forward_diff_no_inf.jl

@@ -28,6 +28,11 @@ module forward_diff_no_inf
         Diffractor.forward_diff_no_inf!(ir, [Core.SSAValue(1) => 1]; transform! = identity_transform!)
         ir2 = Core.Compiler.compact!(ir)
         Core.Compiler.verify_ir(ir2)  # This would error if we were not handling nonconst globals correctly
+        # Assert that the reference to `Main._coeff` is properly typed
+        stmt_idx = findfirst(stmt -> isa(stmt[:inst], GlobalRef), collect(ir2.stmts))
+        stmt = ir2.stmts[stmt_idx]
+        @test stmt[:inst].name == :_coeff
+        @test stmt[:type] == Float64
         f = Core.OpaqueClosure(ir2; do_compile=false)
         @test f(3.5) == 28.0
     end