[llvm] [CodeGenPrepare] Unmerging GEPs across indirect branches must respect types (PR #68587)

Maurice Heumann via llvm-commits llvm-commits at lists.llvm.org
Mon Oct 9 07:09:59 PDT 2023


https://github.com/momo5502 created https://github.com/llvm/llvm-project/pull/68587

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/fvYvvGGjE

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 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 adjusted the optimization and added a simple test case that asserts the correct index is preserved. 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.

>From 7b756c0515bcea238a56d24af26a10f592eb0fff Mon Sep 17 00:00:00 2001
From: Maurice Heumann <maurice.heumann at wibu.com>
Date: Fri, 22 Sep 2023 07:45:05 +0200
Subject: [PATCH] [CodeGenPrepare] Unmerging GEPs across indirect branches must
 respect types

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/fvYvvGGjE

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 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 adjusted the optimization and added a simple test case that asserts the correct index is preserved.
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.
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp           | 58 +++++++++++++++++--
 .../Generic/indirect-br-gep-unmerge.ll        | 29 ++++++++++
 2 files changed, 82 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 371f6598e6b2b35..286b1637ecbc7c5 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -7883,6 +7883,51 @@ static bool GEPSequentialConstIndexed(GetElementPtrInst *GEP) {
          isa<ConstantInt>(GEP->getOperand(1));
 }
 
+// Returns the size of the GEP type.
+// GEP must meet GEPSequentialConstIndexed requirements.
+static TypeSize GEPSequentialConstIndexedTypeSize(GetElementPtrInst &GEP) {
+  const DataLayout &DL = GEP.getModule()->getDataLayout();
+  gep_type_iterator GTI = gep_type_begin(GEP);
+  return DL.getTypeAllocSize(GTI.getIndexedType());
+}
+
+// Returns true if the Target can be addressed from Source
+// in case types mismatch.
+// Both GEPs must meet GEPSequentialConstIndexed requirements.
+static bool GEPIsAddressableFromSource(GetElementPtrInst &Target,
+                                       GetElementPtrInst &Source) {
+  if (Source.getSourceElementType() == Target.getSourceElementType()) {
+    return true;
+  }
+
+  TypeSize SourceTypeSize = GEPSequentialConstIndexedTypeSize(Source);
+  TypeSize TargetTypeSize = GEPSequentialConstIndexedTypeSize(Target);
+
+  return (TargetTypeSize % SourceTypeSize) == 0;
+}
+
+// Calculates the relative index required to address Target
+// from Source with respect to their element types.
+// Both GEPs must meet GEPSequentialConstIndexed requirements.
+static APInt GEPCalculateRelativeIndex(GetElementPtrInst &Target,
+                                       GetElementPtrInst &Source) {
+  uint64_t Scale = 1;
+
+  if (Source.getSourceElementType() != Target.getSourceElementType()) {
+    TypeSize SourceTypeSize = GEPSequentialConstIndexedTypeSize(Source);
+    TypeSize TargetTypeSize = GEPSequentialConstIndexedTypeSize(Target);
+
+    Scale = TargetTypeSize / SourceTypeSize;
+    assert(Scale * SourceTypeSize == TargetTypeSize &&
+           "Target GEP must be addressable from Source GEP");
+  }
+
+  ConstantInt *SourceIdx = cast<ConstantInt>(Source.getOperand(1));
+  ConstantInt *TargetIdx = cast<ConstantInt>(Target.getOperand(1));
+
+  return (TargetIdx->getValue() * Scale) - SourceIdx->getValue();
+}
+
 // Try unmerging GEPs to reduce liveness interference (register pressure) across
 // IndirectBr edges. Since IndirectBr edges tend to touch on many blocks,
 // reducing liveness interference across those edges benefits global register
@@ -7997,6 +8042,9 @@ static bool tryUnmergingGEPsAcrossIndirectBr(GetElementPtrInst *GEPI,
     // up.
     if (!GEPSequentialConstIndexed(UGEPI))
       return false;
+    // Check if GEP Types match or if types are compatible
+    if (!GEPIsAddressableFromSource(*UGEPI, *GEPI))
+      return false;
     if (UGEPI->getOperand(0) != GEPIOp)
       return false;
     if (GEPIIdx->getType() !=
@@ -8013,8 +8061,7 @@ static bool tryUnmergingGEPsAcrossIndirectBr(GetElementPtrInst *GEPI,
     return false;
   // Check the materializing cost of (Uidx-Idx).
   for (GetElementPtrInst *UGEPI : UGEPIs) {
-    ConstantInt *UGEPIIdx = cast<ConstantInt>(UGEPI->getOperand(1));
-    APInt NewIdx = UGEPIIdx->getValue() - GEPIIdx->getValue();
+    APInt NewIdx = GEPCalculateRelativeIndex(*UGEPI, *GEPI);
     InstructionCost ImmCost = TTI->getIntImmCost(
         NewIdx, GEPIIdx->getType(), TargetTransformInfo::TCK_SizeAndLatency);
     if (ImmCost > TargetTransformInfo::TCC_Basic)
@@ -8022,11 +8069,12 @@ static bool tryUnmergingGEPsAcrossIndirectBr(GetElementPtrInst *GEPI,
   }
   // Now unmerge between GEPI and UGEPIs.
   for (GetElementPtrInst *UGEPI : UGEPIs) {
+    APInt NewIdx = GEPCalculateRelativeIndex(*UGEPI, *GEPI);
+    Constant *NewUGEPIIdx = ConstantInt::get(GEPIIdx->getType(), NewIdx);
     UGEPI->setOperand(0, GEPI);
-    ConstantInt *UGEPIIdx = cast<ConstantInt>(UGEPI->getOperand(1));
-    Constant *NewUGEPIIdx = ConstantInt::get(
-        GEPIIdx->getType(), UGEPIIdx->getValue() - GEPIIdx->getValue());
     UGEPI->setOperand(1, NewUGEPIIdx);
+    UGEPI->setSourceElementType(GEPI->getSourceElementType());
+    UGEPI->setResultElementType(GEPI->getResultElementType());
     // If GEPI is not inbounds but UGEPI is inbounds, change UGEPI to not
     // inbounds to avoid UB.
     if (!GEPI->isInBounds()) {
diff --git a/llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll b/llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll
new file mode 100644
index 000000000000000..0a4cf1ae07d461e
--- /dev/null
+++ b/llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll
@@ -0,0 +1,29 @@
+; RUN: llc %s -stop-after=codegenprepare -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 }
+
+ at indirectBrPtr = external hidden global ptr
+
+define dso_local noundef ptr @testFunc(ptr noundef readonly %array, i1 %skip) {
+entry:
+  br i1 %skip, label %loopHeader, label %endBlock
+
+loopHeader:                                                ; preds = %2, %1
+  %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
+  ; CHECK: %nextArrayElement = getelementptr inbounds i8, ptr %elementValuePtr, i64 15
+  %nextArrayElement = getelementptr inbounds %struct.Blub, ptr %currentArrayElement, i64 1
+  br i1 %isGoodValue, label %loopHeader, label %endBlock
+
+endBlock:                                                ; preds = %2
+  %retVal = phi ptr [ %array, %entry ], [ %elementValuePtr, %loopFooter ], [ %elementValuePtr, %loopHeader ]
+  ret ptr %retVal
+}



More information about the llvm-commits mailing list