[LSR][TTI][RISCV] Enable terminator folding for RISC-V

If looking for a miscompile revert candidate, look here!

The transform being enabled prefers comparing to a loop invariant
exit value for a secondary IV over using an otherwise dead primary
IV.  This increases register pressure (by requiring the exit value
to be live through the loop), but reduces the number of instructions
within the loop by one.

On RISC-V which has a large number of scalar registers, this is
generally a profitable transform.  We loose the ability to use a beqz
on what is typically a count down IV, and pay the cost of computing
the exit value on the secondary IV in the loop preheader, but save
an add or sub in the loop body.  For anything except an extremely
short running loop, or one with extreme register pressure, this is
profitable.  On spec2017, we see a 0.42% geomean improvement in
dynamic icount, with no individual workload regressing by more than
0.25%.

Code size wise, we trade a (possibly compressible) beqz and a (possibly
compressible) addi for a uncompressible beq.  We also add instructions
in the preheader.  Net result is a slight regression overall, but
neutral or better inside the loop.

Previous versions of this transform had numerous cornercase correctness
bugs.  All of them ones I can spot by inspection have been fixed, and I
have run this through all of spec2017, but there may be further issues
lurking.  Adding uses to an IV is a fraught thing to do given poison
semantics, so this transform is somewhat inherently risky.

This patch is a reworked version of D134893 by @eop.  That patch has
been abandoned since May, so I picked it up, reworked it a bit, and
am landing it.

Author: preames

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -715,6 +715,11 @@ class TargetTransformInfo {
   /// cost should return false, otherwise return true.
   bool isNumRegsMajorCostOfLSR() const;
 
+  /// Return true if LSR should attempts to replace a use of an otherwise dead
+  /// primary IV in the latch condition with another IV available in the loop.
+  /// When successful, makes the primary IV dead.
+  bool shouldFoldTerminatingConditionAfterLSR() const;
+
   /// \returns true if LSR should not optimize a chain that includes \p I.
   bool isProfitableLSRChainElement(Instruction *I) const;
 
@@ -1760,6 +1765,7 @@ class TargetTransformInfo::Concept {
   virtual bool isLSRCostLess(const TargetTransformInfo::LSRCost &C1,
                              const TargetTransformInfo::LSRCost &C2) = 0;
   virtual bool isNumRegsMajorCostOfLSR() = 0;
+  virtual bool shouldFoldTerminatingConditionAfterLSR() const = 0;
   virtual bool isProfitableLSRChainElement(Instruction *I) = 0;
   virtual bool canMacroFuseCmp() = 0;
   virtual bool canSaveCmp(Loop *L, BranchInst **BI, ScalarEvolution *SE,
@@ -2205,6 +2211,9 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   bool isNumRegsMajorCostOfLSR() override {
     return Impl.isNumRegsMajorCostOfLSR();
   }
+  bool shouldFoldTerminatingConditionAfterLSR() const override {
+    return Impl.shouldFoldTerminatingConditionAfterLSR();
+  }
   bool isProfitableLSRChainElement(Instruction *I) override {
     return Impl.isProfitableLSRChainElement(I);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -231,6 +231,8 @@ class TargetTransformInfoImplBase {
 
   bool isNumRegsMajorCostOfLSR() const { return true; }
 
+  bool shouldFoldTerminatingConditionAfterLSR() const { return false; }
+
   bool isProfitableLSRChainElement(Instruction *I) const { return false; }
 
   bool canMacroFuseCmp() const { return false; }

llvm/include/llvm/CodeGen/BasicTTIImpl.h

@@ -382,6 +382,11 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
     return TargetTransformInfoImplBase::isNumRegsMajorCostOfLSR();
   }
 
+  bool shouldFoldTerminatingConditionAfterLSR() const {
+    return TargetTransformInfoImplBase::
+        shouldFoldTerminatingConditionAfterLSR();
+  }
+
   bool isProfitableLSRChainElement(Instruction *I) {
     return TargetTransformInfoImplBase::isProfitableLSRChainElement(I);
   }

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -402,6 +402,10 @@ bool TargetTransformInfo::isNumRegsMajorCostOfLSR() const {
   return TTIImpl->isNumRegsMajorCostOfLSR();
 }
 
+bool TargetTransformInfo::shouldFoldTerminatingConditionAfterLSR() const {
+  return TTIImpl->shouldFoldTerminatingConditionAfterLSR();
+}
+
 bool TargetTransformInfo::isProfitableLSRChainElement(Instruction *I) const {
   return TTIImpl->isProfitableLSRChainElement(I);
 }

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.h

@@ -359,6 +359,10 @@ class RISCVTTIImpl : public BasicTTIImplBase<RISCVTTIImpl> {
 
   bool isLSRCostLess(const TargetTransformInfo::LSRCost &C1,
                      const TargetTransformInfo::LSRCost &C2);
+
+  bool shouldFoldTerminatingConditionAfterLSR() const {
+    return true;
+  }
 };
 
 } // end namespace llvm

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

@@ -188,8 +188,8 @@ static cl::opt<unsigned> SetupCostDepthLimit(
     "lsr-setupcost-depth-limit", cl::Hidden, cl::init(7),
     cl::desc("The limit on recursion depth for LSRs setup cost"));
 
-static cl::opt<bool> AllowTerminatingConditionFoldingAfterLSR(
-    "lsr-term-fold", cl::Hidden, cl::init(false),
+static cl::opt<cl::boolOrDefault> AllowTerminatingConditionFoldingAfterLSR(
+    "lsr-term-fold", cl::Hidden,
     cl::desc("Attempt to replace primary IV with other IV."));
 
 static cl::opt<bool> AllowDropSolutionIfLessProfitable(
@@ -6949,7 +6949,18 @@ static bool ReduceLoopStrength(Loop *L, IVUsers &IU, ScalarEvolution &SE,
     }
   }
 
-  if (AllowTerminatingConditionFoldingAfterLSR) {
+  const bool EnableFormTerm = [&] {
+    switch (AllowTerminatingConditionFoldingAfterLSR) {
+    case cl::BOU_TRUE:
+      return true;
+    case cl::BOU_FALSE:
+      return false;
+    case cl::BOU_UNSET:
+      return TTI.shouldFoldTerminatingConditionAfterLSR();
+    }
+  }();
+
+  if (EnableFormTerm) {
     if (auto Opt = canFoldTermCondOfLoop(L, SE, DT, LI)) {
       auto [ToFold, ToHelpFold, TermValueS, MustDrop] = *Opt;
 

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

@@ -120,36 +120,45 @@ define i32 @test_lshr2(ptr nocapture %x, ptr nocapture readonly %y, i32 %n) {
 ; RV32-LABEL: test_lshr2:
 ; RV32:       # %bb.0: # %entry
 ; RV32-NEXT:    srli a2, a2, 2
-; RV32-NEXT:    beqz a2, .LBB3_2
-; RV32-NEXT:  .LBB3_1: # %while.body
+; RV32-NEXT:    beqz a2, .LBB3_3
+; RV32-NEXT:  # %bb.1: # %while.body.preheader
+; RV32-NEXT:    slli a2, a2, 2
+; RV32-NEXT:    add a2, a1, a2
+; RV32-NEXT:  .LBB3_2: # %while.body
 ; RV32-NEXT:    # =>This Inner Loop Header: Depth=1
 ; RV32-NEXT:    lw a3, 0(a1)
-; RV32-NEXT:    addi a1, a1, 4
+; RV32-NEXT:    addi a4, a1, 4
 ; RV32-NEXT:    slli a3, a3, 1
-; RV32-NEXT:    addi a4, a0, 4
-; RV32-NEXT:    addi a2, a2, -1
+; RV32-NEXT:    addi a1, a0, 4
 ; RV32-NEXT:    sw a3, 0(a0)
-; RV32-NEXT:    mv a0, a4
-; RV32-NEXT:    bnez a2, .LBB3_1
-; RV32-NEXT:  .LBB3_2: # %while.end
+; RV32-NEXT:    mv a0, a1
+; RV32-NEXT:    mv a1, a4
+; RV32-NEXT:    bne a4, a2, .LBB3_2
+; RV32-NEXT:  .LBB3_3: # %while.end
 ; RV32-NEXT:    li a0, 0
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: test_lshr2:
 ; RV64:       # %bb.0: # %entry
 ; RV64-NEXT:    srliw a2, a2, 2
-; RV64-NEXT:    beqz a2, .LBB3_2
-; RV64-NEXT:  .LBB3_1: # %while.body
+; RV64-NEXT:    beqz a2, .LBB3_3
+; RV64-NEXT:  # %bb.1: # %while.body.preheader
+; RV64-NEXT:    addi a2, a2, -1
+; RV64-NEXT:    slli a2, a2, 32
+; RV64-NEXT:    srli a2, a2, 30
+; RV64-NEXT:    add a2, a2, a1
+; RV64-NEXT:    addi a2, a2, 4
+; RV64-NEXT:  .LBB3_2: # %while.body
 ; RV64-NEXT:    # =>This Inner Loop Header: Depth=1
 ; RV64-NEXT:    lw a3, 0(a1)
-; RV64-NEXT:    addi a1, a1, 4
+; RV64-NEXT:    addi a4, a1, 4
 ; RV64-NEXT:    slli a3, a3, 1
-; RV64-NEXT:    addi a4, a0, 4
-; RV64-NEXT:    addiw a2, a2, -1
+; RV64-NEXT:    addi a1, a0, 4
 ; RV64-NEXT:    sw a3, 0(a0)
-; RV64-NEXT:    mv a0, a4
-; RV64-NEXT:    bnez a2, .LBB3_1
-; RV64-NEXT:  .LBB3_2: # %while.end
+; RV64-NEXT:    mv a0, a1
+; RV64-NEXT:    mv a1, a4
+; RV64-NEXT:    bne a4, a2, .LBB3_2
+; RV64-NEXT:  .LBB3_3: # %while.end
 ; RV64-NEXT:    li a0, 0
 ; RV64-NEXT:    ret
 entry:

llvm/test/CodeGen/RISCV/riscv-codegenprepare-asm.ll

@@ -8,16 +8,19 @@
 define void @test1(ptr nocapture noundef %a, i32 noundef signext %n) {
 ; CHECK-LABEL: test1:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    blez a1, .LBB0_2
-; CHECK-NEXT:  .LBB0_1: # %for.body
+; CHECK-NEXT:    blez a1, .LBB0_3
+; CHECK-NEXT:  # %bb.1: # %for.body.preheader
+; CHECK-NEXT:    slli a1, a1, 32
+; CHECK-NEXT:    srli a1, a1, 30
+; CHECK-NEXT:    add a1, a0, a1
+; CHECK-NEXT:  .LBB0_2: # %for.body
 ; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    lw a2, 0(a0)
 ; CHECK-NEXT:    addiw a2, a2, 4
 ; CHECK-NEXT:    sw a2, 0(a0)
-; CHECK-NEXT:    addi a1, a1, -1
 ; CHECK-NEXT:    addi a0, a0, 4
-; CHECK-NEXT:    bnez a1, .LBB0_1
-; CHECK-NEXT:  .LBB0_2: # %for.cond.cleanup
+; CHECK-NEXT:    bne a0, a1, .LBB0_2
+; CHECK-NEXT:  .LBB0_3: # %for.cond.cleanup
 ; CHECK-NEXT:    ret
 entry:
   %cmp3 = icmp sgt i32 %n, 0