[ExpandMemCmp] Improve memcmp optimisation for boolean results (#71221)

This patch enhances the optimization of memcmp calls when only two
outcomes
are needed and comparison fits into one block, for example:

	bool result = memcmp(a, b, 6) > 0;

Previously, LLVM would generate unnecessary operations even when the
user of
memcmp was only interested in a binary outcome.

llvm/lib/CodeGen/ExpandMemCmp.cpp

@@ -23,6 +23,7 @@
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
@@ -31,6 +32,7 @@
 #include <optional>
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 namespace llvm {
 class TargetLowering;
@@ -636,6 +638,9 @@ Value *MemCmpExpansion::getMemCmpEqZeroOneBlock() {
 
 /// A memcmp expansion that only has one block of load and compare can bypass
 /// the compare, branch, and phi IR that is required in the general case.
+/// This function also analyses users of memcmp, and if there is only one user
+/// from which we can conclude that only 2 out of 3 memcmp outcomes really
+/// matter, then it generates more efficient code with only one comparison.
 Value *MemCmpExpansion::getMemCmpOneBlock() {
   bool NeedsBSwap = DL.isLittleEndian() && Size != 1;
   Type *LoadSizeType = IntegerType::get(CI->getContext(), Size * 8);
@@ -656,6 +661,40 @@ Value *MemCmpExpansion::getMemCmpOneBlock() {
 
   const LoadPair Loads = getLoadPair(LoadSizeType, BSwapSizeType, MaxLoadType,
                                      /*Offset*/ 0);
+
+  // If a user of memcmp cares only about two outcomes, for example:
+  //    bool result = memcmp(a, b, NBYTES) > 0;
+  // We can generate more optimal code with a smaller number of operations
+  if (CI->hasOneUser()) {
+    auto *UI = cast<Instruction>(*CI->user_begin());
+    ICmpInst::Predicate Pred = ICmpInst::Predicate::BAD_ICMP_PREDICATE;
+    uint64_t Shift;
+    bool NeedsZExt = false;
+    // This is a special case because instead of checking if the result is less
+    // than zero:
+    //    bool result = memcmp(a, b, NBYTES) < 0;
+    // Compiler is clever enough to generate the following code:
+    //    bool result = memcmp(a, b, NBYTES) >> 31;
+    if (match(UI, m_LShr(m_Value(), m_ConstantInt(Shift))) &&
+        Shift == (CI->getType()->getIntegerBitWidth() - 1)) {
+      Pred = ICmpInst::ICMP_SLT;
+      NeedsZExt = true;
+    } else {
+      // In case of a successful match this call will set `Pred` variable
+      match(UI, m_ICmp(Pred, m_Specific(CI), m_Zero()));
+    }
+    // Generate new code and remove the original memcmp call and the user
+    if (ICmpInst::isSigned(Pred)) {
+      Value *Cmp = Builder.CreateICmp(CmpInst::getUnsignedPredicate(Pred),
+                                      Loads.Lhs, Loads.Rhs);
+      auto *Result = NeedsZExt ? Builder.CreateZExt(Cmp, UI->getType()) : Cmp;
+      UI->replaceAllUsesWith(Result);
+      UI->eraseFromParent();
+      CI->eraseFromParent();
+      return nullptr;
+    }
+  }
+
   // The result of memcmp is negative, zero, or positive, so produce that by
   // subtracting 2 extended compare bits: sub (ugt, ult).
   // If a target prefers to use selects to get -1/0/1, they should be able
@@ -670,7 +709,7 @@ Value *MemCmpExpansion::getMemCmpOneBlock() {
 }
 
 // This function expands the memcmp call into an inline expansion and returns
-// the memcmp result.
+// the memcmp result. Returns nullptr if the memcmp is already replaced.
 Value *MemCmpExpansion::getMemCmpExpansion() {
   // Create the basic block framework for a multi-block expansion.
   if (getNumBlocks() != 1) {
@@ -838,11 +877,11 @@ static bool expandMemCmp(CallInst *CI, const TargetTransformInfo *TTI,
 
   NumMemCmpInlined++;
 
-  Value *Res = Expansion.getMemCmpExpansion();
-
-  // Replace call with result of expansion and erase call.
-  CI->replaceAllUsesWith(Res);
-  CI->eraseFromParent();
+  if (Value *Res = Expansion.getMemCmpExpansion()) {
+    // Replace call with result of expansion and erase call.
+    CI->replaceAllUsesWith(Res);
+    CI->eraseFromParent();
+  }
 
   return true;
 }