[CodeGenPrepare][RISCV] Reverse transform in CGP to use zero-compare branch

[dtcxzyw]

Summary

This patch handles signed icmp cases of Add and Sub based on D101778 if it guarantees no overflow after hoisting. For example:
https://alive2.llvm.org/ce/z/wbCO7s

define i32 @src(i32 %0) {
  %x = icmp sge i32 %0, 0
  br i1 %x, label %nonnegative, label %5

nonnegative:                                      ; preds = %1
  %2 = icmp slt i32 %0, 13
  br i1 %2, label %5, label %3

3:                                                ; preds = %nonnegative
  %4 = add i32 %0, -13
  ret i32 %4

5:                                                ; preds = %nonnegative, %1
  ret i32 0
}

define i32 @tgt(i32 %0) {
  %x = icmp sge i32 %0, 0
  br i1 %x, label %nonnegative, label %5

nonnegative:                                      ; preds = %1
  %2 = add i32 %0, -13
  %3 = icmp slt i32 %2, 0
  br i1 %3, label %5, label %4

4:                                                ; preds = %nonnegative
  ret i32 %2

5:                                                ; preds = %nonnegative, %1
  ret i32 0
}

This patch addresses performance regressions in RISC-V benchmarks reported by our CI. I think this patch can protect against regressions from patch D147568.

[nikic]

Can you base this on nowrap flags instead of re-deriving them?

[dtcxzyw]

In D147789#4251449, @nikic wrote:

Can you base this on nowrap flags instead of re-deriving them?

We cannot deduce that Add/Sub instructions with NSW flags do not overflow after hoisting.
For example:
https://alive2.llvm.org/ce/z/GVZbvc

define i32 @src(i32 %0) {
  %2 = add nsw i32 %0, -13
  %3 = icmp slt i32 %2, 0
  br i1 %3, label %5, label %4

4:                                                
  ret i32 %2

5:                                               
  ret i32 0
}

define i32 @tgt(i32 %0) {
  %2 = icmp slt i32 %0, 13
  br i1 %2, label %5, label %3

3:                                               
  %4 = add nsw i32 %0, -13
  ret i32 %4

5:                                                
  ret i32 0
}

This transformation in InstCombine is safe but loses the information that add i32 %0, -13 does not overflow before the branch.
Thus, the reverse transformation in CodeGenPrepare depending on the NSW flag will fail:
https://alive2.llvm.org/ce/z/woAW5v

Therefore the re-derivation is necessary as a "mitigation" for the loss of information.

llvm/lib/CodeGen/CodeGenPrepare.cpp

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
class CodeGenPrepare : public FunctionPass {
   const TargetMachine *TM = nullptr;
   const TargetSubtargetInfo *SubtargetInfo;
   const TargetLowering *TLI = nullptr;
   const TargetRegisterInfo *TRI;
   const TargetTransformInfo *TTI = nullptr;
   const BasicBlockSectionsProfileReader *BBSectionsProfileReader = nullptr;
   const TargetLibraryInfo *TLInfo;
   const LoopInfo *LI;
+  AssumptionCache *AC;
   std::unique_ptr<BlockFrequencyInfo> BFI;
   std::unique_ptr<BranchProbabilityInfo> BPI;
   ProfileSummaryInfo *PSI;
 
   /// As we scan instructions optimizing them, this is the next instruction
   /// to optimize. Transforms that can invalidate this should update it.
   BasicBlock::iterator CurInstIterator;
 
public:
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     // FIXME: When we can selectively preserve passes, preserve the domtree.
     AU.addRequired<ProfileSummaryInfoWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addRequired<TargetPassConfig>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
+    AU.addRequired<AssumptionCacheTracker>();
     AU.addUsedIfAvailable<BasicBlockSectionsProfileReader>();
   }
 
 private:
   template <typename F>
   void resetIteratorIfInvalidatedWhileCalling(BasicBlock *BB, F f) {
     // Substituting can cause recursive simplifications, which can invalidate
     // our iterator.  Use a WeakTrackingVH to hold onto it in case this

 } // end anonymous namespace
 
 char CodeGenPrepare::ID = 0;
 
 INITIALIZE_PASS_BEGIN(CodeGenPrepare, DEBUG_TYPE,
                       "Optimize for code generation", false, false)
 INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReader)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(CodeGenPrepare, DEBUG_TYPE, "Optimize for code generation",
                     false, false)
 
 FunctionPass *llvm::createCodeGenPreparePass() { return new CodeGenPrepare(); }
bool CodeGenPrepare::runOnFunction(Function &F) {
 
   TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
   SubtargetInfo = TM->getSubtargetImpl(F);
   TLI = SubtargetInfo->getTargetLowering();
   TRI = SubtargetInfo->getRegisterInfo();
   TLInfo = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
   LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
+  AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   BPI.reset(new BranchProbabilityInfo(F, *LI));
   BFI.reset(new BlockFrequencyInfo(F, *BPI, *LI));
   PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   BBSectionsProfileReader =
       getAnalysisIfAvailable<BasicBlockSectionsProfileReader>();
   OptSize = F.hasOptSize();
   // Use the basic-block-sections profile to promote hot functions to .text.hot
   // if requested.
assert(llvm::none_of(GEPIOp->users(),
                        [&](User *Usr) {
                          return cast<Instruction>(Usr)->getParent() != SrcBlock;
                        }) &&
          "GEPIOp is used outside SrcBlock");
   return true;
 }
 
 static bool optimizeBranch(BranchInst *Branch, const TargetLowering &TLI,
+                           const DataLayout &DL, AssumptionCache *AC,
+                           const DominatorTree *DT,
                            SmallSet<BasicBlock *, 32> &FreshBBs,
                            bool IsHugeFunc) {
   // Try and convert
   //  %c = icmp ult %x, 8
   //  br %c, bla, blb
   //  %tc = lshr %x, 3
   // to
   //  %tc = lshr %x, 3
   //  %c = icmp eq %tc, 0
   //  br %c, bla, blb
   // Creating the cmp to zero can be better for the backend, especially if the
   // lshr produces flags that can be used automatically.
   if (!TLI.preferZeroCompareBranch() || !Branch->isConditional())
     return false;
 
   ICmpInst *Cmp = dyn_cast<ICmpInst>(Branch->getCondition());
   if (!Cmp || !isa<ConstantInt>(Cmp->getOperand(1)) || !Cmp->hasOneUse())
     return false;
 
   Value *X = Cmp->getOperand(0);
-  APInt CmpC = cast<ConstantInt>(Cmp->getOperand(1))->getValue();
+  Value *C = Cmp->getOperand(1);
+  APInt CmpC = cast<ConstantInt>(C)->getValue();
+  if (CmpC.isZero())
+    return false;
 
   for (auto *U : X->users()) {
     Instruction *UI = dyn_cast<Instruction>(U);
     // A quick dominance check
     if (!UI ||
         (UI->getParent() != Branch->getParent() &&
          UI->getParent() != Branch->getSuccessor(0) &&
          UI->getParent() != Branch->getSuccessor(1)) ||
         (UI->getParent() != Branch->getParent() &&
          !UI->getParent()->getSinglePredecessor()))
       continue;
 
     if (CmpC.isPowerOf2() && Cmp->getPredicate() == ICmpInst::ICMP_ULT &&
         match(UI, m_Shr(m_Specific(X), m_SpecificInt(CmpC.logBase2())))) {
       IRBuilder<> Builder(Branch);
       if (UI->getParent() != Branch->getParent())
         UI->moveBefore(Branch);
       Value *NewCmp = Builder.CreateCmp(ICmpInst::ICMP_EQ, UI,
                                         ConstantInt::get(UI->getType(), 0));
       LLVM_DEBUG(dbgs() << "Converting " << *Cmp << "\n");
       LLVM_DEBUG(dbgs() << " to compare on zero: " << *NewCmp << "\n");
       replaceAllUsesWith(Cmp, NewCmp, FreshBBs, IsHugeFunc);
       return true;
     }
-    if (Cmp->isEquality() &&
+    if ((Cmp->isEquality() || Cmp->isSigned()) &&
         (match(UI, m_Add(m_Specific(X), m_SpecificInt(-CmpC))) ||
          match(UI, m_Sub(m_Specific(X), m_SpecificInt(CmpC))))) {
+      if (Cmp->isRelational()) {
+        // A cheap check
+        if (computeOverflowForSignedSub(X, C, DL, AC, Branch, DT) !=
+            OverflowResult::NeverOverflows) {
+          // x = X, c = CmpC
+          //               c < 0             c > 0
+          // add/sub   x - c <= INT_MAX  x - c >= INT_MIN
+          bool IsUpperBound = CmpC.slt(0);
+          APInt Bound =
+              (IsUpperBound ? APInt::getSignedMaxValue(CmpC.getBitWidth())
+                            : APInt::getSignedMinValue(CmpC.getBitWidth())) +
+              CmpC;
+          std::optional<bool> NoOverflow = isImpliedByDomCondition(
+              IsUpperBound ? ICmpInst::ICMP_SLE : ICmpInst::ICMP_SGE, X,
+              ConstantInt::get(UI->getType(), Bound), Branch, DL);
+          if (!NoOverflow || !NoOverflow.value())
+            continue;
+        }
+      }
+
       IRBuilder<> Builder(Branch);
-      if (UI->getParent() != Branch->getParent())
+      if (UI->getParent() != Branch->getParent()) {
         UI->moveBefore(Branch);
+        // clear nsw/nuw flags
+        UI->setHasNoUnsignedWrap(false);
+        UI->setHasNoSignedWrap(false);
+      }
       Value *NewCmp = Builder.CreateCmp(Cmp->getPredicate(), UI,
                                         ConstantInt::get(UI->getType(), 0));
       LLVM_DEBUG(dbgs() << "Converting " << *Cmp << "\n");
       LLVM_DEBUG(dbgs() << " to compare on zero: " << *NewCmp << "\n");
       replaceAllUsesWith(Cmp, NewCmp, FreshBBs, IsHugeFunc);
       return true;
     }
   }
case Instruction::Select:
     return optimizeSelectInst(cast<SelectInst>(I));
   case Instruction::ShuffleVector:
     return optimizeShuffleVectorInst(cast<ShuffleVectorInst>(I));
   case Instruction::Switch:
     return optimizeSwitchInst(cast<SwitchInst>(I));
   case Instruction::ExtractElement:
     return optimizeExtractElementInst(cast<ExtractElementInst>(I));
   case Instruction::Br:
-    return optimizeBranch(cast<BranchInst>(I), *TLI, FreshBBs, IsHugeFunc);
+    return optimizeBranch(cast<BranchInst>(I), *TLI, *DL, AC,
+                          &getDT(*I->getParent()->getParent()), FreshBBs,
+                          IsHugeFunc);
   }
 
   return false;
 }
 
 /// Given an OR instruction, check to see if this is a bitreverse
 /// idiom. If so, insert the new intrinsic and return true.
 bool CodeGenPrepare::makeBitReverse(Instruction &I) {

llvm/test/CodeGen/RISCV/branch-on-zero.ll

while.body:                                       ; preds = %while.body.preheader, %while.body
   store i32 %mul, ptr %x.addr.06, align 4
   %dec = add nsw i32 %c.07, -1
   %tobool.not = icmp eq i32 %dec, 0
   br i1 %tobool.not, label %while.end, label %while.body
 
 while.end:                                        ; preds = %while.body, %entry
   ret i32 0
 }
+
+define i32 @test_nsw_add_may_overflow(i32 %0) {
+; RV32-LABEL: test_nsw_add_may_overflow:
+; RV32:       # %bb.0:
+; RV32-NEXT:    li a1, 13
+; RV32-NEXT:    bge a0, a1, .LBB4_2
+; RV32-NEXT:  # %bb.1:
+; RV32-NEXT:    li a0, 0
+; RV32-NEXT:    ret
+; RV32-NEXT:  .LBB4_2:
+; RV32-NEXT:    addi a0, a0, -13
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test_nsw_add_may_overflow:
+; RV64:       # %bb.0:
+; RV64-NEXT:    sext.w a1, a0
+; RV64-NEXT:    li a2, 13
+; RV64-NEXT:    bge a1, a2, .LBB4_2
+; RV64-NEXT:  # %bb.1:
+; RV64-NEXT:    li a0, 0
+; RV64-NEXT:    ret
+; RV64-NEXT:  .LBB4_2:
+; RV64-NEXT:    addiw a0, a0, -13
+; RV64-NEXT:    ret
+  %2 = icmp slt i32 %0, 13
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = add i32 %0, -13
+  ret i32 %4
+
+5:                                                ; preds = %1
+  ret i32 0
+}
+
+define i32 @test_nsw_add_no_overflow(i32 %0) {
+; RV32-LABEL: test_nsw_add_no_overflow:
+; RV32:       # %bb.0:
+; RV32-NEXT:    bltz a0, .LBB5_3
+; RV32-NEXT:  # %bb.1: # %nonnegative
+; RV32-NEXT:    addi a0, a0, -13
+; RV32-NEXT:    bltz a0, .LBB5_3
+; RV32-NEXT:  # %bb.2:
+; RV32-NEXT:    ret
+; RV32-NEXT:  .LBB5_3:
+; RV32-NEXT:    li a0, 0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test_nsw_add_no_overflow:
+; RV64:       # %bb.0:
+; RV64-NEXT:    sext.w a1, a0
+; RV64-NEXT:    bltz a1, .LBB5_3
+; RV64-NEXT:  # %bb.1: # %nonnegative
+; RV64-NEXT:    addiw a0, a0, -13
+; RV64-NEXT:    bltz a0, .LBB5_3
+; RV64-NEXT:  # %bb.2:
+; RV64-NEXT:    ret
+; RV64-NEXT:  .LBB5_3:
+; RV64-NEXT:    li a0, 0
+; RV64-NEXT:    ret
+  %x = icmp sge i32 %0, 0
+  br i1 %x, label %nonnegative, label %5
+
+nonnegative:                                      ; preds = %1
+  %2 = icmp slt i32 %0, 13
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %nonnegative
+  %4 = add i32 %0, -13
+  ret i32 %4
+
+5:                                                ; preds = %nonnegative, %1
+  ret i32 0
+}