[LSV] Fix the ContextInst for computeKnownBits.

[jlebar]

Summary

Previously we used the later of GEPA or GEPB. This is hacky because
really we should be using the later of the two load/store instructions
being considered. But also it's flat-out incorrect, because GEPA and
GEPB might be in different BBs, in which case we cannot ask which one
comes last (assertion failure,
https://reviews.llvm.org/D149893#4378332).

Fixed, now we use the correct context instruction.

[ronlieb]

thanks for the very rapid fix, i will try it out , and should know in about 1.5 hours. will update here.
in tghe meantime if you think its fine , land it.

[jlebar]

Fix a typo while we're here.

[jlebar]

in the meantime if you think its fine , land it.

OK, will do. I think this is demonstrably more correct than it was before. :) Hopefully we don't have two bugs here.

[ronlieb]

looks like there is a 2nd assert :
llvm/include/llvm/ADT/APInt.h:1039: bool llvm::APInt::operator==(const llvm::APInt&) const: Assertion `BitWidth == RHS.BitWidth && "Comparison requires equal bit widths"' failed.

[jlebar]

Hahaha all right. I am headed out for the day, but will look when I'm back.

[ronlieb]

thanks for looking into it. enjoy rest of your day

[jlebar]

In D151630#4378433, @ronlieb wrote:

thanks for looking into it. enjoy rest of your day

I'm not seeing how this can happen. :(

I see two places where I call APInt::operator==:

• bool AreContiguous = It->OffsetFromLeader == PrevReadEnd;
• if (IOffset == ChainElemOffset ...)

In the first case, the two APInts are from the OffsetFromLeader values from the same Chain. There are two places where a ChainElem is created (search for ChainElem{), and the OffsetFromLeader values explicitly have the same bitwidth.

In the second case, IOffset and ChainElemOffset come from the ChainOffsets hashtable, which is created using the ChainElem::OffsetFromLeader field, so here too they should have the same bitwidth.

Perhaps the assertion is happening in another file (still is probably my bug!), or perhaps I'm missing the call to APInt::operator==. Or perhaps I'm misunderstanding my invariants somehow.

A reproducer, or at least a stacktrace, would be very helpful.

Sorry for the trouble, and thanks for helping me hunt this down.

[ronlieb]

Stack dump:
0. Program arguments: /opt/rocm-5.6.0-12074/llvm/bin/clang-17 -cc1 -triple amdgcn-amd-amdhsa -aux-triple x86_64-unknown-linux-gnu -emit-obj -disable-free -clear-ast-before-backend -main-file-name large_types.cu -mrelocation-model pic -pic-level 2 -fhalf-no-semantic-interposition -mframe-pointer=none -fno-rounding-math -mconstructor-aliases -aux-target-cpu x86-64 -fcuda-is-device -mllvm -amdgpu-internalize-symbols -fcuda-allow-variadic-functions -fvisibility=hidden -fapply-global-visibility-to-externs -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/hip.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/ocml.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/ockl.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_daz_opt_off.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_unsafe_math_off.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_finite_only_off.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_correctly_rounded_sqrt_on.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_wavefrontsize64_on.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_isa_version_906.bc -mlink-builtin-bitcode /opt/rocm-5.6.0-12074/amdgcn/bitcode/oclc_abi_version_400.bc -target-cpu gfx906 -target-feature -xnack -debugger-tuning=gdb -resource-dir /opt/rocm-5.6.0-12074/llvm/lib/clang/17.0.0 -dependency-file CMakeFiles/test_thrust_exclusive_large_types.dir/large_types.cu.o.d -MT testing/async/exclusive_scan/CMakeFiles/test_thrust_exclusive_large_types.dir/large_types.cu.o -sys-header-deps -canonical-system-headers -internal-isystem /opt/rocm-5.6.0-12074/llvm/lib/clang/17.0.0/include/cuda_wrappers -idirafter /opt/rocm-5.6.0-12074/include -include clang_hip_runtime_wrapper.h -c-isystem /opt/rocm-5.6.0-12074/llvm/include/gpu-none-llvm -isystem /opt/rocm-5.6.0-12074/include -isystem /jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/testing -D USE_PROF_API=1 -D HIP_PLATFORM_AMD=1 -D HIP_PLATFORM_HCC=1 -I /jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/build/release/thrust/include -I /jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/thrust/.. -D NDEBUG -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/x86_64-linux-gnu/c++/9 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/backward -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/x86_64-linux-gnu/c++/9 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../include/c++/9/backward -internal-isystem /opt/rocm-5.6.0-12074/llvm/lib/clang/17.0.0/include -internal-isystem /usr/local/include -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../x86_64-linux-gnu/include -internal-externc-isystem /usr/include/x86_64-linux-gnu -internal-externc-isystem /include -internal-externc-isystem /usr/include -internal-isystem /opt/rocm-5.6.0-12074/llvm/lib/clang/17.0.0/include -internal-isystem /usr/local/include -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/9/../../../../x86_64-linux-gnu/include -internal-externc-isystem /usr/include/x86_64-linux-gnu -internal-externc-isystem /include -internal-externc-isystem /usr/include -O3 -Wno-unused-command-line-argument -Wall -Wextra -std=c++14 -fdeprecated-macro -fno-autolink -fdebug-compilation-dir=/jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/build/release/testing/async/exclusive_scan -ferror-limit 19 -pthread -fhip-new-launch-api -fgnuc-version=4.2.1 -fcxx-exceptions -fexceptions -vectorize-loops -vectorize-slp -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -cuid=c55cff91efbd2d1d -fcuda-allow-variadic-functions -faddrsig -DGCC_HAVE_DWARF2_CFI_ASM=1 -o /tmp/large_types-gfx906:xnack--f0faeb.o -x hip /jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/testing/async/exclusive_sca

1. <eof> parser at end of file
2. Code generation
3. Running pass 'CallGraph Pass Manager' on module '/jenkins/workspace/compiler-psdb-amd-stg-open_4/Libs/rocThrust/testing/async/exclusive_scan/large_types.cu'.
4. Running pass 'GPU Load and Store Vectorizer' on function '@_ZN7rocprim6detail20lookback_scan_kernelILb1ENS0_19default_scan_configILj0E11FixedVectorIiLj8EEEEN6thrust6detail15normal_iteratorINS6_10device_ptrIKS4_EEEENS8_INS9_IS4_EEEENS6_7maximumIvEES4_NS0_19lookback_scan_stateIS4_Lb1ELb0EEEEEvT1_T2_mT4_T3_T5_jNS0_16ordered_block_idIjEEPNS0_18input_value_traitsISL_E10value_typeEST_bb'

#0 0x0000563c3056d52f llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x26d652f)
#1 0x0000563c3056ac54 SignalHandler(int) Signals.cpp:0:0
#2 0x00007f02c29a3420 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x14420)
#3 0x00007f02c244000b raise (/lib/x86_64-linux-gnu/libc.so.6+0x4300b)
#4 0x00007f02c241f859 abort (/lib/x86_64-linux-gnu/libc.so.6+0x22859)
#5 0x00007f02c241f729 (/lib/x86_64-linux-gnu/libc.so.6+0x22729)
#6 0x00007f02c2430fd6 (/lib/x86_64-linux-gnu/libc.so.6+0x33fd6)
#7 0x0000563c2f1d5fca (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x133efca)
#8 0x0000563c307592b3 (anonymous namespace)::Vectorizer::getConstantOffset(llvm::Value*, llvm::Value*, llvm::Instruction*, unsigned int) LoadStoreVectorizer.cpp:0:0
#9 0x0000563c30759ca6 (anonymous namespace)::Vectorizer::gatherChains(llvm::ArrayRef<llvm::Instruction*>) LoadStoreVectorizer.cpp:0:0

#10 0x0000563c3075d5b8 (anonymous namespace)::Vectorizer::run() LoadStoreVectorizer.cpp:0:0
#11 0x0000563c30760362 (anonymous namespace)::LoadStoreVectorizerLegacyPass::runOnFunction(llvm::Function&) (.part.0) LoadStoreVectorizer.cpp:0:0
#12 0x0000563c2fecea81 llvm::FPPassManager::runOnFunction(llvm::Function&) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x2037a81)
#13 0x0000563c2f4fa8b7 (anonymous namespace)::CGPassManager::runOnModule(llvm::Module&) CallGraphSCCPass.cpp:0:0
#14 0x0000563c2fecf542 llvm::legacy::PassManagerImpl::run(llvm::Module&) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x2038542)
#15 0x0000563c307ea1c6 clang::EmitBackendOutput(clang::DiagnosticsEngine&, clang::HeaderSearchOptions const&, clang::CodeGenOptions const&, clang::TargetOptions const&, clang::LangOptions const&, llvm::StringRef, llvm::Module*, clang::BackendAction, llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>, std::unique_ptr<llvm::raw_pwrite_stream, std::default_delete<llvm::raw_pwrite_stream>>) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x29531c6)
#16 0x0000563c31809912 clang::BackendConsumer::HandleTranslationUnit(clang::ASTContext&) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x3972912)
#17 0x0000563c32989bd9 clang::ParseAST(clang::Sema&, bool, bool) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x4af2bd9)
#18 0x0000563c318082b8 clang::CodeGenAction::ExecuteAction() (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x39712b8)
#19 0x0000563c3104c789 clang::FrontendAction::Execute() (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x31b5789)
#20 0x0000563c30fcfb16 clang::CompilerInstance::ExecuteAction(clang::FrontendAction&) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x3138b16)
#21 0x0000563c311312e7 clang::ExecuteCompilerInvocation(clang::CompilerInstance*) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0x329a2e7)
#22 0x0000563c2eb78a04 cc1_main(llvm::ArrayRef<char const*>, char const*, void*) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0xce1a04)
#23 0x0000563c2eb748aa ExecuteCC1Tool(llvm::SmallVectorImpl<char const*>&, llvm::ToolContext const&) driver.cpp:0:0
#24 0x0000563c2eb75d2c clang_main(int, char**, llvm::ToolContext const&) (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0xcded2c)
#25 0x0000563c2eac0a45 main (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0xc29a45)
#26 0x00007f02c2421083 __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x24083)
#27 0x0000563c2eb6ef3e _start (/opt/rocm-5.6.0-12074/llvm/bin/clang-17+0xcd7f3e)

[jlebar]

I still have no idea how this could be happening, but I hypothesize 50/50 chance https://reviews.llvm.org/D151640 will fix it for you. I have not been able to write a testcase. :(

[ronlieb]

subsequent testing has shown following assert in our builds of qmcpack and spec cpu2017 538.imagick.

llvm/lib/IR/Constants.cpp:3070: llvm::APInt llvm::ConstantDataSequential::getElementAsAPInt(unsigned int) const: Assertion `isa<IntegerType>(getElementType()) && "Accessor can only be used when element is an integer"' failed.

looking for a small reproducer ...

[ronlieb]

small.ll3 KBDownload

small test case
llc < small.ll

[jlebar]

In D151630#4380613, @ronlieb wrote:

small.ll3 KBDownload

small test case
llc < small.ll

This is working on my machine at revision f81f32adc9a8 (that's the newest I have access to at the moment), and also it doesn't seem to invoke LSV.

LMK if I'm holding it wrong?

[ronlieb]

sorry, i meant to try my trunk build before posting it, forgot. you can disregard, its on my end.
thx for confiorming,

[ronlieb]

subsequent merge cleaned up our downstream issue. all good.

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

private:
   std::vector<Chain> splitChainByAlignment(Chain &C);
 
   /// Converts the instrs in the chain into a single vectorized load or store.
   /// Adds the old scalar loads/stores to ToErase.
   bool vectorizeChain(Chain &C);
 
   /// Tries to compute the offset in bytes PtrB - PtrA.
   std::optional<APInt> getConstantOffset(Value *PtrA, Value *PtrB,
+                                         Instruction *ContextInst,
                                          unsigned Depth = 0);
-  std::optional<APInt> gtConstantOffsetComplexAddrs(Value *PtrA, Value *PtrB,
+  std::optional<APInt> getConstantOffsetComplexAddrs(Value *PtrA, Value *PtrB,
+                                                     Instruction *ContextInst,
                                                      unsigned Depth);
   std::optional<APInt> getConstantOffsetSelects(Value *PtrA, Value *PtrB,
+                                                Instruction *ContextInst,
                                                 unsigned Depth);
 
   /// Gets the element type of the vector that the chain will load or store.
   /// This is nontrivial because the chain may contain elements of different
   /// types; e.g. it's legal to have a chain that contains both i32 and float.
   Type *getChainElemTy(const Chain &C);
 
   /// Determines whether ChainElem can be moved up (if IsLoad) or down (if
if (OtherInstrA && OtherInstrB &&
       if (OtherInstrA->getOperand(0) == OtherInstrB->getOperand(0) &&
           IdxDiff.getSExtValue() == (CstValB - CstValA))
         return true;
     }
   }
   return false;
 }
 
-std::optional<APInt> Vectorizer::gtConstantOffsetComplexAddrs(Value *PtrA,
-                                                              Value *PtrB,
-                                                              unsigned Depth) {
-  LLVM_DEBUG(dbgs() << "LSV: gtConstantOffsetComplexAddrs PtrA=" << *PtrA
-                    << " PtrB=" << *PtrB << " Depth=" << Depth << "\n");
+std::optional<APInt> Vectorizer::getConstantOffsetComplexAddrs(
+    Value *PtrA, Value *PtrB, Instruction *ContextInst, unsigned Depth) {
+  LLVM_DEBUG(dbgs() << "LSV: getConstantOffsetComplexAddrs PtrA=" << *PtrA
+                    << " PtrB=" << *PtrB << " ContextInst=" << *ContextInst
+                    << " Depth=" << Depth << "\n");
   auto *GEPA = dyn_cast<GetElementPtrInst>(PtrA);
   auto *GEPB = dyn_cast<GetElementPtrInst>(PtrB);
   if (!GEPA || !GEPB)
-    return getConstantOffsetSelects(PtrA, PtrB, Depth);
+    return getConstantOffsetSelects(PtrA, PtrB, ContextInst, Depth);
 
   // Look through GEPs after checking they're the same except for the last
   // index.
   if (GEPA->getNumOperands() != GEPB->getNumOperands() ||
       GEPA->getPointerOperand() != GEPB->getPointerOperand())
     return std::nullopt;
   gep_type_iterator GTIA = gep_type_begin(GEPA);
   gep_type_iterator GTIB = gep_type_begin(GEPB);
std::optional<APInt> Vectorizer::getConstantOffsetComplexAddrs(
   if (IdxDiffSCEV == SE.getCouldNotCompute())
     return std::nullopt;
 
   ConstantRange IdxDiffRange = SE.getSignedRange(IdxDiffSCEV);
   if (!IdxDiffRange.isSingleElement())
     return std::nullopt;
   APInt IdxDiff = *IdxDiffRange.getSingleElement();
 
-  LLVM_DEBUG(dbgs() << "LSV: gtConstantOffsetComplexAddrs IdxDiff=" << IdxDiff
+  LLVM_DEBUG(dbgs() << "LSV: getConstantOffsetComplexAddrs IdxDiff=" << IdxDiff
                     << "\n");
 
   // Now we need to prove that adding IdxDiff to ValA won't overflow.
   bool Safe = false;
 
   // First attempt: if OpB is an add with NSW/NUW, and OpB is IdxDiff added to
   // ValA, we're okay.
   if (OpB->getOpcode() == Instruction::Add &&
std::optional<APInt> Vectorizer::getConstantOffsetComplexAddrs(
   // Assuming IdxDiff is positive: If all set bits of IdxDiff or any higher
   // order bit other than the sign bit are known to be zero in ValA, we can add
   // Diff to it while guaranteeing no overflow of any sort.
   //
   // If IdxDiff is negative, do the same, but swap ValA and ValB.
   if (!Safe) {
     // When computing known bits, use the GEPs as context instructions, since
     // they likely are in the same BB as the load/store.
-    Instruction *ContextInst = GEPA->comesBefore(GEPB) ? GEPB : GEPA;
     KnownBits Known(BitWidth);
     computeKnownBits((IdxDiff.sge(0) ? ValA : OpB), Known, DL, 0, &AC,
                      ContextInst, &DT);
     APInt BitsAllowedToBeSet = Known.Zero.zext(IdxDiff.getBitWidth());
     if (Signed)
       BitsAllowedToBeSet.clearBit(BitWidth - 1);
     if (BitsAllowedToBeSet.ult(IdxDiff.abs()))
       return std::nullopt;
     Safe = true;
   }
 
   if (Safe)
     return IdxDiff * Stride;
   return std::nullopt;
 }
 
-std::optional<APInt>
-Vectorizer::getConstantOffsetSelects(Value *PtrA, Value *PtrB, unsigned Depth) {
+std::optional<APInt> Vectorizer::getConstantOffsetSelects(
+    Value *PtrA, Value *PtrB, Instruction *ContextInst, unsigned Depth) {
   if (Depth++ == MaxDepth)
     return std::nullopt;
 
   if (auto *SelectA = dyn_cast<SelectInst>(PtrA)) {
     if (auto *SelectB = dyn_cast<SelectInst>(PtrB)) {
       if (SelectA->getCondition() != SelectB->getCondition())
         return std::nullopt;
       LLVM_DEBUG(dbgs() << "LSV: getConstantOffsetSelects, PtrA=" << *PtrA
-                        << ", PtrB=" << *PtrB << ", Depth=" << Depth << "\n");
+                        << ", PtrB=" << *PtrB << ", ContextInst="
+                        << *ContextInst << ", Depth=" << Depth << "\n");
       std::optional<APInt> TrueDiff = getConstantOffset(
-          SelectA->getTrueValue(), SelectB->getTrueValue(), Depth);
+          SelectA->getTrueValue(), SelectB->getTrueValue(), ContextInst, Depth);
       if (!TrueDiff.has_value())
         return std::nullopt;
-      std::optional<APInt> FalseDiff = getConstantOffset(
-          SelectA->getFalseValue(), SelectB->getFalseValue(), Depth);
+      std::optional<APInt> FalseDiff =
+          getConstantOffset(SelectA->getFalseValue(), SelectB->getFalseValue(),
+                            ContextInst, Depth);
       if (TrueDiff == FalseDiff)
         return TrueDiff;
     }
   }
   return std::nullopt;
 }
 
 EquivalenceClassMap
#endif
   // us to build arbitrarily long chains.
   for (Instruction *I : Instrs) {
     constexpr int MaxChainsToTry = 64;
 
     bool MatchFound = false;
     auto ChainIter = MRU.begin();
     for (size_t J = 0; J < MaxChainsToTry && ChainIter != MRU.end();
          ++J, ++ChainIter) {
-      std::optional<APInt> Offset =
-          getConstantOffset(getLoadStorePointerOperand(ChainIter->first),
-                            getLoadStorePointerOperand(I));
+      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)});
         // Move ChainIter to the front of the MRU list.
         MRU.remove(*ChainIter);
         MRU.push_front(*ChainIter);
#endif
   // Iterate over MRU rather than Chains so the order is deterministic.
   for (auto &E : MRU)
     if (E.second.size() > 1)
       Ret.push_back(std::move(E.second));
   return Ret;
 }
 
 std::optional<APInt> Vectorizer::getConstantOffset(Value *PtrA, Value *PtrB,
+                                                   Instruction *ContextInst,
                                                    unsigned Depth) {
   LLVM_DEBUG(dbgs() << "LSV: getConstantOffset, PtrA=" << *PtrA
-                    << ", PtrB=" << *PtrB << ", Depth=" << Depth << "\n");
+                    << ", PtrB=" << *PtrB << ", ContextInst= " << *ContextInst
+                    << ", Depth=" << Depth << "\n");
   unsigned OffsetBitWidth = DL.getIndexTypeSizeInBits(PtrA->getType());
   APInt OffsetA(OffsetBitWidth, 0);
   APInt OffsetB(OffsetBitWidth, 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;
std::optional<APInt> Vectorizer::getConstantOffset(Value *PtrA, Value *PtrB,
   // 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();
   }
-  std::optional<APInt> Diff = gtConstantOffsetComplexAddrs(PtrA, PtrB, Depth);
+  std::optional<APInt> Diff =
+      getConstantOffsetComplexAddrs(PtrA, PtrB, ContextInst, Depth);
   if (Diff.has_value())
     return OffsetB - OffsetA + Diff->sext(OffsetB.getBitWidth());
   return std::nullopt;
 }

llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/gep-bitcast.ll

define void @sexted_i1_gep_index(ptr addrspace(1) %p, i32 %val) {
   %index.0 = sext i1 %selector to i64
   %index.1 = sext i1 %flipped to i64
   %gep.0 = getelementptr inbounds i32, ptr addrspace(1) %p, i64 %index.0
   %gep.1 = getelementptr inbounds i32, ptr addrspace(1) %p, i64 %index.1
   %val0 = load i32, ptr addrspace(1) %gep.0
   %val1 = load i32, ptr addrspace(1) %gep.1
   ret void
 }
+
+; CHECK-LABEL: @zexted_i1_gep_index_different_bbs
+; CHECK: load i32
+; CHECK: load i32
+define void @zexted_i1_gep_index_different_bbs(ptr addrspace(1) %p, i32 %val) {
+entry:
+  %selector = icmp eq i32 %val, 0
+  %flipped = xor i1 %selector, 1
+  %index.0 = zext i1 %selector to i64
+  %index.1 = zext i1 %flipped to i64
+  %gep.0 = getelementptr inbounds i32, ptr addrspace(1) %p, i64 %index.0
+  br label %next
+
+next:
+  %gep.1 = getelementptr inbounds i32, ptr addrspace(1) %p, i64 %index.1
+  %val0 = load i32, ptr addrspace(1) %gep.0
+  %val1 = load i32, ptr addrspace(1) %gep.1
+  ret void
+}