[NaryReassociate] Fix crash from pointer width / index width confusion

[krzysz00]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Krzysztof Drewniak (krzysz00)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/125923.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/NaryReassociate.cpp (+9-5)
• (modified) llvm/test/Transforms/NaryReassociate/nary-gep.ll (+11)

diff --git a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
index 39720672c202d7..434451abbb84cb 100644
--- a/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/NaryReassociate.cpp
@@ -402,16 +402,20 @@ NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
     IndexExprs.push_back(SE->getSCEV(Index));
   // Replace the I-th index with LHS.
   IndexExprs[I] = SE->getSCEV(LHS);
+  Type *GEPArgType = GEP->getOperand(I)->getType();
+  Type *LHSType = LHS->getType();
+  size_t LHSSize = DL->getTypeSizeInBits(LHSType).getFixedValue();
+  size_t GEPArgSize = DL->getTypeSizeInBits(GEPArgType).getFixedValue();
+  // For pointers, we need to look at the index size, not the total type size.
+  if (isa<PointerType>(GEPArgType))
+    GEPArgSize = DL->getIndexTypeSizeInBits(GEPArgType);
   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 d0ece1e11de5a3..a2cd32f16b88ae 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() {
+; CHECK-LABEL: @zext_fat_pointer_crash(
+; CHECK-NEXT:    [[C:%.*]] = add i32 0, 0
+; CHECK-NEXT:    [[Q:%.*]] = getelementptr double, ptr addrspace(2) null, i32 [[C]]
+; CHECK-NEXT:    ret ptr addrspace(2) [[Q]]
+;
+  %c = add i32 0, 0
+  %q = getelementptr double, ptr addrspace(2) null, 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:%.*]]

[nikic]

Given that this is working with SE, it would be better to use APIs like SE.getTypeSizeInBits or SE.getEffectiveSCEVType.

[krzysz00]

Huh, yeah, I'll swap it over to that

[nikic]

LGTM