[llvm] Change DSOLocalEquivalent type if the underlying global value type changes

[leonardchan]

Summary

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.

[phosek]

Is it feasible to cover this by a test?

llvm/lib/Transforms/Utils/ValueMapper.cpp

435

Super minor nit, I think that E would be more idiomatic here.

[leonardchan]

In D97978#2607908, @phosek wrote:

Is it feasible to cover this by a test?

Was going to add those and some more reviewers. Attempting to make a minimal reproducer for a test.

[leonardchan]

@dexonsmith @mehdi_amini Would you be the right folks to add for reviews on ValueMapper?

[mehdi_amini]

In D97978#2608176, @leonardchan wrote:

@dexonsmith @mehdi_amini Would you be the right folks to add for reviews on ValueMapper?

@dexonsmith is the expert :)

[phosek]

I'm not very familiar with this part of the codebase, but looking through the change, this looks reasonable to me.

llvm/lib/IR/Constants.cpp

Value *DSOLocalEquivalent::handleOperandChangeImpl(Value *From, Value *To) {
   DSOLocalEquivalent *&NewEquiv = getContext().pImpl->DSOLocalEquivalents[Func];
   if (NewEquiv)
     return llvm::ConstantExpr::getBitCast(NewEquiv, getType());
 
   // Replace this with the new one.
   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;
 }
 
 //---- ConstantExpr::get() implementations.
 //
 
 /// This is a utility function to handle folding of casts and lookup of the
 /// cast in the ExprConstants map. It is used by the various get* methods below.

llvm/lib/Transforms/Utils/ValueMapper.cpp

Value *Mapper::mapValue(const Value *V) {
   // is something that is not in the mapping table.
   Constant *C = const_cast<Constant*>(dyn_cast<Constant>(V));
   if (!C)
     return nullptr;
 
   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)) &&
            "Unexpected null mapping for constant operand without "
            "NullMapMissingGlobalValues flag");
     return Mapped;
   };
 

llvm/test/Linker/DSOLocalEquivalent.ll

+; 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
+}

llvm/test/Linker/Inputs/DSOLocalEquivalent.ll

+define void @defined_extern_func() {
+  ret void
+}
+
+define hidden void @hidden_func() {
+  ret void
+}

llvm/test/ThinLTO/X86/DSOLocalEquivalent.ll

+; 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
+}