[SeparateConstOffsetFromGEP] Support multiple indices in reorderGEP

[nikic]

The code was essentially already ready to handle multiple indices -- we only need to adjust the non-negative index check to check all indices, instead of only the first one.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-backend-amdgpu

Author: Nikita Popov (nikic)

Changes

The code was essentially already ready to handle multiple indices -- we only need to adjust the non-negative index check to check all indices, instead of only the first one.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp (+6-13)
• (modified) llvm/test/Transforms/SeparateConstOffsetFromGEP/AMDGPU/reorder-gep.ll (+4-4)

diff --git a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
index 7ac1f43b7b6ac..6ca7d9d54a9b2 100644
--- a/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
+++ b/llvm/lib/Transforms/Scalar/SeparateConstOffsetFromGEP.cpp
@@ -972,14 +972,9 @@ SeparateConstOffsetFromGEP::lowerToArithmetics(GetElementPtrInst *Variadic,
 
 bool SeparateConstOffsetFromGEP::reorderGEP(GetElementPtrInst *GEP,
                                             TargetTransformInfo &TTI) {
-  if (GEP->getNumIndices() != 1)
-    return false;
-
   auto PtrGEP = dyn_cast<GetElementPtrInst>(GEP->getPointerOperand());
   if (!PtrGEP)
     return false;
-  if (PtrGEP->getNumIndices() != 1)
-    return false;
 
   bool NestedNeedsExtraction;
   int64_t NestedByteOffset =
@@ -997,14 +992,12 @@ bool SeparateConstOffsetFromGEP::reorderGEP(GetElementPtrInst *GEP,
   bool PtrGEPInBounds = PtrGEP->isInBounds();
   bool IsChainInBounds = GEPInBounds && PtrGEPInBounds;
   if (IsChainInBounds) {
-    auto GEPIdx = GEP->indices().begin();
-    auto KnownGEPIdx = computeKnownBits(GEPIdx->get(), *DL);
-    IsChainInBounds &= KnownGEPIdx.isNonNegative();
-    if (IsChainInBounds) {
-      auto PtrGEPIdx = PtrGEP->indices().begin();
-      auto KnownPtrGEPIdx = computeKnownBits(PtrGEPIdx->get(), *DL);
-      IsChainInBounds &= KnownPtrGEPIdx.isNonNegative();
-    }
+    auto IsKnownNonNegative = [&](Value *V) {
+      return isKnownNonNegative(V, *DL);
+    };
+    IsChainInBounds &= all_of(GEP->indices(), IsKnownNonNegative);
+    if (IsChainInBounds)
+      IsChainInBounds &= all_of(PtrGEP->indices(), IsKnownNonNegative);
   }
 
   IRBuilder<> Builder(GEP);
diff --git a/llvm/test/Transforms/SeparateConstOffsetFromGEP/AMDGPU/reorder-gep.ll b/llvm/test/Transforms/SeparateConstOffsetFromGEP/AMDGPU/reorder-gep.ll
index a7ca5b93c361d..dd12c98af696d 100644
--- a/llvm/test/Transforms/SeparateConstOffsetFromGEP/AMDGPU/reorder-gep.ll
+++ b/llvm/test/Transforms/SeparateConstOffsetFromGEP/AMDGPU/reorder-gep.ll
@@ -288,8 +288,8 @@ entry:
 define void @multiple_index_maybe_neg(ptr %in.ptr, i64 %in.idx1) {
 ; CHECK-LABEL: define void @multiple_index_maybe_neg(
 ; CHECK-SAME: ptr [[IN_PTR:%.*]], i64 [[IN_IDX1:%.*]]) #[[ATTR0]] {
-; CHECK-NEXT:    [[CONST1:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[IN_PTR]], i64 0, i64 1
-; CHECK-NEXT:    [[IDX1:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[CONST1]], i64 0, i64 [[IN_IDX1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [2 x <2 x i8>], ptr [[IN_PTR]], i64 0, i64 [[IN_IDX1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr [2 x <2 x i8>], ptr [[TMP1]], i64 0, i64 1
 ; CHECK-NEXT:    ret void
 ;
   %const1 = getelementptr inbounds [2 x <2 x i8>], ptr %in.ptr, i64 0, i64 1
@@ -301,8 +301,8 @@ define void @multiple_index_nonneg(ptr %in.ptr, i64 %in.idx1) {
 ; CHECK-LABEL: define void @multiple_index_nonneg(
 ; CHECK-SAME: ptr [[IN_PTR:%.*]], i64 [[IN_IDX1:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    [[IN_IDX1_NNEG:%.*]] = and i64 [[IN_IDX1]], 9223372036854775807
-; CHECK-NEXT:    [[CONST1:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[IN_PTR]], i64 0, i64 1
-; CHECK-NEXT:    [[IDX1:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[CONST1]], i64 0, i64 [[IN_IDX1_NNEG]]
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[IN_PTR]], i64 0, i64 [[IN_IDX1_NNEG]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds [2 x <2 x i8>], ptr [[TMP1]], i64 0, i64 1
 ; CHECK-NEXT:    ret void
 ;
   %in.idx1.nneg = and i64 %in.idx1, 9223372036854775807

[nikic]

Putting this up in the hope that it helps with #68882 (comment).

I've fixed a bug in the inbounds propagation I noticed while writing this in b4d1a60.

[arsenm]

DL is the only capture, so capture by value? Is it worth providing the context instruction to the SimplifyQuery?

[jayfoad]

Does this work? I'm never sure how flexible concat is.

Suggested change

	IsChainInBounds &= all_of(GEP->indices(), IsKnownNonNegative);
	if (IsChainInBounds)
	IsChainInBounds &= all_of(PtrGEP->indices(), IsKnownNonNegative);
	IsChainInBounds &= all_of(concat(GEP->indices(), PtrGEP->indices()), IsKnownNonNegative);

[nikic]

That doesn't work, and I couldn't figure out how to make it to work.

[sgundapa]

Putting this up in the hope that it helps with #68882 (comment).

I was about to push a similar change.
It did solve the regression I have seen. Thanks for the patch.