Add inline cache for binary add

Summary:
This adds the infrastructure for inline caching of binary ops and adds support for inline caches of +. The goal here is to get more precise information about the inline caching then we get from the interpreter before we recompile a function.

The cache is represented by a single 1 byte value which follows transitions of initializing -> known types -> degraded types.  The types tracked are lhs/rhs/return type. So we may start in initializing, go to compact + compact == compact, degrade to long + long == long, and then degrade to fully generic.

Right now we support compact longs, longs, floats, strs, lists, tuples, and complex numbers.

There's a lot of template and macro machinery here but the goal is to make adding new ops easy.

Reviewed By: mpage

Differential Revision: D109088703

fbshipit-source-id: da6fd74d1af882f6dc1f419d74bf35358e74d02f

Author: DinoV

cinderx/Jit/config.h

@@ -150,6 +150,10 @@ struct Config {
   bool attr_caches{!kFreeThreadedBuild};
   // Collect stats information about attribute caches.
   bool collect_attr_cache_stats{false};
+  // Use inline caches for binary operations (currently the add variant only),
+  // dispatching through a self-modifying function pointer that specializes on
+  // the operand types observed at runtime.  Opt-in / off by default.
+  bool binary_op_caches{false};
   // Use type annotations to create runtime checks.
   bool emit_type_annotation_guards{false};
   // Whether or not to JIT specialized opcodes or to fall back to their generic

cinderx/Jit/context.cpp

@@ -373,6 +373,10 @@ LoadTypeMethodCache* Context::allocateLoadTypeMethodCache() {
   return load_type_method_caches_.allocate();
 }
 
+BinaryOpCache* Context::allocateBinaryOpCache(hir::BinaryOpKind op) {
+  return binary_op_caches_.allocate(op);
+}
+
 StoreAttrCache* Context::allocateStoreAttrCache() {
   return store_attr_caches_.allocate();
 }

cinderx/Jit/context.h

@@ -398,6 +398,7 @@ class Context : public IJitContext, public CompiledFunctionOwner {
   LoadModuleMethodCache* allocateLoadModuleMethodCache();
   LoadTypeMethodCache* allocateLoadTypeMethodCache();
   StoreAttrCache* allocateStoreAttrCache();
+  BinaryOpCache* allocateBinaryOpCache(hir::BinaryOpKind op);
 
   const Builtins& builtins();
 
@@ -498,6 +499,7 @@ class Context : public IJitContext, public CompiledFunctionOwner {
   SlabArena<LoadModuleMethodCache> load_module_method_caches_;
   SlabArena<LoadTypeMethodCache> load_type_method_caches_;
   SlabArena<StoreAttrCache, AttributeCacheSizeTrait> store_attr_caches_;
+  SlabArena<BinaryOpCache> binary_op_caches_;
   SlabArena<void*> pointer_caches_;
 
   FunctionEntryCacheMap function_entry_caches_;

cinderx/Jit/hir/builder.cpp

@@ -2027,6 +2027,16 @@ void HIRBuilder::emitBinaryOp(
   if (getConfig().specialized_opcodes) {
     switch (bc_instr.specializedOpcode()) {
       case BINARY_OP_ADD_INT:
+        // If we have inline caches for binary ops enabled then we don't want
+        // to specialize on the last seen interpreter type. The binary cache
+        // ops perform no backoff so once a cache is installed it persists.
+        // We want more accurate tracking of types to not perform deopts when
+        // the caches are enabled.
+        if (!getConfig().binary_op_caches) {
+          tc.emit<GuardType>(left, TLongExact, left, tc.frame);
+          tc.emit<GuardType>(right, TLongExact, right, tc.frame);
+        }
+        break;
       case BINARY_OP_MULTIPLY_INT:
       case BINARY_OP_SUBTRACT_INT:
         tc.emit<GuardType>(left, TLongExact, left, tc.frame);

cinderx/Jit/hir/hir.cpp

@@ -449,6 +449,7 @@ bool Instr::isReplayable() const {
     case Opcode::kBatchDecref:
     case Opcode::kBeginInlinedFunction:
     case Opcode::kBinaryOp:
+    case Opcode::kBinaryOpCached:
     case Opcode::kBranch:
     case Opcode::kBuildSlice:
     case Opcode::kBuildInterpolation:
@@ -733,6 +734,7 @@ bool isPassthrough(const Instr& instr) {
       return (static_cast<const Cast*>(&instr))->pytype() != &PyFloat_Type;
 
     case Opcode::kBinaryOp:
+    case Opcode::kBinaryOpCached:
     case Opcode::kBuildSlice:
     case Opcode::kBuildString:
     case Opcode::kBuildInterpolation:

cinderx/Jit/hir/hir.h

@@ -659,6 +659,40 @@ class INSTR_CLASS(
   BinaryOpKind op_;
 };
 
+// Variant of BinaryOp that dispatches through a per-instruction inline cache
+// (BinaryOpCache).  Used to specialize the operation based on the operand types
+// observed at runtime.  Currently only emitted for the add (kAdd) variant.
+class INSTR_CLASS(
+    BinaryOpCached,
+    (TObject, TObject),
+    HasOutput,
+    Operands<2>,
+    DeoptBase) {
+ public:
+  BinaryOpCached(
+      Register* dst,
+      BinaryOpKind op,
+      Register* left,
+      Register* right,
+      const FrameState& frame)
+      : InstrT(dst, left, right, frame), op_(op) {}
+
+  BinaryOpKind op() const {
+    return op_;
+  }
+
+  Register* left() const {
+    return getOperand(0);
+  }
+
+  Register* right() const {
+    return getOperand(1);
+  }
+
+ private:
+  BinaryOpKind op_;
+};
+
 #define FOREACH_UNARY_OP_KIND(V) \
   V(Not)                         \
   V(Negate)                      \

cinderx/Jit/hir/instr_effects.cpp

@@ -111,6 +111,7 @@ MemoryEffects memoryEffects(const Instr& inst) {
     // inputs, and may write all memory locations (usually from invoking
     // arbitrary user code).
     case Opcode::kBinaryOp:
+    case Opcode::kBinaryOpCached:
     case Opcode::kCallEx:
     case Opcode::kCallInd:
     case Opcode::kCallIntrinsic:
@@ -526,6 +527,7 @@ bool hasArbitraryExecution(const Instr& inst) {
     case Opcode::kLoadSpecial:
     case Opcode::kLongBinaryOp:
     case Opcode::kLongInPlaceOp:
+    case Opcode::kBinaryOpCached:
     case Opcode::kMakeFunction:
     case Opcode::kMergeSetUnpack:
     case Opcode::kMatchClass:

cinderx/Jit/hir/ops.h

@@ -13,6 +13,7 @@ namespace cinderx::jit::hir {
   V(BatchDecref)                   \
   V(BeginInlinedFunction)          \
   V(BinaryOp)                      \
+  V(BinaryOpCached)                \
   V(BitCast)                       \
   V(Branch)                        \
   V(BuildSlice)                    \

cinderx/Jit/hir/parser.cpp

@@ -543,6 +543,15 @@ HIRParser::parseInstr(std::string_view opcode, Register* dst, int bb_index) {
       instruction = newInstr<BinaryOp>(dst, op, left, right);
       break;
     }
+    case Opcode::kBinaryOpCached: {
+      expect("<");
+      BinaryOpKind op = ParseBinaryOpName(getNextToken());
+      expect(">");
+      auto left = parseRegister();
+      auto right = parseRegister();
+      instruction = newInstr<BinaryOpCached>(dst, op, left, right);
+      break;
+    }
     case Opcode::kLongBinaryOp: {
       expect("<");
       BinaryOpKind op = ParseBinaryOpName(getNextToken());

cinderx/Jit/hir/pass.cpp

@@ -188,6 +188,7 @@ Type outputType(
       return TObject;
     }
 
+    case Opcode::kBinaryOpCached:
     case Opcode::kBuildInterpolation:
     case Opcode::kBuildTemplate:
     case Opcode::kCallIntrinsic: