[CodeGenPrepare] Check types when unmerging GEPs across indirect branches

[momo5502]

The optimization in CodeGenPrepare, where GEPs are unmerged across indirect branches must respect the types of both GEPs and their sizes when adjusting the indices.

The sample here shows the bug:

https://godbolt.org/z/8e9o5sYPP

The value %elementValuePtr addresses the second field of the %struct.Blub. It is therefore a GEP with index 1 and type i8.
The value %nextArrayElement addresses the next array element. It is therefore a GEP with index 1 and type %struct.Blub.

Both values point to completely different addresses, even if the indices are the same, due to the types being different.
However, after CodeGenPrepare has run, %nextArrayElement is a bitcast from %elementValuePtr, meaning both were treated as equal.

The cause for this is that the unmerging optimization does not take types into consideration.
It sees both GEPs have %currentArrayElement as source operand and therefore tries to rewrite %nextArrayElement in terms of %elementValuePtr.
It changes the index to the difference of the two GEPs. As both indices are 1, the difference is 0. As the indices are 0 the GEP is later replaced with a simple bitcast in CodeGenPrepare.

Before adjusting the indices, the types of the GEPs would have to be aligned and the indices scaled accordingly for the optimization to be correct.
Due to the size of the struct being 16 and the %elementValuePtr pointing to offset 1, the correct index for the unmerged %nextArrayElement would be 15.

I assume this bug emerged from the opaque pointer change as GEPs like %elementValuePtr that access the struct field based of type i8 did not naturally occur before.

In light of future migration to ptradd, simply not performing the optimization if the types mismatch should be sufficient.

[nikic]

High level comment: Would bailing out of the transform if the source element types don't match be "good enough"?

I'd prefer not to add the extra logic if we don't particularly care about making it work with mismatching GEP types. (Longer term this will be automatically fixed by migration to ptradd.)

[momo5502]

High level comment: Would bailing out of the transform if the source element types don't match be "good enough"?

I'd prefer not to add the extra logic if we don't particularly care about making it work with mismatching GEP types. (Longer term this will be automatically fixed by migration to ptradd.)

I agree. I initially chose to bail out, but thought it could be interpreted as the "lazy way" and feared the PR might get rejected, hence the fix.
However, under the aspect of future migration to ptradd, it doesn't make sense to introduce a fix for something that is going to become obsolete sooner or later.
I will adjust it and revert to bailing out tomorrow.

[nikic]

This looks good, just some test nits.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Maurice Heumann (momo5502)

Changes

The optimization in CodeGenPrepare, where GEPs are unmerged across indirect branches must respect the types of both GEPs and their sizes when adjusting the indices.

The sample here shows the bug:

https://godbolt.org/z/8e9o5sYPP

The value %elementValuePtr addresses the second field of the %struct.Blub. It is therefore a GEP with index 1 and type i8.
The value %nextArrayElement addresses the next array element. It is therefore a GEP with index 1 and type %struct.Blub.

Both values point to completely different addresses, even if the indices are the same, due to the types being different.
However, after CodeGenPrepare has run, %nextArrayElement is a bitcast from %elementValuePtr, meaning both were treated as equal.

The cause for this is that the unmerging optimization does not take types into consideration.
It sees both GEPs have %currentArrayElement as source operand and therefore tries to rewrite %nextArrayElement in terms of %elementValuePtr.
It changes the index to the difference of the two GEPs. As both indices are 1, the difference is 0. As the indices are 0 the GEP is later replaced with a simple bitcast in CodeGenPrepare.

Before adjusting the indices, the types of the GEPs would have to be aligned and the indices scaled accordingly for the optimization to be correct.
Due to the size of the struct being 16 and the %elementValuePtr pointing to offset 1, the correct index for the unmerged %nextArrayElement would be 15.

I assume this bug emerged from the opaque pointer change as GEPs like %elementValuePtr that access the struct field based of type i8 did not naturally occur before.

In light of future migration to ptradd, simply not performing the optimization if the types mismatch should be sufficient.

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/CodeGenPrepare.cpp (+2)
• (added) llvm/test/CodeGen/X86/indirect-br-gep-unmerge.ll (+51)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 371f6598e6b2b35..187820717b6fd5c 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -7999,6 +7999,8 @@ static bool tryUnmergingGEPsAcrossIndirectBr(GetElementPtrInst *GEPI,
       return false;
     if (UGEPI->getOperand(0) != GEPIOp)
       return false;
+    if (UGEPI->getSourceElementType() != GEPI->getSourceElementType())
+      return false;
     if (GEPIIdx->getType() !=
         cast<ConstantInt>(UGEPI->getOperand(1))->getType())
       return false;
diff --git a/llvm/test/CodeGen/X86/indirect-br-gep-unmerge.ll b/llvm/test/CodeGen/X86/indirect-br-gep-unmerge.ll
new file mode 100644
index 000000000000000..655971631bc384f
--- /dev/null
+++ b/llvm/test/CodeGen/X86/indirect-br-gep-unmerge.ll
@@ -0,0 +1,51 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -S -codegenprepare %s -o - | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.Blub = type { i8, i8, ptr }
+
+@indirectBrPtr = external hidden global ptr
+
+define ptr @testFunc(ptr noundef readonly %array, i1 %skip) {
+; CHECK-LABEL: define ptr @testFunc(
+; CHECK-SAME: ptr noundef readonly [[ARRAY:%.*]], i1 [[SKIP:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[SKIP]], label [[LOOPHEADER:%.*]], label [[ENDBLOCK_CLONE:%.*]]
+; CHECK:       loopHeader:
+; CHECK-NEXT:    [[CURRENTARRAYELEMENT:%.*]] = phi ptr [ [[ARRAY]], [[ENTRY:%.*]] ], [ [[NEXTARRAYELEMENT:%.*]], [[LOOPFOOTER:%.*]] ]
+; CHECK-NEXT:    [[ELEMENTVALUEPTR:%.*]] = getelementptr inbounds i8, ptr [[CURRENTARRAYELEMENT]], i64 1
+; CHECK-NEXT:    [[ELEMENTVALUE:%.*]] = load i8, ptr [[ELEMENTVALUEPTR]], align 1
+; CHECK-NEXT:    indirectbr ptr @indirectBrPtr, [label [[LOOPFOOTER]], label %endBlock]
+; CHECK:       loopFooter:
+; CHECK-NEXT:    [[ISGOODVALUE:%.*]] = icmp eq i8 [[ELEMENTVALUE]], 0
+; CHECK-NEXT:    [[NEXTARRAYELEMENT]] = getelementptr inbounds [[STRUCT_BLUB:%.*]], ptr [[CURRENTARRAYELEMENT]], i64 1
+; CHECK-NEXT:    br i1 [[ISGOODVALUE]], label [[LOOPHEADER]], label [[ENDBLOCK_CLONE]]
+; CHECK:       endBlock:
+; CHECK-NEXT:    br label [[DOTSPLIT:%.*]]
+; CHECK:       .split:
+; CHECK-NEXT:    [[MERGE:%.*]] = phi ptr [ [[ELEMENTVALUEPTR]], [[ENDBLOCK:%.*]] ], [ [[RETVAL_CLONE:%.*]], [[ENDBLOCK_CLONE]] ]
+; CHECK-NEXT:    ret ptr [[MERGE]]
+; CHECK:       endBlock.clone:
+; CHECK-NEXT:    [[RETVAL_CLONE]] = phi ptr [ [[ARRAY]], [[ENTRY]] ], [ [[ELEMENTVALUEPTR]], [[LOOPFOOTER]] ]
+; CHECK-NEXT:    br label [[DOTSPLIT]]
+;
+entry:
+  br i1 %skip, label %loopHeader, label %endBlock
+
+loopHeader:
+  %currentArrayElement = phi ptr [ %array, %entry ], [ %nextArrayElement, %loopFooter ]
+  %elementValuePtr = getelementptr inbounds i8, ptr %currentArrayElement, i64 1
+  %elementValue = load i8, ptr %elementValuePtr, align 1
+  indirectbr ptr @indirectBrPtr, [label %loopFooter, label %endBlock]
+
+loopFooter:
+  %isGoodValue = icmp eq i8 %elementValue, 0
+  %nextArrayElement = getelementptr inbounds %struct.Blub, ptr %currentArrayElement, i64 1
+  br i1 %isGoodValue, label %loopHeader, label %endBlock
+
+endBlock:
+  %retVal = phi ptr [ %array, %entry ], [ %elementValuePtr, %loopFooter ], [ %elementValuePtr, %loopHeader ]
+  ret ptr %retVal
+}

[momo5502]

Thank you :)