[IVUser] Add a trunc if used only by a single IV user

[junbuml]

Summary

Add a trunc as an IV user if it is used only by a single IV user that is added
directly through the trunc and there are other IV users that has the same SCEV
recurrence as the trunc. By adding such trunc, we expect that LSR chooses a
better solution in case where there are other IV users which have the same SCEV
recurrence as the trunc.

Motivating case in AArch64 is added as a testcase where extra IV increments are
added in LSR which picked suboptimal formulae.

For the loop :

for.body.preheader:
  %n_sext = sext i32 %n to i64
  br label %for.body 
for.body:
  %K.in = phi i64 [ %n_sext, %for.body.preheader ], [ %K, %for.body ]
  %K = add i64 %K.in, 1
  %StoredAddr = getelementptr inbounds [12 x i32], [12 x i32]* @gvarray, i64 0, i64 %K
  %StoredValue = trunc i64 %K to i32
  store volatile i32 %StoredValue, i32* %StoredAddr
  %cmp = icmp sgt i64 %K, 1
  br i1 %cmp, label %for.body, label %for.end

Before this change :

.LBB0_2:
  add x11, x9, x10, lsl #2
  add x10, x10, #1
  str w8, [x11, #4]
  add w8, w8, #1
  cmp x10, #1
  b.gt .LBB0_2

After this change :

.LBB0_2:
add	x10, x8, #1
cmp		x8, #1
str	w8, [x9, x8, lsl #2]
mov	 x8, x10
b.gt	.LBB0_2

[atrick]

This is definitely a case worth fixing.

The "problem" with this case from LSR point of view is that SCEV is
too smart and hoists the 'trunc' out of the recurrence. This makes
it look like the expression that feeds the store and the compare are
based on different initial values.

This patch looks like it goes through a fairly involved analysis
before LSR to undo the normalization that happens during IVUsers. I
think that only obscures the input to LSR. It happens to generate
better code in this case, but LSR still thinks that two registers are
needed for the loop counter. So LSR hasn't really been fixed, and I'm
afraid this will prevent LSR from optimizing other cases.

Look at the output of LSR *after* applying this patch:

The chosen solution requires 3 regs, with addrec cost 2, plus 1 scale cost, plus 1 setup cost:

LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i64
  reg({(sext i32 %n to i64),+,1}<%for.body>)
LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i32
  reg({%n,+,1}<%for.body>)
LSR Use: Kind=Address of i32 in addrspace(0), Offsets={0}, widest fixup type: i32*
  reg(@gvarray) + 4*reg({(sext i32 %n to i64),+,1}<%for.body>)

So the problem has not been fixed in LSR. We just get lucky and "replaceCongruentPhis" happens to cleanup the extra IV:

INDVARS: Eliminated congruent iv.inc: %lsr.iv.next = add i32 %lsr.iv, 1
INDVARS: Eliminated congruent iv: %lsr.iv = phi i32 [ %lsr, %for.body ], [ %n, %for.body.preheader ]

I have two suggestions. I'm not sure either will work:

(1) Fix IVUsers so that the truncate is itself the IV users. Maybe do this only if the trunc has a single user and/or the use it itself not an "interesting" SCEV. Then LSR will see the same SCEV recurrence for all users. This will be easy, but I don't know if it will mess up any other cases.

(2) Improve replaceCongruentPhis to cleanup these sort of off-by-one-with-postinc-user phis. This requires a fair amount of pattern matching.

In any case, the important, and time-consuming part will be evaluating AArch64 performance on the LLVM test suite after the change. Although if there are few changes, it may be sufficient just to look at the disassembly diffs to see that they're reasonable.

[atrick]

Actually, I'm sure fix #1 works in this case, I'm just not sure it won't regress anything else.

[junbuml]

(1) Fix IVUsers so that the truncate is itself the IV users. Maybe do this only if the trunc has a single user and/or the use it itself not an "interesting" SCEV. Then LSR will see the same SCEV recurrence for all users. This will be easy, but I don't know if it will mess up any other cases.

Based on Andrew's suggestion #1, updated the patch and summary.

Looked at the disassembly diffs in spec2000/2006. This patch is not applied in spec2000, but applied in 6 benchmarks (hmmer,h264ref, omnetpp, gobmk, sphinx3, and gcc) in Spec2006. Except gcc in which one more add instruction is added in a loop, all diffs looks reasonable to me as it remove more instructions. I will continue running performance tests and update with them.

[atrick]

Does this patch still need review, or has it been subsumed by other changes?

[junbuml]

Does this patch still need review, or has it been subsumed by other changes?

I deprioritized this because I didn't see clear performance gains in benchmarks I was targeting. However, there was no other changes covering this issue, and this change is still good to be reviewed.

[atrick]

It would be great to fix the example shown in the summary. However, I'm concerned that the proposed heuristic could hurt codegen in other cases.The cases where this heuristic generates worse code should be identified to determine whether this is the right strategy.

Maybe this should just be a bug report since the patch is in limbo?

[junbuml]

Thanks Andrew for reviewing this old change. Let me try to check any code change with/without this patch in llvmTS and spec2000/2006. I added a comment in Bugzilla (https://llvm.org/bugs/show_bug.cgi?id=26913), mentioning this change.

include/llvm/Analysis/IVUsers.h

class IVUsers : public LoopPass {
 
   /// IVUses - A list of all tracked IV uses of induction variable expressions
   /// we are interested in.
   ilist<IVStrideUse> IVUses;
 
   // Ephemeral values used by @llvm.assume in this function.
   SmallPtrSet<const Value *, 32> EphValues;
 
+  // A list of truncs which is used only by a single IV user added directly
+  // through the trunc.
+  SmallPtrSet<Instruction *, 2> TruncsToSingleIVUser;
+
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override;
 
   void releaseMemory() override;
 
 public:
   static char ID; // Pass ID, replacement for typeid
   IVUsers();
 
   Loop *getLoop() const { return L; }
 
   /// AddUsersIfInteresting - Inspect the specified Instruction.  If it is a
   /// reducible SCEV, recursively add its users to the IVUsesByStride set and
   /// return true.  Otherwise, return false.
   bool AddUsersIfInteresting(Instruction *I);
 
   IVStrideUse &AddUser(Instruction *User, Value *Operand);
 
+  void FilterOutTruncIVUsers();
+
   /// getReplacementExpr - Return a SCEV expression which computes the
   /// value of the OperandValToReplace of the given IVStrideUse.
   const SCEV *getReplacementExpr(const IVStrideUse &IU) const;
 
   /// getExpr - Return the expression for the use.
   const SCEV *getExpr(const IVStrideUse &IU) const;
 
   const SCEV *getStride(const IVStrideUse &IU, const Loop *L) const;

lib/Analysis/IVUsers.cpp

 #include "llvm/IR/Type.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 using namespace llvm;
 
 #define DEBUG_TYPE "iv-users"
 
+
 char IVUsers::ID = 0;
 INITIALIZE_PASS_BEGIN(IVUsers, "iv-users",
                       "Induction Variable Users", false, true)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_END(IVUsers, "iv-users",
if (LI->getLoopFor(User->getParent()) != L) {
         AddUserToIVUsers = true;
       }
     } else if (Processed.count(User) || !AddUsersImpl(User, SimpleLoopNests)) {
       DEBUG(dbgs() << "FOUND USER: " << *User << '\n'
                    << "   OF SCEV: " << *ISE << '\n');
       AddUserToIVUsers = true;
     }
 
+    // See if User is a trunc which is used only by a single IV user just added
+    // directly through the trunc. By adding such trunc, we expect that LSR
+    // chooses a better solution in case there are other IV users which have the
+    // same SCEV recurrence as the trunc. Later in the pass, truncs added here
+    // will be removed if none of IV users has the same SCEV recurrence.
+    if (isa<TruncInst>(User) && !AddUserToIVUsers && User->hasOneUse()) {
+      Instruction *OneTruncUse = User->user_back();
+      if (LI->getLoopFor(User->getParent()) == L &&
+          LI->getLoopFor(OneTruncUse->getParent()) == L && !IVUses.empty() &&
+          (IVUses.back()).getOperandValToReplace() == User) {
+        AddUserToIVUsers = true;
+        TruncsToSingleIVUser.insert(User);
+      }
+    }
+
     if (AddUserToIVUsers) {
       // Okay, we found a user that we cannot reduce.
       IVStrideUse &NewUse = AddUser(User, I);
       // Autodetect the post-inc loop set, populating NewUse.PostIncLoops.
       // The regular return value here is discarded; instead of recording
       // it, we just recompute it when we need it.
       const SCEV *OriginalISE = ISE;
       ISE = TransformForPostIncUse(NormalizeAutodetect,
IVStrideUse &IVUsers::AddUser(Instruction *User, Value *Operand) {
   return IVUses.back();
 }
 
 IVUsers::IVUsers()
     : LoopPass(ID) {
   initializeIVUsersPass(*PassRegistry::getPassRegistry());
 }
 
+/// Remove truncs in TruncsToSingleIVUser if none of IV users has the same SCEV
+/// recurrence.
+void IVUsers::FilterOutTruncIVUsers() {
+  if (TruncsToSingleIVUser.empty())
+    return;
+  for (IVUsers::iterator U = begin(), E = end(); U != E;) {
+    if (!TruncsToSingleIVUser.count(U->getUser())) {
+      ++U;
+      continue;
+    }
+    Instruction *TruncI = U->getUser();
+    assert(isa<TruncInst>(TruncI) && TruncI->hasOneUse() &&
+           "Unexpected TruncInst to remove");
+    const SCEV *TruncS = SE->getSCEV(U->getOperandValToReplace());
+    bool HasSameRec = false;
+    for (IVStrideUse &UU : IVUses) {
+      if (&UU == &*U)
+        continue;
+      if (SE->getSCEV(UU.getOperandValToReplace()) == TruncS) {
+        HasSameRec = true;
+        break;
+      }
+    }
+    if (!HasSameRec)
+      U = IVUses.erase(U);
+    else
+      ++U;
+  }
+}
+
 void IVUsers::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<AssumptionCacheTracker>();
   AU.addRequired<LoopInfoWrapperPass>();
   AU.addRequired<DominatorTreeWrapperPass>();
   AU.addRequired<ScalarEvolutionWrapperPass>();
   AU.setPreservesAll();
 }
 
bool IVUsers::runOnLoop(Loop *l, LPPassManager &LPM) {
   CodeMetrics::collectEphemeralValues(L, AC, EphValues);
 
   // Find all uses of induction variables in this loop, and categorize
   // them by stride.  Start by finding all of the PHI nodes in the header for
   // this loop.  If they are induction variables, inspect their uses.
   for (BasicBlock::iterator I = L->getHeader()->begin(); isa<PHINode>(I); ++I)
     (void)AddUsersIfInteresting(&*I);
 
+  FilterOutTruncIVUsers();
   return false;
 }
 
 void IVUsers::print(raw_ostream &OS, const Module *M) const {
   OS << "IV Users for loop ";
   L->getHeader()->printAsOperand(OS, false);
   if (SE->hasLoopInvariantBackedgeTakenCount(L)) {
     OS << " with backedge-taken count "

 LLVM_DUMP_METHOD void IVUsers::dump() const {
   print(dbgs());
 }
 #endif
 
 void IVUsers::releaseMemory() {
   Processed.clear();
   IVUses.clear();
+  TruncsToSingleIVUser.clear();
 }
 
 /// getReplacementExpr - Return a SCEV expression which computes the
 /// value of the OperandValToReplace.
 const SCEV *IVUsers::getReplacementExpr(const IVStrideUse &IU) const {
   return SE->getSCEV(IU.getOperandValToReplace());
 }
 

test/Transforms/LoopStrengthReduce/AArch64/lsr-trunc.ll

+; RUN: llc  < %s | FileCheck %s
+
+@gvarray = common global [12 x i32] zeroinitializer, align 4
+
+target triple = "arm64-unknown-unknown"
+
+; Check if the same register is used in the stored value, indexing, and compare.
+
+define void @test(i32 %n) {
+entry:
+  %cmp7 = icmp sgt i32 %n, 1
+  br i1 %cmp7, label %for.body.preheader, label %for.end
+
+for.body.preheader:
+  %n_sext = sext i32 %n to i64
+  br label %for.body
+
+for.body:
+; CHECK: cmp x[[REG:[0-9]+]]
+; CHECK: str w[[REG]], [x{{[0-9]+}}, x[[REG]], lsl #2]
+  %K.in = phi i64 [ %n_sext, %for.body.preheader ], [ %K, %for.body ]
+  %K = add i64 %K.in, 1
+  %StoredAddr = getelementptr inbounds [12 x i32], [12 x i32]* @gvarray, i64 0, i64 %K
+  %StoredValue = trunc i64 %K to i32
+  store volatile i32 %StoredValue, i32* %StoredAddr
+  %cmp = icmp sgt i64 %K, 1
+  br i1 %cmp, label %for.body, label %for.end
+for.end:
+  ret void
+}