[llvm] [X86] Extend `combinei64TruncSrlAdd` to handle patterns with `or` and `xor` (PR #128435)

[Phoebe Wang via llvm-commits]

=?utf-8?q?João?= Gouveia <jtalonegouveia at gmail.com>,
=?utf-8?q?João?= Gouveia <jtalonegouveia at gmail.com>,
=?utf-8?q?João?= Gouveia <jtalonegouveia at gmail.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/128435 at github.com>


================
@@ -53733,36 +53733,53 @@ static SDValue combineLRINT_LLRINT(SDNode *N, SelectionDAG &DAG,
   return DAG.getNode(X86ISD::CVTP2SI, DL, VT, Src);
 }
 
-// Attempt to fold some (truncate (srl (add X, C1), C2)) patterns to
-// (add (truncate (srl X, C2)), C1'). C1' will be smaller than C1 so we are able
-// to avoid generating code with MOVABS and large constants in certain cases.
-static SDValue combinei64TruncSrlAdd(SDValue N, EVT VT, SelectionDAG &DAG,
-                                     const SDLoc &DL) {
-  using namespace llvm::SDPatternMatch;
-
-  SDValue AddLhs;
-  APInt AddConst, SrlConst;
-  if (VT != MVT::i32 ||
-      !sd_match(N, m_AllOf(m_SpecificVT(MVT::i64),
-                           m_Srl(m_OneUse(m_Add(m_Value(AddLhs),
-                                                m_ConstInt(AddConst))),
-                                 m_ConstInt(SrlConst)))))
-    return SDValue();
+// Attempt to fold some (truncate (srl (binop X, C1), C2)) patterns to
+// (binop (truncate (srl X, C2)), C1'). C1' will be smaller than C1 so we are
+// able to avoid generating code with MOVABS and large constants in certain
+// cases.
+static SDValue combinei64TruncSrlBinop(SDValue N, EVT VT, SelectionDAG &DAG,
+                                       const SDLoc &DL) {
+
+  SDValue Binop = N.getOperand(0);
+  APInt BinopConst = Binop.getConstantOperandAPInt(1);
+  APInt SrlConst = N.getConstantOperandAPInt(1);
+  unsigned Opcode = Binop.getOpcode();
+
+  auto CleanUpFn = +[](SDValue Op, EVT CleanUpVT, EVT VT, SelectionDAG &DAG,
+                       const SDLoc &DL) {
+    SDValue CleanUp = DAG.getAnyExtOrTrunc(Op, DL, CleanUpVT);
+    return DAG.getAnyExtOrTrunc(CleanUp, DL, VT);
+  };
+  auto ZeroExtCleanUp = +[](SDValue Op, EVT CleanUpVT, EVT VT,
+                            SelectionDAG &DAG, const SDLoc &DL) {
+    return DAG.getZeroExtendInReg(Op, DL, CleanUpVT);
+  };
 
-  if (SrlConst.ule(32) || AddConst.countr_zero() < SrlConst.getZExtValue())
+  switch (Opcode) {
+  default:
     return SDValue();
+  case ISD::ADD:
+    if (BinopConst.countr_zero() < SrlConst.getZExtValue())
+      return SDValue();
+    CleanUpFn = ZeroExtCleanUp;
+    [[fallthrough]];
+  case ISD::OR:
+  case ISD::XOR:
+    if (SrlConst.ule(32))
+      return SDValue();
+    break;
+  }
 
-  SDValue AddLHSSrl =
-      DAG.getNode(ISD::SRL, DL, MVT::i64, AddLhs, N.getOperand(1));
-  SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, AddLHSSrl);
-
-  APInt NewAddConstVal = AddConst.lshr(SrlConst).trunc(VT.getSizeInBits());
-  SDValue NewAddConst = DAG.getConstant(NewAddConstVal, DL, VT);
-  SDValue NewAddNode = DAG.getNode(ISD::ADD, DL, VT, Trunc, NewAddConst);
+  SDValue BinopLHSSrl =
+      DAG.getNode(ISD::SRL, DL, MVT::i64, Binop.getOperand(0), N.getOperand(1));
+  SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, BinopLHSSrl);
 
+  APInt NewBinopConstVal = BinopConst.lshr(SrlConst).trunc(VT.getSizeInBits());
+  SDValue NewBinopConst = DAG.getConstant(NewBinopConstVal, DL, VT);
+  SDValue NewBinopNode = DAG.getNode(Opcode, DL, VT, Trunc, NewBinopConst);
----------------
phoebewang wrote:

I think we can remove Binop, it's a bit misleading that works for all binary operators.

https://github.com/llvm/llvm-project/pull/128435