[llvm] 10223c7 - [ConstraintElim] Use nusw flag for GEP decomposition

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2024-12-03T15:56:29+01:00
New Revision: 10223c72a9c2afc3ae86eb7995f0dc47f5f9ad44

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

LOG: [ConstraintElim] Use nusw flag for GEP decomposition

Check for nusw instead of inbounds when decomposing GEPs.

In this particular case, we can also look through multiple nusw
flags, because we will ultimately be working in the unsigned
constraint system.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
    llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 4884c23f16e12a..944be38cb94bc7 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -387,12 +387,12 @@ struct OffsetResult {
   Value *BasePtr;
   APInt ConstantOffset;
   SmallMapVector<Value *, APInt, 4> VariableOffsets;
-  bool AllInbounds;
+  GEPNoWrapFlags NW;
 
   OffsetResult() : BasePtr(nullptr), ConstantOffset(0, uint64_t(0)) {}
 
   OffsetResult(GEPOperator &GEP, const DataLayout &DL)
-      : BasePtr(GEP.getPointerOperand()), AllInbounds(GEP.isInBounds()) {
+      : BasePtr(GEP.getPointerOperand()), NW(GEP.getNoWrapFlags()) {
     ConstantOffset = APInt(DL.getIndexTypeSizeInBits(BasePtr->getType()), 0);
   }
 };
@@ -426,7 +426,7 @@ static OffsetResult collectOffsets(GEPOperator &GEP, const DataLayout &DL) {
     Result.ConstantOffset += ConstantOffset2;
     if (Result.VariableOffsets.size() == 0 && VariableOffsets2.size() == 1)
       Result.VariableOffsets = VariableOffsets2;
-    Result.AllInbounds &= InnerGEP->isInBounds();
+    Result.NW &= InnerGEP->getNoWrapFlags();
   }
   return Result;
 }
@@ -450,9 +450,9 @@ static Decomposition decomposeGEP(GEPOperator &GEP,
 
   assert(!IsSigned && "The logic below only supports decomposition for "
                       "unsigned predicates at the moment.");
-  const auto &[BasePtr, ConstantOffset, VariableOffsets, AllInbounds] =
+  const auto &[BasePtr, ConstantOffset, VariableOffsets, NW] =
       collectOffsets(GEP, DL);
-  if (!BasePtr || !AllInbounds)
+  if (!BasePtr || !NW.hasNoUnsignedSignedWrap())
     return &GEP;
 
   Decomposition Result(ConstantOffset.getSExtValue(), DecompEntry(1, BasePtr));

diff  --git a/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll b/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
index a4d825b327969c..5e2bab28807f54 100644
--- a/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
+++ b/llvm/test/Transforms/ConstraintElimination/gep-arithmetic.ll
@@ -630,3 +630,68 @@ ptr.check:
 exit:
   ret i4 3
 }
+
+define i1 @test_nusw(ptr %p, i32 %x, i32 %y) {
+; CHECK-LABEL: @test_nusw(
+; CHECK-NEXT:    [[X_EXT:%.*]] = zext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[Y_EXT:%.*]] = zext i32 [[Y:%.*]] to i64
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[X_EXT]], [[Y_EXT]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
+; CHECK-NEXT:    [[GEP_X:%.*]] = getelementptr nusw i8, ptr [[P:%.*]], i64 [[X_EXT]]
+; CHECK-NEXT:    [[GEP_Y:%.*]] = getelementptr nusw i8, ptr [[P]], i64 [[Y_EXT]]
+; CHECK-NEXT:    ret i1 true
+;
+  %x.ext = zext i32 %x to i64
+  %y.ext = zext i32 %y to i64
+  %cmp1 = icmp ugt i64 %x.ext, %y.ext
+  call void @llvm.assume(i1 %cmp1)
+  %gep.x = getelementptr nusw i8, ptr %p, i64 %x.ext
+  %gep.y = getelementptr nusw i8, ptr %p, i64 %y.ext
+  %cmp2 = icmp ugt ptr %gep.x, %gep.y
+  ret i1 %cmp2
+}
+
+define i1 @test_nusw_nested(ptr %p, i32 %x, i32 %y) {
+; CHECK-LABEL: @test_nusw_nested(
+; CHECK-NEXT:    [[X_EXT:%.*]] = zext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[Y_EXT:%.*]] = zext i32 [[Y:%.*]] to i64
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[X_EXT]], [[Y_EXT]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
+; CHECK-NEXT:    [[GEP_X:%.*]] = getelementptr nusw i8, ptr [[P:%.*]], i64 [[X_EXT]]
+; CHECK-NEXT:    [[GEP_X1:%.*]] = getelementptr nusw i8, ptr [[GEP_X]], i64 1
+; CHECK-NEXT:    [[GEP_Y:%.*]] = getelementptr nusw i8, ptr [[P]], i64 [[Y_EXT]]
+; CHECK-NEXT:    ret i1 true
+;
+  %x.ext = zext i32 %x to i64
+  %y.ext = zext i32 %y to i64
+  %cmp1 = icmp ugt i64 %x.ext, %y.ext
+  call void @llvm.assume(i1 %cmp1)
+  %gep.x = getelementptr nusw i8, ptr %p, i64 %x.ext
+  %gep.x1 = getelementptr nusw i8, ptr %gep.x, i64 1
+  %gep.y = getelementptr nusw i8, ptr %p, i64 %y.ext
+  %cmp2 = icmp ugt ptr %gep.x1, %gep.y
+  ret i1 %cmp2
+}
+
+define i1 @test_missing_nusw(ptr %p, i32 %x, i32 %y) {
+; CHECK-LABEL: @test_missing_nusw(
+; CHECK-NEXT:    [[X_EXT:%.*]] = zext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[Y_EXT:%.*]] = zext i32 [[Y:%.*]] to i64
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[X_EXT]], [[Y_EXT]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP1]])
+; CHECK-NEXT:    [[GEP_X:%.*]] = getelementptr nusw i8, ptr [[P:%.*]], i64 [[X_EXT]]
+; CHECK-NEXT:    [[GEP_X1:%.*]] = getelementptr i8, ptr [[GEP_X]], i64 1
+; CHECK-NEXT:    [[GEP_Y:%.*]] = getelementptr nusw i8, ptr [[P]], i64 [[Y_EXT]]
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp ugt ptr [[GEP_X1]], [[GEP_Y]]
+; CHECK-NEXT:    ret i1 [[CMP2]]
+;
+  %x.ext = zext i32 %x to i64
+  %y.ext = zext i32 %y to i64
+  %cmp1 = icmp ugt i64 %x.ext, %y.ext
+  call void @llvm.assume(i1 %cmp1)
+  %gep.x = getelementptr nusw i8, ptr %p, i64 %x.ext
+  %gep.x1 = getelementptr i8, ptr %gep.x, i64 1
+  %gep.y = getelementptr nusw i8, ptr %p, i64 %y.ext
+  %cmp2 = icmp ugt ptr %gep.x1, %gep.y
+  ret i1 %cmp2
+}