Add prevent_primitive_inlining flag for forward-over-reverse mode

Author: sunxd3

src/interpreter/forward_mode.jl

@@ -29,7 +29,9 @@ end
 function build_frule(args...; debug_mode=false, silence_debug_messages=true)
     sig = _typeof(TestUtils.__get_primals(args))
     interp = get_interpreter(ForwardMode)
-    return build_frule(interp, sig; debug_mode, silence_debug_messages)
+    return build_frule(
+        interp, sig; debug_mode, silence_debug_messages, prevent_primitive_inlining=false
+    )
 end
 
 struct DualRuleInfo
@@ -45,20 +47,25 @@ end
         debug_mode=false,
         silence_debug_messages=true,
         skip_world_age_check=false,
+        prevent_primitive_inlining=false,
     ) where {C}
 
 Returns a function which performs forward-mode AD for `sig_or_mi`. Will derive a rule if
 `sig_or_mi` is not a primitive.
 
 Set `skip_world_age_check=true` when the interpreter's world age is intentionally older
 than the current world (e.g., when building rules for MistyClosure which uses its own world).
+
+Set `prevent_primitive_inlining=true` for forward-over-reverse mode. This prevents primitives
+from being inlined during optimization, which would expose foreigncalls that lack frule!!.
 """
 function build_frule(
     interp::MooncakeInterpreter{C},
     sig_or_mi;
     debug_mode=false,
     silence_debug_messages=true,
     skip_world_age_check=false,
+    prevent_primitive_inlining=false,
 ) where {C}
     @nospecialize sig_or_mi
 
@@ -89,12 +96,16 @@ function build_frule(
     try
         # If we've already derived the OpaqueClosures and info, do not re-derive, just
         # create a copy and pass in new shared data.
-        oc_cache_key = ClosureCacheKey(interp.world, (sig_or_mi, debug_mode, :forward))
+        oc_cache_key = ClosureCacheKey(
+            interp.world, (sig_or_mi, debug_mode, :forward, prevent_primitive_inlining)
+        )
         if haskey(interp.oc_cache, oc_cache_key)
             return interp.oc_cache[oc_cache_key]
         else
             # Derive forward-pass IR, and shove in a `MistyClosure`.
-            dual_ir, captures, info = generate_dual_ir(interp, sig_or_mi; debug_mode)
+            dual_ir, captures, info = generate_dual_ir(
+                interp, sig_or_mi; debug_mode, prevent_primitive_inlining
+            )
             dual_oc = misty_closure(
                 info.dual_ret_type, dual_ir, captures...; do_compile=true
             )
@@ -162,7 +173,11 @@ struct DualInfo
 end
 
 function generate_dual_ir(
-    interp::MooncakeInterpreter, sig_or_mi; debug_mode=false, do_inline=true
+    interp::MooncakeInterpreter,
+    sig_or_mi;
+    debug_mode=false,
+    do_inline=true,
+    prevent_primitive_inlining=false,
 )
     # Reset id count. This ensures that the IDs generated are the same each time this
     # function runs.
@@ -214,11 +229,15 @@ function generate_dual_ir(
     captures_tuple = (captures...,)
     dual_ir.argtypes[1] = _typeof(captures_tuple)
 
-    # Optimize dual IR using ForwardOverReverseInterpreter to prevent inlining primitives.
-    # This is needed for forward-over-reverse: if primitives are inlined, foreigncalls
-    # get exposed and we fail when trying to differentiate through them.
-    fo_r_interp = ForwardOverReverseInterpreter(interp)
-    dual_ir_opt = optimise_ir!(dual_ir; interp=fo_r_interp, do_inline)
+    # Use ForwardOverReverseInterpreter when prevent_primitive_inlining is true.
+    # This prevents primitives from being inlined, which is needed for forward-over-reverse
+    # mode where inlining exposes foreigncalls that lack frule!!.
+    if prevent_primitive_inlining
+        fo_r_interp = ForwardOverReverseInterpreter(interp)
+        dual_ir_opt = optimise_ir!(dual_ir; interp=fo_r_interp, do_inline)
+    else
+        dual_ir_opt = optimise_ir!(dual_ir; do_inline)
+    end
     return dual_ir_opt, captures_tuple, DualRuleInfo(isva, nargs, dual_ret_type(primal_ir))
 end
 

src/interpreter/reverse_mode.jl

@@ -1216,12 +1216,8 @@ function generate_ir(
     rvs_ir = pullback_ir(
         primal_ir, Treturn, ad_stmts_blocks, pb_comms_insts, info, _typeof(shared_data)
     )
-    # Use ForwardOverReverseInterpreter to prevent inlining primitives.
-    # This is needed for forward-over-reverse: if primitives are inlined, foreigncalls
-    # get exposed and we fail when trying to differentiate through them.
-    fo_r_interp = ForwardOverReverseInterpreter(interp)
-    opt_fwd_ir = optimise_ir!(IRCode(fwd_ir); interp=fo_r_interp, do_inline)
-    opt_rvs_ir = optimise_ir!(IRCode(rvs_ir); interp=fo_r_interp, do_inline)
+    opt_fwd_ir = optimise_ir!(IRCode(fwd_ir); do_inline)
+    opt_rvs_ir = optimise_ir!(IRCode(rvs_ir); do_inline)
     return DerivedRuleInfo(
         ir, opt_fwd_ir, fwd_ret_type, opt_rvs_ir, rvs_ret_type, shared_data, info, isva
     )

src/rrules/misty_closures.jl

@@ -24,10 +24,15 @@ end
 #    bypassing method table lookups that would require a current-world interpreter
 # 3. Any nested non-primitive calls use LazyFRule/DynamicFRule which obtain a current-world
 #    interpreter via get_interpreter() at runtime
+# We set prevent_primitive_inlining=true because MistyClosure is used in forward-over-reverse
+# mode. Without this, primitives get inlined during optimization, exposing foreigncalls that
+# lack frule!! definitions.
 function _dual_mc(p::MistyClosure)
     mc_world = UInt(p.oc.world)
     interp = MooncakeInterpreter(DefaultCtx, ForwardMode; world=mc_world)
-    return build_frule(interp, p; skip_world_age_check=true)
+    return build_frule(
+        interp, p; skip_world_age_check=true, prevent_primitive_inlining=true
+    )
 end
 
 tangent_type(::Type{<:MistyClosure}) = MistyClosureTangent