[InstructionSimplify][FIX] Handle bit widths in the presence of AS-casts

[jdoerfert]

Summary

While D64768 fixed some problems, it did not make sure we always keep
matching types for constant offsets in the presence of AS-casts.

The new approach is mostly keeping the constant offset bit width in sync
with the original value type (in the spirit of D118727). That said, if
we perform multiple strip operations we still need to adjust the types
if the first pointer strip causes a bit width change.

NOTE: Test is WIP.

[nikic]

Does https://github.com/llvm/llvm-project/commit/f35af77573d9e80bf6e61b3fdd20fe55191e962f fix your original problem?

[jdoerfert]

In D119740#3322389, @nikic wrote:

Does https://github.com/llvm/llvm-project/commit/f35af77573d9e80bf6e61b3fdd20fe55191e962f fix your original problem?

It makes it easier to fix but without testing it I would not assume the issue is gone.
The strip calls can still yield constants with different bit-width as they are matched to the bit-width of the resulting value.
So LHS/RHS that come out of the strip call can have different types and address spaces, and the constant offsets can have different bit-widths.

If we always return the constants with bit-widths matching the inputs (which I do here), the only remaining problem is that we call stripPointerCasts first, which causes the LHS/RHS to be "not in sync" when we call stripAndAccumulate. (see my comment in the code).

llvm/lib/Analysis/InstructionSimplify.cpp

2577

^Here

[nikic]

In D119740#3323013, @jdoerfert wrote:

In D119740#3322389, @nikic wrote:

Does https://github.com/llvm/llvm-project/commit/f35af77573d9e80bf6e61b3fdd20fe55191e962f fix your original problem?

It makes it easier to fix but without testing it I would not assume the issue is gone.
The strip calls can still yield constants with different bit-width as they are matched to the bit-width of the resulting value.
So LHS/RHS that come out of the strip call can have different types and address spaces, and the constant offsets can have different bit-widths.

The reason why I'd expect it to be fixed is that RHSOffset and LHSOffset only get used together if RHS == LHS now, in which case the type of the (offset-stripped) pointers match and as such should also have the same index size.

If we always return the constants with bit-widths matching the inputs (which I do here), the only remaining problem is that we call stripPointerCasts first, which causes the LHS/RHS to be "not in sync" when we call stripAndAccumulate. (see my comment in the code).

Not just that, but computePointerICmp() itself may be called on different types -- apparently, it looks through ptrtoint.

[jdoerfert]

In D119740#3325545, @nikic wrote:

In D119740#3323013, @jdoerfert wrote:

In D119740#3322389, @nikic wrote:

Does https://github.com/llvm/llvm-project/commit/f35af77573d9e80bf6e61b3fdd20fe55191e962f fix your original problem?

It makes it easier to fix but without testing it I would not assume the issue is gone.
The strip calls can still yield constants with different bit-width as they are matched to the bit-width of the resulting value.
So LHS/RHS that come out of the strip call can have different types and address spaces, and the constant offsets can have different bit-widths.

The reason why I'd expect it to be fixed is that RHSOffset and LHSOffset only get used together if RHS == LHS now, in which case the type of the (offset-stripped) pointers match and as such should also have the same index size.

I think you are right.

llvm/lib/Analysis/InstructionSimplify.cpp

 #include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstSimplifyFolder.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/OverflowInstAnalysis.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
+#include "llvm/IR/Constant.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Operator.h"
static Value *SimplifyAddInst(Value *Op0, Value *Op1, bool IsNSW, bool IsNUW,
   return nullptr;
 }
 
 Value *llvm::SimplifyAddInst(Value *Op0, Value *Op1, bool IsNSW, bool IsNUW,
                              const SimplifyQuery &Query) {
   return ::SimplifyAddInst(Op0, Op1, IsNSW, IsNUW, Query, RecursionLimit);
 }
 
+static Type *getBiggestTypeIfDifferent(const DataLayout &DL,
+                                       ArrayRef<Constant *> Constants) {
+  assert(Constants.size() > 1 && "Expect at least two constant values!");
+  bool AllEqual = true;
+  Type *MaxTy = Constants.front()->getType();
+  TypeSize MaxTS = DL.getTypeSizeInBits(MaxTy);
+  for (Constant *C : Constants.drop_front()) {
+    Type *Ty = C->getType();
+    TypeSize TS = DL.getTypeSizeInBits(Ty);
+    AllEqual &= Ty == MaxTy;
+    if (TS <= MaxTS)
+      continue;
+    MaxTy = Ty;
+    MaxTS = TS;
+  }
+  return AllEqual ? nullptr : MaxTy;
+}
+
 /// Compute the base pointer and cumulative constant offsets for V.
 ///
 /// This strips all constant offsets off of V, leaving it the base pointer, and
 /// accumulates the total constant offset applied in the returned constant. It
 /// returns 0 if V is not a pointer, and returns the constant '0' if there are
 /// no constant offsets applied.
 ///
 /// This is very similar to GetPointerBaseWithConstantOffset except it doesn't
 /// follow non-inbounds geps. This allows it to remain usable for icmp ult/etc.
 /// folding.
+///
+/// Note that this may trace through `addrspacecast` and the offset bit width
+/// might not match the stripped value type. Instead, it will match the input
+/// value type index size.
 static Constant *stripAndComputeConstantOffsets(const DataLayout &DL, Value *&V,
                                                 bool AllowNonInbounds = false) {
   assert(V->getType()->isPtrOrPtrVectorTy());
 
   APInt Offset = APInt::getZero(DL.getIndexTypeSizeInBits(V->getType()));
 
   V = V->stripAndAccumulateConstantOffsets(DL, Offset, AllowNonInbounds);
-  // As that strip may trace through `addrspacecast`, need to sext or trunc
-  // the offset calculated.
-  Type *IntIdxTy = DL.getIndexType(V->getType())->getScalarType();
-  Offset = Offset.sextOrTrunc(IntIdxTy->getIntegerBitWidth());
 
-  Constant *OffsetIntPtr = ConstantInt::get(IntIdxTy, Offset);
+  Constant *OffsetIntPtr = ConstantInt::get(V->getContext(), Offset);
   if (VectorType *VecTy = dyn_cast<VectorType>(V->getType()))
     return ConstantVector::getSplat(VecTy->getElementCount(), OffsetIntPtr);
   return OffsetIntPtr;
 }
 
 /// Compute the constant difference between two pointer values.
 /// If the difference is not a constant, returns zero.
 static Constant *computePointerDifference(const DataLayout &DL, Value *LHS,
                                           Value *RHS) {
   Constant *LHSOffset = stripAndComputeConstantOffsets(DL, LHS);
   Constant *RHSOffset = stripAndComputeConstantOffsets(DL, RHS);
+  // Note that we do not need to adjust the offset bit widths as they match
+  // LHS/RHS which are expected to have equal type.
+  assert(LHSOffset->getType() == RHSOffset->getType() &&
+         "Stripped offset types should match!");
 
   // If LHS and RHS are not related via constant offsets to the same base
   // value, there is nothing we can do here.
   if (LHS != RHS)
     return nullptr;
 
   // Otherwise, the difference of LHS - RHS can be computed as:
   //    LHS - RHS
computePointerICmp(CmpInst::Predicate Pred, Value *LHS, Value *RHS,
   const DataLayout &DL = Q.DL;
   const TargetLibraryInfo *TLI = Q.TLI;
   const DominatorTree *DT = Q.DT;
   const Instruction *CxtI = Q.CxtI;
   const InstrInfoQuery &IIQ = Q.IIQ;
 
   // First, skip past any trivial no-ops.
   LHS = LHS->stripPointerCasts();
   RHS = RHS->stripPointerCasts();
 
   // A non-null pointer is not equal to a null pointer.
   if (isa<ConstantPointerNull>(RHS) && ICmpInst::isEquality(Pred) &&
       llvm::isKnownNonZero(LHS, DL, 0, nullptr, nullptr, nullptr,
                            IIQ.UseInstrInfo))
     return ConstantInt::get(GetCompareTy(LHS),
                             !CmpInst::isTrueWhenEqual(Pred));
 
computePointerICmp(CmpInst::Predicate Pred, Value *LHS, Value *RHS,
   // Strip off any constant offsets so that we can reason about them.
   // It's tempting to use getUnderlyingObject or even just stripInBoundsOffsets
   // here and compare base addresses like AliasAnalysis does, however there are
   // numerous hazards. AliasAnalysis and its utilities rely on special rules
   // governing loads and stores which don't apply to icmps. Also, AliasAnalysis
   // doesn't need to guarantee pointer inequality when it says NoAlias.
   Constant *LHSOffset = stripAndComputeConstantOffsets(DL, LHS);
   Constant *RHSOffset = stripAndComputeConstantOffsets(DL, RHS);
+  // Since we started with stripped versions of LHS/RHS the inbounds offsets
+  // might not have the same bit widths.
+  if (Type *MaxTy = getBiggestTypeIfDifferent(DL, {LHSOffset, RHSOffset})) {
+    if (LHSOffset->getType() != MaxTy)
+      LHSOffset = ConstantExpr::getZExt(LHSOffset, MaxTy);
+    if (RHSOffset->getType() != MaxTy)
+      RHSOffset = ConstantExpr::getZExt(RHSOffset, MaxTy);
+  }
+  assert(LHSOffset->getType() == RHSOffset->getType() &&
+         "Stripped offset types should match!");
 
   // If LHS and RHS are related via constant offsets to the same base
   // value, we can replace it with an icmp which just compares the offsets.
   if (LHS == RHS)
     return ConstantExpr::getICmp(Pred, LHSOffset, RHSOffset);
 
   // Various optimizations for (in)equality comparisons.
   if (Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE) {
if (Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE) {
     }
 
     // Even if an non-inbounds GEP occurs along the path we can still optimize
     // equality comparisons concerning the result. We avoid walking the whole
     // chain again by starting where the last calls to
     // stripAndComputeConstantOffsets left off and accumulate the offsets.
     Constant *LHSNoBound = stripAndComputeConstantOffsets(DL, LHS, true);
     Constant *RHSNoBound = stripAndComputeConstantOffsets(DL, RHS, true);
-    if (LHS == RHS)
+    if (LHS == RHS) {
+      // Since we started with stripped versions of LHS/RHS the no-inbounds
+      // offsets might not have the same bit widths. We adjust all constants to
+      // the biggest type among them to ensure constant expressions below use
+      // equal-typed operands.
+      if (Type *MaxTy = getBiggestTypeIfDifferent(
+              DL, {LHSOffset, LHSNoBound, RHSNoBound})) {
+        if (LHSOffset->getType() != MaxTy)
+          LHSOffset = ConstantExpr::getZExt(LHSOffset, MaxTy);
+        if (RHSOffset->getType() != MaxTy)
+          RHSOffset = ConstantExpr::getZExt(RHSOffset, MaxTy);
+        if (LHSNoBound->getType() != MaxTy)
+          LHSNoBound = ConstantExpr::getZExt(LHSNoBound, MaxTy);
+        if (RHSNoBound->getType() != MaxTy)
+          RHSNoBound = ConstantExpr::getZExt(RHSNoBound, MaxTy);
+      }
       return ConstantExpr::getICmp(Pred,
                                    ConstantExpr::getAdd(LHSOffset, LHSNoBound),
                                    ConstantExpr::getAdd(RHSOffset, RHSNoBound));
+    }
 
     // If one side of the equality comparison must come from a noalias call
     // (meaning a system memory allocation function), and the other side must
     // come from a pointer that cannot overlap with dynamically-allocated
     // memory within the lifetime of the current function (allocas, byval
     // arguments, globals), then determine the comparison result here.
     SmallVector<const Value *, 8> LHSUObjs, RHSUObjs;
     getUnderlyingObjects(LHS, LHSUObjs);