Sink non-escaping PrimitiveBox onto deopt paths

Summary:
Chained primitive float arithmetic already runs unboxed on the fast path
(`simplifyFloatBinaryOp` emits `PrimitiveUnbox`/`DoubleBinaryOp`/`PrimitiveBox`,
and `simplifyUnbox` fuses `unbox(box(x))`).  But the intermediate `PrimitiveBox`
instructions often survive: the boxed value has no data uses, yet it is held in
deopt frame state (`Locals`/`Stack`) so the interpreter can recover it if a
guard deopts.  The box is then pure fast-path overhead, an allocation that
exists only for a path that usually does not run.

This adds a `SinkPrimitiveBox` HIR pass.  For a `PrimitiveBox` whose result
never escapes (no data-operand use, only deopt frame-state references and
`UseType` assertions) it rewrites the frame-state references to the unboxed
source value and deletes the box (and the now-meaningless `UseType`s).  The
deopt machinery already re-boxes unboxed live values from their
`LiveValue::value_kind` (`kDouble` -> `PyFloat_FromDouble`), so deopt stays
correct while the fast path stays unboxed.  `UseType` operands are excluded when
deciding whether a value escapes, since `UseType` is a no-op type assertion, not
a real consumer.

Scoped to `CDouble` (floats) for now; the pass generalizes to other primitive
kinds (e.g. `CInt64`/`CUInt64`) with a single type-check change, which is the
higher-value follow-up for integer-arithmetic-heavy code where boxing is
expensive.

Reviewed By: mpage

Differential Revision: D110384610

fbshipit-source-id: 502d54a7199abbda046af2f663f9caa6ee626c55

Author: alexmalyshev

cinderx/Jit/compiler.cpp

@@ -20,6 +20,7 @@
 #include "cinderx/Jit/hir/printer.h"
 #include "cinderx/Jit/hir/refcount_insertion.h"
 #include "cinderx/Jit/hir/simplify.h"
+#include "cinderx/Jit/hir/sink_primitive_box.h"
 #include "cinderx/Jit/hir/ssa.h"
 #include "cinderx/Jit/hir/stats.h"
 #include "cinderx/Jit/jit_time_log.h"
@@ -107,6 +108,7 @@ void Compiler::runPasses(
       hir::BuiltinLoadMethodElimination{}, PassConfig::kBuiltinLoadMethodElim);
   runPassIf(hir::Simplify{}, PassConfig::kSimplify);
   runPassIf(hir::CleanCFG{}, PassConfig::kCleanCFG);
+  runPassIf(hir::SinkPrimitiveBox{}, PassConfig::kSinkPrimitiveBox);
   runPassIf(hir::DeadCodeElimination{}, PassConfig::kDeadCodeElim);
   runPassIf(hir::CleanCFG{}, PassConfig::kCleanCFG);
 
@@ -164,6 +166,7 @@ PassConfig createConfig() {
   set(hir_opts.insert_update_prev_instr, PassConfig::kInsertUpdatePrevInstr);
   set(hir_opts.phi_elim, PassConfig::kPhiElim);
   set(hir_opts.simplify, PassConfig::kSimplify);
+  set(hir_opts.sink_primitive_box, PassConfig::kSinkPrimitiveBox);
 
   return static_cast<PassConfig>(result);
 }

cinderx/Jit/compiler.h

@@ -38,6 +38,7 @@ enum PassConfig : uint64_t {
   kPhiElim = 1 << 7,
   kSimplify = 1 << 8,
   kInsertUpdatePrevInstr = 1 << 9,
+  kSinkPrimitiveBox = 1 << 10,
 
   // Run all the passes.
   kAll = ~uint64_t{0},

cinderx/Jit/config.h

@@ -43,6 +43,7 @@ struct HIROptimizations {
   bool insert_update_prev_instr{true};
   bool phi_elim{true};
   bool simplify{true};
+  bool sink_primitive_box{true};
 };
 
 // List of LIR optimization passes to run.

cinderx/Jit/hir/sink_primitive_box.cpp

@@ -0,0 +1,88 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include "cinderx/Jit/hir/sink_primitive_box.h"
+
+#include "cinderx/Jit/hir/hir.h"
+
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace cinderx::jit::hir {
+
+void SinkPrimitiveBox::run(Function& func) {
+  // Registers that escape as a real PyObject, i.e. appear as a data operand of
+  // some consuming instruction.  UseType is excluded: it is a no-op type
+  // assertion (it keeps a GuardType alive), not a real consumer, so a box used
+  // only by UseType and deopt frame state does not actually escape.
+  std::unordered_set<Register*> escapes;
+  for (auto& block : func.cfg.blocks) {
+    for (Instr& instr : block) {
+      if (instr.isUseType()) {
+        continue;
+      }
+      for (std::size_t i = 0, n = instr.numOperands(); i < n; ++i) {
+        escapes.insert(instr.getOperand(i));
+      }
+    }
+  }
+
+  // Map each sinkable box's result to its unboxed source value, remember the
+  // box instructions, and collect the UseType assertions on those boxes (which
+  // become meaningless once the box is gone).
+  std::unordered_map<Register*, Register*> sink_map;
+  std::vector<Instr*> dead_instrs;
+  for (auto& block : func.cfg.blocks) {
+    for (Instr& instr : block) {
+      if (!instr.isPrimitiveBox()) {
+        continue;
+      }
+      auto& box = static_cast<PrimitiveBox&>(instr);
+      // Limited to floats for now: deopt re-boxes a CDouble via PyFloat.
+      if (!(box.type() <= TCDouble)) {
+        continue;
+      }
+      if (!escapes.contains(box.output())) {
+        sink_map.emplace(box.output(), box.value());
+        dead_instrs.push_back(&instr);
+      }
+    }
+  }
+
+  if (sink_map.empty()) {
+    return;
+  }
+
+  // Drop UseType assertions on sunk boxes before rewriting, so we don't rewrite
+  // them to a primitive-typed register (which would be type-inconsistent).
+  for (auto& block : func.cfg.blocks) {
+    for (Instr& instr : block) {
+      if (instr.isUseType() && sink_map.contains(instr.getOperand(0))) {
+        dead_instrs.push_back(&instr);
+      }
+    }
+  }
+
+  // Rewrite the boxes' remaining (frame-state only) uses to the unboxed value.
+  // The deopt machinery records the unboxed value's kind and re-boxes it if a
+  // deopt fires, so the box is no longer needed on the fast path.
+  for (auto& block : func.cfg.blocks) {
+    for (Instr& instr : block) {
+      instr.visitUses([&](Register*& reg) {
+        auto it = sink_map.find(reg);
+        if (it != sink_map.end()) {
+          reg = it->second;
+        }
+        return true;
+      });
+    }
+  }
+
+  // The boxes (and their UseType assertions) now have no uses.
+  for (Instr* instr : dead_instrs) {
+    instr->unlink();
+    delete instr;
+  }
+}
+
+} // namespace cinderx::jit::hir

cinderx/Jit/hir/sink_primitive_box.h

@@ -0,0 +1,29 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#pragma once
+
+#include "cinderx/Jit/hir/pass.h"
+
+namespace cinderx::jit::hir {
+
+// Sink PrimitiveBox instructions onto deopt paths.
+//
+// When a boxed primitive's result has no data-operand uses (it is referenced
+// only by deopt frame state) the box exists solely so the value can be
+// materialized for the interpreter if a deopt occurs.  The deopt machinery can
+// re-box an unboxed primitive from its LiveValue (via value_kind), so we
+// rewrite those frame-state references to the unboxed source value, leaving the
+// box dead. This keeps chained primitive arithmetic unboxed on the fast path
+// while staying correct on deopt.  Currently limited to CDouble (floats).
+class SinkPrimitiveBox final : public Pass {
+ public:
+  SinkPrimitiveBox() : Pass("SinkPrimitiveBox") {}
+
+  void run(Function& irfunc) override;
+
+  static std::unique_ptr<SinkPrimitiveBox> factory() {
+    return std::make_unique<SinkPrimitiveBox>();
+  }
+};
+
+} // namespace cinderx::jit::hir

cinderx/Jit/pyjit.cpp

@@ -564,6 +564,11 @@ FlagProcessor initFlagProcessor() {
       "CINDERX_JIT_PHI_ELIM");
   HIR_OPTIMIZATION_OPTION(
       "simplify", simplify, "cinderx-jit-simplify", "CINDERX_JIT_SIMPLIFY");
+  HIR_OPTIMIZATION_OPTION(
+      "sink primitive box",
+      sink_primitive_box,
+      "cinderx-jit-sink-primitive-box",
+      "CINDERX_JIT_SINK_PRIMITIVE_BOX");
 
   flag_processor.addOption(
       "cinderx-jit-simplify-iteration-limit",