Implement float comparisons with primitive operations

Summary:
Lower float comparisons to unboxed `PrimitiveCompare` on `CDouble` instead of
falling back to `PyFloat_Type.tp_richcompare`.  This avoids boxing, a
type-generic dispatch, and a call per comparison on the hot path.

The subtlety is NaN.  Python requires every ordering comparison involving a NaN
to be false, `nan == nan` to be false, and `nan != nan` to be true, but the
hardware float-compare instructions expose different flags for unordered (NaN)
operands on x86-64 (`comisd`) and ARM64 (`fcmp`).  The HIR stays
architecture-independent, it emits the natural comparison
(`PrimitiveCompare<LessThan> x, y` and friends, and
`PrimitiveCompare<Equal>`/`<NotEqual>` for `==`/`!=`), and the NaN-correct
condition is chosen when the compare is lowered to LIR.

The choice must be made at the LIR-opcode level (not by, say, swapping operands
inside the compare's asm emission) because a compare that feeds a branch is
fused into a conditional jump whose condition is derived from the LIR opcode via
`compareToBranchCC`, reusing the compare's flags.  If the compare's asm silently
swapped operands, the fused jump, which is unaware of that swap, would test the
wrong condition.  So the LIR generator picks the opcode and operand order so
that the opcode's standard condition is NaN-correct for both the standalone
`setcc`/`cset` and the fused `jcc`/`b.cc`:

- x86-64: `comisd` sets the flags in an unsigned sense and marks unordered
  operands with CF=ZF=PF=1, so only the above / above-equal conditions are false
  for NaN.  Every ordering is expressed as `>` / `>=`
  (`kGreaterThanUnsigned`/`kGreaterThanEqualUnsigned`), swapping the operands
  for `<` / `<=`. `comisd` folds unordered into ZF (the equality flag), so
  `==`/`!=` have no NaN-correct single-instruction form; they are decomposed
  into the ordering comparisons: `a == b` becomes `(a <= b) && (a >= b)` (both
  false when a NaN is involved), and `a != b` negates that.  `(a < b) || (a > b)`
  would be wrong, it yields false for NaN.

- ARM64: `fcmp` sets NZCV and marks unordered operands with C=1, V=1 while
  leaving Z=0.  `>` / `>=` use the signed `GT`/`GE` conditions and `<` / `<=` the
  unsigned `LO`/`LS` conditions (all false for unordered); `==`/`!=` use
  `EQ`/`NE` (which key off Z alone and are already NaN-correct).

`TranslateCompare` is unchanged in spirit: it emits `comisd`/`fcmp` for double
operands and falls through to the shared per-opcode condition switch, so the
standalone and fused forms stay consistent.  Floats are neither signed nor
unsigned, so both the signed and unsigned `PrimitiveCompareOp` variants are
handled together in the generator (Static Python `double` emits the unsigned
variants; the Python-float path emits the natural ones).  A side benefit is that
Static Python `double` `<`/`<=` are now NaN-correct too (they previously used
`setb`/`setbe`, which are true for unordered operands).

Unboxing the nms float comparisons speeds it up ~7%; the other workloads are
within noise.

Reviewed By: mpage

Differential Revision: D110238069

fbshipit-source-id: 24952346336c44ea6d7225add39596c255a53f8b

Author: alexmalyshev

cinderx/Jit/codegen/autogen.cpp

@@ -531,6 +531,13 @@ void TranslateCompare(Environ* env, const Instruction* instr) {
   } else if (!inp1->isVecD()) {
     as->cmp(AutoTranslator::getGp(inp0), AutoTranslator::getGp(inp1));
   } else {
+    // Floating-point comparison; both operands are in XMM registers.  `comisd`
+    // sets the flags in the unsigned sense (CF/ZF) and reports unordered (NaN)
+    // operands as CF=ZF=PF=1; the shared condition switch below then reads
+    // those flags.  NaN-correctness and the comparison direction are chosen
+    // when the compare is lowered to LIR, so a compare fused into a branch,
+    // which reuses these flags via compareToBranchCC on the LIR opcode, stays
+    // consistent with the standalone setcc emitted here.
     as->comisd(AutoTranslator::getVecD(inp0), AutoTranslator::getVecD(inp1));
   }
   auto output = AutoTranslator::getGp(instr->output());
@@ -566,7 +573,7 @@ void TranslateCompare(Environ* env, const Instruction* instr) {
       as->setbe(output);
       break;
     default:
-      JIT_ABORT("bad instruction for TranslateCompare");
+      JIT_ABORT("Bad instruction for TranslateCompare {}", instr->opname());
   }
   if (instr->output()->dataType() != lir::Operand::k8bit) {
     as->movzx(
@@ -593,6 +600,11 @@ void TranslateCompare(Environ* env, const Instruction* instr) {
   } else if (!inp1->isVecD()) {
     as->cmp(AutoTranslator::getGpWiden(inp0), AutoTranslator::getGpWiden(inp1));
   } else {
+    // Floating-point comparison, see the note in the x86-64 path.  `fcmp` sets
+    // NZCV (unordered/NaN operands set C=1, V=1 while leaving Z=0), the shared
+    // condition switch below picks the cset. NaN-correctness and the comparison
+    // direction are chosen when the compare is lowered to LIR, keeping the
+    // standalone cset and any fused b.cc consistent.
     as->fcmp(AutoTranslator::getVecD(inp0), AutoTranslator::getVecD(inp1));
   }
 
@@ -629,7 +641,7 @@ void TranslateCompare(Environ* env, const Instruction* instr) {
       as->cset(output, arm::CondCode::kLS);
       break;
     default:
-      JIT_ABORT("bad instruction for TranslateCompare");
+      JIT_ABORT("Bad instruction for TranslateCompare {}", instr->opname());
   }
 #else
   CINDER_UNSUPPORTED

cinderx/Jit/hir/hir.cpp

@@ -394,7 +394,6 @@ bool Instr::isReplayable() const {
     case Opcode::kCIntToCBool:
     case Opcode::kCompactLongUnbox:
     case Opcode::kDoubleBinaryOp:
-    case Opcode::kFloatCompare:
     case Opcode::kFormatValue:
     case Opcode::kFormatWithSpec:
     case Opcode::kGetSecondOutput:
@@ -760,7 +759,6 @@ bool isPassthrough(const Instr& instr) {
     case Opcode::kFillTypeAttrCache:
     case Opcode::kFillTypeMethodCache:
     case Opcode::kFloatBinaryOp:
-    case Opcode::kFloatCompare:
     case Opcode::kFormatValue:
     case Opcode::kFormatWithSpec:
     case Opcode::kGetAIter:

cinderx/Jit/hir/hir.h

@@ -1785,32 +1785,6 @@ class INSTR_CLASS(
   CompareOp op_;
 };
 
-// Perform the comparison indicated by op between two floats
-class INSTR_CLASS(
-    FloatCompare,
-    (TFloatExact, TFloatExact),
-    HasOutput,
-    Operands<2>) {
- public:
-  FloatCompare(Register* dst, CompareOp op, Register* left, Register* right)
-      : InstrT(dst, left, right), op_(op) {}
-
-  CompareOp op() const {
-    return op_;
-  }
-
-  Register* left() const {
-    return getOperand(0);
-  }
-
-  Register* right() const {
-    return getOperand(1);
-  }
-
- private:
-  CompareOp op_;
-};
-
 // Perform the comparison indicated by op between two longs
 class INSTR_CLASS(
     LongCompare,

cinderx/Jit/hir/instr_effects.cpp

@@ -41,7 +41,6 @@ MemoryEffects memoryEffects(const Instr& inst) {
     case Opcode::kDeopt:
     case Opcode::kDeoptPatchpoint:
     case Opcode::kDoubleBinaryOp:
-    case Opcode::kFloatCompare:
     case Opcode::kGetSecondOutput:
     case Opcode::kHintType:
     case Opcode::kIndexUnbox:
@@ -400,7 +399,6 @@ bool hasArbitraryExecution(const Instr& inst) {
     case Opcode::kDeoptPatchpoint:
     case Opcode::kDoubleBinaryOp:
     case Opcode::kEndInlinedFunction:
-    case Opcode::kFloatCompare:
     case Opcode::kGetSecondOutput:
     case Opcode::kGuardIs:
     case Opcode::kHintType:

cinderx/Jit/hir/ops.h

@@ -57,7 +57,6 @@ namespace cinderx::jit::hir {
   V(FillTypeAttrCache)             \
   V(FillTypeMethodCache)           \
   V(FloatBinaryOp)                 \
-  V(FloatCompare)                  \
   V(FormatValue)                   \
   V(FormatWithSpec)                \
   V(GetAIter)                      \

cinderx/Jit/hir/parser.cpp

@@ -588,15 +588,6 @@ HIRParser::parseInstr(std::string_view opcode, Register* dst, int bb_index) {
       instruction = newInstr<Compare>(dst, op, left, right);
       break;
     }
-    case Opcode::kFloatCompare: {
-      expect("<");
-      CompareOp op = ParseCompareOpName(getNextToken());
-      expect(">");
-      auto left = parseRegister();
-      auto right = parseRegister();
-      NEW_INSTR(FloatCompare, dst, op, left, right);
-      break;
-    }
     case Opcode::kLongCompare: {
       expect("<");
       CompareOp op = ParseCompareOpName(getNextToken());

cinderx/Jit/hir/pass.cpp

@@ -354,7 +354,6 @@ Type outputType(
     }
     case Opcode::kFloatBinaryOp:
       return TFloatExact;
-    case Opcode::kFloatCompare:
     case Opcode::kLongCompare:
     case Opcode::kUnicodeCompare:
       return TImmortalBool;

cinderx/Jit/hir/printer.cpp

@@ -474,10 +474,6 @@ static std::string format_immediates(const Function* func, const Instr& instr) {
       const auto& cmp = static_cast<const Compare&>(instr);
       return std::string{GetCompareOpName(cmp.op())};
     }
-    case Opcode::kFloatCompare: {
-      const auto& cmp = static_cast<const FloatCompare&>(instr);
-      return std::string{GetCompareOpName(cmp.op())};
-    }
     case Opcode::kLongCompare: {
       const auto& cmp = static_cast<const LongCompare&>(instr);
       return std::string{GetCompareOpName(cmp.op())};

cinderx/Jit/hir/simplify.cpp

@@ -430,12 +430,22 @@ Register* simplifyCompare(Env& env, const Compare* instr) {
     }
   }
 
-  // Emit FloatCompare if both args are FloatExact and the op is supported
-  // between two longs.
+  // Emit primitive comparisons for floats: unbox and compare as CDouble.  The
+  // op is emitted naturally (PrimitiveCompare<LessThan> and friends). NaNs are
+  // handled using Python's rules (`NaN == NaN` is false, `NaN != NaN` is true,
+  // all other comparison types with NaN are false).
   if (left->isA(TFloatExact) && right->isA(TFloatExact) &&
-      !(op == CompareOp::kIn || op == CompareOp::kNotIn ||
-        op == CompareOp::kExcMatch)) {
-    return env.emit<FloatCompare>(instr->op(), left, right);
+      (op == CompareOp::kLessThan || op == CompareOp::kLessThanEqual ||
+       op == CompareOp::kGreaterThan || op == CompareOp::kGreaterThanEqual ||
+       op == CompareOp::kEqual || op == CompareOp::kNotEqual)) {
+    std::optional<PrimitiveCompareOp> prim_op = toPrimitiveCompareOp(op);
+    env.emit<UseType>(left, TFloatExact);
+    env.emit<UseType>(right, TFloatExact);
+    Register* unboxed_left = env.emit<PrimitiveUnbox>(left, TCDouble);
+    Register* unboxed_right = env.emit<PrimitiveUnbox>(right, TCDouble);
+    Register* result =
+        env.emit<PrimitiveCompare>(*prim_op, unboxed_left, unboxed_right);
+    return env.emit<PrimitiveBoxBool>(result);
   }
 
   // Emit LongCompare if both args are LongExact and the op is supported between

cinderx/Jit/lir/generator.cpp

@@ -2592,6 +2592,126 @@ LIRGenerator::TranslatedBlock LIRGenerator::translateOneBasicBlock(
       }
       case Opcode::kPrimitiveCompare: {
         auto instr = static_cast<const PrimitiveCompare*>(&i);
+        // Float comparisons need NaN-correct condition codes, and the choice is
+        // architecture-specific because comisd (x86-64) and fcmp (ARM64) expose
+        // different flags for unordered (NaN) operands.  Pick the LIR opcode
+        // (and operand order) so the opcode's standard condition, used both by
+        // the standalone setcc/cset and by a compare fused into a branch
+        // through compareToBranchCC, yields Python's result: every ordering
+        // comparison involving a NaN is false, `NaN == NaN` is False, `Nan !=
+        // NaN` is True.
+        if (instr->left()->type() <= TCDouble) {
+          Register* lhs = instr->left();
+          Register* rhs = instr->right();
+#if defined(CINDER_X86_64)
+          // comisd sets unsigned-sense flags and reports unordered as CF=1, so
+          // only the above / above-equal conditions are false for NaN.  Express
+          // every ordering as > / >=, swapping operands for < / <=.  == / !=
+          // have no NaN-correct single-instruction form (comisd folds unordered
+          // into ZF), so build them from ordering comparisons:
+          //   a == b == (a <= b) && (a >= b);  a != b == !(a == b)
+          // Floats are neither signed nor unsigned, so the signed and unsigned
+          // PrimitiveCompareOp variants denote the same ordering (Static Python
+          // `double` emits the unsigned variants; the Python-float path emits
+          // the natural ones).
+          switch (instr->op()) {
+            case PrimitiveCompareOp::kGreaterThan:
+            case PrimitiveCompareOp::kGreaterThanUnsigned:
+              bbb.appendInstr(
+                  instr->output(), Instruction::kGreaterThanUnsigned, lhs, rhs);
+              break;
+            case PrimitiveCompareOp::kGreaterThanEqual:
+            case PrimitiveCompareOp::kGreaterThanEqualUnsigned:
+              bbb.appendInstr(
+                  instr->output(),
+                  Instruction::kGreaterThanEqualUnsigned,
+                  lhs,
+                  rhs);
+              break;
+            case PrimitiveCompareOp::kLessThan: // a < b == b > a
+            case PrimitiveCompareOp::kLessThanUnsigned:
+              bbb.appendInstr(
+                  instr->output(), Instruction::kGreaterThanUnsigned, rhs, lhs);
+              break;
+            case PrimitiveCompareOp::kLessThanEqual: // a <= b == b >= a
+            case PrimitiveCompareOp::kLessThanEqualUnsigned:
+              bbb.appendInstr(
+                  instr->output(),
+                  Instruction::kGreaterThanEqualUnsigned,
+                  rhs,
+                  lhs);
+              break;
+            case PrimitiveCompareOp::kEqual:
+            case PrimitiveCompareOp::kNotEqual: {
+              Instruction* le = bbb.appendInstr(
+                  OutVReg{Operand::k8bit},
+                  Instruction::kGreaterThanEqualUnsigned,
+                  rhs,
+                  lhs); // a <= b
+              Instruction* ge = bbb.appendInstr(
+                  OutVReg{Operand::k8bit},
+                  Instruction::kGreaterThanEqualUnsigned,
+                  lhs,
+                  rhs); // a >= b
+              if (instr->op() == PrimitiveCompareOp::kEqual) {
+                bbb.appendInstr(instr->output(), Instruction::kAnd, le, ge);
+              } else {
+                Instruction* eq = bbb.appendInstr(
+                    OutVReg{Operand::k8bit}, Instruction::kAnd, le, ge);
+                bbb.appendInstr(
+                    instr->output(),
+                    Instruction::kXor,
+                    eq,
+                    Imm{1, DataType::k8bit});
+              }
+              break;
+            }
+            default:
+              JIT_ABORT(
+                  "Not a float comparison {}", static_cast<int>(instr->op()));
+          }
+#elif defined(CINDER_AARCH64)
+          // fcmp leaves Z=0 for unordered operands and sets C=1, V=1.  Pick the
+          // condition that is false for NaN on each ordering: GT/GE for > / >=
+          // (signed opcodes), LO/LS for < / <= (unsigned opcodes).  EQ/NE key
+          // off Z alone and are already NaN-correct.  Floats are neither signed
+          // nor unsigned, so the signed and unsigned PrimitiveCompareOp
+          // variants denote the same ordering (Static Python `double` emits the
+          // unsigned variants; the Python-float path emits the natural ones).
+          Instruction::Opcode op;
+          switch (instr->op()) {
+            case PrimitiveCompareOp::kEqual:
+              op = Instruction::kEqual;
+              break;
+            case PrimitiveCompareOp::kNotEqual:
+              op = Instruction::kNotEqual;
+              break;
+            case PrimitiveCompareOp::kGreaterThan:
+            case PrimitiveCompareOp::kGreaterThanUnsigned:
+              op = Instruction::kGreaterThanSigned;
+              break;
+            case PrimitiveCompareOp::kGreaterThanEqual:
+            case PrimitiveCompareOp::kGreaterThanEqualUnsigned:
+              op = Instruction::kGreaterThanEqualSigned;
+              break;
+            case PrimitiveCompareOp::kLessThan:
+            case PrimitiveCompareOp::kLessThanUnsigned:
+              op = Instruction::kLessThanUnsigned;
+              break;
+            case PrimitiveCompareOp::kLessThanEqual:
+            case PrimitiveCompareOp::kLessThanEqualUnsigned:
+              op = Instruction::kLessThanEqualUnsigned;
+              break;
+            default:
+              JIT_ABORT(
+                  "Not a float comparison {}", static_cast<int>(instr->op()));
+          }
+          bbb.appendInstr(instr->output(), op, lhs, rhs);
+#else
+          CINDER_UNSUPPORTED
+#endif
+          break;
+        }
         Instruction::Opcode op;
         switch (instr->op()) {
           case PrimitiveCompareOp::kEqual:
@@ -3341,17 +3461,6 @@ LIRGenerator::TranslatedBlock LIRGenerator::translateOneBasicBlock(
             op);
         break;
       }
-      case Opcode::kFloatCompare: {
-        auto instr = static_cast<const FloatCompare*>(&i);
-
-        bbb.appendCallInstruction(
-            instr->output(),
-            PyFloat_Type.tp_richcompare,
-            instr->left(),
-            instr->right(),
-            static_cast<int>(instr->op()));
-        break;
-      }
       case Opcode::kLongCompare: {
         auto instr = static_cast<const LongCompare*>(&i);