[memtag] Plug in stack safety analysis.

[eugenis]

Summary

Run StackSafetyAnalysis at the end of the IR pipeline and annotate
proven safe allocas with !stack-safe metadata. Do not instrument such
allocas in the AArch64StackTagging pass.

[eugenis]

To explain some design decisions:

The analysis is added to the IR pipeline and not to the AArch64 codegen pipeline mainly because we plan to extend it with ThinLTO support in the future. To do that, module summary builder will need to depend on the (function-)local stack safety analysis pass, and the global analysis result will be provided by the combined summary reader, and all of that must happen in the target-independent IR pipeline.

Another reason is that the legacy pass manager, which is still used for codegen, is not good at mixing function and module analyses. AFAIK a module pass that depends on function analysis will always produce an on-the-fly pass manager which can not reuse earlier analysis results; and a function pass can not depend on a module analysis at all (at least I could not make it work).

For these two reasons I opted not to use analysis manager framework to pass the safety information to AArch64StackTaggingPass, and chose to use alloca metadata instead.

[merge_guards_bot]

Unit tests: pass. 62250 tests passed, 0 failed and 816 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[vitalybuka]

llvm/lib/Analysis/StackSafetyAnalysis.cpp

670

Could you please add basic test in StackSafetyAnalysis to checks "stack-safe" in IR?

[eugenis]

+ a test

clang/lib/CodeGen/BackendUtil.cpp

 #include "clang/Basic/TargetOptions.h"
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Lex/HeaderSearchOptions.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/Bitcode/BitcodeWriterPass.h"
 #include "llvm/CodeGen/RegAllocRegistry.h"
 #include "llvm/CodeGen/SchedulerRegistry.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"

 static void addDataFlowSanitizerPass(const PassManagerBuilder &Builder,
                                      legacy::PassManagerBase &PM) {
   const PassManagerBuilderWrapper &BuilderWrapper =
       static_cast<const PassManagerBuilderWrapper&>(Builder);
   const LangOptions &LangOpts = BuilderWrapper.getLangOpts();
   PM.add(createDataFlowSanitizerPass(LangOpts.SanitizerBlacklistFiles));
 }
 
+static void addMemTagOptimizationPasses(const PassManagerBuilder &Builder,
+                                        legacy::PassManagerBase &PM) {
+  PM.add(createStackSafetyGlobalInfoWrapperPass(/*SetMetadata=*/true));
+}
+
 static TargetLibraryInfoImpl *createTLII(llvm::Triple &TargetTriple,
                                          const CodeGenOptions &CodeGenOpts) {
   TargetLibraryInfoImpl *TLII = new TargetLibraryInfoImpl(TargetTriple);
 
   switch (CodeGenOpts.getVecLib()) {
   case CodeGenOptions::Accelerate:
     TLII->addVectorizableFunctionsFromVecLib(TargetLibraryInfoImpl::Accelerate);
     break;
void EmitAssemblyHelper::CreatePasses(legacy::PassManager &MPM,
 
   if (LangOpts.Sanitize.has(SanitizerKind::DataFlow)) {
     PMBuilder.addExtension(PassManagerBuilder::EP_OptimizerLast,
                            addDataFlowSanitizerPass);
     PMBuilder.addExtension(PassManagerBuilder::EP_EnabledOnOptLevel0,
                            addDataFlowSanitizerPass);
   }
 
+  if (LangOpts.Sanitize.has(SanitizerKind::MemTag)) {
+    PMBuilder.addExtension(PassManagerBuilder::EP_OptimizerLast,
+                           addMemTagOptimizationPasses);
+  }
+
   // Set up the per-function pass manager.
   FPM.add(new TargetLibraryInfoWrapperPass(*TLII));
   if (CodeGenOpts.VerifyModule)
     FPM.add(createVerifierPass());
 
   // Set up the per-module pass manager.
   if (!CodeGenOpts.RewriteMapFiles.empty())
     addSymbolRewriterPass(CodeGenOpts, &MPM);
if (LangOpts.Sanitize.has(SanitizerKind::HWAddress)) {
       MPM.addPass(HWAddressSanitizerPass(
           /*CompileKernel=*/false, Recover));
     }
     if (LangOpts.Sanitize.has(SanitizerKind::KernelHWAddress)) {
       MPM.addPass(HWAddressSanitizerPass(
           /*CompileKernel=*/true, /*Recover=*/true));
     }
 
+    if (CodeGenOpts.OptimizationLevel > 0 &&
+        LangOpts.Sanitize.has(SanitizerKind::MemTag)) {
+      MPM.addPass(StackSafetyGlobalAnnotatorPass());
+    }
+
     if (CodeGenOpts.OptimizationLevel == 0) {
       addCoroutinePassesAtO0(MPM, LangOpts, CodeGenOpts);
       addSanitizersAtO0(MPM, TargetTriple, LangOpts, CodeGenOpts);
     }
   }
 
   // FIXME: We still use the legacy pass manager to do code generation. We
   // create that pass manager here and use it as needed below.

clang/test/Driver/memtag.c

+// REQUIRES: aarch64-registered-target
+
+// Old pass manager.
+// RUN: %clang     -fno-experimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-NO-SAFETY
+// RUN: %clang -O1 -fno-experimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+// RUN: %clang -O2 -fno-experimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+// RUN: %clang -O3 -fno-experimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+
+// New pass manager.
+// RUN: %clang     -fexperimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-NO-SAFETY
+// RUN: %clang -O1 -fexperimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+// RUN: %clang -O2 -fexperimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+// RUN: %clang -O3 -fexperimental-new-pass-manager -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag %s -S -emit-llvm -o - | FileCheck %s --check-prefix=CHECK-SAFETY
+
+int z;
+__attribute__((noinline)) void use(int *p) { *p = z; }
+int foo() { int x; use(&x); return x; }
+
+// CHECK-NO-SAFETY: define dso_local i32 @foo()
+// CHECK-NO-SAFETY: %x = alloca i32, align 4{{$}}
+
+// CHECK-SAFETY: define dso_local i32 @foo()
+// CHECK-SAFETY: %x = alloca i32, align 4, !stack-safe

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

 
 public:
   StackSafetyInfo();
   StackSafetyInfo(FunctionInfo &&Info);
   StackSafetyInfo(StackSafetyInfo &&);
   StackSafetyInfo &operator=(StackSafetyInfo &&);
   ~StackSafetyInfo();
 
+  FunctionInfo *getInfo() const { return Info.get(); }
+
   // TODO: Add useful for client methods.
   void print(raw_ostream &O) const;
 };
 
 /// StackSafetyInfo wrapper for the new pass manager.
 class StackSafetyAnalysis : public AnalysisInfoMixin<StackSafetyAnalysis> {
   friend AnalysisInfoMixin<StackSafetyAnalysis>;
   static AnalysisKey Key;
class StackSafetyGlobalPrinterPass
     : public PassInfoMixin<StackSafetyGlobalPrinterPass> {
   raw_ostream &OS;
 
 public:
   explicit StackSafetyGlobalPrinterPass(raw_ostream &OS) : OS(OS) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+class StackSafetyGlobalAnnotatorPass
+    : public PassInfoMixin<StackSafetyGlobalAnnotatorPass> {
+
+public:
+  explicit StackSafetyGlobalAnnotatorPass() {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 /// This pass performs the global (interprocedural) stack safety analysis
 /// (legacy pass manager).
 class StackSafetyGlobalInfoWrapperPass : public ModulePass {
-  StackSafetyGlobalInfo SSI;
+  StackSafetyGlobalInfo SSGI;
+  bool SetMetadata;
 
 public:
   static char ID;
 
-  StackSafetyGlobalInfoWrapperPass();
+  StackSafetyGlobalInfoWrapperPass(bool SetMetadata = false);
 
-  const StackSafetyGlobalInfo &getResult() const { return SSI; }
+  const StackSafetyGlobalInfo &getResult() const { return SSGI; }
 
   void print(raw_ostream &O, const Module *M) const override;
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   bool runOnModule(Module &M) override;
 };
 
+ModulePass *createStackSafetyGlobalInfoWrapperPass(bool SetMetadata);
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_STACKSAFETYANALYSIS_H

llvm/lib/Analysis/StackSafetyAnalysis.cpp

 raw_ostream &operator<<(raw_ostream &OS, const UseInfo &U) {
   OS << U.Range;
   for (auto &Call : U.Calls)
     OS << ", " << Call;
   return OS;
 }
 
 struct AllocaInfo {
-  const AllocaInst *AI = nullptr;
+  AllocaInst *AI = nullptr;
   uint64_t Size = 0;
   UseInfo Use;
 
-  AllocaInfo(unsigned PointerSize, const AllocaInst *AI, uint64_t Size)
+  AllocaInfo(unsigned PointerSize, AllocaInst *AI, uint64_t Size)
       : AI(AI), Size(Size), Use(PointerSize) {}
 
   StringRef getName() const { return AI->getName(); }
 };
 
 raw_ostream &operator<<(raw_ostream &OS, const AllocaInfo &A) {
   return OS << A.getName() << "[" << A.Size << "]: " << A.Use;
 }
for (unsigned ArgNo = 0; ArgNo < Type->getNumParams(); ArgNo++) {
     UseInfo &US = Params.back().Use;
     US.Calls.emplace_back(Aliasee, ArgNo, ConstantRange(APInt(PointerSize, 0)));
   }
 }
 
 namespace {
 
 class StackSafetyLocalAnalysis {
-  const Function &F;
+  Function &F;
   const DataLayout &DL;
   ScalarEvolution &SE;
   unsigned PointerSize = 0;
 
   const ConstantRange UnknownRange;
 
   ConstantRange offsetFromAlloca(Value *Addr, const Value *AllocaPtr);
   ConstantRange getAccessRange(Value *Addr, const Value *AllocaPtr,
                                ConstantRange SizeRange);
   ConstantRange getAccessRange(Value *Addr, const Value *AllocaPtr,
                                TypeSize Size);
   ConstantRange getMemIntrinsicAccessRange(const MemIntrinsic *MI, const Use &U,
                                            const Value *AllocaPtr);
 
   bool analyzeAllUses(const Value *Ptr, UseInfo &AS);
 
   ConstantRange getRange(uint64_t Lower, uint64_t Upper) const {
     return ConstantRange(APInt(PointerSize, Lower), APInt(PointerSize, Upper));
   }
 
 public:
-  StackSafetyLocalAnalysis(const Function &F, ScalarEvolution &SE)
+  StackSafetyLocalAnalysis(Function &F, ScalarEvolution &SE)
       : F(F), DL(F.getParent()->getDataLayout()), SE(SE),
         PointerSize(DL.getPointerSizeInBits()),
         UnknownRange(PointerSize, true) {}
 
   // Run the transformation on the associated function.
   StackSafetyInfo run();
 };
 

 
 PreservedAnalyses StackSafetyGlobalPrinterPass::run(Module &M,
                                                     ModuleAnalysisManager &AM) {
   OS << "'Stack Safety Analysis' for module '" << M.getName() << "'\n";
   print(AM.getResult<StackSafetyGlobalAnalysis>(M), OS, M);
   return PreservedAnalyses::all();
 }
 
+static bool SetStackSafetyMetadata(Module &M,
+                                   const StackSafetyGlobalInfo &SSGI) {
+  bool Changed = false;
+  unsigned Width = M.getDataLayout().getPointerSizeInBits();
+  for (auto &F : M.functions()) {
+    if (F.isDeclaration() || F.hasOptNone())
+      continue;
+    auto Iter = SSGI.find(&F);
+    if (Iter == SSGI.end())
+      continue;
+    StackSafetyInfo::FunctionInfo *Summary = Iter->second.getInfo();
+    for (auto &AS : Summary->Allocas) {
+      ConstantRange AllocaRange{APInt(Width, 0), APInt(Width, AS.Size)};
+      if (AllocaRange.contains(AS.Use.Range)) {
+        AS.AI->setMetadata(M.getMDKindID("stack-safe"),
+                           MDNode::get(M.getContext(), None));
+        Changed = true;
+      }
+    }
+  }
+  return Changed;
+}
+
+PreservedAnalyses
+StackSafetyGlobalAnnotatorPass::run(Module &M, ModuleAnalysisManager &AM) {
+  auto &SSGI = AM.getResult<StackSafetyGlobalAnalysis>(M);
+  (void)SetStackSafetyMetadata(M, SSGI);
+  return PreservedAnalyses::all();
+}
+
 char StackSafetyGlobalInfoWrapperPass::ID = 0;
 
-StackSafetyGlobalInfoWrapperPass::StackSafetyGlobalInfoWrapperPass()
-    : ModulePass(ID) {
+StackSafetyGlobalInfoWrapperPass::StackSafetyGlobalInfoWrapperPass(bool SetMetadata)
+    : ModulePass(ID), SetMetadata(SetMetadata) {
   initializeStackSafetyGlobalInfoWrapperPassPass(
       *PassRegistry::getPassRegistry());
 }
 
 void StackSafetyGlobalInfoWrapperPass::print(raw_ostream &O,
                                              const Module *M) const {
-  ::print(SSI, O, *M);
+  ::print(SSGI, O, *M);
 }
 
 void StackSafetyGlobalInfoWrapperPass::getAnalysisUsage(
     AnalysisUsage &AU) const {
   AU.addRequired<StackSafetyInfoWrapperPass>();
 }
 
 bool StackSafetyGlobalInfoWrapperPass::runOnModule(Module &M) {
   StackSafetyDataFlowAnalysis SSDFA(
       M, [this](Function &F) -> const StackSafetyInfo & {
         return getAnalysis<StackSafetyInfoWrapperPass>(F).getResult();
       });
-  SSI = SSDFA.run();
-  return false;
+  SSGI = SSDFA.run();
+  return SetMetadata ? SetStackSafetyMetadata(M, SSGI) : false;
+}
+
+ModulePass *llvm::createStackSafetyGlobalInfoWrapperPass(bool SetMetadata) {
+  return new StackSafetyGlobalInfoWrapperPass(SetMetadata);
 }
 
 static const char LocalPassArg[] = "stack-safety-local";
 static const char LocalPassName[] = "Stack Safety Local Analysis";
 INITIALIZE_PASS_BEGIN(StackSafetyInfoWrapperPass, LocalPassArg, LocalPassName,
                       false, true)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_END(StackSafetyInfoWrapperPass, LocalPassArg, LocalPassName,
                     false, true)
 
 static const char GlobalPassName[] = "Stack Safety Analysis";
 INITIALIZE_PASS_BEGIN(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
                       GlobalPassName, false, false)
 INITIALIZE_PASS_DEPENDENCY(StackSafetyInfoWrapperPass)
 INITIALIZE_PASS_END(StackSafetyGlobalInfoWrapperPass, DEBUG_TYPE,
                     GlobalPassName, false, false)

llvm/lib/Passes/PassRegistry.def

 MODULE_PASS("wholeprogramdevirt", WholeProgramDevirtPass(nullptr, nullptr))
 MODULE_PASS("verify", VerifierPass())
 MODULE_PASS("asan-module", ModuleAddressSanitizerPass(/*CompileKernel=*/false, false, true, false))
 MODULE_PASS("msan-module", MemorySanitizerPass({}))
 MODULE_PASS("tsan-module", ThreadSanitizerPass())
 MODULE_PASS("kasan-module", ModuleAddressSanitizerPass(/*CompileKernel=*/true, false, true, false))
 MODULE_PASS("sancov-module", ModuleSanitizerCoveragePass())
 MODULE_PASS("poison-checking", PoisonCheckingPass())
+MODULE_PASS("stack-safety-annotator", StackSafetyGlobalAnnotatorPass())
 #undef MODULE_PASS
 
 #ifndef CGSCC_ANALYSIS
 #define CGSCC_ANALYSIS(NAME, CREATE_PASS)
 #endif
 CGSCC_ANALYSIS("no-op-cgscc", NoOpCGSCCAnalysis())
 CGSCC_ANALYSIS("fam-proxy", FunctionAnalysisManagerCGSCCProxy())
 CGSCC_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

bool AArch64StackTagging::isInterestingAlloca(const AllocaInst &AI) {
   bool IsInteresting =
       AI.getAllocatedType()->isSized() && AI.isStaticAlloca() &&
       // alloca() may be called with 0 size, ignore it.
       AI.getAllocationSizeInBits(*DL).getValue() > 0 &&
       // inalloca allocas are not treated as static, and we don't want
       // dynamic alloca instrumentation for them as well.
       !AI.isUsedWithInAlloca() &&
       // swifterror allocas are register promoted by ISel
-      !AI.isSwiftError();
+      !AI.isSwiftError() &&
+      // safe allocas are not interesting
+      !AI.getMetadata("stack-safe");
   return IsInteresting;
 }
 
 void AArch64StackTagging::tagAlloca(AllocaInst *AI, Instruction *InsertBefore,
                                     Value *Ptr, uint64_t Size) {
   auto SetTagZeroFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag_zero);
   auto StgpFunc =

llvm/test/Analysis/StackSafetyAnalysis/ipa-attr.ll

+; RUN: llvm-as %s -o %t0.bc
+; RUN: llvm-as %S/Inputs/ipa.ll -o %t1.bc
+; RUN: llvm-link -disable-lazy-loading %t0.bc %t1.bc -o %t.combined.bc
+; RUN: opt -S -passes="stack-safety-annotator" %t.combined.bc -o - 2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @Write1(i8* %p)
+declare void @Write8(i8* %p)
+
+; Basic out-of-bounds.
+define void @f1() {
+; CHECK-LABEL: define void @f1() {
+; CHECK: alloca i32, align 4{{$}}
+entry:
+  %x = alloca i32, align 4
+  %x1 = bitcast i32* %x to i8*
+  call void @Write8(i8* %x1)
+  ret void
+}
+
+; Basic in-bounds.
+define void @f2() {
+; CHECK-LABEL: define void @f2() {
+; CHECK: alloca i32, align 4, !stack-safe ![[A:[0-9]+]]{{$}}
+entry:
+  %x = alloca i32, align 4
+  %x1 = bitcast i32* %x to i8*
+  call void @Write1(i8* %x1)
+  ret void
+}
+
+; CHECK: ![[A]] = !{}

llvm/test/CodeGen/AArch64/stack-tagging.ll

 ; CHECK:  ret void
 
 
 define void @ManyVars() sanitize_memtag {
 entry:
   %x1 = alloca i32, align 4
   %x2 = alloca i8, align 4
   %x3 = alloca i32, i32 11, align 4
+  %x4 = alloca i32, align 4, !stack-safe !0
   call void @use32(i32* %x1)
   call void @use8(i8* %x2)
   call void @use32(i32* %x3)
   ret void
 }
 
 ; CHECK-LABEL: define void @ManyVars(
 ; CHECK:  alloca { i32, [12 x i8] }, align 16
 ; CHECK:  call { i32, [12 x i8] }* @llvm.aarch64.tagp.{{.*}}({ i32, [12 x i8] }* {{.*}}, i64 0)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  alloca { i8, [15 x i8] }, align 16
 ; CHECK:  call { i8, [15 x i8] }* @llvm.aarch64.tagp.{{.*}}({ i8, [15 x i8] }* {{.*}}, i64 1)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  alloca { [11 x i32], [4 x i8] }, align 16
 ; CHECK:  call { [11 x i32], [4 x i8] }* @llvm.aarch64.tagp.{{.*}}({ [11 x i32], [4 x i8] }* {{.*}}, i64 2)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)
+; CHECK:  alloca i32, align 4
+; CHECK-NOT: @llvm.aarch64.tagp
+; CHECK-NOT: @llvm.aarch64.settag
 
 ; CHECK:  call void @use32(
 ; CHECK:  call void @use8(
 ; CHECK:  call void @use32(
 
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 16)
 ; CHECK:  call void @llvm.aarch64.settag(i8* {{.*}}, i64 48)

 ; CHECK: call void @llvm.aarch64.settag(
 ; CHECK: store i32
 ; CHECK: store i32
 ; CHECK: store i32
 ; CHECK: call void @llvm.aarch64.settag(
 ; CHECK: call void @llvm.aarch64.settag(
 ; CHECK: call void @llvm.aarch64.settag(
 ; CHECK: ret void
+
+!0 = !{}
 No newline at end of file