[llvm] e41ffd3 - [NaryReassociate] Fix crash from pointer width / index width confusion (#125923)

[via llvm-commits]

Author: Krzysztof Drewniak
Date: 2025-02-06T12:48:52-06:00
New Revision: e41ffd3420d7c00c4c8f241b8f7f9ba210eea68f

URL: https://github.com/llvm/llvm-project/commit/e41ffd3420d7c00c4c8f241b8f7f9ba210eea68f
DIFF: https://github.com/llvm/llvm-project/commit/e41ffd3420d7c00c4c8f241b8f7f9ba210eea68f.diff

LOG: [NaryReassociate] Fix crash from pointer width / index width confusion (#125923)

NaryReassociate would crash on expressions like the one in the added
test that involved pointers where the size of the type was greater than
the index width of the pointer, causing calls to SCEV's zext expression
on types that didn't need to be zero-extended.

This commit fixes the issue.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/NaryReassociate.cpp
    llvm/test/Transforms/NaryReassociate/nary-gep.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index 39720672c202d7f..3b3f32a0ea591fb 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -402,16 +402,17 @@ NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
     IndexExprs.push_back(SE->getSCEV(Index));
   // Replace the I-th index with LHS.
   IndexExprs[I] = SE->getSCEV(LHS);
+  Type *GEPArgType = SE->getEffectiveSCEVType(GEP->getOperand(I)->getType());
+  Type *LHSType = SE->getEffectiveSCEVType(LHS->getType());
+  size_t LHSSize = DL->getTypeSizeInBits(LHSType).getFixedValue();
+  size_t GEPArgSize = DL->getTypeSizeInBits(GEPArgType).getFixedValue();
   if (isKnownNonNegative(LHS, SimplifyQuery(*DL, DT, AC, GEP)) &&
-      DL->getTypeSizeInBits(LHS->getType()).getFixedValue() <
-          DL->getTypeSizeInBits(GEP->getOperand(I)->getType())
-              .getFixedValue()) {
+      LHSSize < GEPArgSize) {
     // Zero-extend LHS if it is non-negative. InstCombine canonicalizes sext to
     // zext if the source operand is proved non-negative. We should do that
     // consistently so that CandidateExpr more likely appears before. See
     // @reassociate_gep_assume for an example of this canonicalization.
-    IndexExprs[I] =
-        SE->getZeroExtendExpr(IndexExprs[I], GEP->getOperand(I)->getType());
+    IndexExprs[I] = SE->getZeroExtendExpr(IndexExprs[I], GEPArgType);
   }
   const SCEV *CandidateExpr = SE->getGEPExpr(cast<GEPOperator>(GEP),
                                              IndexExprs);

diff  --git a/llvm/test/Transforms/NaryReassociate/nary-gep.ll b/llvm/test/Transforms/NaryReassociate/nary-gep.ll
index d0ece1e11de5a3c..b56fdbe4a0669d4 100644
--- a/llvm/test/Transforms/NaryReassociate/nary-gep.ll
+++ b/llvm/test/Transforms/NaryReassociate/nary-gep.ll
@@ -21,6 +21,17 @@ define void @no_sext_fat_pointer(ptr addrspace(2) %a, i32 %i, i32 %j) {
   ret void
 }
 
+define ptr addrspace(2) @zext_fat_pointer_crash(ptr addrspace(2) %p, i32 %a) {
+; CHECK-LABEL: @zext_fat_pointer_crash(
+; CHECK-NEXT:    [[C:%.*]] = add i32 [[A:%.*]], 1
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr double, ptr addrspace(2) [[P:%.*]], i32 [[C]]
+; CHECK-NEXT:    ret ptr addrspace(2) [[Q]]
+;
+  %c = add i32 %a, 1
+  %q = getelementptr double, ptr addrspace(2) %p, i32 %c
+  ret ptr addrspace(2) %q
+}
+
 define void @or_disjoint(ptr addrspace(2) %a, i32 %i, i32 %j, i32 %k) {
 ; CHECK-LABEL: @or_disjoint(
 ; CHECK-NEXT:    [[OR:%.*]] = or disjoint i32 [[I:%.*]], [[J:%.*]]