[PATCH] D75937: Add Support to X86 for Load Hardening to Mitigate Load Value Injection (LVI) [5/6]

[Matthew Riley via Phabricator via llvm-commits]

mattdr accepted this revision.
mattdr added a comment.
This revision is now accepted and ready to land.

Still some comments to address, but I think this is substantially close to where it needs to be. Thanks so much for your work making the code more straightforward and readable!



================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:351
+
+    Graph = GraphBuilder::trim(
+        *ElimGraph, MachineGadgetGraph::NodeSet{*ElimGraph}, CutEdges);
----------------
sconstab wrote:
> mattdr wrote:
> > I know there was a lot of effort to make `ImmutableGraph` efficient, but this is an `O(N)` reallocation on every loop, right?
> > 
> > In practice, how many times do we end up looping before we hit a fixed point?
> The answers to your questions are "yes" and "until all of the gadgets have been mitigated". The intent here is to optimize for the MILP plugin:
> https://github.com/intel/lvi-llvm-optimization-plugin
> The plugin eliminates lots of gadgets all at once, and therefore rebuilding the graph is a relatively rare event. I'm not sure it would be worthwhile to use an entirely different data structure and algorithm for the greedy mitigation strategy.
> 
> What I have done is rewrite the Greedy heuristic to keep track of which edges have been cut and eliminated, thus never actually having to rebuild the graph.
Did you mean to remove this call to `GraphBuilder::trim()`? As it stands right now the loop actually reallocates the graph twice per iteration -- once in `trimMitigatedEdges`, once in this call to `trim()`.


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:567
+        [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
+          if (!FirstNode) {
+            Visited.insert(*N);
----------------
mattdr wrote:
> We already capture `RootN` later on. We can remove the boolean conditional and make this more readable by writing `if (*N != RootN)` (maybe `if (N != &RootN)`? I've sort of lost track of the API for nodes and edges.)
> 
> Pushing further: maybe we can just remove the tiny optimization we get by special-casing the root node in the interest of simplicity?
Still think it makes sense to add the root node to the reachable set (as it's trivially reachable from itself)


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:643
+          for (const auto &EE : N.edges()) {
+            if (MachineGadgetGraph::isCFGEdge(EE)) {
+              if (!CheapestSoFar || EE.getValue() < CheapestSoFar->getValue())
----------------
sconstab wrote:
> mattdr wrote:
> > Every edge touched by this loop will //also// be touched by the other leg of this loop at one point or another ("E is a CFG edge"). Can we find a way to avoid the `O(E^2)` inner loop?
> > 
> > For example: what if we made the outermost loop edge-major? Keep sets of mitigated source-nodes and sink-nodes. Create a list of all edges by weight, sort it, then go through the edges from lowest to highest weight and use them to try to mitigate a new source or sink node.
> > 
> > There's a good chance I'm missing something fundamental there, and I thank you in advance for your patience explaining it to me. But if I can I'd really like to avoid repeating the `CheapestSoFar` comparison on edge weights in two places.
> Yes... in hindsight this whole algorithm was more than a bit sloppy on my part. I completely revamped the algorithm and I think that now it is O(N + E). Please check and make sure you agree.
Each iteration loops over all edges and removes exactly one... so this is probably still `O(E^2)`.

Seems like we can get it down closer to `O(E * lg E)` if we:
1. Compute `GadgetSources` and `GadgetSinks` once
2. Put the edges into a vector, then sorting them by weight
3. Iterate through that vector and, for each edge, add it as an edge to cut if it is still relevant (i.e. not yet otherwise mitigated)

The insight here is that edge weights don't change, so mitigating an edge doesn't change the ranking of other edges.

That said, I'm happy with the readability of the current implementation and would be satisfied for now if we just add `//FIXME: this is O(E^2), it could probably be O(E * lg E) with some work`


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:170-171
                  const MachineDominanceFrontier &MDF, bool FixedLoads) const;
+  int elimEdges(MachineGadgetGraph &G, EdgeSet &CutEdges /* in, out */,
+                NodeSet &Nodedges /* in, out */) const;
+  std::unique_ptr<MachineGadgetGraph>
----------------
In the implementation the latter two parameters are called `ElimEdges` and `ElimNodes`, which I think are good names


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:534
+// Returns the number of remaining gadgets after edge elimination
+int X86LoadValueInjectionLoadHardeningPass::elimEdges(
+    MachineGadgetGraph &G, MachineGadgetGraph::EdgeSet &ElimEdges /* in, out */,
----------------
This is a lot clearer now, thanks!

Given what this is doing I think it makes sense to call it `findMitigatedEdges`.


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:537
+    MachineGadgetGraph::NodeSet &ElimNodes /* in, out */) const {
+  if (G.NumFences > 0) { // eliminate fences
+    for (const auto &E : G.edges()) {
----------------
"Eliminate CFG edges that target a fence, as they're trivially mitigated"


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:541
+      if (isFence(Dest->getValue())) {
+        ElimNodes.insert(*Dest);
+        ElimEdges.insert(E);
----------------
How are we sure we've removed all edges that pointed to this same destination node?


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:553
+  int MitigatedGadgets = 0, RemainingGadgets = 0;
+  MachineGadgetGraph::NodeSet Visited{G};
+  for (const auto &RootN : G.nodes()) {
----------------
To make it easier to follow, let's call this something like `ReachableGadgetSinkNodes`


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:555-556
+  for (const auto &RootN : G.nodes()) {
+    if (llvm::none_of(RootN.edges(), MachineGadgetGraph::isGadgetEdge))
+      continue; // skip this node if it isn't a gadget source
+
----------------
I think this is just intended as an optimization -- it's not necessary for correctness. Assuming that's right, suggest removing it since it doesn't really make things faster but it does add some extra complexity.


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:559
+    // Find all of the nodes that are reachable from RootN
+    std::function<void(const MachineGadgetGraph::Node *, bool)> TraverseDFS =
+        [&](const MachineGadgetGraph::Node *N, bool FirstNode) {
----------------
Maybe `FindCFGReachableGadgetSinksDFS`?


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:583
+    }
+    Visited.clear();
+  }
----------------
Suggest putting this at the top of the loop so it's obvious from the beginning that results aren't intended to accrue from iteration to iteration.


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:636
+        CutEdges.set(I);
+  } else { // Use the default greedy heuristic
+    MachineGadgetGraph::NodeSet ElimNodes{G}, GadgetSinks{G};
----------------
I really like that the refinement loop (trying to get to a fixed point) for the greedy algorithm is now entirely within this function. I'd suggest going one step further: take the loop out of `runOnMachineFunction` entirely and add one around the calls to the plugin in the lines above. That way we keep the implementation detail that "this needs to be run in a loop" as close to the actual algorithm as possible.


================
Comment at: llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp:715-716
+          for (const auto &E : N.edges()) {
+            if (MachineGadgetGraph::isCFGEdge(E) && !CutEdges.contains(E)) {
+              CutEdges.insert(E);
+              ++AdditionalEdgesCut;
----------------
You can skip the double-lookup and just do:

```
if (MachineGadgetGraph::isCFGEdge(E)) {
  if (CutEdges.insert(E)) {  // returns true if not already present
    ++AdditionalEdgesCut;
  }
}


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

https://reviews.llvm.org/D75937