[SeparateConstOffsetFromGEP] Fix bug handling negative offsets

[tstellar]

Summary

Fix bug constants and sub instructions

When finding constants in a chain starting with the RHS operator of
sub instructions, we were negating the constant before zero extending
it, which is incorrect.

Unfortunately, I was unable to find a simple way to implement this
transformation correctly, so for now I just disabled this optimization
for constants that feed into the RHS of a sub.

Resolves #62379

Transformation from alive2.llvm.org:

define i16 @src(i8 %a, i8 %b, i8 %c) {
entry:
%0 = sub nuw nsw i8 %c, %a
%1 = sub nuw nsw i8 %b, %0
%2 = zext i8 %1 to i16
ret i16 %2
}

Before/Bad:

define i16 @tgt(i8 %a, i8 %b, i8 %c) {
entry:
%0 = zext i8 %a to i16
%1 = zext i8 %b to i16
%c_neg = sub i8 0, %c
%c_zext = zext i8 %c_neg to i16
%2 = sub i16 0, %0
%3 = sub i16 %1, %2
%4 = add i16 %3, %c_zext
ret i16 %4
}

Correct:

define i16 @tgt(i8 %a, i8 %b, i8 %c) {
entry:
%0 = zext i8 %a to i16
%1 = zext i8 %b to i16
%c_zext = zext i8 %c to i16
%c_neg = sub i16 0, %c_zext
%2 = sub i16 0, %0
%3 = sub i16 %1, %2
%4 = add i16 %3, %c_neg
ret i16 %4
}

[nikic]

Haven't looked at the code change yet, but I'm pretty sure the test can use -passes='separate-const-offset-from-gep<lower-gep>' instead of adding a new opt flag.

[Peakulorain]

Transforming zext(A op B) to zext(A) op zext(B) should guarantee that A and B are non-negative. I think it would be better to do this fix during ConstantOffsetExtractor::CanTraceInto.

[tstellar]

In D149507#4307404, @Peakulorain wrote:

Transforming zext(A op B) to zext(A) op zext(B) should guarantee that A and B are non-negative. I think it would be better to do this fix during ConstantOffsetExtractor::CanTraceInto.

The problem I'm running into is that it's very difficult to to prove a value is non-negative, so adding those kinds of checks into CanTraceInto more or less disables the pass. The pass also makes an incorrect assumption that operands to an inbounds GEP are always non-negative. And it's not just assuming that direct operands to inbounds GEPS are non-negative, it assumes that any operands in the tree of values feeding into the GEP are non-negative.

[tstellar]

With some help from @nikic, I discovered that the bug is much more narrow in scope than I thought initially, so I should be able to do a more targeted fix.

[tstellar]

Use a more targetted fix.

[nikic]

This looks reasonable to me, but maybe wait a bit in case @Peakulorain has a better suggestion.

https://github.com/llvm/llvm-project/blob/1d0ccebcd725309399262af346494242b064e2ed/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp#L590-L593 is the code negating the offset.

[Peakulorain]

Since more knowledgeable people gived professional advice, it is also ok for me. :)

[tstellar]

Delete now dead code that was negating the constant.

[tstellar]

In D149507#4312067, @nikic wrote:

This looks reasonable to me, but maybe wait a bit in case @Peakulorain has a better suggestion.

https://github.com/llvm/llvm-project/blob/1d0ccebcd725309399262af346494242b064e2ed/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp#L590-L593 is the code negating the offset.

I actually meant to delete that code in this patch, so I've updated the patch to take it out.

[nikic]

In D149507#4312220, @tstellar wrote:

In D149507#4312067, @nikic wrote:

This looks reasonable to me, but maybe wait a bit in case @Peakulorain has a better suggestion.

https://github.com/llvm/llvm-project/blob/1d0ccebcd725309399262af346494242b064e2ed/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp#L590-L593 is the code negating the offset.

I actually meant to delete that code in this patch, so I've updated the patch to take it out.

Hm, why is it safe to delete that code? Don't we still need it in the case where we have no extension or a sign extension?

[tstellar]

Add back negation code.

[tstellar]

In D149507#4312221, @nikic wrote:

In D149507#4312220, @tstellar wrote:

In D149507#4312067, @nikic wrote:

This looks reasonable to me, but maybe wait a bit in case @Peakulorain has a better suggestion.

https://github.com/llvm/llvm-project/blob/1d0ccebcd725309399262af346494242b064e2ed/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp#L590-L593 is the code negating the offset.

I actually meant to delete that code in this patch, so I've updated the patch to take it out.

Hm, why is it safe to delete that code? Don't we still need it in the case where we have no extension or a sign extension?

Yes, I've added it back.

llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp

 //===- SeparateConstOffsetFromGEP.cpp -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 // Setting this flag may emit false positives when the input module already
 // contains dead instructions. Therefore, we set it only in unit tests that are
 // free of dead code.
 static cl::opt<bool>
     VerifyNoDeadCode("reassociate-geps-verify-no-dead-code", cl::init(false),
                      cl::desc("Verify this pass produces no dead code"),
                      cl::Hidden);
 
+static cl::opt<bool> ForceEnableLowerGEP(
+    "seperate-const-offset-from-gep-force-enable-lower-gep", cl::init(false),
+    cl::desc("Enable the LowerGEP feature even for targets that disable it."
+       "  This is mainly for testing purposes."),
+    cl::Hidden);
+
 namespace {
 
 /// A helper class for separating a constant offset from a GEP index.
 ///
 /// In real programs, a GEP index may be more complicated than a simple addition
 /// of something and a constant integer which can be trivially splitted. For
 /// example, to split ((a << 3) | 5) + b, we need to search deeper for the
 /// constant offset, so that we can separate the index to (a << 3) + b and 5.

 /// A pass that tries to split every GEP in the function into a variadic
 /// base and a constant offset. It is a FunctionPass because searching for the
 /// constant offset may inspect other basic blocks.
 class SeparateConstOffsetFromGEPLegacyPass : public FunctionPass {
 public:
   static char ID;
 
   SeparateConstOffsetFromGEPLegacyPass(bool LowerGEP = false)
-      : FunctionPass(ID), LowerGEP(LowerGEP) {
+      : FunctionPass(ID), LowerGEP(ForceEnableLowerGEP ? true : LowerGEP) {
     initializeSeparateConstOffsetFromGEPLegacyPassPass(
         *PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();

 /// base and a constant offset. It is a FunctionPass because searching for the
 /// constant offset may inspect other basic blocks.
 class SeparateConstOffsetFromGEP {
 public:
   SeparateConstOffsetFromGEP(
       DominatorTree *DT, ScalarEvolution *SE, LoopInfo *LI,
       TargetLibraryInfo *TLI,
       function_ref<TargetTransformInfo &(Function &)> GetTTI, bool LowerGEP)
-      : DT(DT), SE(SE), LI(LI), TLI(TLI), GetTTI(GetTTI), LowerGEP(LowerGEP) {}
+      : DT(DT), SE(SE), LI(LI), TLI(TLI), GetTTI(GetTTI),
+        LowerGEP(ForceEnableLowerGEP ? true : LowerGEP) {}
 
   bool run(Function &F);
 
 private:
   /// Tries to split the given GEP into a variadic base and a constant offset,
   /// and returns true if the splitting succeeds.
   bool splitGEP(GetElementPtrInst *GEP);
 
if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
     // sext(zext(a)) = zext(a). Verified in @sext_zext in split-gep.ll.
     //
     // Clear the NonNegative flag, because zext(a) >= 0 does not imply a >= 0.
     ConstantOffset =
         find(U->getOperand(0), /* SignExtended */ false,
              /* ZeroExtended */ true, /* NonNegative */ false).zext(BitWidth);
   }
 
+  // If SignExtend and ZeroExtend are both true then the zext and sext
+  // operations are in this order:
+  // zext ( sext ( a op b)).
+  // We know this, because if the sext came first then the zext instruction
+  // would have changed SignExtended to false.
+  //
+  // We can't handle negative offsets that are going to be zero_extended later.
+  // For example:
+  // %1 = add i32 -15, %0
+  // %2 = zext i32 %1 to i64
+  //
+  // Will be transformed to:
+  //
+  // %1 = zext i32 -15 to i64
+  // %2 = zext i32 %0 to i64
+  // %3 = add i64 %1, %2
+  //
+  // Which is not the same.  Instead of -15 + %0, we end up with
+  // 4294967281 + %0
+  // TODO: It should be possible to fix this pass to be able to correctly handle
+  // negative offsets.
+  if (ZeroExtended && !SignExtended && !isa<SExtInst>(V) && ConstantOffset.isNegative())
+    ConstantOffset = APInt(BitWidth, 0);
+
   // If we found a non-zero constant offset, add it to the path for
   // rebuildWithoutConstOffset. Zero is a valid constant offset, but doesn't
   // help this optimization.
   if (ConstantOffset != 0)
     UserChain.push_back(U);
   return ConstantOffset;
 }
 

llvm/test/Transforms/SeparateConstOffsetFromGEP/pr62379-zeroext-negative.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=separate-const-offset-from-gep  -seperate-const-offset-from-gep-force-enable-lower-gep < %s | FileCheck %s
+
+; These tests check that this pass does not incorrectly handle negative offsets.
+
+@c = internal constant [4 x i32] [i32 0, i32 1, i32 2, i32 3]
+
+define i32 @add_negative_rhs(i32 %a, i32 %b, ptr %ptr) {
+; CHECK-LABEL: @add_negative_rhs(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add nuw nsw i32 -15, [[B:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[TMP0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [4 x i32], ptr [[PTR:%.*]], i64 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = load i32, ptr [[TMP2]], align 4
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+entry:
+  %0 = add nuw nsw i32 -15, %b
+  %1 = zext i32 %0 to i64
+  %2 = getelementptr inbounds [4 x i32], ptr %ptr, i64 0, i64 %1
+  %3 = load i32, ptr %2
+  ret i32 %3
+}
+
+define i32 @add_negative_lhs(i32 %a, i32 %b, ptr %ptr) {
+; CHECK-LABEL: @add_negative_lhs(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add nuw nsw i32 [[B:%.*]], -15
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[TMP0]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [4 x i32], ptr [[PTR:%.*]], i64 0, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = load i32, ptr [[TMP2]], align 4
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+entry:
+  %0 = add nuw nsw i32 %b, -15
+  %1 = zext i32 %0 to i64
+  %2 = getelementptr inbounds [4 x i32], ptr %ptr, i64 0, i64 %1
+  %3 = load i32, ptr %2
+  ret i32 %3
+}
+
+
+define i32 @sub_positive(i32 %a, i32 %b, ptr %ptr) {
+; CHECK-LABEL: @sub_positive(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = sub nuw nsw i32 15, [[A:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw nsw i32 [[B:%.*]], [[TMP0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i32 [[TMP1]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds [4 x i32], ptr [[PTR:%.*]], i64 0, i64 [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = load i32, ptr [[TMP3]], align 4
+; CHECK-NEXT:    ret i32 [[TMP4]]
+;
+entry:
+  %0 = sub nuw nsw i32 15, %a
+  %1 = sub nuw nsw i32 %b, %0
+  %2 = zext i32 %1 to i64
+  %3 = getelementptr inbounds [4 x i32], ptr %ptr, i64 0, i64 %2
+  %4 = load i32, ptr %3
+  ret i32 %4
+}