[LoadStoreVectorizer] Fix index width != pointer width case

[krzysz00]

Summary

Fixes https://github.com/llvm/llvm-project/issues/62856

[jlebar]

Yes, this patch looks correct to me.

Thanks, and sorry for the breakage.

[foad]

Thanks!

llvm/lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

std::optional<APInt> Vectorizer::getConstantOffset(Value *PtrA, Value *PtrB,
   if (PtrA == PtrB)
     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())
-          .sextOrTrunc(OrigBitWidth);
+    if (DistRange.isSingleElement()) {
+      // Handle index width (the width of Dist) != pointer width (the width of
+      // the Offset*s at this point).
+      APInt Dist = DistRange.getSingleElement()->sextOrTrunc(NewPtrBitWidth);
+      return (OffsetB - OffsetA + Dist).sextOrTrunc(OrigBitWidth);
+    }
   }
   std::optional<APInt> Diff =
       getConstantOffsetComplexAddrs(PtrA, PtrB, ContextInst, Depth);
   if (Diff.has_value())
     return (OffsetB - OffsetA + Diff->sext(OffsetB.getBitWidth()))
         .sextOrTrunc(OrigBitWidth);
   return std::nullopt;
 }

llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/addrspace-7.ll

-; REQUIRES: asserts
-; RUN: not --crash opt -mtriple=amdgcn-amd-amdhsa -passes=load-store-vectorizer -S -o - %s
-; RUN: not --crash opt -mtriple=amdgcn-amd-amdhsa -aa-pipeline=basic-aa -passes='function(load-store-vectorizer)' -S -o - %s
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -passes=load-store-vectorizer -S -o - %s | FileCheck %s
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -aa-pipeline=basic-aa -passes='function(load-store-vectorizer)' -S -o - %s | FileCheck %s
 
 target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5"
 
 define { float, float } @f() {
+; CHECK-LABEL: define { float, float } @f() {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[L1:%.*]] = load float, ptr addrspace(7) null, align 4
+; CHECK-NEXT:    [[L2:%.*]] = load float, ptr addrspace(7) getelementptr (i8, ptr addrspace(7) null, i64 24), align 4
+; CHECK-NEXT:    [[IV1:%.*]] = insertvalue { float, float } zeroinitializer, float [[L1]], 0
+; CHECK-NEXT:    [[IV2:%.*]] = insertvalue { float, float } [[IV1]], float [[L2]], 1
+; CHECK-NEXT:    ret { float, float } [[IV2]]
+;
 bb:
   %l1 = load float, ptr addrspace(7) null
   %l2 = load float, ptr addrspace(7) getelementptr (i8, ptr addrspace(7) null, i64 24)
   %iv1 = insertvalue { float, float } zeroinitializer, float %l1, 0
   %iv2 = insertvalue { float, float } %iv1, float %l2, 1
   ret { float, float } %iv2
 }