samitolvanen updated this revision to Diff 473321.
samitolvanen marked 6 inline comments as done.
samitolvanen added a comment.

Addressed Fangrui's feedback.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D135411/new/

https://reviews.llvm.org/D135411

Files:
  clang/lib/CodeGen/BackendUtil.cpp
  clang/lib/Driver/ToolChain.cpp
  llvm/include/llvm/InitializePasses.h
  llvm/include/llvm/Transforms/Instrumentation/KCFI.h
  llvm/lib/Passes/PassBuilder.cpp
  llvm/lib/Passes/PassRegistry.def
  llvm/lib/Transforms/Instrumentation/CMakeLists.txt
  llvm/lib/Transforms/Instrumentation/KCFI.cpp
  llvm/test/Transforms/KCFI/kcfi-patchable-function-prefix.ll
  llvm/test/Transforms/KCFI/kcfi-supported.ll
  llvm/test/Transforms/KCFI/kcfi.ll
  llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn

Index: llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
===================================================================
--- llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
+++ llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
@@ -20,6 +20,7 @@
     "InstrOrderFile.cpp",
     "InstrProfiling.cpp",
     "Instrumentation.cpp",
+    "KCFI.cpp",
     "MemProfiler.cpp",
     "MemorySanitizer.cpp",
     "PGOInstrumentation.cpp",
Index: llvm/test/Transforms/KCFI/kcfi.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/KCFI/kcfi.ll
@@ -0,0 +1,21 @@
+; RUN: opt -S -passes=kcfi < %s | FileCheck %s
+
+; CHECK-LABEL: define void @f1(
+define void @f1(ptr noundef %x) {
+  ; CHECK:      %[[#GEPI:]] = getelementptr inbounds i32, ptr %x, i32 -1
+  ; CHECK-NEXT: %[[#LOAD:]] = load i32, ptr %[[#GEPI]], align 4
+  ; CHECK-NEXT: %[[#ICMP:]] = icmp ne i32 %[[#LOAD]], 12345678
+  ; CHECK-NEXT: br i1 %[[#ICMP]], label %[[#TRAP:]], label %[[#CALL:]], !prof ![[#WEIGHTS:]]
+  ; CHECK:      [[#TRAP]]:
+  ; CHECK-NEXT: call void @llvm.trap()
+  ; CHECK-NEXT: br label %[[#CALL]]
+  ; CHECK:      [[#CALL]]:
+  ; CHECK-NEXT: call void %x()
+  ; CHECK-NOT:  [ "kcfi"(i32 12345678) ]
+  call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"kcfi", i32 1}
+; CHECK: ![[#WEIGHTS]] = !{!"branch_weights", i32 1, i32 1048575}
Index: llvm/test/Transforms/KCFI/kcfi-supported.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/KCFI/kcfi-supported.ll
@@ -0,0 +1,18 @@
+; REQUIRES: x86-registered-target
+; RUN: opt -S -passes=kcfi < %s | FileCheck %s
+
+;; If the back-end supports KCFI operand bundle lowering, KCFIPass should be a no-op.
+
+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"
+
+; CHECK-LABEL: define void @f1
+define void @f1(ptr noundef %x) {
+  ; CHECK-NOT: call void @llvm.trap()
+  ; CHECK: call void %x() [ "kcfi"(i32 12345678) ]
+  call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"kcfi", i32 1}
Index: llvm/test/Transforms/KCFI/kcfi-patchable-function-prefix.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/KCFI/kcfi-patchable-function-prefix.ll
@@ -0,0 +1,12 @@
+; RUN: not opt -S -passes=kcfi < %s 2>&1 | FileCheck %s
+
+; CHECK: LLVM ERROR: -fpatchable-function-entry=N,M, where M>0 is not compatible with -fsanitize=kcfi on this target
+define void @f1(ptr noundef %x) #0 {
+  call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+attributes #0 = { "patchable-function-prefix"="1" }
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"kcfi", i32 1}
Index: llvm/lib/Transforms/Instrumentation/KCFI.cpp
===================================================================
--- /dev/null
+++ llvm/lib/Transforms/Instrumentation/KCFI.cpp
@@ -0,0 +1,101 @@
+//===-- KCFI.cpp - Generic KCFI operand bundle lowering ---------*- 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 pass emits generic KCFI indirect call checks for targets that don't
+// support lowering KCFI operand bundles in the back-end.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Instrumentation/KCFI.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalObject.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/Transforms/Instrumentation.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "kcfi"
+
+STATISTIC(NumKCFIChecks, "Number of kcfi operands transformed into checks");
+
+PreservedAnalyses KCFIPass::run(Function &F, FunctionAnalysisManager &AM) {
+  Module &M = *F.getParent();
+  if (!M.getModuleFlag("kcfi") ||
+      (TM &&
+       TM->getSubtargetImpl(F)->getTargetLowering()->supportKCFIBundles()))
+    return PreservedAnalyses::all();
+
+  // Find call instructions with KCFI operand bundles.
+  SmallVector<CallInst *> KCFICalls;
+  for (Instruction &I : instructions(F)) {
+    if (auto *CI = dyn_cast<CallInst>(&I))
+      if (CI->getOperandBundle(LLVMContext::OB_kcfi))
+        KCFICalls.push_back(CI);
+  }
+
+  if (KCFICalls.empty())
+    return PreservedAnalyses::all();
+
+  // patchable-function-prefix emits nops between the KCFI type identifier
+  // and the function start. As we don't know the size of the emitted nops,
+  // don't allow this attribute with generic lowering.
+  if (F.hasFnAttribute("patchable-function-prefix"))
+    report_fatal_error("-fpatchable-function-entry=N,M, where M>0 is not "
+                       "compatible with -fsanitize=kcfi on this target",
+                       /*GenCrashDiag=*/false);
+
+  LLVMContext &Ctx = M.getContext();
+  IntegerType *Int32Ty = Type::getInt32Ty(Ctx);
+  MDNode *VeryUnlikelyWeights =
+      MDBuilder(Ctx).createBranchWeights(1, (1U << 20) - 1);
+
+  for (CallInst *CI : KCFICalls) {
+    // Get the expected hash value.
+    const uint32_t ExpectedHash =
+        cast<ConstantInt>(CI->getOperandBundle(LLVMContext::OB_kcfi)->Inputs[0])
+            ->getZExtValue();
+
+    // Drop the KCFI operand bundle.
+    CallBase *Call =
+        CallBase::removeOperandBundle(CI, LLVMContext::OB_kcfi, CI);
+    assert(Call != CI);
+    Call->copyMetadata(*CI);
+    CI->replaceAllUsesWith(Call);
+    CI->eraseFromParent();
+
+    if (!Call->isIndirectCall())
+      continue;
+
+    // Emit a check and trap if the target hash doesn't match.
+    IRBuilder<> Builder(Call);
+    Value *HashPtr = Builder.CreateConstInBoundsGEP1_32(
+        Int32Ty, Call->getCalledOperand(), -1);
+    Value *Test = Builder.CreateICmpNE(Builder.CreateLoad(Int32Ty, HashPtr),
+                                       ConstantInt::get(Int32Ty, ExpectedHash));
+    Instruction *ThenTerm =
+        SplitBlockAndInsertIfThen(Test, Call, false, VeryUnlikelyWeights);
+    Builder.SetInsertPoint(ThenTerm);
+    Builder.CreateCall(Intrinsic::getDeclaration(&M, Intrinsic::trap));
+    ++NumKCFIChecks;
+  }
+
+  return PreservedAnalyses::none();
+}
Index: llvm/lib/Transforms/Instrumentation/CMakeLists.txt
===================================================================
--- llvm/lib/Transforms/Instrumentation/CMakeLists.txt
+++ llvm/lib/Transforms/Instrumentation/CMakeLists.txt
@@ -11,6 +11,7 @@
   Instrumentation.cpp
   InstrOrderFile.cpp
   InstrProfiling.cpp
+  KCFI.cpp
   PGOInstrumentation.cpp
   PGOMemOPSizeOpt.cpp
   PoisonChecking.cpp
Index: llvm/lib/Passes/PassRegistry.def
===================================================================
--- llvm/lib/Passes/PassRegistry.def
+++ llvm/lib/Passes/PassRegistry.def
@@ -322,6 +322,7 @@
 FUNCTION_PASS("newgvn", NewGVNPass())
 FUNCTION_PASS("jump-threading", JumpThreadingPass())
 FUNCTION_PASS("partially-inline-libcalls", PartiallyInlineLibCallsPass())
+FUNCTION_PASS("kcfi", KCFIPass(TM))
 FUNCTION_PASS("lcssa", LCSSAPass())
 FUNCTION_PASS("loop-data-prefetch", LoopDataPrefetchPass())
 FUNCTION_PASS("loop-load-elim", LoopLoadEliminationPass())
Index: llvm/lib/Passes/PassBuilder.cpp
===================================================================
--- llvm/lib/Passes/PassBuilder.cpp
+++ llvm/lib/Passes/PassBuilder.cpp
@@ -136,6 +136,7 @@
 #include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/InstrOrderFile.h"
 #include "llvm/Transforms/Instrumentation/InstrProfiling.h"
+#include "llvm/Transforms/Instrumentation/KCFI.h"
 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
Index: llvm/include/llvm/Transforms/Instrumentation/KCFI.h
===================================================================
--- /dev/null
+++ llvm/include/llvm/Transforms/Instrumentation/KCFI.h
@@ -0,0 +1,31 @@
+//===-- KCFI.h - Generic KCFI operand bundle lowering -----------*- 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 pass emits generic KCFI indirect call checks for targets that don't
+// support lowering KCFI operand bundles in the back-end.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_KCFI_H
+#define LLVM_TRANSFORMS_INSTRUMENTATION_KCFI_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+class TargetMachine;
+
+class KCFIPass : public PassInfoMixin<KCFIPass> {
+  TargetMachine *TM;
+
+public:
+  KCFIPass(TargetMachine *TM) : TM(TM) {}
+  static bool isRequired() { return true; }
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+} // namespace llvm
+#endif // LLVM_TRANSFORMS_INSTRUMENTATION_KCFI_H
Index: llvm/include/llvm/InitializePasses.h
===================================================================
--- llvm/include/llvm/InitializePasses.h
+++ llvm/include/llvm/InitializePasses.h
@@ -262,6 +262,7 @@
 void initializeLowerSwitchLegacyPassPass(PassRegistry &);
 void initializeLowerMatrixIntrinsicsLegacyPassPass(PassRegistry &);
 void initializeLowerMatrixIntrinsicsMinimalLegacyPassPass(PassRegistry &);
+void initializeKCFIPass(PassRegistry &);
 void initializeMIRAddFSDiscriminatorsPass(PassRegistry &);
 void initializeMIRCanonicalizerPass(PassRegistry &);
 void initializeMIRNamerPass(PassRegistry &);
Index: clang/lib/Driver/ToolChain.cpp
===================================================================
--- clang/lib/Driver/ToolChain.cpp
+++ clang/lib/Driver/ToolChain.cpp
@@ -1089,7 +1089,7 @@
        ~SanitizerKind::Function) |
       (SanitizerKind::CFI & ~SanitizerKind::CFIICall) |
       SanitizerKind::CFICastStrict | SanitizerKind::FloatDivideByZero |
-      SanitizerKind::UnsignedIntegerOverflow |
+      SanitizerKind::KCFI | SanitizerKind::UnsignedIntegerOverflow |
       SanitizerKind::UnsignedShiftBase | SanitizerKind::ImplicitConversion |
       SanitizerKind::Nullability | SanitizerKind::LocalBounds;
   if (getTriple().getArch() == llvm::Triple::x86 ||
@@ -1097,9 +1097,6 @@
       getTriple().getArch() == llvm::Triple::arm || getTriple().isWasm() ||
       getTriple().isAArch64() || getTriple().isRISCV())
     Res |= SanitizerKind::CFIICall;
-  if (getTriple().getArch() == llvm::Triple::x86_64 ||
-      getTriple().isAArch64(64))
-    Res |= SanitizerKind::KCFI;
   if (getTriple().getArch() == llvm::Triple::x86_64 ||
       getTriple().isAArch64(64) || getTriple().isRISCV())
     Res |= SanitizerKind::ShadowCallStack;
Index: clang/lib/CodeGen/BackendUtil.cpp
===================================================================
--- clang/lib/CodeGen/BackendUtil.cpp
+++ clang/lib/CodeGen/BackendUtil.cpp
@@ -71,6 +71,7 @@
 #include "llvm/Transforms/Instrumentation/GCOVProfiler.h"
 #include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/InstrProfiling.h"
+#include "llvm/Transforms/Instrumentation/KCFI.h"
 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h"
@@ -640,6 +641,31 @@
   }
 }
 
+static void addKCFIPass(TargetMachine *TM, const Triple &TargetTriple,
+                        const LangOptions &LangOpts, PassBuilder &PB) {
+  // If the back-end supports KCFI operand bundle lowering, skip KCFIPass.
+  if (TargetTriple.getArch() == llvm::Triple::x86_64 ||
+      TargetTriple.isAArch64(64))
+    return;
+
+  // Ensure we lower KCFI operand bundles with -O0.
+  PB.registerOptimizerLastEPCallback(
+      [&, TM](ModulePassManager &MPM, OptimizationLevel Level) {
+        if (Level == OptimizationLevel::O0 &&
+            LangOpts.Sanitize.has(SanitizerKind::KCFI))
+          MPM.addPass(createModuleToFunctionPassAdaptor(KCFIPass(TM)));
+      });
+
+  // When optimizations are requested, run KCIFPass after InstCombine to
+  // avoid unnecessary checks.
+  PB.registerPeepholeEPCallback(
+      [&, TM](FunctionPassManager &FPM, OptimizationLevel Level) {
+        if (Level != OptimizationLevel::O0 &&
+            LangOpts.Sanitize.has(SanitizerKind::KCFI))
+          FPM.addPass(KCFIPass(TM));
+      });
+}
+
 static void addSanitizers(const Triple &TargetTriple,
                           const CodeGenOptions &CodeGenOpts,
                           const LangOptions &LangOpts, PassBuilder &PB) {
@@ -938,8 +964,10 @@
 
     // Don't add sanitizers if we are here from ThinLTO PostLink. That already
     // done on PreLink stage.
-    if (!IsThinLTOPostLink)
+    if (!IsThinLTOPostLink) {
       addSanitizers(TargetTriple, CodeGenOpts, LangOpts, PB);
+      addKCFIPass(TM.get(), TargetTriple, LangOpts, PB);
+    }
 
     if (Optional<GCOVOptions> Options = getGCOVOptions(CodeGenOpts, LangOpts))
       PB.registerPipelineStartEPCallback(
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to