[llvm] [InstCombine] Lower flag check pattern to use a bitmask-shift (PR #169557)

[Craig Topper via llvm-commits]

================
@@ -3119,6 +3119,127 @@ static Instruction *matchFunnelShift(Instruction &Or, InstCombinerImpl &IC) {
   return nullptr;
 }
 
+static Value *combineOrOfImmCmpToBitExtract(Instruction &Or,
+                                            InstCombiner::BuilderTy &Builder,
+                                            const DataLayout &DL) {
+
+  auto isICmpEqImm = [](Value *N, ConstantInt *&Imm, Value *&X) -> bool {
+    if (X)
+      return match(N, m_OneUse(m_SpecificICmp(ICmpInst::ICMP_EQ, m_Specific(X),
+                                              m_ConstantInt(Imm))));
+
+    return match(N, m_OneUse(m_SpecificICmp(ICmpInst::ICMP_EQ, m_Value(X),
+                                            m_ConstantInt(Imm))));
+  };
+
+  // %srl = lshr %bitmap, %X
+  // %icmp = icmp ult %X, %max_value
+  // %trunc = trunc %srl to i1
+  // %sel = select %icmp, %trunc, false
+  auto CreateBitExtractSeq = [&](APInt BitMap, APInt MaxValue,
+                                 Value *X) -> Value * {
+    LLVMContext &Context = Or.getContext();
+
+    // %srl = lshr %bitmap, %X
+    // It is okay for the shift amount to be truncated because
+    // if information is lost then it is garunteed to fail the bounds
+    // check and the shift result will be discarded
+    ConstantInt *BitMapConst = ConstantInt::get(Context, BitMap);
+    Value *ShiftAmt =
+        Builder.CreateZExtOrTrunc(X, BitMapConst->getIntegerType());
+    Value *LShr = Builder.CreateLShr(BitMapConst, ShiftAmt);
+
+    // %icmp = icmp ult %X, %max_value
+    // Use the type that is the larger of 'X' and the bounds integer
+    // so that no information is lost
+    Value *MaxVal = ConstantInt::get(Context, MaxValue);
+    if (MaxVal->getType()->getIntegerBitWidth() >
+        X->getType()->getIntegerBitWidth())
+      X = Builder.CreateZExt(X, MaxVal->getType());
+    else
+      MaxVal = Builder.CreateZExt(MaxVal, X->getType());
+    Value *BoundsCheck = Builder.CreateICmp(ICmpInst::ICMP_ULT, X, MaxVal);
+
+    // %trunc = trunc %srl to i1
+    // Only care about the low bit
+    Value *ShrTrunc = Builder.CreateTrunc(LShr, IntegerType::get(Context, 1));
+
+    // %sel = select %icmp, %trunc, false
+    return Builder.CreateSelect(BoundsCheck, ShrTrunc,
+                                ConstantInt::getFalse(Context));
+  };
+
+  // Our BitMap should be able to fit into a single arch register
+  // otherwise the tranformation won't be profitable
+  unsigned XLen = DL.getLargestLegalIntTypeSizeInBits();
+  auto validImm = [&](APInt APImm) -> bool {
+    auto Imm = APImm.tryZExtValue();
+    return Imm && (*Imm < XLen);
+  };
+
+  // Match (or (icmp eq X, Imm0), (icmp eq X, Imm1))
+  ConstantInt *LHS, *RHS;
+  Value *X = nullptr;
+  if (isICmpEqImm(Or.getOperand(0), LHS, X) &&
+      isICmpEqImm(Or.getOperand(1), RHS, X)) {
+    // The Shr with become poison when shifted by Undef
+    if (!isGuaranteedNotToBeUndefOrPoison(X))
+      return nullptr;
+
+    APInt LHSAP = LHS->getValue();
+    APInt RHSAP = RHS->getValue();
+    if (!validImm(LHSAP) || !validImm(RHSAP))
+      return nullptr;
+    LHSAP = LHSAP.zextOrTrunc(XLen);
+    RHSAP = RHSAP.zextOrTrunc(XLen);
----------------
topperc wrote:

Can we extract `LHSAP` and `RHSAP` to `unsigned` since we know they represent bit positions and are less than `XLen`?

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