[SeparateConstOffsetFromGEP][PowerPC] Fix: sext(a) + sext(b) -> sext(a + b) matches add and sub instructions with one another

[ajwock]

Summary

During the SeparateConstOffsetFromGEP pass, signed extensions are distributed and then later recombined. The recombination stage is somewhat problematic- it doesn't differ add and sub instructions from another when matching the sext(a) +/- sext(b) -> sext(a +/- b) pattern, although this issue is incredibly rare. It only occurs under the below conditions. The test case only miscompiles on PowerPC as per my current testing; however, the issue is not in PowerPC specific code as far as I can tell, and it appears that the right IR input could theoretically trigger this on multiple platforms.

The IR contains:
%unextendedA
%unextendedB
%subuAuB = unextendedA - unextendedB
%extA = extend A
%extB = extend B
%addeAeB = extA + extB

The problematic code will transform that into:

%unextendedA
%unextendedB
%subuAuB = unextendedA - unextendedB
%extA = extend A
%extB = extend B
%addeAeB = extend subuAuB ; Obviously not semantically equivalent to the IR input.

The operands must be in the same order as above- and other optimization passes must preserve this behavior until this section of the code is reached, which is unlikely with the possibly-unnecessary extensions. My test case was heavily reduced both at the C level and the IR level from cblas_ztrsv, but there are many extraneous instructions creating dependencies necessary to preserve the miscompilation.

I do appreciate your documentation, jingyue.

[ajwock]

Diff did not include my test case- correcting that.

[lebedev.ri]

test coverage missing.

[ajwock]

Ping.

[ajwock]

ping.

[ajwock]

ping

[jlebar]

Hi, I am not familiar with this code but seeing as you're not getting reviews and this seems like a serious problem, I am willing to take a look (and maybe if I say something wrong here, Cunningham's Law will save us).

Without looking at the code too deeply, I am wondering about the testcase. Per the commit message it sounds like you're saying that invalid LLVM IR is generated. If so, can we have an LLVM IR testcase that runs this specific pass, rather than an llc testcase? Based on what's in the commit message, it sounds like such a testcase would be far simpler (and indeed less fragile) than what we have here.

[arsenm]

test/Transforms/SeparateConstOffsetFromGEP/PowerPC/test-nomatch-add-and-sub.ll

8 ↗

(On Diff #214216)

I'm not opposed to including this testcase, but I would like to have a pure IR one that only runs the one pass

[ajwock]

Sorry for the delay- I had to investigate some of the fragility so that my test case reduced to this pass would still trigger the issue. However, I did manage to get a test case that is only a few lines long that would trigger the issue without the update.

[nemanjai]

In D65967#1730540, @ajwock wrote:

Sorry for the delay- I had to investigate some of the fragility so that my test case reduced to this pass would still trigger the issue. However, I did manage to get a test case that is only a few lines long that would trigger the issue without the update.

I suppose now that the test case is not PPC specific, it can be moved to a different directory?

[ajwock]

Changed directory.

[arsenm]

LGTM. I'm not super familiar with this code, but this seems reasonable. There's nothing new going on, and just splitting add and sub handling

[kpn]

Excellent! Thanks Matt!

Drew will be back in the office tomorrow. I'll commit this for him then.

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

private:
   ///   => constant extraction     &a[sext(i) + sext(j)] + 5
   ///   => reunion                 &a[sext(i +nsw j)] + 5
   bool reuniteExts(Function &F);
 
   /// A helper that reunites sexts in an instruction.
   bool reuniteExts(Instruction *I);
 
   /// Find the closest dominator of <Dominatee> that is equivalent to <Key>.
-  Instruction *findClosestMatchingDominator(const SCEV *Key,
-                                            Instruction *Dominatee);
+  Instruction *findClosestMatchingDominator(
+      const SCEV *Key, Instruction *Dominatee,
+      DenseMap<const SCEV *, SmallVector<Instruction *, 2>> &DominatingExprs);
+
   /// Verify F is free of dead code.
   void verifyNoDeadCode(Function &F);
 
   bool hasMoreThanOneUseInLoop(Value *v, Loop *L);
 
   // Swap the index operand of two GEP.
   void swapGEPOperand(GetElementPtrInst *First, GetElementPtrInst *Second);
 
   // Check if it is safe to swap operand of two GEP.
   bool isLegalToSwapOperand(GetElementPtrInst *First, GetElementPtrInst *Second,
                             Loop *CurLoop);
 
   const DataLayout *DL = nullptr;
   DominatorTree *DT = nullptr;
   ScalarEvolution *SE;
 
   LoopInfo *LI;
   TargetLibraryInfo *TLI;
 
   /// Whether to lower a GEP with multiple indices into arithmetic operations or
   /// multiple GEPs with a single index.
   bool LowerGEP;
 
-  DenseMap<const SCEV *, SmallVector<Instruction *, 2>> DominatingExprs;
+  DenseMap<const SCEV *, SmallVector<Instruction *, 2>> DominatingAdds;
+  DenseMap<const SCEV *, SmallVector<Instruction *, 2>> DominatingSubs;
 };
 
 } // end anonymous namespace
 
 char SeparateConstOffsetFromGEP::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     SeparateConstOffsetFromGEP, "separate-const-offset-from-gep",
bool SeparateConstOffsetFromGEP::runOnFunction(Function &F) {
 
   if (VerifyNoDeadCode)
     verifyNoDeadCode(F);
 
   return Changed;
 }
 
 Instruction *SeparateConstOffsetFromGEP::findClosestMatchingDominator(
-    const SCEV *Key, Instruction *Dominatee) {
+    const SCEV *Key, Instruction *Dominatee,
+    DenseMap<const SCEV *, SmallVector<Instruction *, 2>> &DominatingExprs) {
   auto Pos = DominatingExprs.find(Key);
   if (Pos == DominatingExprs.end())
     return nullptr;
 
   auto &Candidates = Pos->second;
   // Because we process the basic blocks in pre-order of the dominator tree, a
   // candidate that doesn't dominate the current instruction won't dominate any
   // future instruction either. Therefore, we pop it out of the stack. This
bool SeparateConstOffsetFromGEP::reuniteExts(Instruction *I) {
   if (!SE->isSCEVable(I->getType()))
     return false;
 
   //   Dom: LHS+RHS
   //   I: sext(LHS)+sext(RHS)
   // If Dom can't sign overflow and Dom dominates I, optimize I to sext(Dom).
   // TODO: handle zext
   Value *LHS = nullptr, *RHS = nullptr;
-  if (match(I, m_Add(m_SExt(m_Value(LHS)), m_SExt(m_Value(RHS)))) ||
-      match(I, m_Sub(m_SExt(m_Value(LHS)), m_SExt(m_Value(RHS))))) {
+  if (match(I, m_Add(m_SExt(m_Value(LHS)), m_SExt(m_Value(RHS))))) {
     if (LHS->getType() == RHS->getType()) {
       const SCEV *Key =
           SE->getAddExpr(SE->getUnknown(LHS), SE->getUnknown(RHS));
-      if (auto *Dom = findClosestMatchingDominator(Key, I)) {
+      if (auto *Dom = findClosestMatchingDominator(Key, I, DominatingAdds)) {
+        Instruction *NewSExt = new SExtInst(Dom, I->getType(), "", I);
+        NewSExt->takeName(I);
+        I->replaceAllUsesWith(NewSExt);
+        RecursivelyDeleteTriviallyDeadInstructions(I);
+        return true;
+      }
+    }
+  } else if (match(I, m_Sub(m_SExt(m_Value(LHS)), m_SExt(m_Value(RHS))))) {
+    if (LHS->getType() == RHS->getType()) {
+      const SCEV *Key =
+          SE->getAddExpr(SE->getUnknown(LHS), SE->getUnknown(RHS));
+      if (auto *Dom = findClosestMatchingDominator(Key, I, DominatingSubs)) {
         Instruction *NewSExt = new SExtInst(Dom, I->getType(), "", I);
         NewSExt->takeName(I);
         I->replaceAllUsesWith(NewSExt);
         RecursivelyDeleteTriviallyDeadInstructions(I);
         return true;
       }
     }
   }
 
   // Add I to DominatingExprs if it's an add/sub that can't sign overflow.
-  if (match(I, m_NSWAdd(m_Value(LHS), m_Value(RHS))) ||
-      match(I, m_NSWSub(m_Value(LHS), m_Value(RHS)))) {
+  if (match(I, m_NSWAdd(m_Value(LHS), m_Value(RHS)))) {
+    if (programUndefinedIfFullPoison(I)) {
+      const SCEV *Key =
+          SE->getAddExpr(SE->getUnknown(LHS), SE->getUnknown(RHS));
+      DominatingAdds[Key].push_back(I);
+    }
+  } else if (match(I, m_NSWSub(m_Value(LHS), m_Value(RHS)))) {
     if (programUndefinedIfFullPoison(I)) {
       const SCEV *Key =
           SE->getAddExpr(SE->getUnknown(LHS), SE->getUnknown(RHS));
-      DominatingExprs[Key].push_back(I);
+      DominatingSubs[Key].push_back(I);
     }
   }
   return false;
 }
 
 bool SeparateConstOffsetFromGEP::reuniteExts(Function &F) {
   bool Changed = false;
-  DominatingExprs.clear();
+  DominatingAdds.clear();
+  DominatingSubs.clear();
   for (const auto Node : depth_first(DT)) {
     BasicBlock *BB = Node->getBlock();
     for (auto I = BB->begin(); I != BB->end(); ) {
       Instruction *Cur = &*I++;
       Changed |= reuniteExts(Cur);
     }
   }
   return Changed;

llvm/test/Transforms/SeparateConstOffsetFromGEP/test-add-sub-separation.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -separate-const-offset-from-gep < %s | FileCheck %s
+
+define void @matchingExtensions(i32* %ap, i32* %bp, i64* %result) {
+; CHECK-LABEL: @matchingExtensions(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = load i32, i32* [[AP:%.*]]
+; CHECK-NEXT:    [[B:%.*]] = load i32, i32* [[BP:%.*]]
+; CHECK-NEXT:    [[EB:%.*]] = sext i32 [[B]] to i64
+; CHECK-NEXT:    [[SUBAB:%.*]] = sub nsw i32 [[A]], [[B]]
+; CHECK-NEXT:    [[EA:%.*]] = sext i32 [[A]] to i64
+; CHECK-NEXT:    [[ADDEAEB:%.*]] = add nsw i64 [[EA]], [[EB]]
+; CHECK-NEXT:    [[EXTSUB:%.*]] = sext i32 [[SUBAB]] to i64
+; CHECK-NEXT:    [[IDX:%.*]] = getelementptr i32, i32* [[AP]], i64 [[EXTSUB]]
+; CHECK-NEXT:    store i64 [[ADDEAEB]], i64* [[RESULT:%.*]]
+; CHECK-NEXT:    store i32 [[SUBAB]], i32* [[IDX]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a = load i32, i32* %ap
+  %b = load i32, i32* %bp
+  %eb = sext i32 %b to i64
+  %subab = sub nsw i32 %a, %b
+  %ea = sext i32 %a to i64
+  %addeaeb = add nsw i64 %ea, %eb
+  %extsub = sext i32 %subab to i64
+  %idx = getelementptr i32, i32* %ap, i64 %extsub
+  store i64 %addeaeb, i64* %result
+  store i32 %subab, i32* %idx
+  ret void
+}