Break kYieldValue into components

Summary:
Breaks kYieldValue into a few smaller instructions:
  * kYieldValue - populates the return value into the return register
  * kStoreGenYieldPoint
  * kBranchToYieldExit
  * kResumeGenYield

This is going to ultimately let us get rid of `kYieldValue` instead - instead we'll be able to flow the return register directly to the end instead of having a hard-coded assignment into it. But before we can do that we need these broken up like this.

Reviewed By: kddnewton

Differential Revision: D95829402

fbshipit-source-id: 5a1d077af1714c41ff5f1ab84c82c94e787206f8

Author: DinoV

cinderx/Jit/codegen/autogen.cpp

@@ -8,6 +8,7 @@
 #include "cinderx/Jit/codegen/gen_asm_utils.h"
 #include "cinderx/Jit/frame.h"
 #include "cinderx/Jit/generators_rt.h"
+#include "cinderx/Jit/hir/hir.h"
 #include "cinderx/Jit/jit_rt.h"
 #include "cinderx/Jit/lir/instruction.h"
 #include "cinderx/Jit/lir/printer.h"
@@ -823,23 +824,64 @@ void translateYieldValue(Environ* env, const Instruction* instr) {
     PhyLocation value_out = instr->getInput(1)->getStackSlot();
     as->mov(x86::rax, x86::ptr(x86::rbp, value_out.loc));
   }
+#elif defined(CINDER_AARCH64)
+  a64::Builder* as = env->as;
 
-  // Arbitrary scratch register for use in emitStoreGenYieldPoint()
+  // Value to send goes to x0 so it can be yielded (returned) by epilogue.
+  if (instr->getInput(1)->isImm()) {
+    as->mov(a64::x0, instr->getInput(1)->getConstant());
+  } else {
+    PhyLocation value_out = instr->getInput(1)->getStackSlot();
+    as->ldr(
+        a64::x0,
+        arch::ptr_resolve(as, arch::fp, value_out.loc, arch::reg_scratch_0));
+  }
+#else
+  CINDER_UNSUPPORTED
+#endif
+}
+
+void translateStoreGenYieldPoint(Environ* env, const Instruction* instr) {
+#if defined(CINDER_X86_64)
+  arch::Builder* as = env->as;
   auto scratch_r = x86::r9;
-  auto resume_label = as->newLabel();
-  emitStoreGenYieldPoint(as, env, instr, resume_label, x86::rbp, scratch_r);
+  env->pending_yield_resume_label = as->newLabel();
+  emitStoreGenYieldPoint(
+      as, env, instr, env->pending_yield_resume_label, x86::rbp, scratch_r);
+#elif defined(CINDER_AARCH64)
+  a64::Builder* as = env->as;
+  auto scratch_r = arch::reg_scratch_0;
+  env->pending_yield_resume_label = as->newLabel();
+  emitStoreGenYieldPoint(
+      as, env, instr, env->pending_yield_resume_label, arch::fp, scratch_r);
+#else
+  CINDER_UNSUPPORTED
+#endif
+}
 
-  // Jump to epilogue
-  as->jmp(env->exit_for_yield_label);
+void translateBranchToYieldExit(Environ* env, const Instruction*) {
+#if defined(CINDER_X86_64)
+  env->as->jmp(env->exit_for_yield_label);
+#elif defined(CINDER_AARCH64)
+  env->as->b(env->exit_for_yield_label);
+#else
+  CINDER_UNSUPPORTED
+#endif
+}
+
+void translateResumeGenYield(Environ* env, const Instruction* instr) {
+#if defined(CINDER_X86_64)
+  arch::Builder* as = env->as;
 
   // Resumed execution in this generator begins here
-  as->bind(resume_label);
+  as->bind(env->pending_yield_resume_label);
 
 #if PY_VERSION_HEX < 0x030C0000
   // On 3.10, for yield-from yield points, store the finish_yield_from arg
   // (RDX from resume entry) into GenDataFooter so the subsequent Send
   // instruction can load it.
-  if (instr->isInYieldFromContext()) {
+  if (instr->origin() &&
+      static_cast<const hir::YieldValue*>(instr->origin())->isYieldFrom()) {
     auto fyf_offset = offsetof(GenDataFooter, finishYieldFrom);
     as->mov(x86::qword_ptr(x86::rbp, fyf_offset), x86::rdx);
   }
@@ -850,32 +892,15 @@ void translateYieldValue(Environ* env, const Instruction* instr) {
 #elif defined(CINDER_AARCH64)
   a64::Builder* as = env->as;
 
-  // Value to send goes to x0 so it can be yielded (returned) by epilogue.
-  if (instr->getInput(1)->isImm()) {
-    as->mov(a64::x0, instr->getInput(1)->getConstant());
-  } else {
-    PhyLocation value_out = instr->getInput(1)->getStackSlot();
-    as->ldr(
-        a64::x0,
-        arch::ptr_resolve(as, arch::fp, value_out.loc, arch::reg_scratch_0));
-  }
-
-  // Arbitrary scratch register for use in emitStoreGenYieldPoint()
-  auto scratch_r = arch::reg_scratch_0;
-  auto resume_label = as->newLabel();
-  emitStoreGenYieldPoint(as, env, instr, resume_label, arch::fp, scratch_r);
-
-  // Jump to epilogue
-  as->b(env->exit_for_yield_label);
-
   // Resumed execution in this generator begins here
-  as->bind(resume_label);
+  as->bind(env->pending_yield_resume_label);
 
 #if PY_VERSION_HEX < 0x030C0000
   // On 3.10, for yield-from yield points, store the finish_yield_from arg
   // (X2 from resume entry) into GenDataFooter so the subsequent Send
   // instruction can load it.
-  if (instr->isInYieldFromContext()) {
+  if (instr->origin() &&
+      static_cast<const hir::YieldValue*>(instr->origin())->isYieldFrom()) {
     auto fyf_offset = offsetof(GenDataFooter, finishYieldFrom);
     as->str(
         a64::x2,
@@ -1702,6 +1727,18 @@ BEGIN_RULES(Instruction::kYieldValue)
   GEN(ANY, CALL_C(translateYieldValue))
 END_RULES
 
+BEGIN_RULES(Instruction::kStoreGenYieldPoint)
+  GEN(ANY, CALL_C(translateStoreGenYieldPoint))
+END_RULES
+
+BEGIN_RULES(Instruction::kBranchToYieldExit)
+  GEN(ANY, CALL_C(translateBranchToYieldExit))
+END_RULES
+
+BEGIN_RULES(Instruction::kResumeGenYield)
+  GEN(ANY, CALL_C(translateResumeGenYield))
+END_RULES
+
 BEGIN_RULES(Instruction::kYieldExitPoint)
   GEN(ANY, CALL_C(translateYieldExitPoint))
 END_RULES
@@ -2993,6 +3030,18 @@ BEGIN_RULES(Instruction::kYieldValue)
   GEN(ANY, CALL_C(translateYieldValue))
 END_RULES
 
+BEGIN_RULES(Instruction::kStoreGenYieldPoint)
+  GEN(ANY, CALL_C(translateStoreGenYieldPoint))
+END_RULES
+
+BEGIN_RULES(Instruction::kBranchToYieldExit)
+  GEN(ANY, CALL_C(translateBranchToYieldExit))
+END_RULES
+
+BEGIN_RULES(Instruction::kResumeGenYield)
+  GEN(ANY, CALL_C(translateResumeGenYield))
+END_RULES
+
 BEGIN_RULES(Instruction::kYieldExitPoint)
   GEN(ANY, CALL_C(translateYieldExitPoint))
 END_RULES

cinderx/Jit/codegen/environ.h

@@ -42,6 +42,10 @@ struct Environ {
   asmjit::Label exit_for_yield_label;
   asmjit::Label gen_resume_entry_label;
 
+  // Resume label shared between StoreGenYieldPoint and ResumeGenYield.
+  // Created by translateStoreGenYieldPoint, bound by translateResumeGenYield.
+  asmjit::Label pending_yield_resume_label;
+
   // Deopt exits. One per guard.
   struct DeoptExit {
     DeoptExit(size_t idx, asmjit::Label lbl, const jit::lir::Instruction* ins)

cinderx/Jit/lir/generator.cpp

@@ -1309,18 +1309,34 @@ LIRGenerator::TranslatedBlock LIRGenerator::TranslateOneBasicBlock(
       }
       case Opcode::kYieldValue: {
         auto hir_instr = static_cast<const YieldValue*>(&i);
-        Instruction* instr = bbb.appendInstr(
-            hir_instr->output(),
-            Instruction::kYieldValue,
-            env_->asm_tstate,
-            hir_instr->reg());
+
+        // 1. kYieldValue: load tstate and yield value into ABI registers.
+        //    No output — the yielded value goes to the return register in
+        //    codegen.
+        bbb.appendInstr(
+            Instruction::kYieldValue, env_->asm_tstate, hir_instr->reg());
+
+        // 2. kStoreGenYieldPoint: store yield point metadata and live regs.
+        Instruction* store_instr =
+            bbb.appendInstr(Instruction::kStoreGenYieldPoint);
         if (hir_instr->isYieldFrom()) {
           // Add the sub-iterator as an extra input so that
           // emitStoreGenYieldPoint can capture its spill offset.
-          instr->addOperands(VReg{bbb.getDefInstr(hir_instr->yieldFromIter())});
-          instr->setYieldFromInputIdx(instr->getNumInputs() - 1);
+          store_instr->addOperands(
+              VReg{bbb.getDefInstr(hir_instr->yieldFromIter())});
+          store_instr->setYieldFromInputIdx(store_instr->getNumInputs() - 1);
         }
-        finishYield(bbb, instr, hir_instr);
+        finishYield(bbb, store_instr, hir_instr);
+
+        // 3. kBranchToYieldExit: jump to yield exit epilogue.
+        bbb.appendInstr(Instruction::kBranchToYieldExit);
+
+        // 4. kResumeGenYield: bind resume label, load resumed inputs.
+        //    Has output (the sent-in value) and takes tstate as input.
+        bbb.appendInstr(
+            hir_instr->output(),
+            Instruction::kResumeGenYield,
+            env_->asm_tstate);
         break;
       }
       case Opcode::kInitialYield: {

cinderx/Jit/lir/instruction.cpp

@@ -289,6 +289,7 @@ bool Instruction::isAnyYield() const {
   switch (opcode_) {
     case kYieldInitial:
     case kYieldValue:
+    case kStoreGenYieldPoint:
       return true;
     default:
       return false;

cinderx/Jit/lir/instruction.h

@@ -168,6 +168,9 @@ enum OperandSizeType {
   X(IntToBool, false, FlagEffects::kSet, kDefault, 1, {1})                    \
   X(YieldInitial, false, FlagEffects::kInvalidate, kDefault, 0, {}, 1)        \
   X(YieldValue, false, FlagEffects::kInvalidate, kDefault, 0, {}, 1)          \
+  X(StoreGenYieldPoint, false, FlagEffects::kInvalidate, kDefault, 0, {}, 1)  \
+  X(BranchToYieldExit, false, FlagEffects::kNone, kDefault, 0, {}, 1)         \
+  X(ResumeGenYield, false, FlagEffects::kInvalidate, kDefault, 0, {}, 1)      \
   X(YieldExitPoint, false, FlagEffects::kNone, kDefault, 0, {}, 1)            \
   X(EpilogueEnd, false, FlagEffects::kInvalidate, kDefault, 0, {}, 1)