BinaryOpCache: multiply (*) inline-cache support

Summary:
Add multiply (`*`) support to BinaryOpCache on top of the add-only base. The BinaryOpCached HIR node, the switch-dispatch machinery, and codegen are op-agnostic and unchanged here; this commit only adds the multiply operation.

Included:
- FOREACH_MULTIPLY_SPECIALIZATION and its folding into FOREACH_BINARY_OP_SPECIALIZATION, generating the kMul* / kUninitializedMultiply / kMultiplyGeneric Specialization enum values.
- Multiply op helpers: longMul, compactLongMul, floatMul, complexMul, and the sequence*long repeats (listMul, strMul, tupleMul via sequenceRepeat); multiplyGeneric and populateAndInvokeMultiply.
- selectInitialSpecialization / invoke / specializedTypes multiply arms.
- simplify rule extended to also rewrite a generic BinaryOp<Multiply> into BinaryOpCached<Multiply>.
- Multiply tests (inline_cache_test.cpp multiply specialization/compute/deopt cases, hir_simplify_test.cpp GenericMultiplyBecomesBinaryOpCached).

Reviewed By: yoney

Differential Revision: D109363776

fbshipit-source-id: 71f24c4f8620db15aa2ed6285c2bdc98b380fc0e

Author: DinoV

cinderx/Jit/hir/builder.cpp

@@ -2027,6 +2027,7 @@ void HIRBuilder::emitBinaryOp(
   if (getConfig().specialized_opcodes) {
     switch (bc_instr.specializedOpcode()) {
       case BINARY_OP_ADD_INT:
+      case BINARY_OP_MULTIPLY_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.
@@ -2037,7 +2038,6 @@ void HIRBuilder::emitBinaryOp(
           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);
         tc.emit<GuardType>(right, TLongExact, right, tc.frame);

cinderx/Jit/hir/simplify.cpp

@@ -1098,10 +1098,11 @@ Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
     return env.emit<UnicodeConcat>(lhs, rhs, *instr->frameState());
   }
 
-  // Generic add where the operand types aren't statically known: emit an
-  // inline-cached variant that specializes on the operand types seen at
-  // runtime (e.g. a fast path for int + int).
-  if (getConfig().binary_op_caches && op == BinaryOpKind::kAdd) {
+  // Generic add/multiply where the operand types aren't statically known: emit
+  // an inline-cached variant that specializes on the operand types seen at
+  // runtime (e.g. a fast path for int + int or int * int).
+  if (getConfig().binary_op_caches &&
+      (op == BinaryOpKind::kAdd || op == BinaryOpKind::kMultiply)) {
     return env.emit<BinaryOpCached>(op, lhs, rhs, *instr->frameState());
   }
 

cinderx/Jit/inline_cache.cpp

@@ -2103,7 +2103,7 @@ LoadModuleMethodCache::lookupSlowPath(BorrowedRef<> obj, BorrowedRef<> name) {
   return {nullptr, nullptr};
 }
 
-// Single source of truth for BinaryOpCache's add specializations, in priority
+// Single source of truth for BinaryOpCache's specializations, in priority
 // order.  Each row is X(Name, Lhs, Rhs, Ret, Op, Fallback):
 //   - Name           unique identifier; becomes Specialization::k<Name>.
 //   - Lhs, Rhs, Ret  operand and result types (SpecializedType values); their
@@ -2114,9 +2114,15 @@ LoadModuleMethodCache::lookupSlowPath(BorrowedRef<> obj, BorrowedRef<> name) {
 //                    compact/compact/long when a result overflows the compact
 //                    range, which in turn steps down to long/long/long when the
 //                    operands stop being compact.
-//   - Op             fast-path operation (e.g. longAdd, defined below).
-//   - Fallback       the Specialization to step down to once this row stops
-//                    matching.
+//   - Op             fast-path operation (e.g. longAdd, defined in
+//                    inline_cache.cpp).
+//   - Fallback       the per-op Specialization to step down to once this row
+//                    stops matching.
+//
+// These macros are defined here (rather than in inline_cache.cpp) so the
+// Specialization enum below can be generated from the same lists that drive the
+// dispatch switches.  The Op/Fallback columns are only expanded inside
+// inline_cache.cpp, so naming file-local op helpers here is fine.
 #define FOREACH_ADD_SPECIALIZATION(X)                                     \
   X(AddCompactCompactCompact,                                             \
     CompactLong,                                                          \
@@ -2137,12 +2143,45 @@ LoadModuleMethodCache::lookupSlowPath(BorrowedRef<> obj, BorrowedRef<> name) {
   X(AddTuple, Tuple, Tuple, Tuple, tupleAdd, kAddGeneric)                 \
   X(AddComplex, Complex, Complex, Complex, complexAdd, kAddGeneric)
 
+// Specializations for the multiply op.  Note the (sequence, long) rows have
+// distinct lhs/rhs/result types: list/str/tuple repeated by an integer count.
+#define FOREACH_MULTIPLY_SPECIALIZATION(X)                        \
+  X(MulCompactCompactCompact,                                     \
+    CompactLong,                                                  \
+    CompactLong,                                                  \
+    CompactLong,                                                  \
+    compactLongMul,                                               \
+    kMulCompactCompactLong)                                       \
+  X(MulCompactCompactLong,                                        \
+    CompactLong,                                                  \
+    CompactLong,                                                  \
+    Long,                                                         \
+    compactLongMul,                                               \
+    kMulLongLongLong)                                             \
+  X(MulLongLongLong, Long, Long, Long, longMul, kMultiplyGeneric) \
+  X(MulFloat, Float, Float, Float, floatMul, kMultiplyGeneric)    \
+  X(MulList, List, Long, List, listMul, kMultiplyGeneric)         \
+  X(MulUnicode, Unicode, Long, Unicode, strMul, kMultiplyGeneric) \
+  X(MulTuple, Tuple, Long, Tuple, tupleMul, kMultiplyGeneric)     \
+  X(MulComplex, Complex, Long, Complex, complexMul, kMultiplyGeneric)
+
+// The full specialization list (add followed by multiply), used to generate the
+// single Specialization enum and the specializedTypes() switch that covers all
+// values.
+#define FOREACH_BINARY_OP_SPECIALIZATION(X) \
+  FOREACH_ADD_SPECIALIZATION(X)             \
+  FOREACH_MULTIPLY_SPECIALIZATION(X)
+
 enum class BinaryOpCache::Specialization : uint8_t {
-#define DECLARE_ADD_SPECIALIZATION(NAME, LHS, RHS, RET, OP, FALLBACK) k##NAME,
+#define DECLARE_BINARY_OP_SPECIALIZATION(NAME, LHS, RHS, RET, OP, FALLBACK) \
+  k##NAME,
   kUninitializedAdd,
   kAddGeneric,
-  FOREACH_ADD_SPECIALIZATION(DECLARE_ADD_SPECIALIZATION)
-#undef DECLARE_ADD_SPECIALIZATION
+  FOREACH_ADD_SPECIALIZATION(DECLARE_BINARY_OP_SPECIALIZATION)
+      kUninitializedMultiply,
+  kMultiplyGeneric,
+  FOREACH_MULTIPLY_SPECIALIZATION(DECLARE_BINARY_OP_SPECIALIZATION)
+#undef DECLARE_BINARY_OP_SPECIALIZATION
 };
 
 BinaryOpCache::BinaryOpCache(cinderx::jit::hir::BinaryOpKind op)
@@ -2153,6 +2192,8 @@ BinaryOpCache::Specialization BinaryOpCache::selectInitialSpecialization(
   switch (op) {
     case cinderx::jit::hir::BinaryOpKind::kAdd:
       return Specialization::kUninitializedAdd;
+    case cinderx::jit::hir::BinaryOpKind::kMultiply:
+      return Specialization::kUninitializedMultiply;
     default:
       throw std::runtime_error(
           fmt::format(
@@ -2291,6 +2332,8 @@ PyObject* BinaryOpCache::invokeSpecialized(
   }
 #define POPULATE_ADD_SPECIALIZATION(...) \
   POPULATE_BINARY_SPECIALIZATION(add, __VA_ARGS__)
+#define POPULATE_MULTIPLY_SPECIALIZATION(...) \
+  POPULATE_BINARY_SPECIALIZATION(multiply, __VA_ARGS__)
 
 // Emits one dispatch-switch arm that runs the specialization directly via
 // invokeSpecialized<>, threading the Fallback value and the matching
@@ -2307,6 +2350,8 @@ PyObject* BinaryOpCache::invokeSpecialized(
         &BinaryOpCache::DISPATCH>(lhs, rhs, cache);
 #define DISPATCH_ADD_SPECIALIZATION(...) \
   DISPATCH_BINARY_SPECIALIZATION(add, __VA_ARGS__)
+#define DISPATCH_MULTIPLY_SPECIALIZATION(...) \
+  DISPATCH_BINARY_SPECIALIZATION(multiply, __VA_ARGS__)
 
 // Emits one specializedTypes() switch arm mapping a specialization to its
 // (lhs, rhs, return) operand types.  A single enum lets one switch cover both
@@ -2353,13 +2398,78 @@ static inline PyObject* tupleAdd(PyObject* lhs, PyObject* rhs) {
   return PyTuple_Type.tp_as_sequence->sq_concat(lhs, rhs);
 }
 
+static inline PyObject* longMul(PyObject* lhs, PyObject* rhs) {
+#if PY_VERSION_HEX >= 0x030F0000
+  // _PyLong_Multiply was removed in 3.15.  Both operands are exact ints here,
+  // so the public number slot is equivalent and returns a new reference.
+  return PyLong_Type.tp_as_number->nb_multiply(lhs, rhs);
+#else
+  return _PyLong_Multiply(
+      reinterpret_cast<PyLongObject*>(lhs),
+      reinterpret_cast<PyLongObject*>(rhs));
+#endif
+}
+
+// Fast path for two compact ints: multiply their machine-word values
+// directly. Both operands are single-digit (guaranteed by _PyLong_IsCompact),
+// so the product cannot overflow Py_ssize_t.  The result may itself be
+// non-compact; the compact/compact/compact specialization detects that via
+// its return-type check and steps down to compact/compact/long.
+static inline PyObject* compactLongMul(PyObject* lhs, PyObject* rhs) {
+  Py_ssize_t a = _PyLong_CompactValue(reinterpret_cast<PyLongObject*>(lhs));
+  Py_ssize_t b = _PyLong_CompactValue(reinterpret_cast<PyLongObject*>(rhs));
+  return PyLong_FromSsize_t(a * b);
+}
+
+static inline PyObject* floatMul(PyObject* lhs, PyObject* rhs) {
+  double a = reinterpret_cast<PyFloatObject*>(lhs)->ob_fval;
+  double b = reinterpret_cast<PyFloatObject*>(rhs)->ob_fval;
+  return PyFloat_FromDouble(a * b);
+}
+
+static inline PyObject* complexMul(PyObject* lhs, PyObject* rhs) {
+  // complex * long: the complex nb_multiply slot coerces the integer operand.
+  return PyComplex_Type.tp_as_number->nb_multiply(lhs, rhs);
+}
+
+// Sequence-repeat helpers for (sequence, long) multiplication.  The repeat
+// count is the integer rhs; an out-of-range count surfaces as an error from
+// PyLong_AsSsize_t, matching the generic path.
+static inline PyObject*
+sequenceRepeat(PySequenceMethods* methods, PyObject* seq, PyObject* count) {
+  Py_ssize_t n = PyLong_AsSsize_t(count);
+  if (n == -1 && PyErr_Occurred()) {
+    return nullptr;
+  }
+  return methods->sq_repeat(seq, n);
+}
+
+static inline PyObject* listMul(PyObject* lhs, PyObject* rhs) {
+  return sequenceRepeat(PyList_Type.tp_as_sequence, lhs, rhs);
+}
+
+static inline PyObject* strMul(PyObject* lhs, PyObject* rhs) {
+  return sequenceRepeat(PyUnicode_Type.tp_as_sequence, lhs, rhs);
+}
+
+static inline PyObject* tupleMul(PyObject* lhs, PyObject* rhs) {
+  return sequenceRepeat(PyTuple_Type.tp_as_sequence, lhs, rhs);
+}
+
 PyObject* BinaryOpCache::addGeneric(
     PyObject* lhs,
     PyObject* rhs,
     BinaryOpCache* /* cache */) {
   return PyNumber_Add(lhs, rhs);
 }
 
+PyObject* BinaryOpCache::multiplyGeneric(
+    PyObject* lhs,
+    PyObject* rhs,
+    BinaryOpCache* /* cache */) {
+  return PyNumber_Multiply(lhs, rhs);
+}
+
 PyObject* BinaryOpCache::populateAndInvokeAdd(
     PyObject* lhs,
     PyObject* rhs,
@@ -2370,10 +2480,20 @@ PyObject* BinaryOpCache::populateAndInvokeAdd(
   return addGeneric(lhs, rhs, cache);
 }
 
-// Dispatch on the cache's current specialization and run the corresponding add
-// operation directly.  The specialized arms are generated from
-// FOREACH_ADD_SPECIALIZATION so this stays in sync with the Specialization
-// enum. There is no indirect call through a function pointer.
+PyObject* BinaryOpCache::populateAndInvokeMultiply(
+    PyObject* lhs,
+    PyObject* rhs,
+    BinaryOpCache* cache) {
+  FOREACH_MULTIPLY_SPECIALIZATION(POPULATE_MULTIPLY_SPECIALIZATION)
+
+  cache->specialization_ = Specialization::kMultiplyGeneric;
+  return multiplyGeneric(lhs, rhs, cache);
+}
+
+// Dispatch on the cache's current specialization and run the corresponding
+// add operation directly.  The arms cover only the add subset of the single
+// Specialization enum (generated from FOREACH_ADD_SPECIALIZATION); multiply
+// states never reach here because codegen calls add() only for kAdd caches.
 PyObject*
 BinaryOpCache::add(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache) {
   switch (cache->specialization_) {
@@ -2382,23 +2502,41 @@ BinaryOpCache::add(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache) {
     case Specialization::kAddGeneric:
       return addGeneric(lhs, rhs, cache);
       FOREACH_ADD_SPECIALIZATION(DISPATCH_ADD_SPECIALIZATION)
+    default:
+      JIT_ABORT("Unexpected specialization in BinaryOpCache::add");
+  }
+}
+
+// Dispatch on the cache's current specialization.  Mirrors add() but over the
+// multiply subset of the enum (FOREACH_MULTIPLY_SPECIALIZATION).
+PyObject*
+BinaryOpCache::multiply(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache) {
+  switch (cache->specialization_) {
+    case Specialization::kUninitializedMultiply:
+      return populateAndInvokeMultiply(lhs, rhs, cache);
+    case Specialization::kMultiplyGeneric:
+      return multiplyGeneric(lhs, rhs, cache);
+      FOREACH_MULTIPLY_SPECIALIZATION(DISPATCH_MULTIPLY_SPECIALIZATION)
+    default:
+      JIT_ABORT("Unexpected specialization in BinaryOpCache::multiply");
   }
-  JIT_ABORT("Unknown BinaryOpCache specialization");
 }
 
 BinaryOpCache::BinarySpecialization BinaryOpCache::specializedTypes() const {
   switch (specialization_) {
     case Specialization::kUninitializedAdd:
+    case Specialization::kUninitializedMultiply:
       return BinarySpecialization{
           SpecializedType::kUninitialized,
           SpecializedType::kUninitialized,
           SpecializedType::kUninitialized};
     case Specialization::kAddGeneric:
+    case Specialization::kMultiplyGeneric:
       return BinarySpecialization{
           SpecializedType::kGeneric,
           SpecializedType::kGeneric,
           SpecializedType::kGeneric};
-      FOREACH_ADD_SPECIALIZATION(SPECIALIZATION_TYPES_ENTRY)
+      FOREACH_BINARY_OP_SPECIALIZATION(SPECIALIZATION_TYPES_ENTRY)
   }
   JIT_ABORT("Unknown BinaryOpCache specialization");
 }

cinderx/Jit/inline_cache.h

@@ -486,21 +486,29 @@ enum class SpecializedType : uint8_t {
 // A cache for an individual BinaryOpCached instruction.
 //
 // Implements an inline cache for binary operations as a small state machine.
-// A cache starts in the populate state, which checks the inputs for known cache
-// types on the first invocation, then transitions specialization_ to the
-// matching specialized state, or to a generic state when no SpecializedType
-// applies.
+// A single Specialization enum covers both add and multiply states, but add and
+// multiply have separate dispatch entry points (add() / multiply()) that each
+// switch over their op's subset of the enum.  A cache is constructed for a
+// single op; it starts in that op's populate state, which checks the inputs for
+// known cache types on the first invocation, then transitions specialization_
+// to the matching specialized state, or to a generic state when no
+// SpecializedType applies.
 //
-// Codegen emits a direct call to add(), which switches on specialization_ and
-// calls the matching specialized operation directly -- there is no indirect
-// call through a function pointer.
+// Codegen emits a direct call to add() (for kAdd) or multiply() (for
+// kMultiply); each switches on specialization_ and calls the matching
+// specialized operation directly -- there is no indirect call through a
+// function pointer.
 class BinaryOpCache {
  public:
   // Identifies which specialization the cache has settled on, i.e. which
-  // operation add() dispatches to.  The k<Name> values are auto-generated from
-  // FOREACH_ADD_SPECIALIZATION (in inline_cache.cpp); kUninitializedAdd is the
-  // initial (lazily specializing) populate state and kAddGeneric is the
-  // permanent generic fallback.  Defined out-of-line in inline_cache.cpp.
+  // operation add()/multiply() dispatches to.  A single enum holds both ops'
+  // states: the k<Name> values are auto-generated from
+  // FOREACH_BINARY_OP_SPECIALIZATION, the kUninitialized* values are the
+  // initial (lazily specializing) populate states, and
+  // kAddGeneric/kMultiplyGeneric are the permanent generic fallbacks.  add()
+  // only ever observes the add subset and multiply() the multiply subset, but a
+  // single enum lets specializedTypes() switch over all values without a
+  // discriminant.
   enum class Specialization : uint8_t;
 
   // The (lhs, rhs, return) operand/result types a cache has specialized to.
@@ -520,9 +528,11 @@ class BinaryOpCache {
   // op has no inline-cache support.
   explicit BinaryOpCache(cinderx::jit::hir::BinaryOpKind op);
 
-  // Dispatch entry point called directly by codegen for kAdd.  Switches on the
-  // cache's specialization enum and runs the corresponding operation directly.
+  // Dispatch entry points called directly by codegen: add() for kAdd,
+  // multiply() for kMultiply.  Each switches on the cache's per-op
+  // specialization enum and runs the corresponding operation directly.
   static PyObject* add(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
+  static PyObject* multiply(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
 
   // Returns the (lhs, rhs, return) operand types the cache has settled on
   // ({kUninitialized, ...} before the first call).
@@ -541,10 +551,19 @@ class BinaryOpCache {
   static PyObject*
   populateAndInvokeAdd(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
 
+  // Initial entry point for the multiply op: inspects the operand types,
+  // transitions the multiply specialization, and performs the operation.
+  static PyObject*
+  populateAndInvokeMultiply(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
+
   // Permanent generic fallback that just calls PyNumber_Add.
   static PyObject*
   addGeneric(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
 
+  // Permanent generic fallback that just calls PyNumber_Multiply.
+  static PyObject*
+  multiplyGeneric(PyObject* lhs, PyObject* rhs, BinaryOpCache* cache);
+
   // Specialized entry for a (lhs, rhs) -> ret triple.  Guards that lhs passes
   // checkFor(LhsKind) and rhs passes checkFor(RhsKind) and, if so, runs the
   // fast-path Op.  When the return type is refined (returnNeedsCheck), it also

cinderx/Jit/lir/generator.cpp

@@ -3351,16 +3351,28 @@ LIRGenerator::TranslatedBlock LIRGenerator::translateOneBasicBlock(
       case Opcode::kBinaryOpCached: {
         auto instr = static_cast<const BinaryOpCached*>(&i);
         BinaryOpCache* cache = getContext()->allocateBinaryOpCache(instr->op());
-        // Emit a direct call to the dispatch entry point add(), which switches
-        // on the cache's specialization enum -- there is no indirect call
-        // through a function pointer.  allocateBinaryOpCache() already rejected
-        // any unsupported op kind.
-        bbb.appendCallInstruction(
-            instr->output(),
-            BinaryOpCache::add,
-            instr->left(),
-            instr->right(),
-            cache);
+        // Emit a direct call to the op-specific dispatch entry point: add() for
+        // kAdd, multiply() for kMultiply.  Each switches on the cache's per-op
+        // specialization enum -- there is no indirect call through a function
+        // pointer.  allocateBinaryOpCache() already rejected any other op kind.
+        if (instr->op() == BinaryOpKind::kMultiply) {
+          bbb.appendCallInstruction(
+              instr->output(),
+              BinaryOpCache::multiply,
+              instr->left(),
+              instr->right(),
+              cache);
+        } else {
+          JIT_DCHECK(
+              instr->op() == BinaryOpKind::kAdd,
+              "BinaryOpCached only supports add and multiply");
+          bbb.appendCallInstruction(
+              instr->output(),
+              BinaryOpCache::add,
+              instr->left(),
+              instr->right(),
+              cache);
+        }
         break;
       }
       case Opcode::kLongBinaryOp: {