From cf371573b0b841db96df729e5b61c6a6f36c17d2 Mon Sep 17 00:00:00 2001
From: Leonard Chan <leonardchan@google.com>
Date: Fri, 5 Mar 2021 16:11:01 -0800
Subject: [PATCH] [llvm] Change DSOLocalEquivalent type if the underlying
 global value type changes

We encountered an issue where LTO running on IR that used the DSOLocalEquivalent
constant would result in bad codegen. The underlying issue was ValueMapper wasn't
properly handling DSOLocalEquivalent, so this just adds the machinery for handling
it. This code path is triggered by a fix to DSOLocalEquivalent::handleOperandChangeImpl
where DSOLocalEquivalent could potentially not have the same type as its underlying GV.

This updates DSOLocalEquivalent::handleOperandChangeImpl to change the type if
the GV type changes and handles this constant in ValueMapper.

Differential Revision: https://reviews.llvm.org/D97978
---
 llvm/lib/IR/Constants.cpp                     |  6 ++++
 llvm/lib/Transforms/Utils/ValueMapper.cpp     | 14 ++++++++
 llvm/test/Linker/DSOLocalEquivalent.ll        | 34 +++++++++++++++++++
 llvm/test/Linker/Inputs/DSOLocalEquivalent.ll |  7 ++++
 llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll   | 17 ++++++++++
 5 files changed, 78 insertions(+)
 create mode 100644 llvm/test/Linker/DSOLocalEquivalent.ll
 create mode 100644 llvm/test/Linker/Inputs/DSOLocalEquivalent.ll
 create mode 100644 llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll

diff --git a/llvm/lib/IR/Constants.cpp b/llvm/lib/IR/Constants.cpp
index 62760ff7aec2..0f55a535ffe9 100644
--- a/llvm/lib/IR/Constants.cpp
+++ b/llvm/lib/IR/Constants.cpp
@@ -1914,6 +1914,12 @@ Value *DSOLocalEquivalent::handleOperandChangeImpl(Value *From, Value *To) {
   getContext().pImpl->DSOLocalEquivalents.erase(getGlobalValue());
   NewEquiv = this;
   setOperand(0, Func);
+
+  if (Func->getType() != getType()) {
+    // It is ok to mutate the type here because this constant should always
+    // reflect the type of the function it's holding.
+    mutateType(Func->getType());
+  }
   return nullptr;
 }
 
diff --git a/llvm/lib/Transforms/Utils/ValueMapper.cpp b/llvm/lib/Transforms/Utils/ValueMapper.cpp
index d0ec368dd84f..46ba2e7e0296 100644
--- a/llvm/lib/Transforms/Utils/ValueMapper.cpp
+++ b/llvm/lib/Transforms/Utils/ValueMapper.cpp
@@ -432,6 +432,20 @@ Value *Mapper::mapValue(const Value *V) {
   if (BlockAddress *BA = dyn_cast<BlockAddress>(C))
     return mapBlockAddress(*BA);
 
+  if (const auto *E = dyn_cast<DSOLocalEquivalent>(C)) {
+    auto *Val = mapValue(E->getGlobalValue());
+    GlobalValue *GV = dyn_cast<GlobalValue>(Val);
+    if (GV)
+      return getVM()[E] = DSOLocalEquivalent::get(GV);
+
+    auto *Func = cast<Function>(Val->stripPointerCastsAndAliases());
+    Type *NewTy = E->getType();
+    if (TypeMapper)
+      NewTy = TypeMapper->remapType(NewTy);
+    return getVM()[E] = llvm::ConstantExpr::getBitCast(
+               DSOLocalEquivalent::get(Func), NewTy);
+  }
+
   auto mapValueOrNull = [this](Value *V) {
     auto Mapped = mapValue(V);
     assert((Mapped || (Flags & RF_NullMapMissingGlobalValues)) &&
diff --git a/llvm/test/Linker/DSOLocalEquivalent.ll b/llvm/test/Linker/DSOLocalEquivalent.ll
new file mode 100644
index 000000000000..954881bbde22
--- /dev/null
+++ b/llvm/test/Linker/DSOLocalEquivalent.ll
@@ -0,0 +1,34 @@
+; RUN: llvm-link %s %S/Inputs/DSOLocalEquivalent.ll -S | FileCheck %s
+; RUN: llvm-link %S/Inputs/DSOLocalEquivalent.ll %s -S | FileCheck %s
+
+declare void @extern_func()
+declare void @defined_extern_func()
+declare hidden void @hidden_func()
+
+; CHECK:      define void @call_extern_func() {
+; CHECK-NEXT:   call void dso_local_equivalent @extern_func()
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+define void @call_extern_func() {
+
+  call void dso_local_equivalent @extern_func()
+  ret void
+}
+
+; CHECK:      define void @call_defined_extern_func() {
+; CHECK-NEXT:   call void dso_local_equivalent @defined_extern_func()
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+define void @call_defined_extern_func() {
+  call void dso_local_equivalent @defined_extern_func()
+  ret void
+}
+
+; CHECK:      define void @call_hidden_func() {
+; CHECK-NEXT:   call void dso_local_equivalent @hidden_func()
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+define void @call_hidden_func() {
+  call void dso_local_equivalent @hidden_func()
+  ret void
+}
diff --git a/llvm/test/Linker/Inputs/DSOLocalEquivalent.ll b/llvm/test/Linker/Inputs/DSOLocalEquivalent.ll
new file mode 100644
index 000000000000..efd34576861f
--- /dev/null
+++ b/llvm/test/Linker/Inputs/DSOLocalEquivalent.ll
@@ -0,0 +1,7 @@
+define void @defined_extern_func() {
+  ret void
+}
+
+define hidden void @hidden_func() {
+  ret void
+}
diff --git a/llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll b/llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll
new file mode 100644
index 000000000000..04fa0e9e682d
--- /dev/null
+++ b/llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll
@@ -0,0 +1,17 @@
+; RUN: llvm-as %s -o %t1.bc
+; RUN: llvm-lto2 run %t1.bc -o %t2.o -r=%t1.bc,caller,plx -r=%t1.bc,extern_func,plx -save-temps
+; RUN: llvm-dis %t2.o.0.5.precodegen.bc -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-fuchsia"
+
+declare void @extern_func()
+
+; CHECK:      define {{.*}} void @caller() {{.*}}{
+; CHECK-NEXT:   tail call void dso_local_equivalent @extern_func()
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+define void @caller() {
+  call void dso_local_equivalent @extern_func()
+  ret void
+}