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

[Maurice Heumann via llvm-commits]

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

>From 66994cd1656fdd4a889036e19c1dc5e1d30b353c 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] Check types when unmerging GEPs across
 indirect branches

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.
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp           |  2 +
 .../Generic/indirect-br-gep-unmerge.ll        | 51 +++++++++++++++++++
 2 files changed, 53 insertions(+)
 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..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/Generic/indirect-br-gep-unmerge.ll b/llvm/test/CodeGen/Generic/indirect-br-gep-unmerge.ll
new file mode 100644
index 000000000000000..703476143ee0fcf
--- /dev/null
+++ b/llvm/test/CodeGen/Generic/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 }
+
+ at indirectBrPtr = external hidden global ptr
+
+define dso_local noundef ptr @testFunc(ptr noundef readonly %array, i1 %skip) {
+; CHECK-LABEL: define dso_local noundef 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:                                                ; 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
+  %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
+}