Hook up eras mode in stage 2

Author: oxinabox

src/codegen/forward_demand.jl

@@ -40,13 +40,13 @@ function forward_diff!(ir::IRCode, interp::AbstractInterpreter, irsv::IRInterpre
 end
 function forward_diff!(ir::IRCode, interp::AbstractInterpreter, irsv::IRInterpretationState,
                        val, order::Int;
-                       custom_diff!, diff_cache)
+                       custom_diff!, diff_cache, eras_mode)
     return ChainRulesCore.zero_tangent(val)
 end
 function forward_diff!(ir::IRCode, interp::AbstractInterpreter, irsv::IRInterpretationState,
                        arg::Argument, order::Int;
-                       custom_diff!, diff_cache)
-    recurse(x) = forward_diff!(ir, interp, irsv, x; custom_diff!, diff_cache)
+                       custom_diff!, diff_cache, eras_mode)
+    recurse(x) = forward_diff!(ir, interp, irsv, x; custom_diff!, diff_cache, eras_mode)
     val = custom_diff!(ir, SSAValue(0), arg, recurse)
     if val !== nothing
         return val
@@ -56,9 +56,9 @@ end
 
 function forward_diff_uncached!(ir::IRCode, interp::AbstractInterpreter, irsv::IRInterpretationState,
                                 ssa::SSAValue, inst::Core.Compiler.Instruction, order::Int;
-                                custom_diff!, diff_cache)
+                                custom_diff!, diff_cache, eras_mode)
     stmt = inst[:inst]
-    recurse(x) = forward_diff!(ir, interp, irsv, x, order; custom_diff!, diff_cache)
+    recurse(x) = forward_diff!(ir, interp, irsv, x, order; custom_diff!, diff_cache, eras_mode)
     if (val = custom_diff!(ir, ssa, stmt, recurse)) !== nothing
         return val
     elseif isa(stmt, PiNode)
@@ -212,8 +212,10 @@ Internal method which generates the code for forward mode diffentiation
 		decides if the custom `transform!` should be applied to a `stmt` or not
 		Default: `false` for all statements
  - `transform!(ir::IRCode, ssa::SSAValue, order::Int)` mutates `ir` to do a custom tranformation.
+ - `eras_mode`: determines if to error if not all derivatives are taylor
 """
 function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
+                              eras_mode = false,
                               visit_custom! = (@nospecialize args...)->false,
                               transform! = (@nospecialize args...)->error())
     # Step 1: For each SSAValue in the IR, keep track of the differentiation order needed
@@ -286,12 +288,12 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 newargs = map(stmt.args[2:end]) do @nospecialize arg
                     maparg(arg, SSAValue(ssa), order)
                 end
-                replace_call!(ir, SSAValue(ssa), Expr(:call, ∂☆{order}(), newargs...))
+                replace_call!(ir, SSAValue(ssa), Expr(:call, ∂☆{order, eras_mode}(), newargs...))
             elseif isexpr(stmt, :call) || isexpr(stmt, :new)
                 newargs = map(stmt.args) do @nospecialize arg
                     maparg(arg, SSAValue(ssa), order)
                 end
-                f = isexpr(stmt, :call) ? ∂☆{order}() : ∂☆new{order}()
+                f = isexpr(stmt, :call) ? ∂☆{order, eras_mode}() : ∂☆new{order, eras_mode}()
                 replace_call!(ir, SSAValue(ssa), Expr(:call, f, newargs...))
             elseif isa(stmt, PiNode)
                 # TODO: New PiNode that discriminates based on primal?

src/stage2/forward.jl

@@ -19,7 +19,7 @@ end
 
 # Engineering entry point for the 2nd-order forward AD functionality. This is
 # unlikely to be the actual interface. For now, it is used for testing.
-function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1)
+function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1; eras_mode = false)
     interp = ADInterpreter(; forward=true, backward=false)
     match = Base._which(tt)
     frame = Core.Compiler.typeinf_frame(interp, match.method, match.spec_types, match.sparams, #=run_optimizer=#true)
@@ -82,7 +82,7 @@ function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1)
         end
     end
 
-    ir = forward_diff!(interp, ir, src, mi, vals; visit_custom!, transform!)
+    ir = forward_diff!(interp, ir, src, mi, vals; visit_custom!, transform!, eras_mode)
 
     return OpaqueClosure(ir)
 end

test/forward_diff_no_inf.jl

@@ -1,21 +1,44 @@
 
 module forward_diff_no_inf
+    using Core.Compiler: SSAValue
+    const CC = Core.Compiler
+
     using Diffractor, Test
     # 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]
+    identity_transform!(ir, ssa::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))
+        return CC.insert_node!(ir, SSAValue(1), CC.NewInstruction(Expr(:call, Diffractor.zero_bundle{1}(), arg), Any))
     end
 
+
+    function infer_ir!(ir)
+        interp = CC.NativeInterpreter()
+        mi = ccall(:jl_new_method_instance_uninit, Ref{Core.MethodInstance}, ());
+        mi.specTypes = Tuple{map(CC.widenconst, ir.argtypes)...}
+        mi.def = @__MODULE__
+
+        for i in 1:length(ir.stmts)  # For testuing purposes we are going to refine everything
+            ir[SSAValue(i)][:flag] |= CC.IR_FLAG_REFINED
+        end
+    
+        method_info = CC.MethodInfo(#=propagate_inbounds=#true, nothing)
+        min_world = world = (interp).world
+        max_world = Diffractor.get_world_counter()
+        irsv = CC.IRInterpretationState(interp, method_info, ir, mi, ir.argtypes, world, min_world, max_world)
+        (rt, nothrow) = CC._ir_abstract_constant_propagation(interp, irsv)
+        return rt
+    end
+    
+   
     @testset "Constructors in forward_diff_no_inf!" begin
         struct Bar148
             v
         end
         foo_148(x) = Bar148(x)
 
         ir = first(only(Base.code_ircode(foo_148, Tuple{Float64})))
-        Diffractor.forward_diff_no_inf!(ir, [Core.SSAValue(1) => 1]; transform! = identity_transform!)
-        ir2 = Core.Compiler.compact!(ir)
+        Diffractor.forward_diff_no_inf!(ir, [SSAValue(1) => 1]; transform! = identity_transform!)
+        ir2 = CC.compact!(ir)
         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
@@ -25,9 +48,9 @@ module forward_diff_no_inf
         plus_a_global(x) = x + _coeff
 
         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
+        Diffractor.forward_diff_no_inf!(ir, [SSAValue(1) => 1]; transform! = identity_transform!)
+        ir2 = CC.compact!(ir)
+        CC.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]
@@ -51,17 +74,40 @@ module forward_diff_no_inf
         input_ir = first(only(Base.code_ircode(phi_run, Tuple{Float64})))
         ir = copy(input_ir)
         #Workout where to diff to trigger error
-        diff_ssa = Core.SSAValue[]
+        diff_ssa = SSAValue[]
         for idx in 1:length(ir.stmts)
             if ir.stmts[idx][:inst] isa Core.PhiNode
-                push!(diff_ssa, Core.SSAValue(idx))
+                push!(diff_ssa, SSAValue(idx))
             end
         end
 
         Diffractor.forward_diff_no_inf!(ir, diff_ssa .=> 1; transform! = identity_transform!)
-        ir2 = Core.Compiler.compact!(ir)
-        Core.Compiler.verify_ir(ir2)  # This would error if we were not handling nonconst phi nodes correctly (after https://github.com/JuliaLang/julia/pull/50158)
+        ir2 = CC.compact!(ir)
+        CC.verify_ir(ir2)  # This would error if we were not handling nonconst phi nodes correctly (after https://github.com/JuliaLang/julia/pull/50158)
         f = Core.OpaqueClosure(ir2; do_compile=false)
         @test f(3.5) == 3.5  # this will segfault if we are not handling phi nodes correctly
     end
+
+    @testset "Eras mode" begin
+        foo(x, y) = x*x + y*y
+        # Eras mode should make this all inferable
+        eras_mode = true
+        ir = first(only(Base.code_ircode(foo, Tuple{Any, Any})))
+        #@assert  ir[SSAValue(1)][:inst].args[1].name == :literal_pow
+        @assert  ir[SSAValue(3)][:inst].args[1].name == :+
+        Diffractor.forward_diff_no_inf!(ir, [SSAValue(1)] .=> 1; transform! = identity_transform!, eras_mode)
+        ir = CC.compact!(ir)
+        #@assert ir[SSAValue(5)][:inst].args[1] == Diffractor.∂☆{1, eras_mode}()
+        #@assert ir[SSAValue(5)][:inst].args[2].primal == *
+        ir.argtypes[2:end] .= Float64
+        @assert infer_ir!(ir) == Float64
+        
+        Diffractor.forward_diff_no_inf!(ir, [SSAValue(6)] .=> 1; transform! = identity_transform!, eras_mode=eras_mode)
+        ir = CC.compact!(ir)
+        CC.verify_ir(ir)
+        infer_ir!(ir)
+        
+
+
+    end
 end