[llvm] [InstCombine] Improve `(icmp pred (and X, Y), ...)` fold. (PR #66787)

[Simon Pilgrim via llvm-commits]

================
@@ -230,49 +230,114 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner {
                                                 PatternMatch::m_Value()));
   }
 
-  /// Return true if the specified value is free to invert (apply ~ to).
-  /// This happens in cases where the ~ can be eliminated.  If WillInvertAllUses
-  /// is true, work under the assumption that the caller intends to remove all
-  /// uses of V and only keep uses of ~V.
-  ///
-  /// See also: canFreelyInvertAllUsersOf()
-  static bool isFreeToInvert(Value *V, bool WillInvertAllUses) {
+  /// Return nonnull value if V is free to invert (with condition) regarding
+  /// WillInvertAllUses.
+  /// If Builder is nonnull, it will return a simplified ~V
+  /// If builder is null, it will return an arbitrary nonnull value (not
+  /// dereferenceable).
+  static Value *getFreeInverted(Value *V, bool WillInvertAllUses,
+                                BuilderTy *Builder, unsigned Depth = 0) {
+    static Value *const kNonNull = reinterpret_cast<Value *>(uintptr_t(1));
     // ~(~(X)) -> X.
-    if (match(V, m_Not(PatternMatch::m_Value())))
-      return true;
+    Value *A, *B;
+    if (match(V, m_Not(PatternMatch::m_Value(A))))
+      return A;
 
+    Constant *C;
     // Constants can be considered to be not'ed values.
-    if (match(V, PatternMatch::m_AnyIntegralConstant()))
-      return true;
+    if (match(V, PatternMatch::m_ImmConstant(C)))
+      return ConstantExpr::getNot(C);
+
+    if (Depth++ >= MaxAnalysisRecursionDepth)
+      return nullptr;
+
+    // The rest of the cases require that we invert all uses so don't bother
+    // doing the analysis if we know we can't use the result.
+    if (!WillInvertAllUses)
+      return nullptr;
 
     // Compares can be inverted if all of their uses are being modified to use
     // the ~V.
-    if (isa<CmpInst>(V))
-      return WillInvertAllUses;
+    if (auto *I = dyn_cast<CmpInst>(V)) {
+      if (Builder != nullptr)
+        return Builder->CreateCmp(I->getInversePredicate(), I->getOperand(0),
+                                  I->getOperand(1));
+      return kNonNull;
+    }
 
     // If `V` is of the form `A + Constant` then `-1 - V` can be folded into
     // `(-1 - Constant) - A` if we are willing to invert all of the uses.
-    if (match(V, m_Add(PatternMatch::m_Value(), PatternMatch::m_ImmConstant())))
-      return WillInvertAllUses;
+    if (match(V, m_Add(PatternMatch::m_Value(A), PatternMatch::m_Value(B)))) {
+      if (auto *BV = getFreeInverted(B, B->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateSub(BV, A) : kNonNull;
+      if (auto *AV = getFreeInverted(A, A->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateSub(AV, B) : kNonNull;
+      return nullptr;
+    }
+
+    // If `V` is of the form `A ^ ~B` then `~(A ^ ~B)` can be folded
+    // into `A ^ B` if we are willing to invert all of the uses.
+    if (match(V, m_Xor(PatternMatch::m_Value(A), PatternMatch::m_Value(B)))) {
+      if (auto *BV = getFreeInverted(B, B->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateXor(A, BV) : kNonNull;
+      if (auto *AV = getFreeInverted(A, A->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateXor(AV, B) : kNonNull;
+      return nullptr;
+    }
 
     // If `V` is of the form `Constant - A` then `-1 - V` can be folded into
     // `A + (-1 - Constant)` if we are willing to invert all of the uses.
-    if (match(V, m_Sub(PatternMatch::m_ImmConstant(), PatternMatch::m_Value())))
-      return WillInvertAllUses;
-
-    // Selects with invertible operands are freely invertible
-    if (match(V,
-              m_Select(PatternMatch::m_Value(), m_Not(PatternMatch::m_Value()),
-                       m_Not(PatternMatch::m_Value()))))
-      return WillInvertAllUses;
-
-    // Min/max may be in the form of intrinsics, so handle those identically
-    // to select patterns.
-    if (match(V, m_MaxOrMin(m_Not(PatternMatch::m_Value()),
-                            m_Not(PatternMatch::m_Value()))))
-      return WillInvertAllUses;
-
-    return false;
+    if (match(V, m_Sub(PatternMatch::m_Value(A), PatternMatch::m_Value(B)))) {
+      if (auto *AV = getFreeInverted(A, A->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateAdd(AV, B) : kNonNull;
+      return nullptr;
+    }
+
+    // If `V` is of the form `Constant - A` then `-1 - V` can be folded into
+    // `A + (-1 - Constant)` if we are willing to invert all of the uses.
+    if (match(V, m_AShr(PatternMatch::m_Value(A), PatternMatch::m_Value(B)))) {
+      if (auto *AV = getFreeInverted(A, A->hasOneUse(), Builder, Depth))
+        return Builder ? Builder->CreateAShr(AV, B) : kNonNull;
+      return nullptr;
+    }
+
+    Value *Cond;
+    // LogicOps are special in that we canonicalize them at the cost of an
+    // instruction.
+    bool IsSelect =
+        match(V, m_Select(PatternMatch::m_Value(Cond), PatternMatch::m_Value(A),
+                          PatternMatch::m_Value(B))) &&
+        !match(V, PatternMatch::m_LogicalOp(PatternMatch::m_Value(),
+                                            PatternMatch::m_Value()));
+    // Selects/min/max with invertible operands are freely invertible
+    if (IsSelect || match(V, m_MaxOrMin(PatternMatch::m_Value(A),
+                                        PatternMatch::m_Value(B)))) {
+      if (!getFreeInverted(A, A->hasOneUse(), /*Builder*/ nullptr, Depth) ||
+          !getFreeInverted(B, B->hasOneUse(), /*Builder*/ nullptr, Depth))
+        return nullptr;
+      if (Builder != nullptr) {
+        A = getFreeInverted(A, A->hasOneUse(), Builder, Depth);
+        B = getFreeInverted(B, B->hasOneUse(), Builder, Depth);
+        assert(A != nullptr && B != nullptr);
----------------
RKSimon wrote:

(style) missing assert message

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