[Metadata] Use temporary MD nodes when appending module flags during module linking

[wolfgangp]

Summary

This is a proposal to fix issue 51893.

As described there, we're seeing excessive memory usage at link time when using LTO + IPGO, and we traced it back to the appending of module flags during module linking.

This patch suggests to use temporary MD nodes for the list nodes that are newly created when module flags with SrcBehavior "append" are handled. Since we have to turn them into permanent nodes at some point, we keep track of the temp nodes in the module, and make them permanent when we know we're done with the linking.

If anyone has any better suggestions, please let me know.

@dexonsmith Apologies if you are not the right reviewer. The code has been implemented by Rafael originally, but IIRC he's no longer working on llvm.

[dexonsmith]

Seems like a good problem solve. IIUC, the root problem is that there are O(N) "versions" of the module flags being created, something like:

• Flags1 = !{!1}
• Flags2 = !{!1, !2}
• Flags3 = !{!1, !2, !3}
• ...
• FlagsN = !{!1, !2, !3, ... !N}

IIUC, the status quo (without this patch) is that all these iterations of Flags are "leaked" onto the LLVMContext.

• With this patch, each iteration is instead made a temporary node so it can be deleted after. That seems reasonable (although the delete operation seems a bit error-prone).
• This patch does not fix the root problem, which is that "append" is O(N) instead of amortized O(1). As a result, O(N^2) work is still being done here.

I think it'd be better to fix the workload to use a "vector" data structure.

One approach would be to (carefully) add a capacity and resize operation to MDNode operands:

• Change MDNode to allow hung-off operands; e.g., steal a bit from NumOperands to indicate whether hung-off, and store the pointer to operands and a capacity where co-allocated operands usually go.
• Add a function MDNode::resize() that only appropriate subclasses enable (only some MDNode subclasses would want to enable this feature). Asserts if the subclass doesn't allow it, and asserts if the node is Uniqued (i.e., only legal for Temporary or Distinct, which have a real "identity" and are not stored in uniquing sets).
• Ensure that an empty temporary or distinct node can be resized later by co-allocating enough space; as a side effect, even an empty node would start with some small storage, but that's probably okay. (An empty uniqued node wouldn't need anything co-allocated.)
• (The capacity of a distinct node would not be serialized in textual IR or bitcode...)
• (Probably there are other implementation strategies.)

With MDNode::resize() in place:

• Change module flags in "append" or "append-unique" mode to use a distinct metadata node to store the list of nodes. The flag tuple that references it should be distinct as well. (distinct means "don't store in the uniquing table"; appropriate here since the content can change.)
• Ideally, add an IR verifier check that these nodes are distinct. Fix textual IR testcases to use distinct in these places and add a bitcode upgrade to apply distinct retroactively.
• During module linking, call resize() (and/or a derived push_back()) to extend the existing metadata node in-place rather than creating a new one. (If the "ideal" verifier/upgrade step above is skipped, then the append operation would need to check "are the right nodes distinct?" and if not create new distinct ones that can be modified in place.)

A second approach would be to add first-class support for module flags. E.g., could create an MDModuleFlag metadata type to use instead of the 3-part tuple. This node type could support the necessary operations efficiently.

A third approach would be to not fix the IR at all, but instead would be to change llvm::Module to have a "lazy module flags" mode. When enabled, !llvm.module.flags is NOT kept up-to-date with changes. Append and AppendUnique flags that have been modified are stored in staging data structures (maybe, SmallVector and SmallSetVector, respectively). LTO could put the module into this "lazy" mode before linking, and put it back into "normal" mode after it was finished linking to commit the changes.

IMO, the first approach I mention above would be best (possibly with the second approach as a follow up!) since the resize() operation would be generally useful; there are lots of metadata nodes that are really just lists of arbitrary size, and that get extended later on (e.g., @dblaikie @aprantl, I think there are places in debug info IR where new nodes are allocated for this sort of thing, is that right?).

[dblaikie]

IMO, the first approach I mention above would be best (possibly with the second approach as a follow up!) since the resize() operation would be generally useful; there are lots of metadata nodes that are really just lists of arbitrary size, and that get extended later on (e.g., @dblaikie @aprantl, I think there are places in debug info IR where new nodes are allocated for this sort of thing, is that right?).

I'm actually not sure how many debug info nodes there are that get incrementally appended to... - the CU list itself, yes, but otherwise mostly the lists are made by frontends and made the right size the first time, I think? yeah, DIBuilder keeping a SmallVector of retained type nodes, then creating the necessary vector to create an MDNode from that in one shot, without incrementally appending.

But I could be wrong - just my rough estimate/recollection.

[dexonsmith]

In D118416#3286076, @dblaikie wrote:

IMO, the first approach I mention above would be best (possibly with the second approach as a follow up!) since the resize() operation would be generally useful; there are lots of metadata nodes that are really just lists of arbitrary size, and that get extended later on (e.g., @dblaikie @aprantl, I think there are places in debug info IR where new nodes are allocated for this sort of thing, is that right?).

I'm actually not sure how many debug info nodes there are that get incrementally appended to... - the CU list itself, yes, but otherwise mostly the lists are made by frontends and made the right size the first time, I think? yeah, DIBuilder keeping a SmallVector of retained type nodes, then creating the necessary vector to create an MDNode from that in one shot, without incrementally appending.

But I could be wrong - just my rough estimate/recollection.

Yeah, I think it was the CU list (hits this kind of problem in LTO, probably has some custom code now) and the retained types (delays creation in IRBuilder to work around this) that I was thinking of.

Since distinct and temporary metadata are non-const anyway, it seems better to me to allow the size to change dynamically, and teaching the clients to take advantage of that rather than having extra side data, or temporarily invalid state (note that the verifier fails if it finds a temporary node).

[dblaikie]

In D118416#3289247, @dexonsmith wrote:

In D118416#3286076, @dblaikie wrote:

IMO, the first approach I mention above would be best (possibly with the second approach as a follow up!) since the resize() operation would be generally useful; there are lots of metadata nodes that are really just lists of arbitrary size, and that get extended later on (e.g., @dblaikie @aprantl, I think there are places in debug info IR where new nodes are allocated for this sort of thing, is that right?).

I'm actually not sure how many debug info nodes there are that get incrementally appended to... - the CU list itself, yes, but otherwise mostly the lists are made by frontends and made the right size the first time, I think? yeah, DIBuilder keeping a SmallVector of retained type nodes, then creating the necessary vector to create an MDNode from that in one shot, without incrementally appending.

But I could be wrong - just my rough estimate/recollection.

Yeah, I think it was the CU list (hits this kind of problem in LTO, probably has some custom code now) and the retained types (delays creation in IRBuilder to work around this) that I was thinking of.

Since distinct and temporary metadata are non-const anyway, it seems better to me to allow the size to change dynamically, and teaching the clients to take advantage of that rather than having extra side data, or temporarily invalid state (note that the verifier fails if it finds a temporary node).

Sounds OK to me, if it makes sense to you!

[wolfgangp]

Thanks for looking into this.

It seems best to abandon the proposed change and introduce a resize() capability for MDNodes.

Not being all that intimately familiar with MD, I might have some silly questions later on.

[dexonsmith]

In D118416#3291434, @wolfgangp wrote:

Thanks for looking into this.

It seems best to abandon the proposed change and introduce a resize() capability for MDNodes.

Not being all that intimately familiar with MD, I might have some silly questions later on.

(A bit behind on reviews); sounds great! I think the main gotcha is that an MDNode in a "uniqued" state shouldn't be resizable since it's constant-like.

[wolfgangp]

I'm finally getting around to take a stab at this. One thing in particular is giving me trouble:

In D118416#3285239, @dexonsmith wrote:

• Ensure that an empty temporary or distinct node can be resized later by co-allocating enough space; as a side effect, even an empty node would start with some small storage, but that's probably okay. (An empty uniqued node wouldn't need anything co-allocated.)

I'm largely following your implementation strategy. On first allocation I'm creating all nodes they way they are currently, i.e. with co-allocated operands. For temporary and distinct nodes I'm allocating extra space for a potentially needed hung-off-storage descriptor, if the number of operands is small (i.e. the space allocated for them is not enough for what's needed for the descriptor). For uniqued nodes I don't allocate anything extra.

The resize operation is fairly straightforward to implement, but the trouble I'm having is at deallocation. Temporary MDnodes can be made unique, and later at deallocation time we can't distinguish between uniqued nodes that have originally been allocated as temporary nodes and the ones that were allocated as unique from the start. A possible solution would be to steal another bit from NumOperands (reducing it to 30 bits) and capture the original allocation state there, but I'm still looking for alternatives.

Cloning these small temporary nodes (without the extra size) at uniquing time would be another option, though it seems like the infrastructure is not available for this operation. I assume we don't want to saddle all MDnodes (even the uniqued ones) with the extra space.

At the moment I'm trying the additional bit approach, which seems to be working, but I wouldn't mind hearing your thoughts on this.

[dexonsmith]

In D118416#3408550, @wolfgangp wrote:

I'm finally getting around to take a stab at this. One thing in particular is giving me trouble:

In D118416#3285239, @dexonsmith wrote:

• Ensure that an empty temporary or distinct node can be resized later by co-allocating enough space; as a side effect, even an empty node would start with some small storage, but that's probably okay. (An empty uniqued node wouldn't need anything co-allocated.)

I'm largely following your implementation strategy. On first allocation I'm creating all nodes they way they are currently, i.e. with co-allocated operands. For temporary and distinct nodes I'm allocating extra space for a potentially needed hung-off-storage descriptor, if the number of operands is small (i.e. the space allocated for them is not enough for what's needed for the descriptor). For uniqued nodes I don't allocate anything extra.

The resize operation is fairly straightforward to implement, but the trouble I'm having is at deallocation. Temporary MDnodes can be made unique, and later at deallocation time we can't distinguish between uniqued nodes that have originally been allocated as temporary nodes and the ones that were allocated as unique from the start. A possible solution would be to steal another bit from NumOperands (reducing it to 30 bits) and capture the original allocation state there, but I'm still looking for alternatives.

Cloning these small temporary nodes (without the extra size) at uniquing time would be another option, though it seems like the infrastructure is not available for this operation. I assume we don't want to saddle all MDnodes (even the uniqued ones) with the extra space.

At the moment I'm trying the additional bit approach, which seems to be working, but I wouldn't mind hearing your thoughts on this.

Stealing a bit down to NumOperands=30 seems pretty reasonable. I think an MDNode with more than a billion operands probably doesn't need to be supported.

I think it's okay to saddle anything that came through the "temporary" or "distinct" path with one extra word; not that bitcode layout of metadata is designed to avoid metadata temporaries in most cases for uniqued nodes.

But I'm not sure you need it. I think you can do something like this pseudo-code for the co-allocation:

struct HungOffType {
  uint32_t Capacity;
  uint32_t NumOperands;
  MDOperand Ops[Capacity];
};
union CoallocationType {
  MDOperand Small[SmallCapacity];
  std::unique_ptr<HungOffType> Large;
};

You can cut up NumOperands without loss of generality once you have that setup:

uint32_t IsSmall : 1;          // Whether it's small.
uint32_t SmallCapacity : 8;    // Co-allocate up to 256 operands.
uint32_t SmallNumOperands : 8; // How much small storage is used?

size_t getNumOperands() const {
  return IsSmall ? SmallNumOperands : getLargeOps().NumOperands;
}

This limits the co-allocation to 256 operands, but if the initial number of operands is too big for that it can hang them off (maybe we want to limit small size further; a histogram might tell us). It also buries getNumOperands() behind a an extra pointer, but anyone checking the size is probably iterating through anyway so it's a pointer they were already going to traverse.

[wolfgangp]

This is a breakdown of MDNodes by number of operands using a bootstrap and a couple of other internal builds (-g -O2). The second column gives the percentage of MDnodes with the number of operands in the first column.
99.9 percent of all nodes have 15 operands or fewer. Almost half have 2 operands.

Operands	Pct	Cumulative
0	3.7	3.7
1	20.2	23.8
2	47.0	70.8
3	2.8	73.6
4	5.6	79.2
5	4.2	83.5
6	8.9	92.3
7	0.1	92.4
8	5.2	97.6
9-15	2.3	99.9
16-31	0.1	100.0
32-64	0.0	100.0
> 64	0.0	100.0

[dexonsmith]

Great; seems like maybe we'd be okay to max out at 16 operands in "small" mode.

llvm/include/llvm/IR/Module.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_IR_MODULE_H
 #define LLVM_IR_MODULE_H
 
 #include "llvm-c/Types.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Comdat.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
private:
   Materializer;                   ///< Used to materialize GlobalValues
   std::string ModuleID;           ///< Human readable identifier for the module
   std::string SourceFileName;     ///< Original source file name for module,
                                   ///< recorded in bitcode.
   std::string TargetTriple;       ///< Platform target triple Module compiled on
                                   ///< Format: (arch)(sub)-(vendor)-(sys0-(abi)
   NamedMDSymTabType NamedMDSymTab;  ///< NamedMDNode names.
   DataLayout DL;                  ///< DataLayout associated with the module
+  SmallPtrSet<MDNode *, 1> TemporaryMDNodes; ///< Holds the temporary MD nodes
+                                             ///< that are used as list nodes.
+
   StringMap<unsigned>
       CurrentIntrinsicIds; ///< Keep track of the current unique id count for
                            ///< the specified intrinsic basename.
   DenseMap<std::pair<Intrinsic::ID, const FunctionType *>, unsigned>
       UniquedIntrinsicNames; ///< Keep track of uniqued names of intrinsics
                              ///< based on unnamed types. The combination of
                              ///< ID and FunctionType maps to the extension that
                              ///< is used to make the intrinsic name unique.
/// @{
   /// equivalent to getDataLayout()->getStringRepresentation().
   const std::string &getDataLayoutStr() const {
     return DL.getStringRepresentation();
   }
 
   /// Get the data layout for the module's target platform.
   const DataLayout &getDataLayout() const;
 
+  /// Get the set of temporary MD nodes.
+  SmallPtrSet<MDNode *, 1> &getTempMDNodes() { return TemporaryMDNodes; }
+
   /// Get the target triple which is a string describing the target host.
   /// @returns a string containing the target triple.
   const std::string &getTargetTriple() const { return TargetTriple; }
 
   /// Get the global data context.
   /// @returns LLVMContext - a container for LLVM's global information
   LLVMContext &getContext() const { return Context; }
 
/// @}
   /// target triple for a macOS target.
   /// @returns a string containing the target variant triple.
   StringRef getDarwinTargetVariantTriple() const;
 
   /// Get the target variant version build SDK version metadata.
   ///
   /// An empty version is returned if no such metadata is attached.
   VersionTuple getDarwinTargetVariantSDKVersion() const;
+
+  /// Turn temporary MD nodes that may have been created during linking into
+  /// permanent ones.
+  void makeTempMDNodesPermanent();
 };
 
 /// Given "llvm.used" or "llvm.compiler.used" as a global name, collect the
 /// initializer elements of that global in a SmallVector and return the global
 /// itself.
 GlobalVariable *collectUsedGlobalVariables(const Module &M,
                                            SmallVectorImpl<GlobalValue *> &Vec,
                                            bool CompilerUsed);

llvm/lib/IR/Module.cpp

StringRef Module::getDarwinTargetVariantTriple() const {
   if (const auto *MD = getModuleFlag("darwin.target_variant.triple"))
     return cast<MDString>(MD)->getString();
   return "";
 }
 
 VersionTuple Module::getDarwinTargetVariantSDKVersion() const {
   return getSDKVersionMD(getModuleFlag("darwin.target_variant.SDK Version"));
 }
+
+void Module::makeTempMDNodesPermanent() {
+  for (auto *TNode : TemporaryMDNodes) {
+    MDNode::replaceWithPermanent(TempMDNode(TNode));
+  }
+  TemporaryMDNodes.clear();
+}

llvm/lib/LTO/LTO.cpp

Error LTO::runRegularLTO(AddStreamFn AddStream) {
   DiagnosticOutputFile = std::move(*DiagFileOrErr);
 
   // Finalize linking of regular LTO modules containing summaries now that
   // we have computed liveness information.
   for (auto &M : RegularLTO.ModsWithSummaries)
     if (Error Err = linkRegularLTO(std::move(M),
                                    /*LivenessFromIndex=*/true))
       return Err;
+  RegularLTO.CombinedModule->makeTempMDNodesPermanent();
 
   // Ensure we don't have inconsistently split LTO units with type tests.
   // FIXME: this checks both LTO and ThinLTO. It happens to work as we take
   // this path both cases but eventually this should be split into two and
   // do the ThinLTO checks in `runThinLTO`.
   if (Error Err = checkPartiallySplit())
     return Err;
 

llvm/lib/Linker/IRMover.cpp

if (SrcBehaviorValue != DstBehaviorValue) {
                          "': IDs have conflicting behaviors in '" +
                          SrcM->getModuleIdentifier() + "' and '" +
                          DstM.getModuleIdentifier() + "'");
     }
 
     auto replaceDstValue = [&](MDNode *New) {
       Metadata *FlagOps[] = {DstOp->getOperand(0), ID, New};
       MDNode *Flag = MDNode::get(DstM.getContext(), FlagOps);
+      // Find the list node we are about to replace.
+      MDNode *OldOp =
+          cast<MDNode>(DstModFlags->getOperand(DstIndex)->getOperand(2));
+      auto &TempSet = DstM.getTempMDNodes();
+      // If the existing list node is a temporary node,
+      // delete it and remove it from the module's set of temporary nodes.
+      if (TempSet.contains(OldOp)) {
+        TempSet.erase(OldOp);
+        MDNode::deleteTemporary(OldOp);
+      }
+      // If the new list node is a temporary MDnode, add it to the set of
+      // temporary nodes.
+      if (New->isTemporary())
+        TempSet.insert(New);
       DstModFlags->setOperand(DstIndex, Flag);
       Flags[ID].first = Flag;
     };
 
     // Emit a warning if the values differ and either source or destination
     // request Warning behavior.
     if ((DstBehaviorValue == Module::Warning ||
          SrcBehaviorValue == Module::Warning) &&
for (unsigned I = 0, E = SrcModFlags->getNumOperands(); I != E; ++I) {
     case Module::Append: {
       MDNode *DstValue = cast<MDNode>(DstOp->getOperand(2));
       MDNode *SrcValue = cast<MDNode>(SrcOp->getOperand(2));
       SmallVector<Metadata *, 8> MDs;
       MDs.reserve(DstValue->getNumOperands() + SrcValue->getNumOperands());
       MDs.append(DstValue->op_begin(), DstValue->op_end());
       MDs.append(SrcValue->op_begin(), SrcValue->op_end());
 
-      replaceDstValue(MDNode::get(DstM.getContext(), MDs));
+      // Use a temporary MDnode because it can be deleted..
+      replaceDstValue(MDNode::getTemporary(DstM.getContext(), MDs).release());
       break;
     }
     case Module::AppendUnique: {
       SmallSetVector<Metadata *, 16> Elts;
       MDNode *DstValue = cast<MDNode>(DstOp->getOperand(2));
       MDNode *SrcValue = cast<MDNode>(SrcOp->getOperand(2));
       Elts.insert(DstValue->op_begin(), DstValue->op_end());
       Elts.insert(SrcValue->op_begin(), SrcValue->op_end());
 
-      replaceDstValue(MDNode::get(DstM.getContext(),
-                                  makeArrayRef(Elts.begin(), Elts.end())));
+      replaceDstValue(
+          MDNode::getTemporary(DstM.getContext(),
+                               makeArrayRef(Elts.begin(), Elts.end()))
+              .release());
       break;
     }
     }
 
   }
 
   // Check all of the requirements.
   for (unsigned I = 0, E = Requirements.size(); I != E; ++I) {

llvm/tools/llvm-link/llvm-link.cpp

int main(int argc, char **argv) {
   if (!linkFiles(argv[0], Context, L, InputFilenames, Flags))
     return 1;
 
   // Next the -override ones.
   if (!linkFiles(argv[0], Context, L, OverridingInputs,
                  Flags | Linker::Flags::OverrideFromSrc))
     return 1;
 
+  Composite->makeTempMDNodesPermanent();
+
   // Import any functions requested via -import
   if (!importFunctions(argv[0], *Composite))
     return 1;
 
   if (DumpAsm)
     errs() << "Here's the assembly:\n" << *Composite;
 
   std::error_code EC;

llvm/unittests/Linker/LinkModulesTest.cpp

TEST_F(LinkModuleTest, RemangleIntrinsics) {
 
   // "struct.rtx_def" from Foo and "struct.rtx_def.0" from Bar are isomorphic
   // types, so they must be uniquified by linker. Check that they use the same
   // intrinsic definition.
   Function *F = Foo->getFunction("llvm.memset.p0s_struct.rtx_defs.i32");
   ASSERT_EQ(F->getNumUses(), (unsigned)2);
 }
 
+TEST_F(LinkModuleTest, UseTempMDNodesForMduleFlags) {
+  LLVMContext C;
+  SMDiagnostic Err;
+
+  // Link 2 modules with metadata nodes that exhibit the behaviors "append" and
+  // "appendUnique". Make sure that the list MD nodes of the linked module are
+  // temporary MD nodes, so that they can be deleted in subsequent links,
+  // preventing memory being wasted on longer list nodes.
+  // FIXME: This test does not check that such a temp MD node is, in fact,
+  // deleted and truly deallocates its memory. Not sure how to write a test
+  // that would do this.
+
+  const char *FooStr = R"IR(
+    declare void @foo(i32, i32)
+    declare i32  @bar(i32)
+    declare i32  @baz(i32, i32)
+
+    !llvm.module.flags = !{!0, !1}
+    !0 = !{i32 5, !"CG Profile", !2}
+    !1 = !{i32 6, !"MyFlag", !3}
+    !2 = !{!4, !5}
+    !3 = !{!6, !7}
+    !4 = !{void (i32, i32)* @foo, i32 (i32)* @bar, i64 1}
+    !5 = !{void (i32, i32)* @foo, i32 (i32, i32)* @baz, i64 1}
+    !6 = !{i32 100, i32 200}
+    !7 = !{i32 300, i32 400}
+    )IR";
+
+  const char *BarStr = R"IR(
+    declare void @bar(i32, i32)
+    declare i32  @baz(i32)
+    declare i32  @fie(i32, i32)
+
+    !llvm.module.flags = !{!0, !1}
+    !0 = !{i32 5, !"CG Profile", !2}
+    !1 = !{i32 6, !"MyFlag", !3}
+    !2 = !{!4, !5}
+    !3 = !{!6, !7}
+    !4 = !{void (i32, i32)* @bar, i32 (i32)* @baz, i64 1}
+    !5 = !{void (i32, i32)* @bar, i32 (i32, i32)* @fie, i64 1}
+    !6 = !{i32 100, i32 200}
+    !7 = !{i32 300, i32 400}
+    )IR";
+
+  std::unique_ptr<Module> Foo = parseAssemblyString(FooStr, Err, C);
+  assert(Foo);
+  std::unique_ptr<Module> Bar = parseAssemblyString(BarStr, Err, C);
+  assert(Bar);
+  bool Failed = Linker::linkModules(*Foo, std::move(Bar));
+  ASSERT_FALSE(Failed);
+
+  const NamedMDNode *ModFlags = Foo->getModuleFlagsMetadata();
+  const MDNode *CGProfileFlag = ModFlags->getOperand(0);
+  ASSERT_TRUE(CGProfileFlag->getNumOperands() == 3);
+  const MDOperand &ListNode = CGProfileFlag->getOperand(2);
+  const MDNode *ListMDNode = cast<MDNode>(ListNode.get());
+  ASSERT_TRUE(ListMDNode->isTemporary());
+  ASSERT_TRUE(ListMDNode->getNumOperands() == 4);
+
+  const MDNode *MyFlag = ModFlags->getOperand(1);
+  ASSERT_TRUE(MyFlag->getNumOperands() == 3);
+  const MDOperand &ListNode2 = MyFlag->getOperand(2);
+  const MDNode *ListMDNode2 = cast<MDNode>(ListNode2.get());
+  ASSERT_TRUE(ListMDNode2->isTemporary());
+  ASSERT_TRUE(ListMDNode2->getNumOperands() == 2);
+}
+
 } // end anonymous namespace