When reinference is required use IR_FLAG_REFINED (#164)

Co-authored-by: Shuhei Kadowaki <[email protected]>

Author: oxinabox

src/codegen/forward_demand.jl

@@ -248,9 +248,9 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                     # identify where to insert. Must be after phi blocks
                     pos = SSAValue(find_end_of_phi_block(ir, arg.id))
                     if order == 0
-                        insert_node!(ir, pos, NewInstruction(Expr(:call, primal, arg), Any), #=attach_after=#true)
+                        insert_node!(ir, pos, NewInstruction(Expr(:call, primal, arg)), #=attach_after=#true)
                     else
-                        insert_node!(ir, pos, NewInstruction(Expr(:call, truncate, arg, Val{order}()), Any), #=attach_after=#true)
+                        insert_node!(ir, pos, NewInstruction(Expr(:call, truncate, arg, Val{order}())), #=attach_after=#true)
                     end
                 end
             end
@@ -262,7 +262,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             return transform!(ir, arg, order, maparg)
         elseif isa(arg, GlobalRef)
             @assert isconst(arg)
-            return insert_node!(ir, ssa, NewInstruction(Expr(:call, ZeroBundle{order}, arg), Any))
+            return insert_node!(ir, ssa, NewInstruction(Expr(:call, ZeroBundle{order}, arg)))
         elseif isa(arg, QuoteNode)
             return ZeroBundle{order}(arg.value)
         end
@@ -300,6 +300,7 @@ 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
+                inst[:flag] |= CC.IR_FLAG_REFINED
             elseif isa(stmt, GlobalRef)
                 if !isconst(stmt)
                     # Non-const GlobalRefs need to need to be accessed as seperate statements
@@ -310,6 +311,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             elseif isa(stmt, SSAValue) || isa(stmt, QuoteNode)
                 inst[:inst] = maparg(stmt, SSAValue(ssa), order)
                 inst[:type] = Any
+                inst[:flag] |= CC.IR_FLAG_REFINED
             elseif isa(stmt, Expr) || isa(stmt, PhiNode) || isa(stmt, PhiCNode) ||
                    isa(stmt, UpsilonNode) || isa(stmt, GotoIfNot) || isa(stmt, Argument)
                 urs = userefs(stmt)
@@ -318,6 +320,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 end
                 inst[:inst] = urs[]
                 inst[:type] = Any
+                inst[:flag] |= CC.IR_FLAG_REFINED
             else
                 val = ZeroBundle{order}(inst[:inst])
                 inst[:inst] = val
@@ -336,10 +339,12 @@ function forward_diff!(interp::ADInterpreter, ir::IRCode, src::CodeInfo, mi::Met
     ir = compact!(ir)
 
     for i = 1:length(ir.stmts)
-        if ir[SSAValue(i)][:type] == Any
-            # TODO: this flag should actually be being set at the insert site
-            # and we should be filtering on if it is present rather than [:type]=Any
-            ir[SSAValue(i)][:flag] |= CC.IR_FLAG_REFINED
+        inst = ir[SSAValue(i)][:inst]
+        if !isa(inst, ReturnNode) && ir[SSAValue(i)][:type] === Any
+            if iszero(ir[SSAValue(i)][:flag] & CC.IR_FLAG_REFINED)
+                @warn "IR_FLAG_REFINED Flag missed on statement" i inst
+                ir[SSAValue(i)][:flag] |= CC.IR_FLAG_REFINED
+            end
         end
     end
 

src/stage1/compiler_utils.jl

@@ -1,5 +1,5 @@
 # Utilities that should probably go into Core.Compiler
-using Core.Compiler: CFG, BasicBlock, BBIdxIter
+using Core.Compiler: IRCode, CFG, BasicBlock, BBIdxIter
 
 function Base.push!(cfg::CFG, bb::BasicBlock)
     @assert cfg.blocks[end].stmts.stop+1 == bb.stmts.start
@@ -12,25 +12,22 @@ Base.getindex(ir::IRCode, ssa::SSAValue) =
 
 Base.copy(ir::IRCode) = Core.Compiler.copy(ir)
 
-function Core.Compiler.NewInstruction(node)
-    Core.Compiler.NewInstruction(node, Any)
-end
+Core.Compiler.NewInstruction(@nospecialize node) =
+    NewInstruction(node, Any, CC.NoCallInfo(), nothing, CC.IR_FLAG_REFINED)
 
-function Base.setproperty!(x::Core.Compiler.Instruction, f::Symbol, v)
+Base.setproperty!(x::Core.Compiler.Instruction, f::Symbol, v) =
     Core.Compiler.setindex!(x, v, f)
-end
 
-function Base.getproperty(x::Core.Compiler.Instruction, f::Symbol)
+Base.getproperty(x::Core.Compiler.Instruction, f::Symbol) =
     Core.Compiler.getindex(x, f)
-end
 
-function Base.setindex!(ir::Core.Compiler.IRCode, ni::NewInstruction, i::Int)
+function Base.setindex!(ir::IRCode, ni::NewInstruction, i::Int)
     stmt = ir.stmts[i]
     stmt.inst = ni.stmt
     stmt.type = ni.type
     stmt.flag = something(ni.flag, 0)  # fixes 1.9?
     stmt.line = something(ni.line, 0)
-    ni
+    return ni
 end
 
 function Base.push!(ir::IRCode, ni::NewInstruction)
@@ -67,14 +64,14 @@ end
 
 
 """
-    find_end_of_phi_block(ir, start_search_idx)
+    find_end_of_phi_block(ir::IRCode, start_search_idx::Int)
 
 Finds the last index within the same basic block, on or after the `start_search_idx` which is not within a phi block.
 A phi-block is a run on PhiNodes or nothings that must be the first statements within the basic block.
 
 If `start_search_idx` is not within a phi block to begin with, then just returns `start_search_idx`
 """
-function find_end_of_phi_block(ir, start_search_idx::Int)
+function find_end_of_phi_block(ir::IRCode, start_search_idx::Int)
     # Short-cut for early exit:
     stmt = ir.stmts[start_search_idx][:inst]
     stmt !== nothing && !isa(stmt, PhiNode) && return start_search_idx
@@ -90,8 +87,9 @@ function find_end_of_phi_block(ir, start_search_idx::Int)
     return end_search_idx
 end
 
-function replace_call!(ir, idx::SSAValue, new_call)
+function replace_call!(ir::IRCode, idx::SSAValue, new_call::Expr)
     ir[idx][:inst] = new_call
     ir[idx][:type] = Any
     ir[idx][:info] = CC.NoCallInfo()
+    ir[idx][:flag] = CC.IR_FLAG_REFINED
 end

src/stage2/forward.jl

@@ -60,7 +60,7 @@ function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1)
             if order == 0
                 return
             end
-            nr = insert_node!(ir, ssa, NewInstruction(Expr(:call, getindex, stmt.val, TaylorTangentIndex(order)), Any))
+            nr = insert_node!(ir, ssa, NewInstruction(Expr(:call, getindex, stmt.val, TaylorTangentIndex(order))))
             inst[:inst] = ReturnNode(nr)
         elseif is_known_invoke_or_call(stmt, dont_use_ddt_intrinsic, ir)
             arg = maparg(stmt.args[end], ssa, order+1)