https://github.com/ilovepi created 
https://github.com/llvm/llvm-project/pull/125911

When using FatLTO, it is common to want to enable certain types of whole
program optimizations (WPD) or security transforms (CFI), so that they
can be made available when performing LTO. However, these transforms
should not be used when compiling the non-LTO object code. Since the
frontend must emit different IR, we cannot simply clone the module and
optimize the LTO section and non-LTO section differently to work around
this. Instead, we need to remove any problematic instruction sequences.

This patch adds a new pass whose responsibility is to clean up the IR
in the FatLTO pipeline after creating the bitcode section, which is
after running the pre-link pipeline but before running module
optimization. This allows us to safely drop any conflicting instructions
or IR constructs that are inappropriate for non-LTO compilation.

>From 955c40f3e8aab7bfb2c7c7ccc1225ed55b1bce6f Mon Sep 17 00:00:00 2001
From: Paul Kirth <paulki...@google.com>
Date: Thu, 30 Jan 2025 16:44:15 +0000
Subject: [PATCH] [llvm][fatlto] Add FatLTOCleanup pass

When using FatLTO, it is common to want to enable certain types of whole
program optimizations (WPD) or security transforms (CFI), so that they
can be made available when performing LTO. However, these transforms
should not be used when compiling the non-LTO object code. Since the
frontend must emit different IR, we cannot simply clone the module and
optimize the LTO section and non-LTO section differently to work around
this. Instead, we need to remove any problematic instruction sequences.

This patch adds a new pass whose responsibility is to clean up the IR
in the FatLTO pipeline after creating the bitcode section, which is
after running the pre-link pipeline but before running module
optimization. This allows us to safely drop any conflicting instructions
or IR constructs that are inappropriate for non-LTO compilation.
---
 .../test/CodeGen/fat-lto-objects-cfi-bug.cpp  |  34 ++++
 .../llvm/Transforms/IPO/FatLTOCleanup.h       |  36 ++++
 llvm/lib/Passes/PassBuilder.cpp               |   1 +
 llvm/lib/Passes/PassBuilderPipelines.cpp      |   7 +
 llvm/lib/Passes/PassRegistry.def              |   1 +
 llvm/lib/Transforms/IPO/CMakeLists.txt        |   1 +
 llvm/lib/Transforms/IPO/FatLTOCleanup.cpp     | 156 ++++++++++++++++++
 llvm/test/Transforms/FatLTOCleanup/basic.ll   |  46 ++++++
 8 files changed, 282 insertions(+)
 create mode 100644 clang/test/CodeGen/fat-lto-objects-cfi-bug.cpp
 create mode 100644 llvm/include/llvm/Transforms/IPO/FatLTOCleanup.h
 create mode 100644 llvm/lib/Transforms/IPO/FatLTOCleanup.cpp
 create mode 100644 llvm/test/Transforms/FatLTOCleanup/basic.ll

diff --git a/clang/test/CodeGen/fat-lto-objects-cfi-bug.cpp 
b/clang/test/CodeGen/fat-lto-objects-cfi-bug.cpp
new file mode 100644
index 00000000000000..2a61c5ae38ecf6
--- /dev/null
+++ b/clang/test/CodeGen/fat-lto-objects-cfi-bug.cpp
@@ -0,0 +1,34 @@
+// 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 -O1 -emit-llvm -o - -ffat-lto-objects \
+// RUN:      -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 %s 2>&1 | FileCheck %s 
--check-prefix=FATLTO
+
+// 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.
+// FATLTO-LABEL: entry:
+// FATLTO-NEXT:   %vtable = load ptr, ptr %p1
+// FATLTO-NOT: llvm.type.checked.load
+// FATLTO-NEXT:   %vfunc = load ptr, ptr %vtable
+// FATLTO-NEXT:   %call = tail call {{.*}} %vfunc(ptr {{.*}} %p1)
+// FATLTO-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 -O1 --emit-codegen-only -ffat-lto-objects \
+// RUN:      -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 %s
+
+class a {
+public:
+  virtual long b();
+};
+void c(a &p1) { p1.b(); }
diff --git a/llvm/include/llvm/Transforms/IPO/FatLTOCleanup.h 
b/llvm/include/llvm/Transforms/IPO/FatLTOCleanup.h
new file mode 100644
index 00000000000000..1c3de6c6cdd5bb
--- /dev/null
+++ 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
+// and 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);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_FATLTOCLEANUP_H
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index f698a3df08ef78..14cadb3cb84cdd 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -177,6 +177,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"
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp 
b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 4ec0fb8fc81ea4..b056c71c830eff 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/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"
@@ -1651,6 +1652,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
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index a93a995655a147..156ed8985d15df 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/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())
diff --git a/llvm/lib/Transforms/IPO/CMakeLists.txt 
b/llvm/lib/Transforms/IPO/CMakeLists.txt
index 15cb57399d2460..eb048fa814dd18 100644
--- a/llvm/lib/Transforms/IPO/CMakeLists.txt
+++ b/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
diff --git a/llvm/lib/Transforms/IPO/FatLTOCleanup.cpp 
b/llvm/lib/Transforms/IPO/FatLTOCleanup.cpp
new file mode 100644
index 00000000000000..96b13ae99508ca
--- /dev/null
+++ b/llvm/lib/Transforms/IPO/FatLTOCleanup.cpp
@@ -0,0 +1,156 @@
+//===- 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
+// and 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/ADT/SetVector.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 Changed = false;
+  // Use SetVector so we can rely on insertion order to drop instructions
+  // in a safe order (e.g. uses before defs), and avoid adding instructions
+  //  to the drop set multiple times.
+  SetVector<Instruction *> ToDrop;
+  for (Use &U : llvm::make_early_inc_range(CheckedLoadFn.uses())) {
+    for (auto *User : U->users()) {
+      Instruction *I = dyn_cast<Instruction>(User);
+      if (!I)
+        continue;
+      LLVM_DEBUG(dbgs() << "Checking candidate instruction" << *I << "\n");
+
+      IRBuilder<> IRB(I);
+      Value *Ptr = I->getOperand(0);
+      Value *Offset = I->getOperand(1);
+      Type *Ty = I->getType()->getContainedType(0);
+      // Use i8 so we can directly use the offset from llvm.type.checked.load.
+      Value *Gep = IRB.CreateGEP(Type::getInt8Ty(M.getContext()), Ptr, Offset);
+      LoadInst *L = IRB.CreateLoad(Ty, Gep, "vfunc");
+
+      bool Replaced = false;
+      for (auto *MaybeEv : I->users()) {
+        auto *EV = dyn_cast<ExtractValueInst>(MaybeEv);
+        if (!EV)
+          continue;
+
+        size_t Index = EV->getIndices()[0];
+        if (Index == 0) {
+          // If EV is extracting a vtable ptr, replace it w/ a direct load.
+          LLVM_DEBUG(llvm::dbgs()
+                     << "Replacing " << *EV << " with " << *L << "\n");
+          EV->replaceAllUsesWith(L);
+          ToDrop.insert(EV);
+          LLVM_DEBUG(dbgs()
+                     << DEBUG_TYPE << ": add " << *EV << "to drop set.\n");
+          Replaced = true;
+        } else if (Index == 1) {
+          // If EV is extracting the boolean, replace it w/ a constant, so
+          // that the branch to trap will be dropped later.
+          ConstantInt *C = ConstantInt::getTrue(M.getContext());
+          LLVM_DEBUG(llvm::dbgs() << DEBUG_TYPE << ": replacing " << *EV
+                                  << " with " << *C << "\n");
+          EV->replaceAllUsesWith(C);
+          ToDrop.insert(EV);
+          LLVM_DEBUG(dbgs()
+                     << DEBUG_TYPE << ": add " << *EV << "to drop set.\n");
+          Replaced = true;
+        }
+      }
+      if (Replaced) {
+        LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": add " << *I << "to drop set.\n");
+        ToDrop.insert(I);
+        Changed = true;
+      }
+    }
+  }
+  for (Instruction *I : ToDrop) {
+    LLVM_DEBUG(dbgs() << "Dropping instruction:" << *I << "\n");
+    I->eraseFromParent();
+  }
+  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);
+  if (TypeCheckedLoadRelFn)
+    Changed |= cleanUpTypeCheckedLoad(M, *TypeCheckedLoadRelFn);
+
+  if (Changed) {
+    // M.dump();
+    return PreservedAnalyses::none();
+  }
+
+  return PreservedAnalyses::all();
+}
diff --git a/llvm/test/Transforms/FatLTOCleanup/basic.ll 
b/llvm/test/Transforms/FatLTOCleanup/basic.ll
new file mode 100644
index 00000000000000..6067f09df0b2eb
--- /dev/null
+++ b/llvm/test/Transforms/FatLTOCleanup/basic.ll
@@ -0,0 +1,46 @@
+
+; RUN: opt -passes="fatlto-cleanup" -mtriple=x86_64-unknown-fuchsia < %s -S | 
FileCheck %s
+
+
+
+define hidden void @foo(ptr %p1) {
+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
+}
+
+; CHECK-LABEL: define hidden void @foo
+;  CHECK-NEXT: entry:
+;  CHECK-NEXT:  %vtable = load ptr, ptr %p1, align 8
+;  CHECK-NEXT:  %0 = getelementptr i8, ptr %vtable, i32 0
+;  CHECK-NEXT:  %vfunc = load ptr, ptr %0, align 8
+;  CHECK-NEXT:  br i1 true, label %cont2, label %trap1
+
+; CHECK-LABEL: trap1:
+;  CHECK-NEXT:  tail call void @llvm.ubsantrap(i8 2)
+;  CHECK-NEXT:  unreachable
+
+; CHECK-LABEL: cont2:
+;  CHECK-NEXT:  %call = tail call noundef i64 %vfunc(ptr noundef nonnull align 
8 dereferenceable(8) %p1)
+;  CHECK-NEXT:  ret void
+;  CHECK-NEXT:}
+
+; Function Attrs: cold noreturn nounwind
+declare void @llvm.ubsantrap(i8 immarg) #0
+
+; Function Attrs: nocallback nofree nosync nounwind willreturn memory(none)
+declare { ptr, i1 } @llvm.type.checked.load(ptr, i32, metadata) #1
+
+attributes #0 = { cold noreturn nounwind }
+attributes #1 = { nocallback nofree nosync nounwind willreturn memory(none) }

_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to