[LSV] Attempt to fix comparison of APInt's with different bit widths.

[jlebar]

Summary

I have not been able to make a testcase for this, but this should be safe, and
it may fix the assertion being hit in https://reviews.llvm.org/D151630#4378572.

(We should absolutely figure out a testcase for this, either before or after
this is checked in.)

[tschuett]

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

1302

Drive-by nit. .has_value()should only be in unit tests. The .has_value() is redundant in this case.

[ronlieb]

I locally applied this to our downstream repos. CI passed. no issues seen on this patch.

[jlebar]

Thank you for checking. This morning I was able to come up with a reproducer.

[jlebar]

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

1302

Didn't end up using this; removed.

[jlebar]

@ronlieb you might want to double-check that it still works, because this patch is different than the first version. I'm pretty confident it should still work for you.

[ronlieb]

thx for fix, will re-apply patch and verify.

[ronlieb]

all good on newest patch

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

if (auto *SelectB = dyn_cast<SelectInst>(PtrB)) {
         return std::nullopt;
       std::optional<APInt> FalseDiff =
           getConstantOffset(SelectA->getFalseValue(), SelectB->getFalseValue(),
                             ContextInst, Depth);
       if (TrueDiff == FalseDiff)
         return TrueDiff;
     }
   }
   return std::nullopt;
 }
 
 EquivalenceClassMap
 Vectorizer::collectEquivalenceClasses(BasicBlock::iterator Begin,
                                       BasicBlock::iterator End) {
   EquivalenceClassMap Ret;
 
   auto getUnderlyingObject = [](const Value *Ptr) -> const Value * {
for (size_t J = 0; J < MaxChainsToTry && ChainIter != MRU.end();
       std::optional<APInt> Offset = getConstantOffset(
           getLoadStorePointerOperand(ChainIter->first),
           getLoadStorePointerOperand(I),
           /*ContextInst=*/
           (ChainIter->first->comesBefore(I) ? I : ChainIter->first));
       if (Offset.has_value()) {
         // `Offset` might not have the expected number of bits, if e.g. AS has a
         // different number of bits than opaque pointers.
-        ChainIter->second.push_back(
-            ChainElem{I, Offset.value().sextOrTrunc(ASPtrBits)});
+        ChainIter->second.push_back(ChainElem{I, Offset.value()});
         // Move ChainIter to the front of the MRU list.
         MRU.remove(*ChainIter);
         MRU.push_front(*ChainIter);
         MatchFound = true;
         break;
       }
     }
 

 }
 
 std::optional<APInt> Vectorizer::getConstantOffset(Value *PtrA, Value *PtrB,
                                                    Instruction *ContextInst,
                                                    unsigned Depth) {
   LLVM_DEBUG(dbgs() << "LSV: getConstantOffset, PtrA=" << *PtrA
                     << ", PtrB=" << *PtrB << ", ContextInst= " << *ContextInst
                     << ", Depth=" << Depth << "\n");
-  unsigned OffsetBitWidth = DL.getIndexTypeSizeInBits(PtrA->getType());
-  APInt OffsetA(OffsetBitWidth, 0);
-  APInt OffsetB(OffsetBitWidth, 0);
+  // We'll ultimately return a value of this bit width, even if computations
+  // happen in a different width.
+  unsigned OrigBitWidth = DL.getIndexTypeSizeInBits(PtrA->getType());
+  APInt OffsetA(OrigBitWidth, 0);
+  APInt OffsetB(OrigBitWidth, 0);
   PtrA = PtrA->stripAndAccumulateInBoundsConstantOffsets(DL, OffsetA);
   PtrB = PtrB->stripAndAccumulateInBoundsConstantOffsets(DL, OffsetB);
   unsigned NewPtrBitWidth = DL.getTypeStoreSizeInBits(PtrA->getType());
   if (NewPtrBitWidth != DL.getTypeStoreSizeInBits(PtrB->getType()))
     return std::nullopt;
 
   // If we have to shrink the pointer, stripAndAccumulateInBoundsConstantOffsets
   // should properly handle a possible overflow and the value should fit into
   // the smallest data type used in the cast/gep chain.
   assert(OffsetA.getSignificantBits() <= NewPtrBitWidth &&
          OffsetB.getSignificantBits() <= NewPtrBitWidth);
 
   OffsetA = OffsetA.sextOrTrunc(NewPtrBitWidth);
   OffsetB = OffsetB.sextOrTrunc(NewPtrBitWidth);
   if (PtrA == PtrB)
-    return OffsetB - OffsetA;
+    return (OffsetB - OffsetA).sextOrTrunc(OrigBitWidth);
 
   // Try to compute B - A.
   const SCEV *DistScev = SE.getMinusSCEV(SE.getSCEV(PtrB), SE.getSCEV(PtrA));
   if (DistScev != SE.getCouldNotCompute()) {
     LLVM_DEBUG(dbgs() << "LSV: SCEV PtrB - PtrA =" << *DistScev << "\n");
     ConstantRange DistRange = SE.getSignedRange(DistScev);
     if (DistRange.isSingleElement())
-      return OffsetB - OffsetA + *DistRange.getSingleElement();
+      return (OffsetB - OffsetA + *DistRange.getSingleElement())
+          .sextOrTrunc(OrigBitWidth);
   }
   std::optional<APInt> Diff =
       getConstantOffsetComplexAddrs(PtrA, PtrB, ContextInst, Depth);
   if (Diff.has_value())
-    return OffsetB - OffsetA + Diff->sext(OffsetB.getBitWidth());
+    return (OffsetB - OffsetA + Diff->sext(OffsetB.getBitWidth()))
+        .sextOrTrunc(OrigBitWidth);
   return std::nullopt;
 }

llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/vect-ptr-ptr-size-mismatch.ll

entry:
   %gep2 = getelementptr inbounds i32, ptr addrspace(5) %p, i64 1
   %a.ascast = addrspacecast ptr addrspace(5) %p to ptr
   %b.ascast = addrspacecast ptr addrspace(5) %gep2 to ptr
   %tmp1 = load i32, ptr %a.ascast, align 8
   %tmp2 = load i32, ptr %b.ascast, align 8
   unreachable
 }
 
+; CHECK-LABEL: @select_different_as
+; CHECK: load <2 x i32>
+define void @select_different_as(ptr addrspace(1) %p0, ptr addrspace(5) %q0, i1 %cond) {
+entry:
+  %p1 = getelementptr inbounds i32, ptr addrspace(1) %p0, i64 1
+  %q1 = getelementptr inbounds i32, ptr addrspace(5) %q0, i64 1
+  %p0.ascast = addrspacecast ptr addrspace(1) %p0 to ptr
+  %p1.ascast = addrspacecast ptr addrspace(1) %p1 to ptr
+  %q0.ascast = addrspacecast ptr addrspace(5) %q0 to ptr
+  %q1.ascast = addrspacecast ptr addrspace(5) %q1 to ptr
+  %sel0 = select i1 %cond, ptr %p0.ascast, ptr %q0.ascast
+  %sel1 = select i1 %cond, ptr %p1.ascast, ptr %q1.ascast
+  %tmp1 = load i32, ptr %sel0, align 8
+  %tmp2 = load i32, ptr %sel1, align 8
+  unreachable
+}
+
 ; CHECK-LABEL: @shrink_ptr
 ; CHECK: load <2 x i32>
 define void @shrink_ptr(ptr %p) {
 entry:
   %gep2 = getelementptr inbounds i32, ptr %p, i64 1
   %a.ascast = addrspacecast ptr %p to ptr addrspace(5)
   %b.ascast = addrspacecast ptr %gep2 to ptr addrspace(5)
   %tmp1 = load i32, ptr addrspace(5) %a.ascast, align 8