Recognize bit-manipulation idioms and emit bzhi/bts/btc/btr/bt on xarch

src/coreclr/jit/codegen.h

@@ -807,6 +807,7 @@ class CodeGen final : public CodeGenInterface
 
     void genCodeForDivMod(GenTreeOp* treeNode);
     void genCodeForMul(GenTreeOp* treeNode);
+    void genCodeForBitOp(GenTreeOp* treeNode);
     void genCodeForIncSaturate(GenTree* treeNode);
     void genCodeForMulHi(GenTreeOp* treeNode);
     void genLeaInstruction(GenTreeAddrMode* lea);

src/coreclr/jit/codegenxarch.cpp

@@ -1129,6 +1129,51 @@ void CodeGen::genCodeForBinary(GenTreeOp* treeNode)
     genProduceReg(treeNode);
 }
 
+//------------------------------------------------------------------------
+// genCodeForBitOp: Generate code for a GT_BIT_SET/GT_BIT_CLEAR/GT_BIT_INVERT operation, i.e. the
+// value-producing `bts`/`btr`/`btc` instructions which set, reset, or complement a single bit of op1
+// selected by op2.
+//
+// Arguments:
+//    treeNode - the node to generate the code for
+//
+void CodeGen::genCodeForBitOp(GenTreeOp* treeNode)
+{
+    assert(treeNode->OperIs(GT_BIT_SET, GT_BIT_CLEAR, GT_BIT_INVERT));
+
+    GenTree* op1 = treeNode->gtGetOp1(); // value (read-modify-write destination)
+    GenTree* op2 = treeNode->gtGetOp2(); // bit index
+
+    genConsumeOperands(treeNode);
+
+    regNumber targetReg  = treeNode->GetRegNum();
+    var_types targetType = genActualType(treeNode);
+    emitAttr  size       = emitTypeSize(targetType);
+    emitter*  emit       = GetEmitter();
+
+    assert((targetType == TYP_INT) || (targetType == TYP_LONG));
+    assert(op1->isUsedFromReg() && op2->isUsedFromReg());
+
+    instruction ins = treeNode->OperIs(GT_BIT_SET) ? INS_bts : treeNode->OperIs(GT_BIT_CLEAR) ? INS_btr : INS_btc;
+
+    // These are read-modify-write: the `mov` below loads op1 (the value) into the destination and
+    // then `bts`/`btr`/`btc` reads the bit index from op2. LSRA marks op2 as delayFree except when
+    // op2 shares op1's interval and it's their last use -- i.e. `x <op> (1 << x)`, where op1 and op2
+    // are the same value (see AddDelayFreeUses). So the destination can only alias op2 when op1 and
+    // op2 hold the same value, in which case the `mov` writes that same value back into op2's
+    // register and nothing is clobbered before the bit-test reads it. When the operands are distinct
+    // values, delayFree guarantees the destination and op2 use different registers.
+
+    // These are read-modify-write: the destination register also supplies the value operand.
+    inst_Mov(targetType, targetReg, op1->GetRegNum(), /* canSkip */ true);
+
+    // The BT-family reg,reg encoding places the destination in the r/m slot and the bit index in
+    // the reg slot, so the operands are passed reversed (see the note in instrsxarch.h).
+    emit->emitIns_R_R(ins, size, op2->GetRegNum(), targetReg);
+
+    genProduceReg(treeNode);
+}
+
 //------------------------------------------------------------------------
 // genCodeForMul: Generate code for a MUL operation.
 //
@@ -1892,6 +1937,12 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genCodeForBinary(treeNode->AsOp());
             break;
 
+        case GT_BIT_SET:
+        case GT_BIT_CLEAR:
+        case GT_BIT_INVERT:
+            genCodeForBitOp(treeNode->AsOp());
+            break;
+
         case GT_MUL:
             if (varTypeIsFloating(treeNode->TypeGet()))
             {

src/coreclr/jit/emitxarch.cpp

@@ -12999,6 +12999,15 @@ void emitter::emitDispIns(
                         break;
                 }
 
+                if ((ins == INS_bts) || (ins == INS_btr) || (ins == INS_btc))
+                {
+                    // The BT-family reg,reg encoding stores its operands reversed. Display them in
+                    // the normal `dest, index` order.
+                    printf("%s", emitRegName(id->idReg2(), tgtAttr));
+                    printf(", %s", emitRegName(id->idReg1(), srcAttr));
+                    break;
+                }
+
                 printf("%s", emitRegName(id->idReg1(), tgtAttr));
                 printf(", %s", emitRegName(id->idReg2(), srcAttr));
                 break;

src/coreclr/jit/gtlist.h

@@ -298,14 +298,16 @@ GTNODE(SH3ADD_UW        , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
 GTNODE(ADD_UW           , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
 // Maps to riscv64 slli.uw instruction. Computes result = zext(op1[31..0]) << imm.
 GTNODE(SLLI_UW          , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
+#endif // TARGET_RISCV64
 
-// Maps to bset/bseti instruction. Computes result = op1 | (1 << op2)
+#if defined(TARGET_RISCV64) || defined(TARGET_XARCH)
+// Maps to riscv64 bset/bseti and xarch bts. Computes result = op1 | (1 << op2)
 GTNODE(BIT_SET          , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
-// Maps to bclr/bclri instruction. Computes result = op1 & ~(1 << op2)
+// Maps to riscv64 bclr/bclri and xarch btr. Computes result = op1 & ~(1 << op2)
 GTNODE(BIT_CLEAR        , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
-// Maps to binv/binvi instruction. Computes result = op1 ^ (1 << op2)
+// Maps to riscv64 binv/binvi and xarch btc. Computes result = op1 ^ (1 << op2)
 GTNODE(BIT_INVERT       , GenTreeOp          ,0,0,GTK_BINOP|DBK_NOTHIR)
-#endif
+#endif // TARGET_RISCV64 || TARGET_XARCH
 
 //-----------------------------------------------------------------------------
 //  Other nodes that look like unary/binary operators:

src/coreclr/jit/instrsxarch.h

@@ -95,6 +95,12 @@ INST4(lea,              "lea",              IUM_WR, BAD_CODE,     BAD_CODE,
 // and the registers need to be reversed to get the correct encoding.
 INST3(bt,               "bt",               IUM_RD, 0x0F00A3,     BAD_CODE,     0x0F00A3,                                1C,                2X,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF | Encoding_REX2)
 
+// BTS/BTR/BTC are only emitted in their reg,reg form (like BT the registers are reversed to get
+// the correct encoding). They read+write the first operand and write the old bit value to CF.
+INST3(bts,              "bts",              IUM_RW, 0x0F00AB,     BAD_CODE,     0x0F00AB,                                1C,                2X,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF | Encoding_REX2)
+INST3(btr,              "btr",              IUM_RW, 0x0F00B3,     BAD_CODE,     0x0F00B3,                                1C,                2X,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF | Encoding_REX2)
+INST3(btc,              "btc",              IUM_RW, 0x0F00BB,     BAD_CODE,     0x0F00BB,                                1C,                2X,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Undefined_ZF  | Undefined_AF  | Undefined_PF  | Writes_CF | Encoding_REX2)
+
 INST3(bsr,              "bsr",              IUM_WR, BAD_CODE,     BAD_CODE,     0x0F00BD,                                3C,                1C,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Writes_ZF     | Undefined_AF  | Undefined_PF  | Undefined_CF | Encoding_REX2)
 INST3(bsf,              "bsf",              IUM_WR, BAD_CODE,     BAD_CODE,     0x0F00BC,                                3C,                1C,         INS_TT_NONE,    Undefined_OF   | Undefined_SF  | Writes_ZF     | Undefined_AF  | Undefined_PF  | Undefined_CF | Encoding_REX2)
 

src/coreclr/jit/lower.cpp

@@ -4237,6 +4237,34 @@ GenTree* Lowering::OptimizeConstCompare(GenTree* cmp)
             test->gtOp2 = bitOp->gtGetOp2();
             return true;
         }
+
+#ifdef TARGET_XARCH
+        // Also recognize the arithmetic form `(x >> y) & 1`, i.e. AND(RSH|RSZ(x, y), 1), which
+        // tests bit `y` of `x` just like `x & (1 << y)`. Only bit 0 of the shifted value is kept so
+        // the shift kind is irrelevant, and `bt` masks the bit index modulo the operand size, which
+        // matches the C# masked-shift semantics even for an out-of-range `y`. Restricted to a
+        // variable index because a constant index keeps the shift, and `bt` has no immediate form
+        // here (a constant mask `test` is already optimal).
+        GenTree* shiftOp = test->gtOp1;
+        GenTree* oneOp   = test->gtOp2;
+        if (!oneOp->IsIntegralConst(1))
+            std::swap(shiftOp, oneOp);
+
+        if (oneOp->IsIntegralConst(1) && shiftOp->OperIs(GT_RSH, GT_RSZ) && varTypeIsIntOrI(shiftOp) &&
+            !shiftOp->gtGetOp2()->IsIntegralConst())
+        {
+            BlockRange().Remove(oneOp);
+            BlockRange().Remove(shiftOp);
+            test->gtOp1 = shiftOp->gtGetOp1();
+            test->gtOp2 = shiftOp->gtGetOp2();
+
+            // ContainCheckCompare is skipped when this transform succeeds, so clear any containment
+            // the value operand picked up from the removed shift (e.g. a `shrx` memory source) --
+            // the reg,reg `bt` form requires it in a register.
+            test->gtOp1->ClearContained();
+            return true;
+        }
+#endif // TARGET_XARCH
         return false;
     };
 
@@ -10034,6 +10062,117 @@ void Lowering::ContainCheckRet(GenTreeUnOp* ret)
 #endif // FEATURE_MULTIREG_RET
 }
 
+#if defined(TARGET_XARCH) || defined(TARGET_RISCV64)
+//------------------------------------------------------------------------
+// TryLowerBitwiseOpToBitOp: Recognizes the single-bit-manipulation idioms with a variable
+//    (non-constant) bit index and rewrites them in place to the shared GT_BIT_* nodes:
+//
+//        OR (X, LSH(1, Y))       -> BIT_SET    (set bit Y of X)
+//        XOR(X, LSH(1, Y))       -> BIT_INVERT (complement bit Y of X)
+//        AND(X, NOT(LSH(1, Y)))  -> BIT_CLEAR  (reset bit Y of X)
+//
+//    where op1 becomes the value (read-modify-write destination) and op2 becomes the bit index Y.
+//    The `1 << Y` sub-tree may appear on either side of the commutative operation.
+//
+// Arguments:
+//    binOp - a GT_OR, GT_XOR, or GT_AND node of TYP_INT or TYP_LONG
+//
+// Return Value:
+//    The rewritten node (== binOp) on success, or nullptr if the pattern did not match.
+//
+// Notes:
+//    Only the variable-index form is handled. A constant index folds to a constant mask that the
+//    plain `or`/`xor`/`and`-with-immediate form already handles optimally.
+//
+//    The bit index is left as-is; callers are responsible for any target-specific masking of the
+//    index. x86's `bts`/`btr`/`btc` reg,reg form masks the index modulo the operand width (matching
+//    the C# masked-shift semantics of `1 << Y`), whereas RISC-V's `Zbs` ops operate on the full
+//    register and need an explicit `& 31` for 32-bit operands.
+GenTree* Lowering::TryLowerBitwiseOpToBitOp(GenTreeOp* binOp)
+{
+    assert(binOp->OperIs(GT_OR, GT_XOR, GT_AND));
+
+    if (!binOp->TypeIs(TYP_INT, TYP_LONG))
+    {
+        return nullptr;
+    }
+
+    GenTree*& op1 = binOp->gtOp1;
+    GenTree*& op2 = binOp->gtOp2;
+
+    bool isOp1Negated = op1->OperIs(GT_NOT);
+    bool isOp2Negated = op2->OperIs(GT_NOT);
+
+    // For AND/`btr` the `1 << Y` must be negated (`~(1 << Y)`); for OR/XOR it must not be.
+    const bool wantNegated = binOp->OperIs(GT_AND);
+    GenTree*   opp1        = isOp1Negated ? op1->AsUnOp()->gtGetOp1() : op1;
+    GenTree*   opp2        = isOp2Negated ? op2->AsUnOp()->gtGetOp1() : op2;
+
+    bool isOp1SingleBit = (isOp1Negated == wantNegated) && opp1->OperIs(GT_LSH) && opp1->gtGetOp1()->IsIntegralConst(1);
+    bool isOp2SingleBit = (isOp2Negated == wantNegated) && opp2->OperIs(GT_LSH) && opp2->gtGetOp1()->IsIntegralConst(1);
+
+    if (!isOp1SingleBit && !isOp2SingleBit)
+    {
+        return nullptr;
+    }
+
+    // Canonicalize so the `1 << Y` sub-tree is op2 and the value is op1.
+    if (isOp1SingleBit)
+    {
+        std::swap(op1, op2);
+        std::swap(isOp1Negated, isOp2Negated);
+    }
+
+    GenTree* notNode = isOp2Negated ? op2 : nullptr;
+    GenTree* lshNode = (notNode != nullptr) ? op2->AsUnOp()->gtGetOp1() : op2;
+
+    // The shifted value must match the width of the operation.
+    if (!lshNode->TypeIs(binOp->TypeGet()))
+    {
+        return nullptr;
+    }
+
+    // A constant index folds to a constant mask, which the plain form already handles optimally.
+    GenTree* indexNode = lshNode->gtGetOp2();
+    if (indexNode->IsIntegralConst())
+    {
+        return nullptr;
+    }
+
+    // Subsequent nodes may rely on CPU flags set by these nodes, in which case we cannot remove them.
+    if (((binOp->gtFlags & GTF_SET_FLAGS) != 0) || ((lshNode->gtFlags & GTF_SET_FLAGS) != 0) ||
+        ((notNode != nullptr) && ((notNode->gtFlags & GTF_SET_FLAGS) != 0)))
+    {
+        return nullptr;
+    }
+
+    static_assert(AreContiguous(GT_OR, GT_XOR, GT_AND), "");
+    constexpr genTreeOps singleBitOpers[] = {GT_BIT_SET, GT_BIT_INVERT, GT_BIT_CLEAR};
+    const genTreeOps     newOper          = singleBitOpers[binOp->OperGet() - GT_OR];
+
+    JITDUMP("Lower: optimize %s(X, %s)\n", GenTree::OpName(binOp->OperGet()),
+            (notNode != nullptr) ? "NOT(LSH(1, Y))" : "LSH(1, Y)");
+    DISPNODE(binOp);
+
+    // Rewrite in place: op1 stays the value, op2 becomes the bit index. Drop the `1`, the shift, and
+    // the optional NOT.
+    if (notNode != nullptr)
+    {
+        BlockRange().Remove(notNode);
+    }
+    BlockRange().Remove(lshNode->gtGetOp1());
+    BlockRange().Remove(lshNode);
+
+    op2 = indexNode;
+    binOp->ChangeOper(newOper);
+
+    JITDUMP("to:\n");
+    DISPNODE(binOp);
+
+    return binOp;
+}
+#endif // TARGET_XARCH || TARGET_RISCV64
+
 //------------------------------------------------------------------------
 // TryRemoveCast:
 //   Try to remove a cast node by changing its operand.

src/coreclr/jit/lower.h

@@ -453,6 +453,10 @@ class Lowering final : public Phase
     bool TryRemoveCast(GenTreeCast* node);
     bool TryRemoveBitCast(GenTreeUnOp* node);
 
+#if defined(TARGET_XARCH) || defined(TARGET_RISCV64)
+    GenTree* TryLowerBitwiseOpToBitOp(GenTreeOp* binOp);
+#endif // TARGET_XARCH || TARGET_RISCV64
+
 #ifdef TARGET_XARCH
     GenTree* TryLowerMulWithConstant(GenTreeOp* node);
 #endif // TARGET_XARCH
@@ -510,6 +514,7 @@ class Lowering final : public Phase
     GenTree* TryLowerAndOpToResetLowestSetBit(GenTreeOp* andNode);
     GenTree* TryLowerAndOpToExtractLowestSetBit(GenTreeOp* andNode);
     GenTree* TryLowerAndOpToAndNot(GenTreeOp* andNode);
+    GenTree* TryLowerAndOpToZeroHighBits(GenTreeOp* andNode);
     GenTree* TryLowerXorOpToGetMaskUpToLowestSetBit(GenTreeOp* xorNode);
     void     LowerBswapOp(GenTreeOp* node);
     GenTree* LowerHWIntrinsicDotInnerMulSum(GenTreeHWIntrinsic* node);

src/coreclr/jit/lowerriscv64.cpp

@@ -329,44 +329,13 @@ GenTree* Lowering::LowerBinaryArithmetic(GenTreeOp* binOp)
         }
         else // op2 is not constant
         {
-            GenTree* opp1 = isOp1Negated ? op1->gtGetOp1() : op1;
-            GenTree* opp2 = isOp2Negated ? op2->gtGetOp1() : op2;
-
-            bool isOp1SingleBit =
-                (isOp1Negated == binOp->OperIs(GT_AND)) && opp1->OperIs(GT_LSH) && opp1->gtGetOp1()->IsIntegralConst(1);
-            bool isOp2SingleBit =
-                (isOp2Negated == binOp->OperIs(GT_AND)) && opp2->OperIs(GT_LSH) && opp2->gtGetOp1()->IsIntegralConst(1);
-
-            if (isOp1SingleBit || isOp2SingleBit)
+            // a | (1 << b),  a ^ (1 << b),  a & ~(1 << b)   =>   BIT_{SET,INVERT,CLEAR}(a, b)
+            if (TryLowerBitwiseOpToBitOp(binOp) != nullptr)
             {
-                // a | (1 << b),  a ^ (1 << b),  a & ~(1 << b)   =>   BIT_{SET,INVERT,CLEAR}(a, b)
-
-                if (isOp1SingleBit)
-                    std::swap(op1, op2);
-
-                if (binOp->OperIs(GT_AND))
-                {
-                    assert(op2->OperIs(GT_NOT));
-                    BlockRange().Remove(op2);
-                    op2 = op2->gtGetOp1();
-                }
-
-                assert(binOp->OperIs(GT_OR, GT_XOR, GT_AND));
-                static_assert(AreContiguous(GT_OR, GT_XOR, GT_AND), "");
-                constexpr genTreeOps singleBitOpers[] = {GT_BIT_SET, GT_BIT_INVERT, GT_BIT_CLEAR};
-                binOp->ChangeOper(singleBitOpers[binOp->OperGet() - GT_OR]);
-
-                assert(op2->OperIs(GT_LSH));
-                assert(op2->gtGetOp1()->IsIntegralConst(1));
-                assert(!op2->gtGetOp2()->IsIntegralConst());
-                BlockRange().Remove(op2->gtGetOp1());
-                BlockRange().Remove(op2);
-                op2 = op2->gtGetOp2(); // shift amount becomes bit index
-
-                assert(op1->TypeIs(TYP_INT, TYP_LONG));
-                if (!op2->IsIntegralConst() && op1->TypeIs(TYP_INT))
+                // binOp is now BIT_{SET,INVERT,CLEAR} with op2 as the (variable) bit index.
+                if (op1->TypeIs(TYP_INT))
                 {
-                    // Zbs instructions don't have *w variants so wrap the bit index / shift amount to 0-31 manually
+                    // Zbs instructions don't have *w variants so wrap the bit index to 0-31 manually.
                     GenTreeIntCon* mask = m_compiler->gtNewIconNode(0x1F);
                     mask->SetContained();
                     BlockRange().InsertAfter(op2, mask);