Merge pull request #160 from JuliaDiff/ox/globals

Don't inline accesses to nonconst globals

Author: oxinabox

src/codegen/forward_demand.jl

@@ -254,6 +254,11 @@ 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
+
             return insert_node!(ir, ssa, NewInstruction(Expr(:call, ZeroBundle{order}, arg), Any))
         elseif isa(arg, QuoteNode)
             return ZeroBundle{order}(arg.value)

test/stage2_fwd.jl

@@ -62,4 +62,22 @@ module stage2_fwd
         f = Core.OpaqueClosure(ir2; do_compile=false)
         @test f(1.0) == Bar148(1.0)  # This would error if we were not handling constructors (%new) right
     end
+
+    @testset "nonconst globals in forward_diff_no_inf!" begin
+        @eval global _coeff::Float64=24.5
+        plus_a_global(x) = x + _coeff
+
+        # this is needed as transform! is *always* called on Arguments regardless of what visit_custom says
+        identity_transform!(ir, ssa::Core.SSAValue, order) = ir[ssa]
+        function identity_transform!(ir, arg::Core.Argument, order)
+            return Core.Compiler.insert_node!(ir, Core.SSAValue(1), Core.Compiler.NewInstruction(Expr(:call, Diffractor.ZeroBundle{1}, arg), Any))
+        end
+
+        ir = first(only(Base.code_ircode(plus_a_global, Tuple{Float64})))
+        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
+        f = Core.OpaqueClosure(ir2; do_compile=false)
+        @test f(3.2) == 28.0
+    end
 end