[StackSafety] Make full LTO to attach metadata if MTE is enabled

[vitalybuka]

Summary

Right now need StackSafety info only for -fsanitize=memtag and
target is arch64. In future we will need that for asan, hwasa on
any target.

I assume the only way LTO knows about nothing about sanitizers,
so we should always to insert the pass. However we'd like skip
actual work if not sanitizers where enabled. To do so we can scan
function attributes, but more efficient is just to add a module flag.

Depends on D80039.

[vitalybuka]

pm test

[eugenis]

I think this is an OK approach.
Perhaps a module flag would be even better, then it can be decoupled from "sanitize_memtag", and it would not require iterating over the entire module. You can check how it is used in CFI: !"CFI Canonical Jump Tables".

It could also be better, instead of iterating over all functions, simply run the global pass, but make sure that it does not do too much work on non-sanitize_memtag functions. I.e. the function pass should bail out on such functions, and then the data flow will have no data to run on, so that should also be pretty fast.

[vitalybuka]

It could also be better, instead of iterating over all functions, simply run the global pass, but make sure that it does not do too much work on non-sanitize_memtag functions. I.e. the function pass should bail out on such functions, and then the data flow will have no data to run on, so that should also be pretty fast.

Actually I don't think it's right this to do. Function pass is analysis. If it's not needed, just don't query it.

[vitalybuka]

module flag

[vitalybuka]

new line

[vitalybuka]

typo

[efriedma]

If I'm understanding this correctly, the stack-safe metadata on allocas is both produced and consumed at LTO-time. So at the point where the stack-safe metadata would be produced, we can compute whether any later passes will query it.

Given that, why do you need a module flag?

[vitalybuka]

In D80046#2040309, @efriedma wrote:

If I'm understanding this correctly, the stack-safe metadata on allocas is both produced and consumed at LTO-time. So at the point where the stack-safe metadata would be produced, we can compute whether any later passes will query it.

Given that, why do you need a module flag?

I assumed a pass knows nothing about consumers. Can you point to some examples how this can be done?

[vitalybuka]

In D80046#2042874, @vitalybuka wrote:

In D80046#2040309, @efriedma wrote:

If I'm understanding this correctly, the stack-safe metadata on allocas is both produced and consumed at LTO-time. So at the point where the stack-safe metadata would be produced, we can compute whether any later passes will query it.

Given that, why do you need a module flag?

I assumed a pass knows nothing about consumers. Can you point to some examples how this can be done?

Oh, I assume the idea to take that from AArch64TargetMachine.
Still AArch64TargetMachine is going to use this info only if sanitizers where enabled. Also right now SSI is used for memtag. but asan, hwasan may follow.
I assume LTO PM knows nothing about if clang was using sanitizers, so we have always to insert the pass. Then we'd like to return from pass if nothing to do there.
The pass may check if there are functions with corresponding sanitizer attributes. However with module flag we can make it more efficient and avoid scanning all functions.

BTW. function attribute was used in the first snapshot, but @eugenis proposed the flag and I agree with him that now it's nicer.

[efriedma]

Okay, I think we're mostly on the same page, then.

I have a few issues here:

1. Whether the backend wants this information is really a per-function decision, not a per-module decision; using module-level metadata is sort of weird.
2. Having weird rules like this makes it harder to write a non-clang frontend producing LLVM IR.
3. All metadata produced and used by in-tree code needs to be documented in LangRef.
4. Is there some reason this analysis needs to run in the middle of the pass pipeline, as opposed to running it closer to where we actually use the information? Along those lines, do we need to encode this as metadata at all?

[vitalybuka]

In D80046#2042937, @efriedma wrote:

Okay, I think we're mostly on the same page, then.

Probably most important is Q no. 4.

I have a few issues here:

1. Whether the backend wants this information is really a per-function decision, not a per-module decision; using module-level metadata is sort of weird.

Correct. Seems like redundancy.
BTW. Even if the function is not sanitized, it's stack safety analysis can be useful for sanitized callers.

2. Having weird rules like this makes it harder to write a non-clang frontend producing LLVM IR.

So it's traidoff , quit pass early vs more computations. I prefer the flag, but not strongly.

3. All metadata produced and used by in-tree code needs to be documented in LangRef.

I'll update if we go this way

4. Is there some reason this analysis needs to run in the middle of the pass pipeline, as opposed to running it closer to where we actually use the information? Along those lines, do we need to encode this as metadata at all?

In the current design StackSafety is ModulePass and AArch64StackTagging is FunctionPass.
We have some difficulty of using ModulePass from FunctionPass without significant reordering passes. But I suspect it's possible to find good solution. I'll try to get more details.
If we go that way then this patch is not needed.

[vitalybuka]

I'll probably abandon this. I've converted the pass into pure analysis and I need to upstream that.

clang/lib/CodeGen/BackendUtil.cpp

static void addDataFlowSanitizerPass(const PassManagerBuilder &Builder,
   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));
+  PM.add(createStackSafetyGlobalInfoWrapperPass(/*SetMetadata=*/true,
+                                                /*Optional=*/true));
 }
 
 static TargetLibraryInfoImpl *createTLII(llvm::Triple &TargetTriple,
                                          const CodeGenOptions &CodeGenOpts) {
   TargetLibraryInfoImpl *TLII = new TargetLibraryInfoImpl(TargetTriple);
 
   switch (CodeGenOpts.getVecLib()) {
   case CodeGenOptions::Accelerate:
if (!CodeGenOpts.DisableLLVMPasses) {
     }
     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());
+      MPM.addPass(StackSafetyGlobalAnnotatorPass(/*Optional=*/true));
     }
 
     if (CodeGenOpts.OptimizationLevel == 0) {
       addCoroutinePassesAtO0(MPM, LangOpts, CodeGenOpts);
       addSanitizersAtO0(MPM, TargetTriple, LangOpts, CodeGenOpts);
     }
   }
 

clang/lib/CodeGen/CodeGenModule.cpp

void CodeGenModule::Release() {
   }
 
   if (LangOpts.Sanitize.has(SanitizerKind::CFIICall)) {
     getModule().addModuleFlag(llvm::Module::Override,
                               "CFI Canonical Jump Tables",
                               CodeGenOpts.SanitizeCfiCanonicalJumpTables);
   }
 
+  if (CodeGenOpts.OptimizationLevel > 0 &&
+      LangOpts.Sanitize.has(SanitizerKind::MemTag)) {
+    // Indicate that we want StackSafety metadata.
+    getModule().addModuleFlag(llvm::Module::Override, "stack-safe", 1);
+  }
+
   if (CodeGenOpts.CFProtectionReturn &&
       Target.checkCFProtectionReturnSupported(getDiags())) {
     // Indicate that we want to instrument return control flow protection.
     getModule().addModuleFlag(llvm::Module::Override, "cf-protection-return",
                               1);
   }
 
   if (CodeGenOpts.CFProtectionBranch &&

clang/test/CodeGen/stack-safe-flag.cpp

+// RUN: %clang_cc1 -O0 -triple aarch64-unknown-linux -fsanitize=memtag -emit-llvm -o - %s | FileCheck %s -check-prefix=NO
+// RUN: %clang_cc1 -O1 -triple aarch64-unknown-linux -emit-llvm -o - %s                   | FileCheck %s -check-prefix=NO
+// RUN: %clang_cc1 -O1 -triple aarch64-unknown-linux -fsanitize=memtag -emit-llvm -o - %s | FileCheck %s
+
+// NO-NOT: "stack-safe"
+// CHECK: !{i32 4, !"stack-safe", i32 1}

clang/test/Driver/memtag_lto.c

 // RUN: %clang -fno-experimental-new-pass-manager -O1 -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag -c -DBUILD2 %s -flto=full -o %t.lto2.bc
 // RUN: llvm-lto2 run -o %t.lto %t.lto1.bc %t.lto2.bc -save-temps -O1 \
 // RUN:  -r %t.lto1.bc,fn,plx \
 // RUN:  -r %t.lto1.bc,use,lx \
 // RUN:  -r %t.lto1.bc,use_local,plx \
 // RUN:  -r %t.lto1.bc,w, \
 // RUN:  -r %t.lto2.bc,use,plx \
 // RUN:  -r %t.lto2.bc,z,
-// FIXME: Must be -check-prefixes=XSAFE,YSAFE
-// RUN: llvm-dis %t.lto.0.5.precodegen.bc -o - | FileCheck %s -check-prefixes=XUNSAFE,YSAFE
+// RUN: llvm-dis %t.lto.0.5.precodegen.bc -o - | FileCheck %s -check-prefixes=XSAFE,YSAFE
 
 // Full LTO, new PM: both are safe.
 // RUN: %clang -fexperimental-new-pass-manager -O1 -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag -c %s -flto=full -o %t.ltonewpm1.bc
 // RUN: %clang -fexperimental-new-pass-manager -O1 -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag -c -DBUILD2 %s -flto=full -o %t.ltonewpm2.bc
 // RUN: llvm-lto2 run -use-new-pm -o %t.ltonewpm %t.ltonewpm1.bc %t.ltonewpm2.bc -save-temps -O1 \
 // RUN:  -r %t.ltonewpm1.bc,fn,plx \
 // RUN:  -r %t.ltonewpm1.bc,use,lx \
 // RUN:  -r %t.ltonewpm1.bc,use_local,plx \
 // RUN:  -r %t.ltonewpm1.bc,w, \
 // RUN:  -r %t.ltonewpm2.bc,use,plx \
 // RUN:  -r %t.ltonewpm2.bc,z,
-// FIXME: Must be -check-prefixes=XSAFE,YSAFE
-// RUN: llvm-dis %t.ltonewpm.0.5.precodegen.bc -o - | FileCheck %s -check-prefixes=XUNSAFE,YSAFE
+// RUN: llvm-dis %t.ltonewpm.0.5.precodegen.bc -o - | FileCheck %s -check-prefixes=XSAFE,YSAFE
 
 // Thin LTO: both are safe.
 // RUN: %clang -fno-experimental-new-pass-manager -O1 -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag -c %s -flto=thin -o %t.thinlto1.bc
 // RUN: %clang -fno-experimental-new-pass-manager -O1 -target aarch64-unknown-linux -march=armv8+memtag -fsanitize=memtag -c -DBUILD2 %s -flto=thin -o %t.thinlto2.bc
 // RUN: llvm-lto2 run -o %t.thinlto %t.thinlto1.bc %t.thinlto2.bc -save-temps -O1 \
 // RUN:  -r %t.thinlto1.bc,fn,plx \
 // RUN:  -r %t.thinlto1.bc,use,lx \
 // RUN:  -r %t.thinlto1.bc,use_local,plx \

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

class StackSafetyGlobalPrinterPass
     : public PassInfoMixin<StackSafetyGlobalPrinterPass> {
   raw_ostream &OS;
 
 public:
   explicit StackSafetyGlobalPrinterPass(raw_ostream &OS) : OS(OS) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
+// Attaches StackSafety metadata to allocas.
 class StackSafetyGlobalAnnotatorPass
     : public PassInfoMixin<StackSafetyGlobalAnnotatorPass> {
 
+  bool Optional = false;
+
 public:
-  explicit StackSafetyGlobalAnnotatorPass() {}
+  // When Optional is true the pass does nothing if no consumers of meta
+  // information is expected.
+  explicit StackSafetyGlobalAnnotatorPass(bool Optional) : Optional(Optional) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
 /// This pass performs the global (interprocedural) stack safety analysis
 /// (legacy pass manager).
 class StackSafetyGlobalInfoWrapperPass : public ModulePass {
   StackSafetyGlobalInfo SSGI;
-  bool SetMetadata;
+  bool SetMetadata = false;
+  bool Optional = false;
 
 public:
   static char ID;
 
-  StackSafetyGlobalInfoWrapperPass(bool SetMetadata = false);
+  explicit StackSafetyGlobalInfoWrapperPass(bool SetMetadata = false,
+                                            bool Optional = false);
 
   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);
+// Attaches StackSafety metadata to allocas.
+// When Optional is true the pass does nothing if no consumers of meta
+// information is expected.
+ModulePass *createStackSafetyGlobalInfoWrapperPass(bool SetMetadata,
+                                                   bool Optional);
 
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_STACKSAFETYANALYSIS_H

llvm/lib/Analysis/StackSafetyAnalysis.cpp

for (auto &AS : Summary->Allocas) {
                            MDNode::get(M.getContext(), None));
         Changed = true;
       }
     }
   }
   return Changed;
 }
 
+bool isMetadataNeeded(const Module &M) {
+  auto *CI =
+      mdconst::extract_or_null<ConstantInt>(M.getModuleFlag("stack-safe"));
+  return CI && CI->getZExtValue();
+}
+
 } // end anonymous namespace
 
 StackSafetyInfo::StackSafetyInfo() = default;
 StackSafetyInfo::StackSafetyInfo(StackSafetyInfo &&) = default;
 StackSafetyInfo &StackSafetyInfo::operator=(StackSafetyInfo &&) = default;
 
 StackSafetyInfo::StackSafetyInfo(FunctionInfo &&Info)
     : Info(new FunctionInfo(std::move(Info))) {}
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();
 }
 
 PreservedAnalyses
 StackSafetyGlobalAnnotatorPass::run(Module &M, ModuleAnalysisManager &AM) {
+  if (Optional && !isMetadataNeeded(M))
+    return PreservedAnalyses::all();
   auto &SSGI = AM.getResult<StackSafetyGlobalAnalysis>(M);
   (void)setStackSafetyMetadata(M, SSGI);
   return PreservedAnalyses::all();
 }
 
 char StackSafetyGlobalInfoWrapperPass::ID = 0;
 
 StackSafetyGlobalInfoWrapperPass::StackSafetyGlobalInfoWrapperPass(
-    bool SetMetadata)
-    : ModulePass(ID), SetMetadata(SetMetadata) {
+    bool SetMetadata, bool Optional)
+    : ModulePass(ID), SetMetadata(SetMetadata), Optional(Optional) {
   initializeStackSafetyGlobalInfoWrapperPassPass(
       *PassRegistry::getPassRegistry());
 }
 
 void StackSafetyGlobalInfoWrapperPass::print(raw_ostream &O,
                                              const Module *M) const {
   ::print(SSGI, O, *M);
 }
 
 void StackSafetyGlobalInfoWrapperPass::getAnalysisUsage(
     AnalysisUsage &AU) const {
   AU.addRequired<StackSafetyInfoWrapperPass>();
 }
 
 bool StackSafetyGlobalInfoWrapperPass::runOnModule(Module &M) {
+  if (Optional && !isMetadataNeeded(M))
+    return false;
+
   StackSafetyDataFlowAnalysis SSDFA(
       M, [this](Function &F) -> const StackSafetyInfo & {
         return getAnalysis<StackSafetyInfoWrapperPass>(F).getResult();
       });
   SSGI = SSDFA.run();
   return SetMetadata ? setStackSafetyMetadata(M, SSGI) : false;
 }
 
-ModulePass *llvm::createStackSafetyGlobalInfoWrapperPass(bool SetMetadata) {
-  return new StackSafetyGlobalInfoWrapperPass(SetMetadata);
+ModulePass *llvm::createStackSafetyGlobalInfoWrapperPass(bool SetMetadata,
+                                                         bool Optional) {
+  return new StackSafetyGlobalInfoWrapperPass(SetMetadata, Optional);
 }
 
 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/PassBuilder.cpp

PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level, bool DebugLogging,
 
   // Use in-range annotations on GEP indices to split globals where beneficial.
   MPM.addPass(GlobalSplitPass());
 
   // Run whole program optimization of virtual call when the list of callees
   // is fixed.
   MPM.addPass(WholeProgramDevirtPass(ExportSummary, nullptr));
 
+  // Add stack safety metatada if needed.
+  MPM.addPass(StackSafetyGlobalAnnotatorPass(true));
+
   // Stop here at -O1.
   if (Level == OptimizationLevel::O1) {
     // The LowerTypeTestsPass needs to run to lower type metadata and the
     // type.test intrinsics. The pass does nothing if CFI is disabled.
     MPM.addPass(LowerTypeTestsPass(ExportSummary, nullptr));
     return MPM;
   }
 

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())
+MODULE_PASS("stack-safety-annotator", StackSafetyGlobalAnnotatorPass(/*Optional=*/false))
 #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/Transforms/IPO/PassManagerBuilder.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/CFLAndersAliasAnalysis.h"
 #include "llvm/Analysis/CFLSteensAliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/Passes.h"
 #include "llvm/Analysis/ScopedNoAliasAA.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
void PassManagerBuilder::addLTOOptimizationPasses(legacy::PassManagerBase &PM) {
   // Split globals using inrange annotations on GEP indices. This can help
   // improve the quality of generated code when virtual constant propagation or
   // control flow integrity are enabled.
   PM.add(createGlobalSplitPass());
 
   // Apply whole-program devirtualization and virtual constant propagation.
   PM.add(createWholeProgramDevirtPass(ExportSummary, nullptr));
 
+  // Add stack safety metatada if needed.
+  PM.add(createStackSafetyGlobalInfoWrapperPass(
+      /*SetMetadata=*/true, /*Optional=*/true));
+
   // That's all we need at opt level 1.
   if (OptLevel == 1)
     return;
 
   // Now that we internalized some globals, see if we can hack on them!
   PM.add(createGlobalOptimizerPass());
   // Promote any localized global vars.
   PM.add(createPromoteMemoryToRegisterPass());

llvm/test/Other/new-pm-lto-defaults.ll

 ; CHECK-O-NEXT: Running analysis: FunctionAnalysisManagerCGSCCProxy
 ; CHECK-O-NEXT: Running analysis: PassInstrumentationAnalysis
 ; CHECK-O-NEXT: Running analysis: OuterAnalysisManagerProxy<{{.*}}LazyCallGraph{{.*}}>
 ; CHECK-O-NEXT: Running analysis: AAManager
 ; CHECK-O-NEXT: Running pass: ReversePostOrderFunctionAttrsPass
 ; CHECK-O-NEXT: Running analysis: CallGraphAnalysis
 ; CHECK-O-NEXT: Running pass: GlobalSplitPass
 ; CHECK-O-NEXT: Running pass: WholeProgramDevirtPass
+; CHECK-O-NEXT: Running pass: StackSafetyGlobalAnnotatorPass
 ; CHECK-O1-NEXT: Running pass: LowerTypeTestsPass
 ; CHECK-O2-NEXT: Running pass: GlobalOptPass
 ; CHECK-O2-NEXT: Running pass: ModuleToFunctionPassAdaptor<{{.*}}PromotePass>
 ; CHECK-O2-NEXT: Running pass: ConstantMergePass
 ; CHECK-O2-NEXT: Running pass: DeadArgumentEliminationPass
 ; CHECK-O2-NEXT: Running pass: ModuleToFunctionPassAdaptor<{{.*}}PassManager{{.*}}>
 ; CHECK-O2-NEXT: Starting llvm::Function pass manager run.
 ; CHECK-O3-NEXT: Running pass: AggressiveInstCombinePass