Skip RunPeriodicTasks for leaf functions in JIT

Summary:
Trivial JIT-compiled functions like `typing.cast` — which just returns its second argument — generate ~70 instructions of prologue (frame setup, callee-save spills, RunPeriodicTasks check) and ~20 of epilogue, while the actual work is only 2 instructions (Incref + mov). This diff reduces that overhead for "leaf" functions.

A function is a leaf if its bytecode contains only opcodes from a conservative allowlist that cannot invoke user Python code (LOAD_FAST, STORE_FAST, LOAD_CONST, RETURN_VALUE, POP_TOP, COPY, SWAP, etc.) and has no backward jumps (loops). The allowlist approach is safe against future bytecode additions since new opcodes are excluded by default.

For leaf functions, two optimizations apply:

1. Skip the RunPeriodicTasks check at RESUME. This eliminates a LoadEvalBreaker, CondBranch, snapshot, and cold-path call to `_Py_HandlePending` from the function entry. This is safe because leaf functions execute a bounded, small number of instructions and return quickly — the caller will check periodic tasks at its next check point. The no-backward-jumps constraint ensures the function cannot loop, bounding its execution time.

2. Use a specialized `JITRT_UnlinkLeafFrame` at function exit. Since the function has no DeoptBase instructions, its interpreter frame can never be materialized at runtime, so we skip the materialization check (`header->rtfs & JIT_FRAME_INITIALIZED`) that `JITRT_UnlinkLightweightFrameFast` performs. A shared `cleanupLightweightFrameExecutable` helper is extracted to deduplicate the reifier cleanup logic between the two unlink functions.

The `can_deopt` flag on Environ (set from `Function::canDeopt()` after all HIR passes) propagates this information from the HIR layer to the LIR exit block generator, which selects between `JITRT_UnlinkLeafFrame`, `JITRT_UnlinkLightweightFrameFast`, and `JITRT_UnlinkFrame` at JIT compile time.

Before (def f(x): return x):
  3 basic blocks, 11 HIR instructions including LoadEvalBreaker + CondBranch + RunPeriodicTasks

After:
  1 basic block, 5 HIR instructions (LoadArg, LoadCurrentFunc, LoadFrame, Snapshot, Return)

Reviewed By: alexmalyshev

Differential Revision: D103112548

fbshipit-source-id: 92eb90ee6935619c1498a0f02f1e21bca95c74fb

Author: kddnewton

cinderx/Jit/bytecode.cpp

@@ -129,6 +129,10 @@ bool BytecodeInstruction::isBranch() const {
   }
 }
 
+bool BytecodeInstruction::isBackwardBranch() const {
+  return isBranch() && getJumpTarget() <= baseOffset();
+}
+
 bool BytecodeInstruction::isReturn() const {
   switch (opcode()) {
     case RETURN_CONST:

cinderx/Jit/bytecode.h

@@ -46,6 +46,7 @@ class BytecodeInstruction {
   // Check if this instruction is a branch, a return, or a general basic block
   // terminator.
   bool isBranch() const;
+  bool isBackwardBranch() const;
   bool isReturn() const;
   bool isTerminator() const;
 

cinderx/Jit/codegen/environ.h

@@ -170,6 +170,11 @@ struct Environ {
   int win_struct_ret_offset{0};
 #endif
 
+  // True if the function has any DeoptBase instructions in its final HIR.
+  // When false, the interpreter frame can never be materialized, enabling a
+  // cheaper inline frame unlink at exit.
+  bool can_deopt{true};
+
 #if defined(CINDER_AARCH64)
   // Constant pool for large immediate values. translateMovConstPool populates
   // these; gen_asm.cpp emits the pool data after deopt exits.

cinderx/Jit/codegen/gen_asm.cpp

@@ -727,6 +727,8 @@ void* NativeGenerator::getVectorcallEntry() {
   env_.changed_regs = lsalloc.getChangedRegs();
   env_.exit_label = as_->newLabel();
   env_.frame_mode = GetFunction()->frameMode;
+  env_.can_deopt = GetFunction()->canDeopt();
+
   JIT_LOGIF(
       getConfig().log.dump_lir,
       "LIR for {} after register allocation:\n{}",

cinderx/Jit/hir/builder.cpp

@@ -515,6 +515,42 @@ BasicBlock* HIRBuilder::getBlockAtOff(BCOffset off) {
   return it->second;
 }
 
+bool HIRBuilder::isSimpleLeafFunction(BorrowedRef<PyCodeObject> code) {
+  if (code->co_flags & kCoFlagsAnyGenerator) {
+    return false;
+  }
+  for (auto& instr : BytecodeInstructionBlock{code}) {
+    switch (instr.opcode()) {
+      case COPY:
+      case LOAD_CONST:
+      case LOAD_FAST:
+      case LOAD_FAST_AND_CLEAR:
+      case LOAD_FAST_BORROW:
+      case LOAD_FAST_BORROW_LOAD_FAST_BORROW:
+      case LOAD_FAST_CHECK:
+      case LOAD_FAST_LOAD_FAST:
+      case NOP:
+      case NOT_TAKEN:
+      case POP_TOP:
+      case PUSH_NULL:
+      case RESUME:
+      case RETURN_CONST:
+      case RETURN_VALUE:
+      case STORE_FAST:
+      case STORE_FAST_LOAD_FAST:
+      case STORE_FAST_STORE_FAST:
+      case SWAP:
+        break;
+      default:
+        return false;
+    }
+    if (instr.isBackwardBranch()) {
+      return false;
+    }
+  }
+  return true;
+}
+
 std::unique_ptr<Function> buildHIR(const Preloader& preloader) {
   return HIRBuilder{preloader}.buildHIR();
 }
@@ -534,6 +570,8 @@ std::unique_ptr<Function> buildHIR(const Preloader& preloader) {
 std::unique_ptr<Function> HIRBuilder::buildHIR() {
   checkTranslate();
 
+  is_simple_leaf_function_ = isSimpleLeafFunction(code_);
+
   std::unique_ptr<Function> irfunc = preloader_.makeFunction();
   buildHIRImpl(irfunc.get(), /*frame_state=*/nullptr);
   // Use removeTrampolineBlocks and removeUnreachableBlocks directly instead of
@@ -1090,19 +1128,17 @@ void HIRBuilder::translate(
         }
         case POP_JUMP_IF_FALSE:
         case POP_JUMP_IF_TRUE: {
-          BCOffset target_off = bc_instr.getJumpTarget();
-          BasicBlock* target = getBlockAtOff(target_off);
-          if (target_off <= bc_instr.baseOffset()) {
+          BasicBlock* target = getBlockAtOff(bc_instr.getJumpTarget());
+          if (bc_instr.isBackwardBranch()) {
             loop_headers.emplace(target);
           }
           emitPopJumpIf(tc, bc_instr);
           break;
         }
         case POP_JUMP_IF_NONE:
         case POP_JUMP_IF_NOT_NONE: {
-          BCOffset target_off = bc_instr.getJumpTarget();
-          BasicBlock* target = getBlockAtOff(target_off);
-          if (target_off <= bc_instr.baseOffset()) {
+          BasicBlock* target = getBlockAtOff(bc_instr.getJumpTarget());
+          if (bc_instr.isBackwardBranch()) {
             loop_headers.emplace(target);
           }
           emitPopJumpIfNone(tc, bc_instr);
@@ -1936,6 +1972,9 @@ void HIRBuilder::emitResume(
   if (bc_instr.oparg() >= 2) {
     return;
   }
+  if (is_simple_leaf_function_) {
+    return;
+  }
   TranslationContext succ(cfg.AllocateBlock(), tc.frame);
   succ.emitSnapshot();
   insertRunPeriodicActivites(cfg, tc.block, succ.block, tc.frame);

cinderx/Jit/hir/builder.h

@@ -554,6 +554,15 @@ class HIRBuilder {
   Register* kwnames_{nullptr};
 
   OperandStack static_method_stack_;
+
+  // True if the function's bytecode contains only opcodes that cannot invoke
+  // user Python code and has no backward jumps (loops). Stricter than the
+  // common "leaf function" definition (no calls) — this also requires no
+  // complex opcodes. Used to skip RunPeriodicTasks at RESUME since the
+  // function returns quickly and the caller will check periodic tasks.
+  bool is_simple_leaf_function_{false};
+
+  static bool isSimpleLeafFunction(BorrowedRef<PyCodeObject> code);
 };
 
 } // namespace jit::hir

cinderx/Jit/jit_rt.cpp

@@ -854,6 +854,30 @@ void JITRT_UnlinkFrame(PyThreadState* tstate) {
   // JIT frames are stack allocated so there's nothing to pop.
 }
 
+// Clean up the reifier and decref the executable for a lightweight frame.
+// Shared by JITRT_UnlinkLightweightFrameFast and JITRT_UnlinkLeafFrame.
+static void cleanupLightweightFrameExecutable(
+    _PyInterpreterFrame* frame,
+    [[maybe_unused]] jit::FrameHeader* header) {
+#if PY_VERSION_HEX >= 0x030E0000
+  PyStackRef_CLOSE(frame->f_executable);
+#else
+  // Replace the reifier in f_funcobj with the actual function so that any
+  // escaped references to the frame see a valid function pointer, not a
+  // dangling reifier callback.
+  if (jit::hasRtfsFunction(frame)) {
+    frame->f_funcobj = jit::jitFrameGetRtfs(frame)->func();
+  } else {
+    PyObject* func = jit::jitFrameGetFunction(frame);
+    frame->f_funcobj = func;
+    Py_XDECREF(func);
+    header->rtfs = JIT_FRAME_INITIALIZED;
+  }
+
+  Py_DECREF(frameExecutable(frame));
+#endif
+}
+
 void JITRT_UnlinkLightweightFrameFast(PyThreadState* tstate) {
   _PyInterpreterFrame* frame = currentFrame(tstate);
   setCurrentFrame(tstate, frame->previous);
@@ -881,22 +905,19 @@ void JITRT_UnlinkLightweightFrameFast(PyThreadState* tstate) {
     Ci_STACK_CLOSE(frame->f_funcobj);
   }
 
-#if PY_VERSION_HEX >= 0x030E0000
-  PyStackRef_CLOSE(frame->f_executable);
-#else
-  // We can't leave our reifier dangling here otherwise we may
-  // continue to get callbacks, instead leave the function dangling.
-  if (jit::hasRtfsFunction(frame)) {
-    frame->f_funcobj = jit::jitFrameGetRtfs(frame)->func();
-  } else {
-    PyObject* func = jit::jitFrameGetFunction(frame);
-    frame->f_funcobj = func;
-    Py_XDECREF(func);
-    header->rtfs = JIT_FRAME_INITIALIZED;
-  }
+  cleanupLightweightFrameExecutable(frame, header);
+}
 
-  Py_DECREF(frameExecutable(frame));
-#endif
+void JITRT_UnlinkLeafFrame(PyThreadState* tstate) {
+  _PyInterpreterFrame* frame = currentFrame(tstate);
+  setCurrentFrame(tstate, frame->previous);
+
+  // No deopts means the frame was never materialized — skip the
+  // materialization check and just close funcobj + executable directly.
+  Ci_STACK_CLOSE(frame->f_funcobj);
+
+  auto* header = reinterpret_cast<jit::FrameHeader*>(frame) - 1;
+  cleanupLightweightFrameExecutable(frame, header);
 }
 
 PyObject*

cinderx/Jit/jit_rt.h

@@ -71,6 +71,11 @@ void JITRT_UnlinkFrame(PyThreadState* tstate);
 // jitFrameClearExceptCode path.
 void JITRT_UnlinkLightweightFrameFast(PyThreadState* tstate);
 
+// Specialized version for non-deopting leaf functions with lightweight frames.
+// Since no deopts can occur, the frame is guaranteed to never be materialized,
+// so we skip the materialization check entirely.
+void JITRT_UnlinkLeafFrame(PyThreadState* tstate);
+
 /*
  * Handles a call that includes kw arguments where the target function has
  * *args, **kwargs, or keyword only args.

cinderx/Jit/lir/generator.cpp

@@ -935,11 +935,18 @@ void LIRGenerator::GenerateExitBlocks() {
         Instruction::kPhi, nullptr, OutVReg{ret_data_type});
 
     // Unlink frame before epilogue. Non-generators always unlink.
+    bool has_freevars = func_->code != nullptr && func_->code->co_nfreevars > 0;
+    bool uses_lw_frames = getConfig().frame_mode == FrameMode::kLightweight;
+    uint64_t helper;
+    if (!env_->can_deopt && uses_lw_frames && !has_freevars) {
+      helper = reinterpret_cast<uint64_t>(JITRT_UnlinkLeafFrame);
+    } else if (!has_freevars && uses_lw_frames) {
+      helper = reinterpret_cast<uint64_t>(JITRT_UnlinkLightweightFrameFast);
+    } else {
+      helper = reinterpret_cast<uint64_t>(JITRT_UnlinkFrame);
+    }
     block->allocateInstr(
-        Instruction::kCall,
-        nullptr,
-        Imm{reinterpret_cast<uint64_t>(JITRT_UnlinkFrame)},
-        VReg{env_->asm_tstate});
+        Instruction::kCall, nullptr, Imm{helper}, VReg{env_->asm_tstate});
 
     block->allocateInstr(Instruction::kEpilogueEnd, nullptr, VReg{exit_phi_});
     return;

cinderx/RuntimeTests/hir_tests/all_passes_test.txt

@@ -11,23 +11,8 @@ def test():
 fun jittestmodule:test {
   bb 0 {
     LoadFrame
-    UpdatePrevInstr<idx:0 line_no:1: no parent>
-    v5:CInt32 = LoadEvalBreaker
-    CondBranch<2, 1> v5
-  }
-
-  bb 2 (preds 0) {
-    v6:CInt32 = RunPeriodicTasks {
-      FrameState {
-        CurInstrOffset 0
-      }
-    }
-    Branch<1>
-  }
-
-  bb 1 (preds 0, 2) {
-    v7:ImmortalLongExact[1] = LoadConst<ImmortalLongExact[1]>
-    Return<ImmortalLongExact[1]> v7
+    v3:ImmortalLongExact[1] = LoadConst<ImmortalLongExact[1]>
+    Return<ImmortalLongExact[1]> v3
   }
 }
 --- Expected 3.14 ---