[RFC] StableHashTree Implementation.

[plotfi]

Summary

This is a first pass at bringing some work that has been done on assisting the Machine Outliner with cross module outlining decisions. A lot of this work is inspired by or directly refactored from the Global Machine Outliner for ThinLTO talk from EuroLLVM 2020 (https://llvm.org/devmtg/2020-04/talks.html#TechTalk_58).

In this diff however, there is no LTO: This diff enables a way to serialize a representation of the MachineOutliner suffix tree as a HashTree to disk. Serialized HashTrees can be read in and used to aid in making better outlining decisions for modules where a Candidate sequence only occurs once, but which have duplicate Candidates off module. This is a first step and I anticipate this will evolve a lot from its current form.

For now I have a single test case to showcase the mechanics, but am working on more test cases.

[plotfi]

@thegameg: @paquette tells be you might have some ideas on a better format than this json business going on here (based on work you've done on remarks). What do you think?

[plotfi]

Cleaning up patch to be easier to understand.

[plotfi]

spelling and grammar

[plotfi]

clang-tidy

[plotfi]

comments added

[paquette]

I think it would make sense to put the StableHashTree implementation in its own patch. Then, in a follow up, you can plumb through the outliner support.

The data structure itself needs tests outside of the outliner, and I think it references the outliner itself a little too much in the comments.

I also feel like the outliner shouldn't be responsible for producing HashTree. That seems like a different thing which may have its own cost model and considerations. It might make sense to adapt the IRSimilarity framework to MIR and use that for the purposes of producing the tree.

Having the outliner consume the tree seems fine to me though.

llvm/include/llvm/CodeGen/MachineOutliner.h

185

Variable name in comment should match the actual variable name.

(Included existing typo for clarity)

188

• Should say "Matches" or "Match", not "Matche"
• Maybe a more succinct name?

190

Typo

194

If you use an Optional, you can differentiate between "it's just empty" and "it's not actually being used" in the type itself.

Also, would it make sense to use a SmallVector here?

http://llvm.org/docs/ProgrammersManual.html#vector

However, SmallVector<T, 0> is often a better option due to the advantages listed [in the SmallVector section]. std::vector is still useful when you need to store more than UINT32_MAX elements or when interfacing with code that expects vectors

llvm/include/llvm/CodeGen/MachineStableHash.h

1

Might be worth fixing the filename here in a NFC commit?

25

remove unnecessary whitespace change

31

Why not both const?

31

Probably should have a documentation comment?

llvm/include/llvm/CodeGen/StableHashTree.h

12

I think that this comment can describe what this actually does in a bit more detail. If this is intended to be a reusable data structure (as the comment implies), I think it'd make sense to address the following questions:

• What does the stable hash tree actually do?
• Why would you use it in a transformation?

Also "Global Machine Outlining" isn't defined anywhere.

In the patch description you have:

A lot of this work is inspired by or directly refactored from the Global Machine Outliner for ThinLTO talk from EuroLLVM 2020 (https://llvm.org/devmtg/2020-04/talks.html#TechTalk_58).

So it'd be nice to include that somewhere, so people curious about that can take a look.

29

• Should be a Doxygen comment
• Try to use class names to make things clear

34

Move member variable documentation into the struct.

35

Move member variable documentation into the struct.

37

If you use a more meaningful name than "Data", it shouldn't be necessary to document this?

38

Could this just be a function that checks if the map is empty?

39

Would it be appropriate to use an IndexedMap here?

http://llvm.org/docs/ProgrammersManual.html#llvm-adt-indexedmap-h

IndexedMap is a specialized container for mapping small dense integers (or values that can be mapped to small dense integers) to some other type. It is internally implemented as a vector with a mapping function that maps the keys to the dense integer range.

I suppose it depends on if stable_hash tends to be dense, by whatever measure of dense is being used here.

42

Match the name of the file?

46

I think that you can drop the part about walkEdges?

General documentation for the data structure and how it works would be better in the \file comment at the top.

50

These type names are pretty long.

Maybe it'd be good to reduce the cognitive overload by doing something like this somewhere:

/// Graph traversal callback types.
///{
using EdgeCallbackFn = std::function<void(const HashNode *, const HashNode *)>;
using NodeCallbackFn = std::function<void(const HashNode *)>;
///}

51

Although vertex and node are interchangeable terms, I think it'd be good to be consistent and just choose one?

57

Should be Doxygen

60

Documentation comments don't need to include implementation info; that can go out of date.

61

Use Doxygen stuff

71

No need to mention where this is called if you document the algorithm somewhere.

75

- Documentation should just say what the function does, not include implementation details.

• Doxygen comment.

80

Probably good to not mention outlining here if this is supposed to be general-purpose?

90

Needs documentation.

llvm/lib/CodeGen/MachineOutliner.cpp

114

This is a long sentence. Split it up?

372

Capitalization?

596

Can this go in a function?

611

Do you have to use llvm:: here?

629

655

I'm not a fan of messing with the found candidates or cost model to make this work.

If you wanted to handle candidates that appear once across many modules, I think it would be best to pre-populate a hash tree with known beneficial candidates versus trying to guess/mess with stuff during outlining?

Since the tree is serialized to JSON, it should be possible to pre-populate it without using the outliner...

Maybe it'd make sense to adapt the IR similarity framework for this portion somehow versus putting all of this in the outliner? It seems like generating the hash tree is really outside the scope of the pass. Consuming and using it seems okay though.

835

remove braces

933

remove braces

[plotfi]

Updated based on @paquette's feedback. This only includes the StableHashTree data structure and a unit test.

[plotfi]

llvm/include/llvm/CodeGen/StableHashTree.h

38

I thought this myself, but you need an IsTerminal flag because you could have some sequence you want to hash like:

ORRWri
ORRWri
ORRWri
RET

as well as

ORRWri
ORRWri
ORRWri

You need the terminal to know that even though you have a node with successors that that know can be the terminal node in a sequence that was added.

39

Can I add this as a post commit NFC commit? I am unsure on the tradeoffs here at the moment.

[paquette]

If this data structure is a trie, is there any reason you can't just improve the existing SuffixTree to do all of this?

A suffix tree is just a compressed trie.

llvm/include/llvm/CodeGen/StableHashTree.h

59

Should IsTerminal ever be modified outside of Insert and readFromBuffer?

Would it make sense to have it be a private member with a getter?

102

Should this be in StableHashTree.cpp?

124

Needs documentation?

What happens if something inserted was already in the tree?

131

Seems like this may be a bit nicer as an Optional which returns where the thing was found?

Optional<HashNode *> find(const StableHashSequence &Sequence) const;

That way you can also access the specific node if you need it.

Or, even better, if you had an iterator type for this data structure you could do something like this:

iterator find(const StableHashSequence &Sequence) const;

I feel like the iterator idea is probably better, since it's more consistent with the rest of the LLVM data types. You could even have an edge_iterator and a vertex_iterator for edge/vertex walks.

137

Should these functions be in StableHashTree.cpp?

145

Missing comment?

llvm/lib/CodeGen/StableHashTree.cpp

69

Typo

118

Should be consistent with how JSON is capitalized in things the a person might see on the command line.

182

Nit: It's a bit nicer to read if you put the simpler situations as the one you continue from. This lessens the indentation level of the more complicated case:

if (I != Current->Successors.end()) {
  Current = I->second.get();
  continue;
}

// Didn't find the hash in the current node's successors. Create a new one.
std::unique_ptr<HashNode> Next = std::make_unique<HashNode>();
// ...

llvm/include/llvm/CodeGen/MachineOutliner.h

 
 #ifndef LLVM_MACHINEOUTLINER_H
 #define LLVM_MACHINEOUTLINER_H
 
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/StableHashing.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 
 namespace llvm {
 namespace outliner {
 
 /// Represents how an instruction should be mapped by the outliner.
 /// \p Legal instructions are those which are safe to outline.
 /// \p LegalTerminator instructions are safe to outline, but only as the
public:
   unsigned SequenceSize = 0;
 
   /// Target-defined overhead of constructing a frame for this function.
   unsigned FrameOverhead = 0;
 
   /// Target-defined identifier for constructing a frame for this function.
   unsigned FrameConstructionID = 0;
 
+  /// Tells if there is an instance of this OutlinedFunction that is outlined in
+  /// another module. DoesSequenceMatchesOffModule helps getBenefit() to
+  /// consider one additional candidate match that may exists outside of the
+  /// current module.
+  bool DoesSequenceMatcheOffModuleCandidate = false;
+
+  /// The sequence of stable_hash'es for a Candidate in Candidates.
+  /// StableHashSequence is empty if computing hashes is disabled or if
+  /// one of the MachineOperands in one of the MachineInstrs in the Candidates
+  /// is not able have a stable_hash computed.
+  std::vector<stable_hash> StableHashSequence;
+
   /// Return the number of candidates for this \p OutlinedFunction.
-  unsigned getOccurrenceCount() const { return Candidates.size(); }
+  unsigned getOccurrenceCount() const {
+    if (DoesSequenceMatcheOffModuleCandidate)
+      return Candidates.size() + 1;
+    return Candidates.size();
+  }
 
   /// Return the number of bytes it would take to outline this
   /// function.
   unsigned getOutliningCost() const {
     unsigned CallOverhead = 0;
     for (const Candidate &C : Candidates)
       CallOverhead += C.getCallOverhead();
     return CallOverhead + SequenceSize + FrameOverhead;

llvm/include/llvm/CodeGen/MachineStableHash.h

 //===------------ MIRVRegNamerUtils.h - MIR VReg Renaming Utilities -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Stable hashing for MachineInstr and MachineOperand. Useful or getting a
 // hash across runs, modules, etc.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_MACHINESTABLEHASH_H
 #define LLVM_CODEGEN_MACHINESTABLEHASH_H
 
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/StableHashing.h"
 
 namespace llvm {
 class MachineInstr;
 class MachineOperand;
 
 stable_hash stableHashValue(const MachineOperand &MO);
+
 stable_hash stableHashValue(const MachineInstr &MI, bool HashVRegs = false,
                             bool HashConstantPoolIndices = false,
                             bool HashMemOperands = false);
 
+std::vector<stable_hash>
+stableHashMachineInstrs(MachineBasicBlock::iterator &Begin,
+                        const MachineBasicBlock::iterator &End);
 } // namespace llvm
 
 #endif

llvm/include/llvm/CodeGen/StableHashTree.h

+//===-- StableHashTree.h ----------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Contains a stable hash tree implementation based on llvm::stable_hash.
+/// Primarily used or Global Machine Outlining but is reusable.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_STABLEHASHTREE_H
+#define LLVM_STABLEHASHTREE_H
+
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+#include "llvm/CodeGen/StableHashing.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+
+// A node in the hash tree might be terminal, i.e. it represents the end
+// of an stable instruction hash sequence that was outlined in some module.
+// Each node may have several successor nodes that can be reached via
+// different stable instruction hashes.
+//
+// Data is the Hash for the current node
+// IsTerminal is true if this node is the last node in a hash sequence
+struct HashNode {
+  stable_hash Data = 0LL;
+  bool IsTerminal{false};
+  std::unordered_map<stable_hash, std::unique_ptr<HashNode>> Successors;
+};
+
+class HashTree {
+public:
+  /// Walks every edge and vertex in the HashTree and calls CallbackEdge for the
+  /// edges and CallbackVertex for the vertices with the stable_hash for the
+  /// source and the stable_hash of the sink for the edge. Using walkEdges it
+  /// should be possible to traverse the HashTree and serialize it, compute its
+  /// depth, compute the number of vertices, etc.
+  void walkGraph(
+      std::function<void(const HashNode *, const HashNode *)> CallbackEdge,
+      std::function<void(const HashNode *)> CallbackVertex) const;
+
+  /// Walks the edges of a HashTree using walkGraph.
+  void walkEdges(
+      std::function<void(const HashNode *, const HashNode *)> Callback) const;
+
+  // Walks the vertices of a HashTree using walkGraph.
+  void walkVertices(std::function<void(const HashNode *)> Callback) const;
+
+  /// Uses HashTree::walkEdges to print the edges of the hash tree.
+  /// If a DebugMap is provided, then it will be used to provide richer output.
+  void dump(raw_ostream &OS = llvm::errs(),
+            std::unordered_map<stable_hash, std::string> DebugMap = {}) const;
+
+  /// Builds a HashTree from a JSON file. Same format as getJsonMap.
+  llvm::Error readHashTreeFromFile(StringRef Filename);
+
+  /// Writes a JSON file representing the current HashTree.
+  llvm::Error writeHashTreeToFile(StringRef Filename) const;
+
+  /// When building a hash tree, insert sequences of stable instruction hashes.
+  void insertIntoHashTree(
+      const std::vector<std::vector<stable_hash>> &StableHashSequences);
+
+  // When using a hash tree, starting from the root, check whether a sequence
+  // of stable instruction hashes ends up at a terminal node.
+  bool findInHashTree(const std::vector<stable_hash> &StableHashSequence) const;
+
+private:
+  // The hash tree is a compact representation of the set of all outlined
+  // instruction sequences across all modules.
+  // This is not a suffix tree, but just represents a set of instruction
+  // sequences, allowing for efficient walking of instruction sequence
+  // prefixes for matching purposes.
+  // We build the tree by inserting stable instruction sequences during the
+  // first first ThinLTO codegen round, and we use the tree by following and
+  // finding stable instruction hashes during the second ThinLTO codegen round.
+  HashNode HashTreeImpl;
+
+  void insertIntoHashTree(const std::vector<stable_hash> &StableHashSequence);
+};
+
+} // namespace llvm
+
+#endif

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachineLoopInfo.cpp
   MachineLoopUtils.cpp
   MachineModuleInfo.cpp
   MachineModuleInfoImpls.cpp
   MachineOperand.cpp
   MachineOptimizationRemarkEmitter.cpp
   MachineOutliner.cpp
   MachinePassManager.cpp
+  StableHashTree.cpp
   MachinePipeliner.cpp
   MachinePostDominators.cpp
   MachineRegionInfo.cpp
   MachineRegisterInfo.cpp
   MachineScheduler.cpp
   MachineSink.cpp
   MachineSizeOpts.cpp
   MachineSSAUpdater.cpp

llvm/lib/CodeGen/MachineOutliner.cpp

 //===----------------------------------------------------------------------===//
 #include "llvm/CodeGen/MachineOutliner.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
+#include "llvm/CodeGen/MachineStableHash.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/StableHashTree.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"

 /// Number of times to re-run the outliner. This is not the total number of runs
 /// as the outliner will run at least one time. The default value is set to 0,
 /// meaning the outliner will run one time and rerun zero times after that.
 static cl::opt<unsigned> OutlinerReruns(
     "machine-outliner-reruns", cl::init(0), cl::Hidden,
     cl::desc(
         "Number of times to rerun the outliner after the initial outline"));
 
+static cl::opt<std::string> OutlinerHashTreeMode(
+    "outliner-hash-tree-mode", cl::init(""), cl::Hidden,
+    cl::desc("Outliner Hash Tree mode < none | write | read >. Anything but "
+             "'read' or 'write' mode will disable all functionality and this "
+             "is the default. In write mode, the outliner will collect hashes "
+             "of the candidate sequences in a HashTree and write the tree to "
+             "disk. In read mode, the outliner will read a provided HashTree "
+             "from disk and use the tree to aid in adjusting the threshold for "
+             "consideration of a candidate when it is present in "
+             "the HashTree loaded from disk but only occurs once in a given "
+             "module."));
+
+static cl::opt<std::string> OutlinerHashTreeFilename(
+    "outliner-hash-tree-filename", cl::init("OutlinerHashTree.out"), cl::Hidden,
+    cl::desc("Outliner Hash Tree file name written to or read from when using "
+             "-outliner-hash-tree-mode."));
+
 namespace {
 
 /// Maps \p MachineInstrs to unsigned integers and stores the mappings.
 struct InstructionMapper {
 
   /// The next available integer to assign to a \p MachineInstr that
   /// cannot be outlined.
   ///
struct MachineOutliner : public ModulePass {
   unsigned OutlineRepeatedNum = 0;
 
   /// Set to true if the outliner should run on all functions in the module
   /// considered safe for outlining.
   /// Set to true by default for compatibility with llc's -run-pass option.
   /// Set when the pass is constructed in TargetPassConfig.
   bool RunOnAllFunctions = true;
 
+  /// stable-hash of the outlined instruction sequence.
+  HashTree OutlinerHashTree;
+
   StringRef getPassName() const override { return "Machine Outliner"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineModuleInfoWrapperPass>();
     AU.addPreserved<MachineModuleInfoWrapperPass>();
     AU.setPreservesAll();
     ModulePass::getAnalysisUsage(AU);
   }
 
   MachineOutliner() : ModulePass(ID) {
     initializeMachineOutlinerPass(*PassRegistry::getPassRegistry());
+
+    if (StringRef(OutlinerHashTreeMode).lower() == "read")
+      if (auto Err =
+              OutlinerHashTree.readHashTreeFromFile(OutlinerHashTreeFilename))
+        consumeError(std::move(Err));
   }
 
   /// Remark output explaining that not outlining a set of candidates would be
   /// better than outlining that set.
   void emitNotOutliningCheaperRemark(
       unsigned StringLen, std::vector<Candidate> &CandidatesForRepeatedSeq,
       OutlinedFunction &OF);
 
for (const unsigned &StartIdx : RS.StartIndices) {
         MachineBasicBlock *MBB = StartIt->getParent();
 
         CandidatesForRepeatedSeq.emplace_back(StartIdx, StringLen, StartIt,
                                               EndIt, MBB, FunctionList.size(),
                                               Mapper.MBBFlagsMap[MBB]);
       }
     }
 
+    std::vector<stable_hash> StableHashSequence;
+    bool IsCandidateInHashTree = false;
+    bool IsSingleCandidateInHashTree = false;
+    // If we are not either populating (write) or using (read) a HashTree
+    // with the Outliner then don't compute a StableHashSequence.
+    if (CandidatesForRepeatedSeq.size() &&
+        (StringRef(OutlinerHashTreeMode).lower() == "write" ||
+         StringRef(OutlinerHashTreeMode).lower() == "read")) {
+      auto &C = CandidatesForRepeatedSeq.front();
+      auto CandidateBegin = C.front();
+      const auto CandidateEnd = std::next(C.back());
+
+      StableHashSequence =
+          stableHashMachineInstrs(CandidateBegin, CandidateEnd);
+
+      IsCandidateInHashTree =
+          !StableHashSequence.empty() &&
+          OutlinerHashTree.findInHashTree(StableHashSequence);
+
+      IsSingleCandidateInHashTree =
+          IsCandidateInHashTree && CandidatesForRepeatedSeq.size() == 1;
+
+      if (IsSingleCandidateInHashTree)
+        CandidatesForRepeatedSeq.push_back(CandidatesForRepeatedSeq.front());
+    }
+
     // We've found something we might want to outline.
     // Create an OutlinedFunction to store it and check if it'd be beneficial
     // to outline.
     if (CandidatesForRepeatedSeq.size() < 2)
       continue;
 
     // Arbitrarily choose a TII from the first candidate.
     // FIXME: Should getOutliningCandidateInfo move to TargetMachine?
     const TargetInstrInfo *TII =
         CandidatesForRepeatedSeq[0].getMF()->getSubtarget().getInstrInfo();
 
     OutlinedFunction OF =
         TII->getOutliningCandidateInfo(CandidatesForRepeatedSeq);
 
     // If we deleted too many candidates, then there's nothing worth outlining.
     // FIXME: This should take target-specified instruction sizes into account.
     if (OF.Candidates.size() < 2)
       continue;
 
+    // Set the StableHashSequence for this OutlinedFunction if it either exists
+    // in the HashTree or if we are building a HashTree in OutlinerHashTree
+    // write mode.
+    if (IsCandidateInHashTree ||
+        StringRef(OutlinerHashTreeMode).lower() == "write") {
+      OF.StableHashSequence = StableHashSequence;
+      OF.DoesSequenceMatcheOffModuleCandidate = true;
+    }
+
+    // NOTE: This is a not so nice way to trick the backend outliner code to
+    // play nice. I really want some feedback here because currently the
+    // threshold for candidate counts in the backends is 2. But in a case where
+    // We have information that an external module could have an additional
+    // candidate this breaks down.
+    if (IsSingleCandidateInHashTree)
+      OF.Candidates.pop_back();
+
     // Is it better to outline this candidate than not?
     if (OF.getBenefit() < 1) {
       emitNotOutliningCheaperRemark(StringLen, CandidatesForRepeatedSeq, OF);
       continue;
     }
 
     FunctionList.push_back(OF);
   }
bool MachineOutliner::outline(Module &M,
   bool OutlinedSomething = false;
 
   // Sort by benefit. The most beneficial functions should be outlined first.
   llvm::stable_sort(FunctionList, [](const OutlinedFunction &LHS,
                                      const OutlinedFunction &RHS) {
     return LHS.getBenefit() > RHS.getBenefit();
   });
 
+  std::vector<std::vector<stable_hash>> ModuleStableHashSequences;
+
   // Walk over each function, outlining them as we go along. Functions are
   // outlined greedily, based off the sort above.
   for (OutlinedFunction &OF : FunctionList) {
     // If we outlined something that overlapped with a candidate in a previous
     // step, then we can't outline from it.
     erase_if(OF.Candidates, [&Mapper](Candidate &C) {
       return std::any_of(
           Mapper.UnsignedVec.begin() + C.getStartIdx(),
           Mapper.UnsignedVec.begin() + C.getEndIdx() + 1,
           [](unsigned I) { return (I == static_cast<unsigned>(-1)); });
     });
 
     // If we made it unbeneficial to outline this function, skip it.
     if (OF.getBenefit() < 1)
       continue;
 
+    if (OF.StableHashSequence.size()) {
+      ModuleStableHashSequences.push_back(OF.StableHashSequence);
+    }
+
     // It's beneficial. Create the function and outline its sequence's
     // occurrences.
     OF.MF = createOutlinedFunction(M, OF, Mapper, OutlinedFunctionNum);
     emitOutlinedFunctionRemark(OF);
     FunctionsCreated++;
     OutlinedFunctionNum++; // Created a function, move to the next name.
     MachineFunction *MF = OF.MF;
     const TargetSubtargetInfo &STI = MF->getSubtarget();
for (Candidate &C : OF.Candidates) {
                     [](unsigned &I) { I = static_cast<unsigned>(-1); });
       OutlinedSomething = true;
 
       // Statistics.
       NumOutlined++;
     }
   }
 
+  // If we are not in HashTree write mode, then we do not want to modify
+  // the current state of the hash tree at this time.
+  if (StringRef(OutlinerHashTreeMode).lower() == "write" &&
+      ModuleStableHashSequences.size()) {
+    OutlinerHashTree.insertIntoHashTree(ModuleStableHashSequences);
+  }
+
   LLVM_DEBUG(dbgs() << "OutlinedSomething = " << OutlinedSomething << "\n";);
   return OutlinedSomething;
 }
 
 void MachineOutliner::populateMapper(InstructionMapper &Mapper, Module &M,
                                      MachineModuleInfo &MMI) {
   // Build instruction mappings for each function in the module. Start by
   // iterating over each Function in M.
MORE.emit([&]() {
                                                     FnCountAfter)
         << "; Delta: "
         << DiagnosticInfoOptimizationBase::Argument("Delta", FnDelta);
       return R;
     });
   }
 }
 
+static std::unordered_map<stable_hash, std::string>
+getStableHashDebugStrings(MachineModuleInfo &MMI, Module &M) {
+  std::unordered_map<stable_hash, std::string> StableHashMIStrings;
+  for (auto &F : M) {
+    const MachineFunction &MF = MMI.getOrCreateMachineFunction(F);
+    for (const auto &BB : MF) {
+      for (const auto &MI : BB) {
+        std::string MIStr;
+        raw_string_ostream OS(MIStr);
+        MI.print(OS, true, false, false, false);
+        StableHashMIStrings[stableHashValue(MI)] = MIStr;
+      }
+    }
+  }
+
+  return StableHashMIStrings;
+}
+
 bool MachineOutliner::runOnModule(Module &M) {
   // Check if there's anything in the module. If it's empty, then there's
   // nothing to outline.
   if (M.empty())
     return false;
 
+  std::unordered_map<stable_hash, std::string> StableHashMIStrings;
+  LLVM_DEBUG({
+    StableHashMIStrings = getStableHashDebugStrings(
+        getAnalysis<MachineModuleInfoWrapperPass>().getMMI(), M);
+  });
+
   // Number to append to the current outlined function.
   unsigned OutlinedFunctionNum = 0;
 
   OutlineRepeatedNum = 0;
   if (!doOutline(M, OutlinedFunctionNum))
     return false;
 
   for (unsigned I = 0; I < OutlinerReruns; ++I) {
     OutlinedFunctionNum = 0;
     OutlineRepeatedNum++;
     if (!doOutline(M, OutlinedFunctionNum)) {
       LLVM_DEBUG({
         dbgs() << "Did not outline on iteration " << I + 2 << " out of "
                << OutlinerReruns + 1 << "\n";
       });
       break;
     }
   }
 
+  LLVM_DEBUG({
+    llvm::dbgs() << "Dump Outliner Hash Tree:\n";
+    OutlinerHashTree.dump(llvm::dbgs(), StableHashMIStrings);
+  });
+
+  if (StringRef(OutlinerHashTreeMode).lower() == "write") {
+    if (auto Err =
+            OutlinerHashTree.writeHashTreeToFile(OutlinerHashTreeFilename))
+      consumeError(std::move(Err));
+  }
+
   return true;
 }
 
 bool MachineOutliner::doOutline(Module &M, unsigned &OutlinedFunctionNum) {
   MachineModuleInfo &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
 
   // If the user passed -enable-machine-outliner=always or
   // -enable-machine-outliner, the pass will run on all functions in the module.

llvm/lib/CodeGen/MachineStableHash.cpp

for (const auto *Op : MI.memoperands()) {
     HashComponents.push_back(static_cast<unsigned>(Op->getSyncScopeID()));
     HashComponents.push_back(static_cast<unsigned>(Op->getBaseAlign().value()));
     HashComponents.push_back(static_cast<unsigned>(Op->getFailureOrdering()));
   }
 
   return stable_hash_combine_range(HashComponents.begin(),
                                    HashComponents.end());
 }
+
+std::vector<stable_hash>
+llvm::stableHashMachineInstrs(MachineBasicBlock::iterator &Begin,
+                              const MachineBasicBlock::iterator &End) {
+  std::vector<stable_hash> Sequence;
+  for (auto I = Begin; I != End; I++) {
+    const MachineInstr &MI = *I;
+    stable_hash Hash = stableHashValue(MI);
+    if (!Hash)
+      return {};
+    Sequence.push_back(Hash);
+  }
+  return Sequence;
+}

llvm/lib/CodeGen/StableHashTree.cpp

+//===---- StableHashTree.cpp ------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/StableHashTree.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/StableHashing.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/VirtualFileSystem.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdlib>
+#include <functional>
+#include <iterator>
+#include <set>
+
+#include <ios>
+#include <stack>
+#include <string>
+#include <system_error>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#define DEBUG_TYPE "stable-hash-tree"
+
+using namespace llvm;
+
+namespace llvm {
+
+void HashTree::walkGraph(
+    std::function<void(const HashNode *, const HashNode *)> CallbackEdge,
+    std::function<void(const HashNode *)> CallbackVertex) const {
+  std::stack<const HashNode *> Stack;
+  Stack.push(&HashTreeImpl);
+
+  while (!Stack.empty()) {
+    const auto *Current = Stack.top();
+    Stack.pop();
+    CallbackVertex(Current);
+    for (const auto &P : Current->Successors) {
+      CallbackEdge(Current, P.second.get());
+      Stack.push(P.second.get());
+    }
+  }
+}
+
+void HashTree::walkVertices(
+    std::function<void(const HashNode *)> Callback) const {
+  walkGraph([](const HashNode *A, const HashNode *B) {}, Callback);
+}
+
+void HashTree::walkEdges(
+    std::function<void(const HashNode *, const HashNode *)> Callback) const {
+  walkGraph(Callback, [](const HashNode *A) {});
+}
+
+void HashTree::dump(
+    llvm::raw_ostream &OS,
+    std::unordered_map<stable_hash, std::string> DebugMap) const {
+
+  std::unordered_map<const HashNode *, unsigned> NodeMap;
+
+  walkVertices([&NodeMap](const HashNode *Current) {
+    size_t Index = NodeMap.size();
+    NodeMap[Current] = Index;
+    assert(Index = NodeMap.size() + 1 &&
+                   "Expected size of ModeMap to increment by 1");
+  });
+
+  bool IsFirstEntry = true;
+  OS << "{";
+  for (const auto &Entry : NodeMap) {
+    if (!IsFirstEntry)
+      OS << ",";
+    OS << "\n";
+    IsFirstEntry = false;
+    OS << "  \"" << Entry.second << "\" : {\n";
+    OS << "    \"hash\" : \"";
+    OS.raw_ostream::write_hex(Entry.first->Data);
+    OS << "\",\n";
+
+    OS << "    \"isTerminal\" : "
+       << "\"" << (Entry.first->IsTerminal ? "true" : "false") << "\",\n";
+
+    // For debugging we want to provide a string representation of the hashing
+    // source, such as a MachineInstr dump, etc. Not intended for production.
+    auto MII = DebugMap.find(Entry.first->Data);
+    if (MII != DebugMap.end())
+      OS << "    \"source\" : \"" << MII->second << "\",\n";
+
+    OS << "    \"neighbors\" : [";
+
+    bool IsFirst = true;
+    for (const auto &Adj : Entry.first->Successors) {
+      if (!IsFirst)
+        OS << ",";
+      IsFirst = false;
+      OS << " \"";
+      OS << NodeMap[Adj.second.get()];
+      OS << "\" ";
+    }
+
+    OS << "]\n  }";
+  }
+  OS << "\n}\n";
+  OS.flush();
+}
+
+llvm::Error HashTree::writeHashTreeToFile(StringRef Filename) const {
+  std::error_code EC;
+  llvm::raw_fd_ostream OS(Filename, EC, llvm::sys::fs::OF_Text);
+  if (EC)
+    return llvm::createStringError(EC, "Unable to open JSON HashTree");
+  dump(OS);
+  OS.flush();
+  return llvm::Error::success();
+}
+
+llvm::Error HashTree::readHashTreeFromFile(StringRef Filename) {
+  llvm::SmallString<256> Filepath(Filename);
+  auto FileOrError = llvm::vfs::getRealFileSystem()->getBufferForFile(Filepath);
+  if (!FileOrError)
+    return llvm::errorCodeToError(FileOrError.getError());
+
+  auto Json = llvm::json::parse(FileOrError.get()->getBuffer());
+  if (!Json)
+    return Json.takeError();
+
+  const json::Object *JO = Json.get().getAsObject();
+  if (!JO)
+    return llvm::createStringError(std::error_code(), "Bad Json");
+
+  std::unordered_map<unsigned, const llvm::json::Value *> JsonMap;
+  for (const auto &E : *JO)
+    JsonMap[std::stoul(E.first.str())] = &E.second;
+
+  assert(JsonMap.find(0x0) != JsonMap.end() && "Expected a root HashTree node");
+
+  // We have a JsonMap and a NodeMap. We walk the JSON form of the HashTree
+  // using the JsonMap by using the stack of JSON IDs. As we walk we used the
+  // IDs to get the currwent JSON Node and the current HashNode.
+  std::unordered_map<unsigned, HashNode *> NodeMap;
+  std::stack<unsigned> Stack;
+  Stack.push(0);
+  NodeMap[0] = &HashTreeImpl;
+
+  while (!Stack.empty()) {
+    unsigned Current = Stack.top();
+    Stack.pop();
+
+    HashNode *CurrentSubtree = NodeMap[Current];
+    const auto *CurrentJson = JsonMap[Current]->getAsObject();
+
+    std::vector<unsigned> Neighbors;
+    llvm::transform(*CurrentJson->get("neighbors")->getAsArray(),
+                    std::back_inserter(Neighbors),
+                    [](const llvm::json::Value &S) {
+                      return std::stoull(S.getAsString()->str());
+                    });
+
+    stable_hash Hash = std::stoull(
+        CurrentJson->get("hash")->getAsString()->str(), nullptr, 16);
+    CurrentSubtree->Data = Hash;
+
+    std::string IsTerminalStr =
+        StringRef(CurrentJson->get("isTerminal")->getAsString()->str()).lower();
+    CurrentSubtree->IsTerminal =
+        IsTerminalStr == "true" || IsTerminalStr == "on";
+
+    for (auto N : Neighbors) {
+      auto I = JsonMap.find(N);
+      if (I == JsonMap.end())
+        return llvm::createStringError(std::error_code(),
+                                       "Missing neighbor in JSON");
+
+      std::unique_ptr<HashNode> Neighbor = std::make_unique<HashNode>();
+      HashNode *NeighborPtr = Neighbor.get();
+      stable_hash StableHash = std::stoull(
+          I->second->getAsObject()->get("hash")->getAsString()->str(), nullptr,
+          16);
+      CurrentSubtree->Successors.emplace(StableHash, std::move(Neighbor));
+      NodeMap[N] = NeighborPtr;
+
+      Stack.push(I->first);
+    }
+  }
+
+  return llvm::Error::success();
+}
+
+void HashTree::insertIntoHashTree(
+    const std::vector<stable_hash> &StableHashSequence) {
+  HashNode *Current = &HashTreeImpl;
+  for (stable_hash StableHash : StableHashSequence) {
+    auto I = Current->Successors.find(StableHash);
+    if (I == Current->Successors.end()) {
+      std::unique_ptr<HashNode> Next = std::make_unique<HashNode>();
+      HashNode *NextPtr = Next.get();
+      NextPtr->Data = StableHash;
+      Current->Successors.emplace(StableHash, std::move(Next));
+      Current = NextPtr;
+      continue;
+    }
+    Current = I->second.get();
+  }
+  Current->IsTerminal = true;
+}
+
+bool HashTree::findInHashTree(
+    const std::vector<stable_hash> &StableHashSequence) const {
+  const HashNode *Current = &HashTreeImpl;
+  for (stable_hash StableHash : StableHashSequence) {
+    const auto I = Current->Successors.find(StableHash);
+    if (I == Current->Successors.end())
+      return false;
+    Current = I->second.get();
+  }
+  // return Current->Successors.empty();
+  return Current->IsTerminal;
+}
+
+void HashTree::insertIntoHashTree(
+    const std::vector<std::vector<stable_hash>> &StableHashSequences) {
+  for (const auto &StableHashSequence : StableHashSequences)
+    insertIntoHashTree(StableHashSequence);
+}
+
+} // namespace llvm

llvm/test/CodeGen/AArch64/machine-outliner-serialize-hashtree.mir

+
+# This is a check to make sure the external module we are sourcing our tree from
+# outlines the a Candidate Sequence.
+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner -verify-machineinstrs \
+# RUN:   %S/../Inputs/machine-outliner-serialize-hashtree-external.mir -o - | \
+# RUN: FileCheck --check-prefix=CHECK-EXTERNAL %s
+
+# This is a check to make sure the current file which has only a single
+# Candidate Sequence (which matches out external module) will outline the
+# Candidate based on the serialized hash tree even though it only exists in the
+# module one time.
+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner --outliner-hash-tree-mode write \
+# RUN:   -outliner-hash-tree-filename %t1.hashtree \
+# RUN:   -verify-machineinstrs %S/../Inputs/machine-outliner-serialize-hashtree-external.mir && \
+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner --outliner-hash-tree-mode read \
+# RUN:   -outliner-hash-tree-filename %t1.hashtree \
+# RUN:   -verify-machineinstrs %s -o - | FileCheck --check-prefix=CHECK-HASHTREE %s
+
+# This is a check to ensure that without use of the serialized hash tree that we
+# are behaving as expected: one candidate in a module means no outlining.
+# RUN: llc -mtriple=aarch64--- -run-pass=machine-outliner \
+# RUN:   -verify-machineinstrs %s -o - | FileCheck --check-prefix=CHECK-DEFAULT %s
+
+# CHECK-EXTERNAL: OUTLINED_FUNCTION_
+# CHECK-HASHTREE: OUTLINED_FUNCTION_
+# CHECK-DEFAULT-NOT: OUTLINED_FUNCTION_
+
+--- |
+
+  @x = common global i32 0, align 4
+
+  define void @bar(i32 %a) #0 {
+    ret void
+  }
+
+  attributes #0 = { noinline noredzone }
+...
+---
+name:            bar
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0, $lr, $w8
+    $sp = frame-setup SUBXri $sp, 32, 0
+    $fp = frame-setup ADDXri $sp, 16, 0
+  bb.2:
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 0
+    RET undef $lr
+
+...
+---

llvm/test/CodeGen/Inputs/machine-outliner-serialize-hashtree-external.mir

+# RUN: llc -mtriple=aarch64--- -run-pass=prologepilog -run-pass=machine-outliner -verify-machineinstrs -frame-pointer=non-leaf %s -o - | FileCheck %s
+--- |
+
+  @x = common global i32 0, align 4
+
+  define void @bar(i32 %a) #0 {
+    ret void
+  }
+
+  attributes #0 = { noinline noredzone }
+...
+---
+name:            bar
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0, $lr, $w8
+    $sp = frame-setup SUBXri $sp, 32, 0
+    $fp = frame-setup ADDXri $sp, 16, 0
+  bb.2:
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 0
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w15 = ORRWri $wzr, 1
+    $w8 = ORRWri $wzr, 1
+    RET undef $lr
+
+...
+---