[SeparateConstOffsetFromGEP] Fix: `b - a` matched `a - b` during reuniteExts

[Peakulorain]

Summary

During the SeparateConstOffsetFromGEP pass, a - b and b - a will be
considered equivalent in some instances.

An example- the IR contains:

BB1:
    %add = add %a, 511
    br label %BB2
BB2:
    %sub2 = sub %b,  %a
    br label %BB3
BB3:
    %sub1 = sub %add, %b
    %gep = getelementptr float, ptr %p, %sub1

Step 1 in the SeparateConstOffsetFromGEP pass, after split constant index:

BB1:
    %add = add %a, 511
    br label %BB2
BB2:
    %sub2 = sub %b,  %a
    br label %BB3
BB3:
    %sub.t = sub %a, %b
    %gep.base = getelementptr float, ptr %p, %sub.t
    %gep = getelementptr float, ptr %gep.base, 511

Step 2, after reuniteExts:

BB1:
    br label %BB2
BB2:
    %sub2 = sub %b,  %a
    br label %BB3
BB3:
    %gep.base = getelementptr float, ptr %p, %sub2
    %gep = getelementptr float, ptr %gep.base, 511

Obviously, reuniteExts treated a - b and b - a as equivalent.
This patch fixes that.

[Peakulorain]

We triggered this bug on the aarch64 target, since aarch64 has enabled LowerGEP.
Pls use command line in the case split-gep-sub.ll of my patch to trigger this problem.

[spatel]

Please upload with more context - use something like "git diff -U9999".

llvm/test/Transforms/SeparateConstOffsetFromGEP/AArch64/split-gep-sub.ll

1 ↗

(On Diff #495691)

It would be better to use a RUN line with opt like existing tests in this and the parent directory.
IIUC, we have a miscompile, so you can create a baseline test that shows the bug, commit that to main, and then regenerate the CHECK lines with this patch, so we show the difference in IR. Please use "utils/update_test_checks.py" to auto-generate the CHECK lines instead of manually writing them.

[Peakulorain]

In D143542#4123076, @spatel wrote:

Please upload with more context - use something like "git diff -U9999".

Your reply is appreciated. First step to show this miscompile, we should add a flag to enable LowerGEP manually.
So, see the first patch pls: https://reviews.llvm.org/D143980
I will update after the above patch is merged.

[Peakulorain]

In D143542#4123076, @spatel wrote:

Please upload with more context - use something like "git diff -U9999".

Thanks, uploaded again.

[nikic]

LGTM. Using SCEV expressions as a deduplication key seems kind of odd in the first place, but I guess for adds it does provide some value, because you get complexity canonicalization for free. Of course, for subs we don't want that, and this effectively prevents it.

llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-sub.ll

3–4

Please include a comment on what this is testing, it's not particularly obvious.

[spatel]

More knowledgeable reviewers have approved, so just making it clear that there's no need to wait for more comments from me. :)

[Peakulorain]

In D143542#4125864, @spatel wrote:

More knowledgeable reviewers have approved, so just making it clear that there's no need to wait for more comments from me. :)

Also thanks a lot for your review. :)

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

if (LHS->getType() == RHS->getType()) {
         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));
+      const SCEV *Key = SE->getAddExpr(
+          SE->getUnknown(LHS), SE->getNegativeSCEV(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)))) {
     if (programUndefinedIfPoison(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 (programUndefinedIfPoison(I)) {
-      const SCEV *Key =
-          SE->getAddExpr(SE->getUnknown(LHS), SE->getUnknown(RHS));
+      const SCEV *Key = SE->getAddExpr(
+          SE->getUnknown(LHS), SE->getNegativeSCEV(SE->getUnknown(RHS)));
       DominatingSubs[Key].push_back(I);
     }
   }
   return false;
 }
 
 bool SeparateConstOffsetFromGEP::reuniteExts(Function &F) {
   bool Changed = false;

llvm/test/Transforms/SeparateConstOffsetFromGEP/split-gep-sub.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes="separate-const-offset-from-gep<lower-gep>" < %s | FileCheck %s
 
 ; Check that GEP with an index 'A - B + [ConstantInt]' will be split into two
 ; GEPs. eg.
 ;     %A = ...
 ;     %B = ...
 ;     %sub = %A - %B
 ;     %idx = %sub + 10
 ;     %gep = getelementptr int, ptr %p, %idx
 ;  will be transformed into:
 ;     %A = ...

 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[K:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[COND_END:%.*]] ]
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw nsw i32 [[K]], 5
 ; CHECK-NEXT:    [[SUB1:%.*]] = sub nsw i32 [[MUL]], [[REM]]
 ; CHECK-NEXT:    [[CMP26:%.*]] = icmp ult i32 [[SUB1]], 512
 ; CHECK-NEXT:    br i1 [[CMP26]], label [[COND_TRUE:%.*]], label [[COND_END]]
 ; CHECK:       cond.true:
-; CHECK-NEXT:    [[SUB22:%.*]] = sext i32 [[SUB1]] to i64
-; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[P:%.*]] to i64
-; CHECK-NEXT:    [[TMP2:%.*]] = shl i64 [[SUB22]], 2
-; CHECK-NEXT:    [[TMP3:%.*]] = add i64 [[TMP1]], [[TMP2]]
-; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP3]], 2044
-; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to ptr
-; CHECK-NEXT:    store float 1.000000e+00, ptr [[TMP5]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[MUL]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i32 [[REM]] to i64
+; CHECK-NEXT:    [[SUB22:%.*]] = sub i64 [[TMP2]], [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = ptrtoint ptr [[P:%.*]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = shl i64 [[SUB22]], 2
+; CHECK-NEXT:    [[TMP5:%.*]] = add i64 [[TMP3]], [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = add i64 [[TMP5]], 2044
+; CHECK-NEXT:    [[TMP7:%.*]] = inttoptr i64 [[TMP6]] to ptr
+; CHECK-NEXT:    store float 1.000000e+00, ptr [[TMP7]], align 4
 ; CHECK-NEXT:    br label [[COND_END]]
 ; CHECK:       cond.end:
 ; CHECK-NEXT:    [[INC]] = add nuw nsw i32 [[K]], 1
 ; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[INC]], 100
 ; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;