[llvm] [InstCombine] Offset both sides of an equality icmp (PR #134086)

[Nikita Popov via llvm-commits]

================
@@ -5808,6 +5808,134 @@ static Instruction *foldICmpPow2Test(ICmpInst &I,
   return nullptr;
 }
 
+/// Find all possible pairs (BinOp, RHS) that BinOp V, RHS can be simplified.
+using OffsetOp = std::pair<Instruction::BinaryOps, Value *>;
+static void collectOffsetOp(Value *V, SmallVectorImpl<OffsetOp> &Offsets,
+                            bool AllowRecursion) {
+  Instruction *Inst = dyn_cast<Instruction>(V);
+  if (!Inst)
+    return;
+
+  switch (Inst->getOpcode()) {
+  case Instruction::Add: {
+    Constant *C;
+    if (match(Inst->getOperand(1), m_ImmConstant(C)) &&
+        !C->containsUndefOrPoisonElement()) {
+      if (Constant *NegC = ConstantExpr::getNeg(C))
+        Offsets.emplace_back(Instruction::Add, NegC);
+    }
+    break;
+  }
+  case Instruction::Xor:
+    if (isGuaranteedNotToBeUndefOrPoison(Inst->getOperand(1)))
+      Offsets.emplace_back(Instruction::Xor, Inst->getOperand(1));
+    if (isGuaranteedNotToBeUndefOrPoison(Inst->getOperand(0)))
+      Offsets.emplace_back(Instruction::Xor, Inst->getOperand(0));
+    break;
+  case Instruction::Select:
+    if (AllowRecursion) {
+      Value *TrueV = Inst->getOperand(1);
+      if (TrueV->hasOneUse())
+        collectOffsetOp(TrueV, Offsets, /*AllowRecursion=*/false);
+      Value *FalseV = Inst->getOperand(2);
+      if (FalseV->hasOneUse())
+        collectOffsetOp(FalseV, Offsets, /*AllowRecursion=*/false);
+    }
+    break;
+  default:
+    break;
+  }
+}
+
+enum class OffsetKind { Invalid, Value, Select };
+
+struct OffsetResult {
+  OffsetKind Kind;
+  Value *V0, *V1, *V2;
+
+  static OffsetResult invalid() {
+    return {OffsetKind::Invalid, nullptr, nullptr, nullptr};
+  }
+  static OffsetResult value(Value *V) {
+    return {OffsetKind::Value, V, nullptr, nullptr};
+  }
+  static OffsetResult select(Value *Cond, Value *TrueV, Value *FalseV) {
+    return {OffsetKind::Select, Cond, TrueV, FalseV};
+  }
+  bool isValid() const { return Kind != OffsetKind::Invalid; }
+  Value *materialize(InstCombiner::BuilderTy &Builder) const {
+    switch (Kind) {
+    case OffsetKind::Invalid:
+      llvm_unreachable("Invalid offset result");
+    case OffsetKind::Value:
+      return V0;
+    case OffsetKind::Select:
+      return Builder.CreateSelect(V0, V1, V2);
+    }
+  }
+};
+
+/// Offset both sides of an equality icmp to see if we can save some
+/// instructions: icmp eq/ne X, Y -> icmp eq/ne X op Z, Y op Z.
+/// Note: This operation should not introduce poison.
+static Instruction *foldICmpEqualityWithOffset(ICmpInst &I,
+                                               InstCombiner::BuilderTy &Builder,
+                                               const SimplifyQuery &SQ) {
+  assert(I.isEquality() && "Expected an equality icmp");
+  Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
+  if (!Op0->getType()->isIntOrIntVectorTy())
+    return nullptr;
+
+  SmallVector<OffsetOp, 4> OffsetOps;
+  if (Op0->hasOneUse())
----------------
nikic wrote:

Move one-use check into the function?

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