[Utils] Identity map module-level debug info on first use in CloneFunction*

Summary:
To avoid cloning module-level debug info (owned by the module rather
than the function), CloneFunction implementation used to eagerly
identity map such debug info into ValueMap's MD map. In larger modules
with meaningful volume of debug info this gets very expensive.

By passing such debug info metadata via an IdentityMD set for the
ValueMapper to map on first use, we get several benefits:

1. Mapping metadata is not cheap, particularly because of tracking. When
   cloning a Function we identity map lots of global module-level
   metadata to avoid cloning it, while only a fraction of it is actually
   used by the function. Mapping on first use is a lot faster for
   modules with meaningful amount of debug info.

2. Eagerly identity mapping metadata makes it harder to cache
   module-level data (e.g. a set of metadata nodes in a \a DICompileUnit).
   With this patch we can cache certain module-level metadata
   calculations to speed things up further.

Anecdata from compiling a sample cpp file with full debug info shows that this moderately speeds up
CoroSplitPass which is one of the heavier users of cloning:

|                 | Baseline | IdentityMD set |
|-----------------|----------|----------------|
| CoroSplitPass   | 306ms    | 221ms          |
| CoroCloner      | 101ms    | 72ms           |
| Speed up        | 1x       | 1.4x           |

Test Plan:
ninja check-llvm-unit
ninja check-llvm

stack-info: PR: #118627, branch: users/artempyanykh/fast-coro-upstream/8

Author: artempyanykh

llvm/include/llvm/Transforms/Utils/Cloning.h

@@ -193,7 +193,8 @@ void CloneFunctionAttributesInto(Function *NewFunc, const Function *OldFunc,
 void CloneFunctionMetadataInto(Function &NewFunc, const Function &OldFunc,
                                ValueToValueMapTy &VMap, RemapFlags RemapFlag,
                                ValueMapTypeRemapper *TypeMapper = nullptr,
-                               ValueMaterializer *Materializer = nullptr);
+                               ValueMaterializer *Materializer = nullptr,
+                               const MetadataSetTy *IdentityMD = nullptr);
 
 /// Clone OldFunc's body into NewFunc.
 void CloneFunctionBodyInto(Function &NewFunc, const Function &OldFunc,
@@ -202,7 +203,8 @@ void CloneFunctionBodyInto(Function &NewFunc, const Function &OldFunc,
                            const char *NameSuffix = "",
                            ClonedCodeInfo *CodeInfo = nullptr,
                            ValueMapTypeRemapper *TypeMapper = nullptr,
-                           ValueMaterializer *Materializer = nullptr);
+                           ValueMaterializer *Materializer = nullptr,
+                           const MetadataSetTy *IdentityMD = nullptr);
 
 void CloneAndPruneIntoFromInst(Function *NewFunc, const Function *OldFunc,
                                const Instruction *StartingInst,
@@ -242,13 +244,12 @@ DISubprogram *CollectDebugInfoForCloning(const Function &F,
                                          CloneFunctionChangeType Changes,
                                          DebugInfoFinder &DIFinder);
 
-/// Build a map of debug info to use during Metadata cloning.
-/// Returns true if cloning would need module level changes and false if there
-/// would only be local changes.
-bool BuildDebugInfoMDMap(DenseMap<const Metadata *, TrackingMDRef> &MD,
-                         CloneFunctionChangeType Changes,
-                         DebugInfoFinder &DIFinder,
-                         DISubprogram *SPClonedWithinModule);
+/// Based on \p Changes and \p DIFinder populate \p MD with debug info that
+/// needs to be identity mapped during Metadata cloning.
+void FindDebugInfoToIdentityMap(MetadataSetTy &MD,
+                                CloneFunctionChangeType Changes,
+                                DebugInfoFinder &DIFinder,
+                                DISubprogram *SPClonedWithinModule);
 
 /// This class captures the data input to the InlineFunction call, and records
 /// the auxiliary results produced by it.

llvm/include/llvm/Transforms/Utils/ValueMapper.h

@@ -15,6 +15,7 @@
 #define LLVM_TRANSFORMS_UTILS_VALUEMAPPER_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/simple_ilist.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/IR/ValueMap.h"
@@ -35,6 +36,7 @@ class Value;
 
 using ValueToValueMapTy = ValueMap<const Value *, WeakTrackingVH>;
 using DbgRecordIterator = simple_ilist<DbgRecord>::iterator;
+using MetadataSetTy = SmallPtrSet<const Metadata *, 16>;
 
 /// This is a class that can be implemented by clients to remap types when
 /// cloning constants and instructions.
@@ -136,6 +138,18 @@ inline RemapFlags operator|(RemapFlags LHS, RemapFlags RHS) {
 /// alternate \a ValueToValueMapTy and \a ValueMaterializer and returns a ID to
 /// pass into the schedule*() functions.
 ///
+/// NOTE: \c IdentityMD is used by CloneFunction* to directly specify metadata
+/// that should be identity mapped (and hence not cloned). The metadata will be
+/// identity mapped in \c VM on first use. There are several reasons for doing
+/// it this way rather than eagerly identity mapping metadata nodes in \c VM:
+/// 1. Mapping metadata is not cheap, particularly because of tracking.
+/// 2. When cloning a Function we identity map lots of global module-level
+///    metadata to avoid cloning it, while only a fraction of it is actually
+///    used by the function. Mapping on first use is a lot faster for modules
+///    with meaningful amount of debug info.
+/// 3. Eagerly identity mapping metadata makes it harder to cache module-level
+///    data (e.g. a set of metadata nodes in a \a DICompileUnit).
+///
 /// TODO: lib/Linker really doesn't need the \a ValueHandle in the \a
 /// ValueToValueMapTy.  We should template \a ValueMapper (and its
 /// implementation classes), and explicitly instantiate on two concrete
@@ -152,7 +166,8 @@ class ValueMapper {
 public:
   ValueMapper(ValueToValueMapTy &VM, RemapFlags Flags = RF_None,
               ValueMapTypeRemapper *TypeMapper = nullptr,
-              ValueMaterializer *Materializer = nullptr);
+              ValueMaterializer *Materializer = nullptr,
+              const MetadataSetTy *IdentityMD = nullptr);
   ValueMapper(ValueMapper &&) = delete;
   ValueMapper(const ValueMapper &) = delete;
   ValueMapper &operator=(ValueMapper &&) = delete;
@@ -218,8 +233,10 @@ class ValueMapper {
 inline Value *MapValue(const Value *V, ValueToValueMapTy &VM,
                        RemapFlags Flags = RF_None,
                        ValueMapTypeRemapper *TypeMapper = nullptr,
-                       ValueMaterializer *Materializer = nullptr) {
-  return ValueMapper(VM, Flags, TypeMapper, Materializer).mapValue(*V);
+                       ValueMaterializer *Materializer = nullptr,
+                       const MetadataSetTy *IdentityMD = nullptr) {
+  return ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .mapValue(*V);
 }
 
 /// Lookup or compute a mapping for a piece of metadata.
@@ -231,7 +248,9 @@ inline Value *MapValue(const Value *V, ValueToValueMapTy &VM,
 ///     \c MD.
 ///  3. Else if \c MD is a \a ConstantAsMetadata, call \a MapValue() and
 ///     re-wrap its return (returning nullptr on nullptr).
-///  4. Else, \c MD is an \a MDNode.  These are remapped, along with their
+///  4. Else if \c MD is in \c IdentityMD then add an identity mapping for it
+///     and return it.
+///  5. Else, \c MD is an \a MDNode.  These are remapped, along with their
 ///     transitive operands.  Distinct nodes are duplicated or moved depending
 ///     on \a RF_MoveDistinctNodes.  Uniqued nodes are remapped like constants.
 ///
@@ -240,16 +259,20 @@ inline Value *MapValue(const Value *V, ValueToValueMapTy &VM,
 inline Metadata *MapMetadata(const Metadata *MD, ValueToValueMapTy &VM,
                              RemapFlags Flags = RF_None,
                              ValueMapTypeRemapper *TypeMapper = nullptr,
-                             ValueMaterializer *Materializer = nullptr) {
-  return ValueMapper(VM, Flags, TypeMapper, Materializer).mapMetadata(*MD);
+                             ValueMaterializer *Materializer = nullptr,
+                             const MetadataSetTy *IdentityMD = nullptr) {
+  return ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .mapMetadata(*MD);
 }
 
 /// Version of MapMetadata with type safety for MDNode.
 inline MDNode *MapMetadata(const MDNode *MD, ValueToValueMapTy &VM,
                            RemapFlags Flags = RF_None,
                            ValueMapTypeRemapper *TypeMapper = nullptr,
-                           ValueMaterializer *Materializer = nullptr) {
-  return ValueMapper(VM, Flags, TypeMapper, Materializer).mapMDNode(*MD);
+                           ValueMaterializer *Materializer = nullptr,
+                           const MetadataSetTy *IdentityMD = nullptr) {
+  return ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .mapMDNode(*MD);
 }
 
 /// Convert the instruction operands from referencing the current values into
@@ -263,17 +286,21 @@ inline MDNode *MapMetadata(const MDNode *MD, ValueToValueMapTy &VM,
 inline void RemapInstruction(Instruction *I, ValueToValueMapTy &VM,
                              RemapFlags Flags = RF_None,
                              ValueMapTypeRemapper *TypeMapper = nullptr,
-                             ValueMaterializer *Materializer = nullptr) {
-  ValueMapper(VM, Flags, TypeMapper, Materializer).remapInstruction(*I);
+                             ValueMaterializer *Materializer = nullptr,
+                             const MetadataSetTy *IdentityMD = nullptr) {
+  ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .remapInstruction(*I);
 }
 
 /// Remap the Values used in the DbgRecord \a DR using the value map \a
 /// VM.
 inline void RemapDbgRecord(Module *M, DbgRecord *DR, ValueToValueMapTy &VM,
                            RemapFlags Flags = RF_None,
                            ValueMapTypeRemapper *TypeMapper = nullptr,
-                           ValueMaterializer *Materializer = nullptr) {
-  ValueMapper(VM, Flags, TypeMapper, Materializer).remapDbgRecord(M, *DR);
+                           ValueMaterializer *Materializer = nullptr,
+                           const MetadataSetTy *IdentityMD = nullptr) {
+  ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .remapDbgRecord(M, *DR);
 }
 
 /// Remap the Values used in the DbgRecords \a Range using the value map \a
@@ -283,8 +310,9 @@ inline void RemapDbgRecordRange(Module *M,
                                 ValueToValueMapTy &VM,
                                 RemapFlags Flags = RF_None,
                                 ValueMapTypeRemapper *TypeMapper = nullptr,
-                                ValueMaterializer *Materializer = nullptr) {
-  ValueMapper(VM, Flags, TypeMapper, Materializer)
+                                ValueMaterializer *Materializer = nullptr,
+                                const MetadataSetTy *IdentityMD = nullptr) {
+  ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
       .remapDbgRecordRange(M, Range);
 }
 
@@ -297,16 +325,19 @@ inline void RemapDbgRecordRange(Module *M,
 inline void RemapFunction(Function &F, ValueToValueMapTy &VM,
                           RemapFlags Flags = RF_None,
                           ValueMapTypeRemapper *TypeMapper = nullptr,
-                          ValueMaterializer *Materializer = nullptr) {
-  ValueMapper(VM, Flags, TypeMapper, Materializer).remapFunction(F);
+                          ValueMaterializer *Materializer = nullptr,
+                          const MetadataSetTy *IdentityMD = nullptr) {
+  ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD).remapFunction(F);
 }
 
 /// Version of MapValue with type safety for Constant.
 inline Constant *MapValue(const Constant *V, ValueToValueMapTy &VM,
                           RemapFlags Flags = RF_None,
                           ValueMapTypeRemapper *TypeMapper = nullptr,
-                          ValueMaterializer *Materializer = nullptr) {
-  return ValueMapper(VM, Flags, TypeMapper, Materializer).mapConstant(*V);
+                          ValueMaterializer *Materializer = nullptr,
+                          const MetadataSetTy *IdentityMD = nullptr) {
+  return ValueMapper(VM, Flags, TypeMapper, Materializer, IdentityMD)
+      .mapConstant(*V);
 }
 
 } // end namespace llvm

llvm/lib/Transforms/Utils/CloneFunction.cpp

@@ -152,61 +152,54 @@ DISubprogram *llvm::CollectDebugInfoForCloning(const Function &F,
   return SPClonedWithinModule;
 }
 
-bool llvm::BuildDebugInfoMDMap(DenseMap<const Metadata *, TrackingMDRef> &MD,
-                               CloneFunctionChangeType Changes,
-                               DebugInfoFinder &DIFinder,
-                               DISubprogram *SPClonedWithinModule) {
-  bool ModuleLevelChanges = Changes > CloneFunctionChangeType::LocalChangesOnly;
+void llvm::FindDebugInfoToIdentityMap(MetadataSetTy &MD,
+                                      CloneFunctionChangeType Changes,
+                                      DebugInfoFinder &DIFinder,
+                                      DISubprogram *SPClonedWithinModule) {
   if (Changes < CloneFunctionChangeType::DifferentModule &&
       DIFinder.subprogram_count() > 0) {
-    // Turn on module-level changes, since we need to clone (some of) the
-    // debug info metadata.
+    // Even if Changes are local only, we turn on module-level changes, since we
+    // need to clone (some of) the debug info metadata.
     //
     // FIXME: Metadata effectively owned by a function should be made
     // local, and only that local metadata should be cloned.
-    ModuleLevelChanges = true;
-
-    auto mapToSelfIfNew = [&MD](MDNode *N) {
-      // Avoid clobbering an existing mapping.
-      (void)MD.try_emplace(N, N);
-    };
 
     // Avoid cloning types, compile units, and (other) subprograms.
     for (DISubprogram *ISP : DIFinder.subprograms()) {
       if (ISP != SPClonedWithinModule)
-        mapToSelfIfNew(ISP);
+        MD.insert(ISP);
     }
 
     // If a subprogram isn't going to be cloned skip its lexical blocks as well.
     for (DIScope *S : DIFinder.scopes()) {
       auto *LScope = dyn_cast<DILocalScope>(S);
       if (LScope && LScope->getSubprogram() != SPClonedWithinModule)
-        mapToSelfIfNew(S);
+        MD.insert(S);
     }
 
     for (DICompileUnit *CU : DIFinder.compile_units())
-      mapToSelfIfNew(CU);
+      MD.insert(CU);
 
     for (DIType *Type : DIFinder.types())
-      mapToSelfIfNew(Type);
+      MD.insert(Type);
   } else {
     assert(!SPClonedWithinModule &&
            "Subprogram should be in DIFinder->subprogram_count()...");
   }
-
-  return ModuleLevelChanges;
 }
 
 void llvm::CloneFunctionMetadataInto(Function &NewFunc, const Function &OldFunc,
                                      ValueToValueMapTy &VMap,
                                      RemapFlags RemapFlag,
                                      ValueMapTypeRemapper *TypeMapper,
-                                     ValueMaterializer *Materializer) {
+                                     ValueMaterializer *Materializer,
+                                     const MetadataSetTy *IdentityMD) {
   SmallVector<std::pair<unsigned, MDNode *>, 1> MDs;
   OldFunc.getAllMetadata(MDs);
   for (auto MD : MDs) {
-    NewFunc.addMetadata(MD.first, *MapMetadata(MD.second, VMap, RemapFlag,
-                                               TypeMapper, Materializer));
+    NewFunc.addMetadata(MD.first,
+                        *MapMetadata(MD.second, VMap, RemapFlag, TypeMapper,
+                                     Materializer, IdentityMD));
   }
 }
 
@@ -216,7 +209,8 @@ void llvm::CloneFunctionBodyInto(Function &NewFunc, const Function &OldFunc,
                                  const char *NameSuffix,
                                  ClonedCodeInfo *CodeInfo,
                                  ValueMapTypeRemapper *TypeMapper,
-                                 ValueMaterializer *Materializer) {
+                                 ValueMaterializer *Materializer,
+                                 const MetadataSetTy *IdentityMD) {
   if (OldFunc.isDeclaration())
     return;
 
@@ -258,9 +252,10 @@ void llvm::CloneFunctionBodyInto(Function &NewFunc, const Function &OldFunc,
     // Loop over all instructions, fixing each one as we find it, and any
     // attached debug-info records.
     for (Instruction &II : *BB) {
-      RemapInstruction(&II, VMap, RemapFlag, TypeMapper, Materializer);
+      RemapInstruction(&II, VMap, RemapFlag, TypeMapper, Materializer,
+                       IdentityMD);
       RemapDbgRecordRange(II.getModule(), II.getDbgRecordRange(), VMap,
-                          RemapFlag, TypeMapper, Materializer);
+                          RemapFlag, TypeMapper, Materializer, IdentityMD);
     }
 }
 
@@ -322,16 +317,21 @@ void llvm::CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
   DISubprogram *SPClonedWithinModule =
       CollectDebugInfoForCloning(*OldFunc, Changes, DIFinder);
 
-  ModuleLevelChanges =
-      BuildDebugInfoMDMap(VMap.MD(), Changes, DIFinder, SPClonedWithinModule);
+  MetadataSetTy IdentityMD;
+  FindDebugInfoToIdentityMap(IdentityMD, Changes, DIFinder,
+                             SPClonedWithinModule);
 
-  const auto RemapFlag = ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges;
+  // Current implementation always upgrades from local changes to module level
+  // changes due to the way metadata cloning is done. See
+  // BuildDebugInfoToIdentityMap for more details.
+  const auto RemapFlag = RF_None;
 
   CloneFunctionMetadataInto(*NewFunc, *OldFunc, VMap, RemapFlag, TypeMapper,
-                            Materializer);
+                            Materializer, &IdentityMD);
 
   CloneFunctionBodyInto(*NewFunc, *OldFunc, VMap, RemapFlag, Returns,
-                        NameSuffix, CodeInfo, TypeMapper, Materializer);
+                        NameSuffix, CodeInfo, TypeMapper, Materializer,
+                        &IdentityMD);
 
   // Only update !llvm.dbg.cu for DifferentModule (not CloneModule). In the
   // same module, the compile unit will already be listed (or not). When

llvm/lib/Transforms/Utils/ValueMapper.cpp

@@ -120,12 +120,14 @@ class Mapper {
   SmallVector<WorklistEntry, 4> Worklist;
   SmallVector<DelayedBasicBlock, 1> DelayedBBs;
   SmallVector<Constant *, 16> AppendingInits;
+  const MetadataSetTy *IdentityMD;
 
 public:
   Mapper(ValueToValueMapTy &VM, RemapFlags Flags,
-         ValueMapTypeRemapper *TypeMapper, ValueMaterializer *Materializer)
+         ValueMapTypeRemapper *TypeMapper, ValueMaterializer *Materializer,
+         const MetadataSetTy *IdentityMD)
       : Flags(Flags), TypeMapper(TypeMapper),
-        MCs(1, MappingContext(VM, Materializer)) {}
+        MCs(1, MappingContext(VM, Materializer)), IdentityMD(IdentityMD) {}
 
   /// ValueMapper should explicitly call \a flush() before destruction.
   ~Mapper() { assert(!hasWorkToDo() && "Expected to be flushed"); }
@@ -899,6 +901,14 @@ std::optional<Metadata *> Mapper::mapSimpleMetadata(const Metadata *MD) {
     return wrapConstantAsMetadata(*CMD, mapValue(CMD->getValue()));
   }
 
+  // Map metadata from IdentityMD on first use. We need to add these nodes to
+  // the mapping as otherwise metadata nodes numbering gets messed up. This is
+  // still economical because the amount of data in IdentityMD may be a lot
+  // larger than what will actually get used.
+  if (IdentityMD && IdentityMD->contains(MD)) {
+    return getVM().MD()[MD] = TrackingMDRef(const_cast<Metadata *>(MD));
+  }
+
   assert(isa<MDNode>(MD) && "Expected a metadata node");
 
   return std::nullopt;
@@ -1198,8 +1208,9 @@ class FlushingMapper {
 
 ValueMapper::ValueMapper(ValueToValueMapTy &VM, RemapFlags Flags,
                          ValueMapTypeRemapper *TypeMapper,
-                         ValueMaterializer *Materializer)
-    : pImpl(new Mapper(VM, Flags, TypeMapper, Materializer)) {}
+                         ValueMaterializer *Materializer,
+                         const MetadataSetTy *IdentityMD)
+    : pImpl(new Mapper(VM, Flags, TypeMapper, Materializer, IdentityMD)) {}
 
 ValueMapper::~ValueMapper() { delete getAsMapper(pImpl); }