[X86] Fix for ballooning compile times due to Load Value Injection (LVI) mitigations

[sconstab]

Summary

Fix for the issue raised in https://github.com/rust-lang/rust/issues/74632.

The current heuristic for inserting LFENCEs uses a quadratic-time algorithm. This can apparently cause substantial compilation slowdowns for building Rust projects, where functions > 5000 LoC are apparently common.

The updated heuristic in this patch implements a linear-time algorithm. On a set of benchmarks, the slowdown factor for the generated code was comparable (2.55x geo mean for the quadratic-time heuristic, vs. 2.58x for the linear-time heuristic). Both heuristics offer the same security properties, namely, mitigating LVI.

This patch also includes some formatting fixes.

[mattdr]

Glad that we could make this simpler without making the resulting code much slower!

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

682–684

nit: I spent some time trying to figure out how these are "implied" by the requirement and I don't think they are. Let's update the comment to make it clear we're describing one approach that happens to meet the requirement at a sweet spot of simplicity and speed.

699–702

Two things:

1. What's the intuition for aggregating with "max" here rather than, say, summing the weights?

2. It seems like we should be taking CutEdges into account somehow here. If we've already decided to cut an expensive edge, the *incremental* cost of cutting it again is zero and we shouldn't be including it here.

706

EdgeSet has |= for union, maybe that would be better here

[sconstab]

Addressed comments by @mattdr

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

699–702

Thanks for pointing this out. I don't know why my intuition led me to use std::max. The minimum multi-cut algorithm would, in essence, also be adding the weights. And I think you're also correct that edges that have already been cut can be treated as having cost 0.

I implemented these changes, and measured (1) a small decrease in the number of LFENCEs inserted, (2) no significant performance difference, and (3) no exposed LVI vulnerabilities.

706

|= is not overloaded for an array-like container on the RHS. Unless you think that EdgesToCut should also be a BitVector?

[mattdr]

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

700–701

I think this is clearer as:

if (!CutEdges.contains(*EgressEdge))
  EgressCutCost += EgressEdge->getValue();

likewise for the loop over IngressEdges.

706

Ah, my mistake -- one too many autos for me to trace through. This is fine.

[sconstab]

Addressed @mattdr 's lone comment on style.

[sconstab]

Made one further optimization.

At line 669, I added a check that will only invoke elimMitigatedEdgesAndNodes() if the current function already has LFENCEs that may mitigate some gadgets. If the function does not have any LFENCEs, then elimMitigatedEdgesAndNodes() will not transform the gadget graph because (a) there are no fences to eliminate, and (b) the new heuristic only performs *one* round of cuts *after* elimMitigatedEdgesAndNodes() runs.

[mattdr]

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

672–673

Can we add this as an early-exit case in trimMitigatedEdges instead?

678–679

Seems like this early-exit case should not have been guarded by the if.

[sconstab]

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

672–673

That would negatively impact the hardenLoadsWithPlugin() flow that may need to trimMitigatedEdges() even if Graph->NumFences == 0.

[craig.topper]

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

668

I'm not a big fan of calling a pointer a "Ref". It makes me think its a reference but you have to dereference it. Do we gain much from omitting the '*' where this is used?

708

I think llvm::for_each is discouraged in this case by the statement here https://llvm.org/docs/CodingStandards.html#use-range-based-for-loops-wherever-possible

[sconstab]

Addressed @craig.topper 's comments, and also changed some autos to Edge/Node wherever possible to improve readability.

[craig.topper]

LGTM

[craig.topper]

@sconstab Do you need someone to commit this?

[sconstab]

@craig.topper Yes.

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

 //===----------------------------------------------------------------------===//
 
 #include "ImmutableGraph.h"
 #include "X86.h"
 #include "X86Subtarget.h"
 #include "X86TargetMachine.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominanceFrontier.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"

 
 static cl::opt<bool> EmitDotVerify(
     PASS_KEY "-dot-verify",
     cl::desc("For each function, emit a dot graph to stdout depicting "
              "potential LVI gadgets, used for testing purposes only"),
     cl::init(false), cl::Hidden);
 
 static llvm::sys::DynamicLibrary OptimizeDL;
-typedef int (*OptimizeCutT)(unsigned int *nodes, unsigned int nodes_size,
-                            unsigned int *edges, int *edge_values,
-                            int *cut_edges /* out */, unsigned int edges_size);
+typedef int (*OptimizeCutT)(unsigned int *Nodes, unsigned int NodesSize,
+                            unsigned int *Edges, int *EdgeValues,
+                            int *CutEdges /* out */, unsigned int EdgesSize);
 static OptimizeCutT OptimizeCut = nullptr;
 
 namespace {
 
 struct MachineGadgetGraph : ImmutableGraph<MachineInstr *, int> {
   static constexpr int GadgetEdgeSentinel = -1;
   static constexpr MachineInstr *const ArgNodeSentinel = nullptr;
 
private:
   const TargetRegisterInfo *TRI;
 
   std::unique_ptr<MachineGadgetGraph>
   getGadgetGraph(MachineFunction &MF, const MachineLoopInfo &MLI,
                  const MachineDominatorTree &MDT,
                  const MachineDominanceFrontier &MDF) const;
   int hardenLoadsWithPlugin(MachineFunction &MF,
                             std::unique_ptr<MachineGadgetGraph> Graph) const;
-  int hardenLoadsWithGreedyHeuristic(
-      MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const;
+  int hardenLoadsWithHeuristic(MachineFunction &MF,
+                               std::unique_ptr<MachineGadgetGraph> Graph) const;
   int elimMitigatedEdgesAndNodes(MachineGadgetGraph &G,
                                  EdgeSet &ElimEdges /* in, out */,
                                  NodeSet &ElimNodes /* in, out */) const;
   std::unique_ptr<MachineGadgetGraph>
   trimMitigatedEdges(std::unique_ptr<MachineGadgetGraph> Graph) const;
   void findAndCutEdges(MachineGadgetGraph &G,
                        EdgeSet &CutEdges /* out */) const;
   int insertFences(MachineFunction &MF, MachineGadgetGraph &G,

 struct DOTGraphTraits<MachineGadgetGraph *> : DefaultDOTGraphTraits {
   using GraphType = MachineGadgetGraph;
   using Traits = llvm::GraphTraits<GraphType *>;
   using NodeRef = typename Traits::NodeRef;
   using EdgeRef = typename Traits::EdgeRef;
   using ChildIteratorType = typename Traits::ChildIteratorType;
   using ChildEdgeIteratorType = typename Traits::ChildEdgeIteratorType;
 
-  DOTGraphTraits(bool isSimple = false) : DefaultDOTGraphTraits(isSimple) {}
+  DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
 
   std::string getNodeLabel(NodeRef Node, GraphType *) {
     if (Node->getValue() == MachineGadgetGraph::ArgNodeSentinel)
       return "ARGS";
 
     std::string Str;
     raw_string_ostream OS(Str);
     OS << *Node->getValue();
void X86LoadValueInjectionLoadHardeningPass::getAnalysisUsage(
     AnalysisUsage &AU) const {
   MachineFunctionPass::getAnalysisUsage(AU);
   AU.addRequired<MachineLoopInfo>();
   AU.addRequired<MachineDominatorTree>();
   AU.addRequired<MachineDominanceFrontier>();
   AU.setPreservesCFG();
 }
 
-static void WriteGadgetGraph(raw_ostream &OS, MachineFunction &MF,
+static void writeGadgetGraph(raw_ostream &OS, MachineFunction &MF,
                              MachineGadgetGraph *G) {
   WriteGraph(OS, G, /*ShortNames*/ false,
              "Speculative gadgets for \"" + MF.getName() + "\" function");
 }
 
 bool X86LoadValueInjectionLoadHardeningPass::runOnMachineFunction(
     MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "***** " << getPassName() << " : " << MF.getName()
bool X86LoadValueInjectionLoadHardeningPass::runOnMachineFunction(
   const auto &MDT = getAnalysis<MachineDominatorTree>();
   const auto &MDF = getAnalysis<MachineDominanceFrontier>();
   std::unique_ptr<MachineGadgetGraph> Graph = getGadgetGraph(MF, MLI, MDT, MDF);
   LLVM_DEBUG(dbgs() << "Building gadget graph... Done\n");
   if (Graph == nullptr)
     return false; // didn't find any gadgets
 
   if (EmitDotVerify) {
-    WriteGadgetGraph(outs(), MF, Graph.get());
+    writeGadgetGraph(outs(), MF, Graph.get());
     return false;
   }
 
   if (EmitDot || EmitDotOnly) {
     LLVM_DEBUG(dbgs() << "Emitting gadget graph...\n");
     std::error_code FileError;
     std::string FileName = "lvi.";
     FileName += MF.getName();
     FileName += ".dot";
     raw_fd_ostream FileOut(FileName, FileError);
     if (FileError)
       errs() << FileError.message();
-    WriteGadgetGraph(FileOut, MF, Graph.get());
+    writeGadgetGraph(FileOut, MF, Graph.get());
     FileOut.close();
     LLVM_DEBUG(dbgs() << "Emitting gadget graph... Done\n");
     if (EmitDotOnly)
       return false;
   }
 
   int FencesInserted;
   if (!OptimizePluginPath.empty()) {
     if (!OptimizeDL.isValid()) {
       std::string ErrorMsg;
       OptimizeDL = llvm::sys::DynamicLibrary::getPermanentLibrary(
           OptimizePluginPath.c_str(), &ErrorMsg);
       if (!ErrorMsg.empty())
         report_fatal_error("Failed to load opt plugin: \"" + ErrorMsg + '\"');
       OptimizeCut = (OptimizeCutT)OptimizeDL.getAddressOfSymbol("optimize_cut");
       if (!OptimizeCut)
         report_fatal_error("Invalid optimization plugin");
     }
     FencesInserted = hardenLoadsWithPlugin(MF, std::move(Graph));
   } else { // Use the default greedy heuristic
-    FencesInserted = hardenLoadsWithGreedyHeuristic(MF, std::move(Graph));
+    FencesInserted = hardenLoadsWithHeuristic(MF, std::move(Graph));
   }
 
   if (FencesInserted > 0)
     ++NumFunctionsMitigated;
   NumFences += FencesInserted;
   return (FencesInserted > 0);
 }
 
do {
 
     Graph = GraphBuilder::trim(*Graph, MachineGadgetGraph::NodeSet{*Graph},
                                CutEdges);
   } while (true);
 
   return FencesInserted;
 }
 
-int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithGreedyHeuristic(
+int X86LoadValueInjectionLoadHardeningPass::hardenLoadsWithHeuristic(
     MachineFunction &MF, std::unique_ptr<MachineGadgetGraph> Graph) const {
+  using NodeRef = const MachineGadgetGraph::Node *;
+  using EdgeRef = const MachineGadgetGraph::Edge *;
   LLVM_DEBUG(dbgs() << "Eliminating mitigated paths...\n");
   Graph = trimMitigatedEdges(std::move(Graph));
   LLVM_DEBUG(dbgs() << "Eliminating mitigated paths... Done\n");
   if (Graph->NumGadgets == 0)
     return 0;
 
   LLVM_DEBUG(dbgs() << "Cutting edges...\n");
-  MachineGadgetGraph::NodeSet ElimNodes{*Graph}, GadgetSinks{*Graph};
-  MachineGadgetGraph::EdgeSet ElimEdges{*Graph}, CutEdges{*Graph};
-  auto IsCFGEdge = [&ElimEdges, &CutEdges](const MachineGadgetGraph::Edge &E) {
-    return !ElimEdges.contains(E) && !CutEdges.contains(E) &&
-           MachineGadgetGraph::isCFGEdge(E);
-  };
-  auto IsGadgetEdge = [&ElimEdges,
-                       &CutEdges](const MachineGadgetGraph::Edge &E) {
-    return !ElimEdges.contains(E) && !CutEdges.contains(E) &&
-           MachineGadgetGraph::isGadgetEdge(E);
-  };
+  MachineGadgetGraph::EdgeSet CutEdges{*Graph};
 
-  // FIXME: this is O(E^2), we could probably do better.
-  do {
-    // Find the cheapest CFG edge that will eliminate a gadget (by being
-    // egress from a SOURCE node or ingress to a SINK node), and cut it.
-    const MachineGadgetGraph::Edge *CheapestSoFar = nullptr;
+  // Begin by collecting all ingress CFG edges for each node
+  DenseMap<NodeRef, SmallVector<EdgeRef, 2>> IngressEdgeMap;
+  for (const auto &E : Graph->edges())
+    if (MachineGadgetGraph::isCFGEdge(E))
+      IngressEdgeMap[E.getDest()].push_back(&E);
 
-    // First, collect all gadget source and sink nodes.
-    MachineGadgetGraph::NodeSet GadgetSources{*Graph}, GadgetSinks{*Graph};
-    for (const auto &N : Graph->nodes()) {
-      if (ElimNodes.contains(N))
-        continue;
-      for (const auto &E : N.edges()) {
-        if (IsGadgetEdge(E)) {
-          GadgetSources.insert(N);
-          GadgetSinks.insert(*E.getDest());
-        }
-      }
-    }
-
-    // Next, look for the cheapest CFG edge which, when cut, is guaranteed to
-    // mitigate at least one gadget by either:
-    // (a) being egress from a gadget source, or
-    // (b) being ingress to a gadget sink.
-    for (const auto &N : Graph->nodes()) {
-      if (ElimNodes.contains(N))
+  // For each gadget edge, make cuts that guarantee the gadget will be
+  // mitigated. This implies either:
+  // (1) cutting all egress CFG edges from the gadget source, or
+  // (2) cutting all ingress CFG edges to the gadget sink.
+  //
+  // Moreover, the algorithm tries not to make a cut into a loop, if possible.
+      if (!MachineGadgetGraph::isGadgetEdge(E))
         continue;
-      for (const auto &E : N.edges()) {
-        if (IsCFGEdge(E)) {
-          if (GadgetSources.contains(N) || GadgetSinks.contains(*E.getDest())) {
-            if (!CheapestSoFar || E.getValue() < CheapestSoFar->getValue())
-              CheapestSoFar = &E;
-          }
-        }
-      }
-    }
-
-    assert(CheapestSoFar && "Failed to cut an edge");
-    CutEdges.insert(*CheapestSoFar);
-    ElimEdges.insert(*CheapestSoFar);
-  } while (elimMitigatedEdgesAndNodes(*Graph, ElimEdges, ElimNodes));
+
+      SmallVector<EdgeRef, 2> EgressEdges;
+      SmallVector<EdgeRef, 2> &IngressEdges = IngressEdgeMap[E.getDest()];
+      for (const auto &EgressEdge : N.edges())
+        if (MachineGadgetGraph::isCFGEdge(EgressEdge))
+          EgressEdges.push_back(&EgressEdge);
+
+      int EgressCutCost = 0, IngressCutCost = 0;
+      for (const auto *EgressEdge : EgressEdges)
+        EgressCutCost = std::max(EgressCutCost, EgressEdge->getValue());
+      for (const auto *IngressEdge : IngressEdges)
+        IngressCutCost = std::max(IngressCutCost, IngressEdge->getValue());
+
+      auto &EdgesToCut =
+          IngressCutCost < EgressCutCost ? IngressEdges : EgressEdges;
+      llvm::for_each(EdgesToCut, [&](EdgeRef E) { CutEdges.insert(*E); });
   LLVM_DEBUG(dbgs() << "Cutting edges... Done\n");
   LLVM_DEBUG(dbgs() << "Cut " << CutEdges.count() << " edges\n");
 
   LLVM_DEBUG(dbgs() << "Inserting LFENCEs...\n");
   int FencesInserted = insertFences(MF, *Graph, CutEdges);
   LLVM_DEBUG(dbgs() << "Inserting LFENCEs... Done\n");
   LLVM_DEBUG(dbgs() << "Inserted " << FencesInserted << " fences\n");