[llvm][fatlto] Add FatLTOCleanup pass

clang/test/CodeGen/fat-lto-objects-cfi.cpp

@@ -1,21 +1,24 @@
 // REQUIRES: x86-registered-target
 
+// RUN: rm -rf %t && split-file %s %t
 // RUN: %clang_cc1 -triple x86_64-unknown-fuchsia -O2 -flto -ffat-lto-objects \
-// RUN:          -fsanitize=cfi-icall -fsanitize-trap=cfi-icall -fvisibility=hidden  -emit-llvm -o - %s \
+// RUN:          -fsanitize=cfi-icall -fsanitize-trap=cfi-icall -fvisibility=hidden \
+// RUN:          -emit-llvm -o - %t/a.cpp \
 // RUN:   | FileCheck %s --check-prefix=TYPE_TEST
 
+//--- a.cpp
 // TYPE_TEST: llvm.embedded.object
 // TYPE_TEST-SAME: section ".llvm.lto"
 
+// COM: The FatLTO pipeline should remove all llvm.type.test instructions.
 // TYPE_TEST-LABEL: define hidden void @foo
-//      TYPE_TEST: entry:
-// TYPE_TEST-NEXT:   %cmp14.not = icmp eq i64 %len, 0
-// TYPE_TEST-NEXT:   br i1 %cmp14.not, label %for.end7, label %for.cond1.preheader.preheader
-//      TYPE_TEST: for.cond1.preheader.preheader:                    ; preds = %entry
-// TYPE_TEST-NEXT:   %arrayidx.1 = getelementptr inbounds nuw i8, ptr %ptr, i64 4
-// TYPE_TEST-NEXT:   br label %for.cond1.preheader
-
-// TYPE_TEST-NOT: @llvm.type.test
+// TYPE_TEST-NOT:   @llvm.type.test
+// TYPE_TEST-NEXT:  entry:
+// TYPE_TEST-NEXT:    %cmp14.not = icmp eq i64 %len, 0
+// TYPE_TEST-NEXT:    br i1 %cmp14.not, label %for.end7, label %for.cond1.preheader.preheader
+// TYPE_TEST-LABEL: for.cond1.preheader.preheader:                    ; preds = %entry
+// TYPE_TEST-NEXT:    %arrayidx.1 = getelementptr inbounds nuw i8, ptr %ptr, i64 4
+// TYPE_TEST-NEXT:    br label %for.cond1.preheader
 
 // The code below is a reduced case from https://github.com/llvm/llvm-project/issues/112053
 #define __PRINTFLIKE(__fmt, __varargs) __attribute__((__format__(__printf__, __fmt, __varargs)))
@@ -44,3 +47,44 @@ void foo(const void* ptr, size_t len, long disp_addr,
   }
 }
 
+//--- b.cpp
+// COM: Prior to the introduction of the FatLTO cleanup pass, this used to cause
+// COM: the backend to crash, either due to an assertion failure, or because
+// COM: the CFI instructions couldn't be correctly generated. So, check to make
+// COM: sure that the FatLTO pipeline used by clang does not regress.
+
+// COM: Check the generated IR doesn't contain llvm.type.checked.load in the final IR.
+// RUN: %clang_cc1 -triple=x86_64-unknown-fuchsia -O1 -emit-llvm -o - \
+// RUN:      -ffat-lto-objects -fvisibility=hidden \
+// RUN:      -fno-rtti -fsanitize=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fsanitize-trap=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fwhole-program-vtables %t/b.cpp 2>&1 | FileCheck %s --check-prefix=NO_CHECKED_LOAD
+
+// RUN: %clang_cc1 -triple=x86_64-unknown-fuchsia -O1 -emit-llvm -o - \
+// RUN:      -ffat-lto-objects -fvisibility=hidden -fexperimental-relative-c++-abi-vtables \
+// RUN:      -fno-rtti -fsanitize=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fsanitize-trap=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fwhole-program-vtables %t/b.cpp 2>&1 | FileCheck %s --check-prefix=NO_CHECKED_LOAD
+
+// COM: Note that the embedded bitcode section will contain references to
+// COM: llvm.type.checked.load, so we need to match the function body first.
+// NO_CHECKED_LOAD-LABEL: entry:
+// NO_CHECKED_LOAD-NEXT:    %vtable = load ptr, ptr %p1
+// NO_CHECKED_LOAD-NOT:   llvm.type.checked.load
+// NO_CHECKED_LOAD-NEXT:    %vfunc = load ptr, ptr %vtable
+// NO_CHECKED_LOAD-NEXT:    %call = tail call {{.*}} %vfunc(ptr {{.*}} %p1)
+// NO_CHECKED_LOAD-NEXT:    ret void
+
+// COM: Ensure that we don't crash in the backend anymore when clang uses
+// COM: CFI checks with -ffat-lto-objects.
+// RUN: %clang_cc1 -triple=x86_64-unknown-fuchsia -O1 -emit-codegen-only \
+// RUN:      -ffat-lto-objects -fvisibility=hidden \
+// RUN:      -fno-rtti -fsanitize=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fsanitize-trap=cfi-icall,cfi-mfcall,cfi-nvcall,cfi-vcall \
+// RUN:      -fwhole-program-vtables %t/b.cpp
+
+class a {
+public:
+  virtual long b();
+};
+void c(a &p1) { p1.b(); }

llvm/include/llvm/Transforms/IPO/FatLTOCleanup.h

@@ -0,0 +1,36 @@
+//===- FatLtoCleanup.h - clean up IR for the FatLTO pipeline ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines operations used to clean up IR for the FatLTO pipeline.
+// Instrumentation that is beneficial for bitcode sections used in LTO may
+// need to be cleaned up to finish non-LTO compilation. llvm.checked.load is
+// an example of an instruction that we want to preserve for LTO, but is
+// incorrect to leave unchanged during the per-TU compilation in FatLTO.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_IPO_FATLTOCLEANUP_H
+#define LLVM_TRANSFORMS_IPO_FATLTOCLEANUP_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class Module;
+class ModuleSummaryIndex;
+
+class FatLtoCleanup : public PassInfoMixin<FatLtoCleanup> {
+public:
+  FatLtoCleanup() {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+  static bool isRequired() { return true; }
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_FATLTOCLEANUP_H

llvm/lib/Passes/PassBuilder.cpp

@@ -184,6 +184,7 @@
 #include "llvm/Transforms/IPO/ElimAvailExtern.h"
 #include "llvm/Transforms/IPO/EmbedBitcodePass.h"
 #include "llvm/Transforms/IPO/ExpandVariadics.h"
+#include "llvm/Transforms/IPO/FatLTOCleanup.h"
 #include "llvm/Transforms/IPO/ForceFunctionAttrs.h"
 #include "llvm/Transforms/IPO/FunctionAttrs.h"
 #include "llvm/Transforms/IPO/FunctionImport.h"

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -53,6 +53,7 @@
 #include "llvm/Transforms/IPO/ElimAvailExtern.h"
 #include "llvm/Transforms/IPO/EmbedBitcodePass.h"
 #include "llvm/Transforms/IPO/ExpandVariadics.h"
+#include "llvm/Transforms/IPO/FatLTOCleanup.h"
 #include "llvm/Transforms/IPO/ForceFunctionAttrs.h"
 #include "llvm/Transforms/IPO/FunctionAttrs.h"
 #include "llvm/Transforms/IPO/GlobalDCE.h"
@@ -1670,6 +1671,12 @@ PassBuilder::buildFatLTODefaultPipeline(OptimizationLevel Level, bool ThinLTO,
     MPM.addPass(buildLTOPreLinkDefaultPipeline(Level));
   MPM.addPass(EmbedBitcodePass(ThinLTO, EmitSummary));
 
+  // Perform any cleanups to the IR that aren't suitable for per TU compilation,
+  // like removing CFI/WPD related instructions. Note, we reuse
+  // LowerTypeTestsPass to clean up type tests rather than duplicate that logic
+  // in FatLtoCleanup.
+  MPM.addPass(FatLtoCleanup());
+
   // If we're doing FatLTO w/ CFI enabled, we don't want the type tests in the
   // object code, only in the bitcode section, so drop it before we run
   // module optimization and generate machine code. If llvm.type.test() isn't in

llvm/lib/Passes/PassRegistry.def

@@ -98,6 +98,7 @@ MODULE_PASS("lower-emutls", LowerEmuTLSPass())
 MODULE_PASS("lower-global-dtors", LowerGlobalDtorsPass())
 MODULE_PASS("lower-ifunc", LowerIFuncPass())
 MODULE_PASS("lowertypetests", LowerTypeTestsPass())
+MODULE_PASS("fatlto-cleanup", FatLtoCleanup())
 MODULE_PASS("pgo-force-function-attrs", PGOForceFunctionAttrsPass(PGOOpt ? PGOOpt->ColdOptType : PGOOptions::ColdFuncOpt::Default))
 MODULE_PASS("memprof-context-disambiguation", MemProfContextDisambiguation())
 MODULE_PASS("memprof-module", ModuleMemProfilerPass())

llvm/lib/Transforms/IPO/CMakeLists.txt

@@ -14,6 +14,7 @@ add_llvm_component_library(LLVMipo
   EmbedBitcodePass.cpp
   ExpandVariadics.cpp
   ExtractGV.cpp
+  FatLTOCleanup.cpp
   ForceFunctionAttrs.cpp
   FunctionAttrs.cpp
   FunctionImport.cpp

llvm/lib/Transforms/IPO/FatLTOCleanup.cpp

@@ -0,0 +1,122 @@
+//===- FatLtoCleanup.cpp - clean up IR for the FatLTO pipeline --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines operations used to clean up IR for the FatLTO pipeline.
+// Instrumentation that is beneficial for bitcode sections used in LTO may
+// need to be cleaned up to finish non-LTO compilation. llvm.checked.load is
+// an example of an instruction that we want to preserve for LTO, but is
+// incorrect to leave unchanged during the per-TU compilation in FatLTO.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO/FatLTOCleanup.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/Use.h"
+#include "llvm/Support/Debug.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "fatlto-cleanup"
+
+namespace {
+// Replaces uses of llvm.type.checked.load instructions with unchecked loads.
+// In essence, we're undoing the frontends instrumentation, since it isn't
+// correct for the non-LTO part of a FatLTO object.
+//
+// llvm.type.checked.load instruction sequences always have a particular form:
+//
+// clang-format off
+//
+//   %0 = tail call { ptr, i1 } @llvm.type.checked.load(ptr %vtable, i32 0, metadata !"foo"), !nosanitize !0
+//   %1 = extractvalue { ptr, i1 } %0, 1, !nosanitize !0
+//   br i1 %1, label %cont2, label %trap1, !nosanitize !0
+//
+// trap1:                                            ; preds = %entry
+//   tail call void @llvm.ubsantrap(i8 2) #3, !nosanitize !0
+//   unreachable, !nosanitize !0
+//
+// cont2:                                            ; preds = %entry
+//   %2 = extractvalue { ptr, i1 } %0, 0, !nosanitize !0
+//   %call = tail call noundef i64 %2(ptr noundef nonnull align 8 dereferenceable(8) %p1) #4
+//
+// clang-format on
+//
+// In this sequence, the vtable pointer is first loaded and checked against some
+// metadata. The result indicates failure, then the program traps. On the
+// success path, the pointer is used to make an indirect call to the function
+// pointer loaded from the vtable.
+//
+// Since we won't be able to lower this correctly later in non-LTO builds, we
+// need to drop the special load and trap, and emit a normal load of the
+// function pointer from the vtable.
+//
+// This is straight forward, since the checked load can be replaced w/ a load
+// of the vtable pointer and a GEP instruction to index into the vtable and get
+// the correct method/function pointer. We replace the "check" with a constant
+// indicating success, which allows later passes to simplify control flow and
+// remove any now dead instructions.
+//
+// This logic holds for both llvm.type.checked.load and
+// llvm.type.checked.load.relative instructions.
+static bool cleanUpTypeCheckedLoad(Module &M, Function &CheckedLoadFn,
+                                   bool IsRelative) {
+  bool Changed = false;
+  for (User *User : llvm::make_early_inc_range(CheckedLoadFn.users())) {
+    Instruction *I = dyn_cast<Instruction>(User);
+    if (!I)
+      continue;
+    IRBuilder<> IRB(I);
+    Value *Ptr = I->getOperand(0);
+    Value *Offset = I->getOperand(1);
+    Type *PtrTy = I->getType()->getStructElementType(0);
+    ConstantInt *True = ConstantInt::getTrue(M.getContext());
+    Instruction *Load;
+    if (IsRelative) {
+      Load =
+          IRB.CreateIntrinsic(Intrinsic::load_relative, {Offset->getType()},
+                              {Ptr, Offset}, /*FMFSource=*/nullptr, "rel_load");
+    } else {
+      Value *PtrAdd = IRB.CreatePtrAdd(Ptr, Offset);
+      Load = IRB.CreateLoad(PtrTy, PtrAdd, "vfunc");
+    }
+
+    Value *Replacement = PoisonValue::get(I->getType());
+    Replacement = IRB.CreateInsertValue(Replacement, True, {1});
+    Replacement = IRB.CreateInsertValue(Replacement, Load, {0});
+    I->replaceAllUsesWith(Replacement);
+
+    LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": erase " << *I << "\n");
+    I->eraseFromParent();
+    Changed = true;
+  }
+  if (Changed)
+    CheckedLoadFn.eraseFromParent();
+  return Changed;
+}
+} // namespace
+
+PreservedAnalyses FatLtoCleanup::run(Module &M, ModuleAnalysisManager &AM) {
+  Function *TypeCheckedLoadFn =
+      Intrinsic::getDeclarationIfExists(&M, Intrinsic::type_checked_load);
+  Function *TypeCheckedLoadRelFn = Intrinsic::getDeclarationIfExists(
+      &M, Intrinsic::type_checked_load_relative);
+
+  bool Changed = false;
+  if (TypeCheckedLoadFn)
+    Changed |= cleanUpTypeCheckedLoad(M, *TypeCheckedLoadFn, false);
+  if (TypeCheckedLoadRelFn)
+    Changed |= cleanUpTypeCheckedLoad(M, *TypeCheckedLoadRelFn, true);
+
+  if (Changed)
+    return PreservedAnalyses::none();
+  return PreservedAnalyses::all();
+}

llvm/test/Transforms/FatLTOCleanup/remove-type-checked-load.ll

@@ -0,0 +1,75 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+
+; RUN: opt -passes="fatlto-cleanup" < %s -S | FileCheck %s
+
+declare void @llvm.ubsantrap(i8 immarg)
+declare { ptr, i1 } @llvm.type.checked.load(ptr, i32, metadata)
+declare { ptr, i1 } @llvm.type.checked.load.relative(ptr, i32, metadata)
+
+define hidden void @foo(ptr %p1) {
+; CHECK-LABEL: define hidden void @foo(
+; CHECK-SAME: ptr [[P1:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[P1]], align 8
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[VTABLE]], i32 0
+; CHECK-NEXT:    [[VFUNC:%.*]] = load ptr, ptr [[TMP0]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { ptr, i1 } { ptr poison, i1 true }, ptr [[VFUNC]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractvalue { ptr, i1 } [[TMP2]], 1
+; CHECK-NEXT:    br i1 [[TMP3]], label %[[CONT2:.*]], label %[[TRAP1:.*]]
+; CHECK:       [[TRAP1]]:
+; CHECK-NEXT:    tail call void @llvm.ubsantrap(i8 2)
+; CHECK-NEXT:    unreachable
+; CHECK:       [[CONT2]]:
+; CHECK-NEXT:    [[TMP4:%.*]] = extractvalue { ptr, i1 } [[TMP2]], 0
+; CHECK-NEXT:    [[CALL:%.*]] = tail call noundef i64 [[TMP4]](ptr noundef nonnull align 8 dereferenceable(8) [[P1]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %vtable = load ptr, ptr %p1, align 8
+  %0 = tail call { ptr, i1 } @llvm.type.checked.load(ptr %vtable, i32 0, metadata !"_ZTS1a")
+  %1 = extractvalue { ptr, i1 } %0, 1
+  br i1 %1, label %cont2, label %trap1
+
+trap1:
+  tail call void @llvm.ubsantrap(i8 2)
+  unreachable
+
+cont2:
+  %2 = extractvalue { ptr, i1 } %0, 0
+  %call = tail call noundef i64 %2(ptr noundef nonnull align 8 dereferenceable(8) %p1)
+  ret void
+}
+
+define hidden void @relative.vtable(ptr %p1) {
+; CHECK-LABEL: define hidden void @relative.vtable(
+; CHECK-SAME: ptr [[P1:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[P1]], align 8
+; CHECK-NEXT:    [[REL_LOAD:%.*]] = call ptr @llvm.load.relative.i32(ptr [[VTABLE]], i32 0)
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { ptr, i1 } { ptr poison, i1 true }, ptr [[REL_LOAD]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractvalue { ptr, i1 } [[TMP2]], 1
+; CHECK-NEXT:    br i1 [[TMP3]], label %[[CONT2:.*]], label %[[TRAP1:.*]]
+; CHECK:       [[TRAP1]]:
+; CHECK-NEXT:    tail call void @llvm.ubsantrap(i8 2)
+; CHECK-NEXT:    unreachable
+; CHECK:       [[CONT2]]:
+; CHECK-NEXT:    [[TMP4:%.*]] = extractvalue { ptr, i1 } [[TMP2]], 0
+; CHECK-NEXT:    [[CALL:%.*]] = tail call noundef i64 [[TMP4]](ptr noundef nonnull align 8 dereferenceable(8) [[P1]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %vtable = load ptr, ptr %p1, align 8
+  %0 = tail call { ptr, i1 } @llvm.type.checked.load.relative(ptr %vtable, i32 0, metadata !"rel.vtable.type")
+  %1 = extractvalue { ptr, i1 } %0, 1
+  br i1 %1, label %cont2, label %trap1
+
+trap1:
+  tail call void @llvm.ubsantrap(i8 2)
+  unreachable
+
+cont2:
+  %2 = extractvalue { ptr, i1 } %0, 0
+  %call = tail call noundef i64 %2(ptr noundef nonnull align 8 dereferenceable(8) %p1)
+  ret void
+}
+