[CGData] Outlined Hash Tree (#89792)

This defines the OutlinedHashTree class.
It contains sequences of stable hash values of instructions that have
been outlined. This OutlinedHashTree can be used to track the outlined
instruction sequences across modules. A trie structure is used in its
implementation, allowing for a compact sharing of common prefixes.

This is a patch for
https://discourse.llvm.org/t/rfc-enhanced-machine-outliner-part-2-thinlto-nolto/78753.

Author: kyulee-com

llvm/include/llvm/CodeGenData/OutlinedHashTree.h

@@ -0,0 +1,99 @@
+//===- OutlinedHashTree.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
+//
+//===---------------------------------------------------------------------===//
+//
+// This defines the OutlinedHashTree class. It contains sequences of stable
+// hash values of instructions that have been outlined. This OutlinedHashTree
+// can be used to track the outlined instruction sequences across modules.
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGENDATA_OUTLINEDHASHTREE_H
+#define LLVM_CODEGENDATA_OUTLINEDHASHTREE_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StableHashing.h"
+#include "llvm/ObjectYAML/YAML.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <unordered_map>
+#include <vector>
+
+namespace llvm {
+
+/// A HashNode is an entry in an OutlinedHashTree, holding a hash value
+/// and a collection of Successors (other HashNodes). If a HashNode has
+/// a positive terminal value (Terminals > 0), it signifies the end of
+/// a hash sequence with that occurrence count.
+struct HashNode {
+  /// The hash value of the node.
+  stable_hash Hash = 0;
+  /// The number of terminals in the sequence ending at this node.
+  std::optional<unsigned> Terminals;
+  /// The successors of this node.
+  /// We don't use DenseMap as a stable_hash value can be tombstone.
+  std::unordered_map<stable_hash, std::unique_ptr<HashNode>> Successors;
+};
+
+class OutlinedHashTree {
+
+  using EdgeCallbackFn =
+      std::function<void(const HashNode *, const HashNode *)>;
+  using NodeCallbackFn = std::function<void(const HashNode *)>;
+
+  using HashSequence = SmallVector<stable_hash>;
+  using HashSequencePair = std::pair<HashSequence, unsigned>;
+
+public:
+  /// Walks every edge and node in the OutlinedHashTree and calls CallbackEdge
+  /// for the edges and CallbackNode for the nodes with the stable_hash for
+  /// the source and the stable_hash of the sink for an edge. These generic
+  /// callbacks can be used to traverse a OutlinedHashTree for the purpose of
+  /// print debugging or serializing it.
+  void walkGraph(NodeCallbackFn CallbackNode,
+                 EdgeCallbackFn CallbackEdge = nullptr,
+                 bool SortedWalk = false) const;
+
+  /// Release all hash nodes except the root hash node.
+  void clear() {
+    assert(getRoot()->Hash == 0 && !getRoot()->Terminals);
+    getRoot()->Successors.clear();
+  }
+
+  /// \returns true if the hash tree has only the root node.
+  bool empty() { return size() == 1; }
+
+  /// \returns the size of a OutlinedHashTree by traversing it. If
+  /// \p GetTerminalCountOnly is true, it only counts the terminal nodes
+  /// (meaning it returns the the number of hash sequences in the
+  /// OutlinedHashTree).
+  size_t size(bool GetTerminalCountOnly = false) const;
+
+  /// \returns the depth of a OutlinedHashTree by traversing it.
+  size_t depth() const;
+
+  /// \returns the root hash node of a OutlinedHashTree.
+  const HashNode *getRoot() const { return &Root; }
+  HashNode *getRoot() { return &Root; }
+
+  /// Inserts a \p Sequence into the this tree. The last node in the sequence
+  /// will increase Terminals.
+  void insert(const HashSequencePair &SequencePair);
+
+  /// Merge a \p OtherTree into this Tree.
+  void merge(const OutlinedHashTree *OtherTree);
+
+  /// \returns the matching count if \p Sequence exists in the OutlinedHashTree.
+  std::optional<unsigned> find(const HashSequence &Sequence) const;
+
+private:
+  HashNode Root;
+};
+
+} // namespace llvm
+
+#endif

llvm/include/llvm/CodeGenData/OutlinedHashTreeRecord.h

@@ -0,0 +1,75 @@
+//===- OutlinedHashTreeRecord.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
+//
+//===---------------------------------------------------------------------===//
+//
+// This defines the OutlinedHashTreeRecord class. This class holds the outlined
+// hash tree for both serialization and deserialization processes. It utilizes
+// two data formats for serialization: raw binary data and YAML.
+// These two formats can be used interchangeably.
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
+#define LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H
+
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+
+namespace llvm {
+
+/// HashNodeStable is the serialized, stable, and compact representation
+/// of a HashNode.
+struct HashNodeStable {
+  llvm::yaml::Hex64 Hash;
+  unsigned Terminals;
+  std::vector<unsigned> SuccessorIds;
+};
+
+using IdHashNodeStableMapTy = std::map<unsigned, HashNodeStable>;
+using IdHashNodeMapTy = DenseMap<unsigned, HashNode *>;
+using HashNodeIdMapTy = DenseMap<const HashNode *, unsigned>;
+
+struct OutlinedHashTreeRecord {
+  std::unique_ptr<OutlinedHashTree> HashTree;
+
+  OutlinedHashTreeRecord() { HashTree = std::make_unique<OutlinedHashTree>(); }
+  OutlinedHashTreeRecord(std::unique_ptr<OutlinedHashTree> HashTree)
+      : HashTree(std::move(HashTree)) {};
+
+  /// Serialize the outlined hash tree to a raw_ostream.
+  void serialize(raw_ostream &OS) const;
+  /// Deserialize the outlined hash tree from a raw_ostream.
+  void deserialize(const unsigned char *&Ptr);
+  /// Serialize the outlined hash tree to a YAML stream.
+  void serializeYAML(yaml::Output &YOS) const;
+  /// Deserialize the outlined hash tree from a YAML stream.
+  void deserializeYAML(yaml::Input &YIS);
+
+  /// Merge the other outlined hash tree into this one.
+  void merge(const OutlinedHashTreeRecord &Other) {
+    HashTree->merge(Other.HashTree.get());
+  }
+
+  /// \returns true if the outlined hash tree is empty.
+  bool empty() const { return HashTree->empty(); }
+
+  /// Print the outlined hash tree in a YAML format.
+  void print(raw_ostream &OS = llvm::errs()) const {
+    yaml::Output YOS(OS);
+    serializeYAML(YOS);
+  }
+
+private:
+  /// Convert the outlined hash tree to stable data.
+  void convertToStableData(IdHashNodeStableMapTy &IdNodeStableMap) const;
+
+  /// Convert the stable data back to the outlined hash tree.
+  void convertFromStableData(const IdHashNodeStableMapTy &IdNodeStableMap);
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGENDATA_OUTLINEDHASHTREERECORD_H

llvm/lib/CMakeLists.txt

@@ -10,6 +10,7 @@ add_subdirectory(InterfaceStub)
 add_subdirectory(IRPrinter)
 add_subdirectory(IRReader)
 add_subdirectory(CodeGen)
+add_subdirectory(CodeGenData)
 add_subdirectory(CodeGenTypes)
 add_subdirectory(BinaryFormat)
 add_subdirectory(Bitcode)

llvm/lib/CodeGenData/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_llvm_component_library(LLVMCodeGenData
+  OutlinedHashTree.cpp
+  OutlinedHashTreeRecord.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${LLVM_MAIN_INCLUDE_DIR}/llvm/CodeGenData
+
+  DEPENDS
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+  Support
+  )

llvm/lib/CodeGenData/OutlinedHashTree.cpp

@@ -0,0 +1,129 @@
+//===-- OutlinedHashTree.cpp ----------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// An OutlinedHashTree is a Trie that contains sequences of stable hash values
+// of instructions that have been outlined. This OutlinedHashTree can be used
+// to understand the outlined instruction sequences collected across modules.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGenData/OutlinedHashTree.h"
+
+#define DEBUG_TYPE "outlined-hash-tree"
+
+using namespace llvm;
+
+void OutlinedHashTree::walkGraph(NodeCallbackFn CallbackNode,
+                                 EdgeCallbackFn CallbackEdge,
+                                 bool SortedWalk) const {
+  SmallVector<const HashNode *> Stack;
+  Stack.emplace_back(getRoot());
+
+  while (!Stack.empty()) {
+    const auto *Current = Stack.pop_back_val();
+    if (CallbackNode)
+      CallbackNode(Current);
+
+    auto HandleNext = [&](const HashNode *Next) {
+      if (CallbackEdge)
+        CallbackEdge(Current, Next);
+      Stack.emplace_back(Next);
+    };
+    if (SortedWalk) {
+      SmallVector<std::pair<stable_hash, const HashNode *>> SortedSuccessors;
+      for (const auto &[Hash, Successor] : Current->Successors)
+        SortedSuccessors.emplace_back(Hash, Successor.get());
+      llvm::sort(SortedSuccessors);
+      for (const auto &P : SortedSuccessors)
+        HandleNext(P.second);
+    } else {
+      for (const auto &P : Current->Successors)
+        HandleNext(P.second.get());
+    }
+  }
+}
+
+size_t OutlinedHashTree::size(bool GetTerminalCountOnly) const {
+  size_t Size = 0;
+  walkGraph([&Size, GetTerminalCountOnly](const HashNode *N) {
+    Size += (N && (!GetTerminalCountOnly || N->Terminals));
+  });
+  return Size;
+}
+
+size_t OutlinedHashTree::depth() const {
+  size_t Size = 0;
+  DenseMap<const HashNode *, size_t> DepthMap;
+  walkGraph([&Size, &DepthMap](
+                const HashNode *N) { Size = std::max(Size, DepthMap[N]); },
+            [&DepthMap](const HashNode *Src, const HashNode *Dst) {
+              size_t Depth = DepthMap[Src];
+              DepthMap[Dst] = Depth + 1;
+            });
+  return Size;
+}
+
+void OutlinedHashTree::insert(const HashSequencePair &SequencePair) {
+  auto &[Sequence, Count] = SequencePair;
+  HashNode *Current = getRoot();
+
+  for (stable_hash StableHash : Sequence) {
+    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->Hash = StableHash;
+      Current->Successors.emplace(StableHash, std::move(Next));
+      Current = NextPtr;
+    } else
+      Current = I->second.get();
+  }
+  if (Count)
+    Current->Terminals = (Current->Terminals ? *Current->Terminals : 0) + Count;
+}
+
+void OutlinedHashTree::merge(const OutlinedHashTree *Tree) {
+  HashNode *Dst = getRoot();
+  const HashNode *Src = Tree->getRoot();
+  SmallVector<std::pair<HashNode *, const HashNode *>> Stack;
+  Stack.emplace_back(Dst, Src);
+
+  while (!Stack.empty()) {
+    auto [DstNode, SrcNode] = Stack.pop_back_val();
+    if (!SrcNode)
+      continue;
+    if (SrcNode->Terminals)
+      DstNode->Terminals =
+          (DstNode->Terminals ? *DstNode->Terminals : 0) + *SrcNode->Terminals;
+    for (auto &[Hash, NextSrcNode] : SrcNode->Successors) {
+      HashNode *NextDstNode;
+      auto I = DstNode->Successors.find(Hash);
+      if (I == DstNode->Successors.end()) {
+        auto NextDst = std::make_unique<HashNode>();
+        NextDstNode = NextDst.get();
+        NextDstNode->Hash = Hash;
+        DstNode->Successors.emplace(Hash, std::move(NextDst));
+      } else
+        NextDstNode = I->second.get();
+
+      Stack.emplace_back(NextDstNode, NextSrcNode.get());
+    }
+  }
+}
+
+std::optional<unsigned>
+OutlinedHashTree::find(const HashSequence &Sequence) const {
+  const HashNode *Current = getRoot();
+  for (stable_hash StableHash : Sequence) {
+    const auto I = Current->Successors.find(StableHash);
+    if (I == Current->Successors.end())
+      return 0;
+    Current = I->second.get();
+  }
+  return Current->Terminals;
+}