[clang][modules] Account for non-affecting inputs in `ASTWriter`

[jansvoboda11]

Summary

In D106876, we stopped serializing module map files that didn't affect compilation of the current module.

However, since each SourceLocation is simply an offset into SourceManager's global buffer of concatenated input files in, these need to be adjusted during serialization. Otherwise, they can incorrectly point after the buffer or into subsequent input file.

This patch starts adjusting SourceLocations, FileIDs and other SourceManager offsets in ASTWriter.

[Bigcheese]

This looks reasonable. Have you measured the performance impact of this change?

[jansvoboda11]

In D136624#3880526, @Bigcheese wrote:

This looks reasonable. Have you measured the performance impact of this change?

I have done comparison between this patch and https://github.com/apple/llvm-project/pull/5451 (that, instead of leaving non-affecting input files out, serializes one extra bit saying whether they are affecting or not). Both versions carried some additional patches that refine the computation of non-affecting module maps, that I plan to upstream shortly. I used Clang's -ftime-trace on a reasonably large project (1680 TUs, 638 implicitly-built modules). The results say that serialization got 7.09% slower, and deserialization 1.01% slower. The serialization slowdown I understand, but I expected deserialization to get faster, since we now have less of SourceManager to look through. Overall, PCM compilation got 4.35% slower.

[jansvoboda11]

I tried optimizing this patch a bit. Instead of creating compact data structure and using binary search to find the preceding non-affecting file, I now store the adjustment information for each FileID in a vector. During deserialization, FileID is simply used as an index into SLocEntryInfos. That didn't yield any measurable improvement in performance, though. I think the regression must be coming from the SourceLocation/Offset to FileID translation.

I don't see any obvious way to work around that. SourceManager::getFileIDLocal() already implements some optimizations that makes accessing nearby offsets fast. A separate SourceManager would avoid this bottleneck, but I'm not sure how much work that would entail (seems substantial).

@Bigcheese LMK if you're fine with the performance implications here.

[jansvoboda11]

New version with flat vector + FileID indices; replacing the previous compact representation & binary search
approach

[dexonsmith]

In D136624#3888387, @jansvoboda11 wrote:

I tried optimizing this patch a bit. Instead of creating compact data structure and using binary search to find the preceding non-affecting file, I now store the adjustment information for each FileID in a vector. During deserialization, FileID is simply used as an index into SLocEntryInfos. That didn't yield any measurable improvement in performance, though. I think the regression must be coming from the SourceLocation/Offset to FileID translation.

I don't see any obvious way to work around that. SourceManager::getFileIDLocal() already implements some optimizations that makes accessing nearby offsets fast. A separate SourceManager would avoid this bottleneck, but I'm not sure how much work that would entail (seems substantial).

@Bigcheese LMK if you're fine with the performance implications here.

I don't think you need to call getFileID() inside getAdjustment(). There is also an opportunity for a peephole, if getAdjustment() remains expensive.

I've left some comments on the previous version of the patch since it's not obvious to me how to avoid the getFileID() call in the new version.

In D136624#3888097, @jansvoboda11 wrote:

The serialization slowdown I understand, but I expected deserialization to get faster, since we now have less of SourceManager to look through.

Seems worth digging into the deserialization regression. Does the PCM actually get smaller and the ranges more condensed?

One quick test would be to manufacture a situation where two output PCMs would previously have different non-affecting inputs, but now should be bit-for-bit identical. Are they, in fact, bit-for-bit identical? If not, maybe there's something funny to look into...

clang/include/clang/Serialization/ASTWriter.h

449–452

Can you collect a histogram for how big these vectors are? Can we avoid pointer chasing in the common case by making them SmallVector of some size during lookup?

clang/lib/Serialization/ASTWriter.cpp

2054

Can we shift this getAdjustedOffset() computation to after deciding whether to skip the record?

5282–5283

How often does getAdjustment() return the same answer in consecutive calls? If at all common, this would likely benefit from a peephole:

Optional<SLocRange> ASTWriter::CachedAdjustmentRange;
Optional<UIntTy> ASTWriter::CachedAdjustment;

SourceLocation::UIntTy
ASTWriter::getAdjustment(SourceLocation::UIntTy Offset) const {
  // Check for 0.
  //
  // How fast is "isLoadedOffset()"? Can/should we add a peephole, or is it just bit
  // manipulation? (I seem to remember it checking the high bit or something, but if
  // it's doing some sort of look up, maybe it should be in the slow path so it can
  // get cached by LastAdjustment.)
  if (PP->getSourceManager().isLoadedOffset(Offset) ||
      NonAffectingInputs.empty())
    return 0;

  // Check CachedAdjustment.
  if (CachedAdjustment && CachedAdjustmentRange->includes(Offset))
    return *CachedAdjustment;

  // Call getAdjustmentSlow, which updates CachedAdjustment and CachedAdjustmentRange.
  // It's out-of-line so that `getAdjustment` can easily be inlined without inlining
  // the slow path.
  //
  // LastAdjustmentRange would be the size of the "gap" between this adjustment
  // level and the next one (end would be UINTMAX if it's after the last
  // non-affecting range).
  return getAdjustmentSlow(Offset);
}

5289–5290

Why do you need to call getFileID() here?

Instead, I would expect this to be a search through a range of offsets (e.g., see my suggestion at https://reviews.llvm.org/D106876#3869247 -- DroppedMMs contains SourceLocations, not FileIDs).

Two benefits:

1. You don't need to call getFileID() to look up an offset.
2. You can merge adjacent non-affecting files (shrinking the search/storage significantly).

clang/test/Modules/non-affecting-module-maps-source-locations.m

32

This is exercising the code, but it could do one better and check if the output PCMs are bit-for-bit identical when we (now) expect them to be.

Maybe you could do this by having two run lines: one that includes -I %t/second and another that doesn't. Then check if the output PCMs are equal.

[dexonsmith]

clang/test/Modules/non-affecting-module-maps-source-locations.m

32

(Or if the PCM isn't bit-for-bit identical yet, maybe at least the AST block should be...)

[jansvoboda11]

Avoid looking up FileID for an offset.

[jansvoboda11]

I tried implementing your suggestion (merging ranges of adjacent non-affecting files and avoiding FileID lookups), but the numbers from -ftime-trace are very noisy. I got more stable data by measuring clock cycles and instruction counts, but nothing conclusive yet.

Compilation of CompilerInvocation.cpp with implicit modules.

• previous approach with vector + FileID lookup: +0.64% cycles and +1.68% instructions,
• current approach with merged SourceRanges: +0.38% cycles and +1.11% instructions.

I'll post here as I experiment more and get more data.

clang/lib/Serialization/ASTWriter.cpp

5289–5290

My reasoning was that if we search through a range of offsets, we're doing conceptually the same thing as getFileID() (which already has some optimizations baked in). Maybe the non-affecting files are indeed adjacent and we'll be able to merge most of them. I'll give it a shot and report back.

[dexonsmith]

In D136624#3893001, @jansvoboda11 wrote:

I tried implementing your suggestion (merging ranges of adjacent non-affecting files and avoiding FileID lookups), but the numbers from -ftime-trace are very noisy. I got more stable data by measuring clock cycles and instruction counts, but nothing conclusive yet.

Compilation of CompilerInvocation.cpp with implicit modules.

• previous approach with vector + FileID lookup: +0.64% cycles and +1.68% instructions,
• current approach with merged SourceRanges: +0.38% cycles and +1.11% instructions.

I'll post here as I experiment more and get more data.

Nice; that seems like a bit of an improvement.

I'm curious; are system modules allowed to be non-affecting yet, or are they still assumed to be affecting? (It's the system modules that I think are most likely to be adjacent.)

My intuition is that there is likely some peephole that would be quite effective, that might not be useful for general getFileID() lookups.

• I already suggested "same as last lookup?"... I'm curious if that'll help. Maybe that's already in getFileID(), but now that you've factored out that call, it could be useful to replicate.
• You could also try: "past the the last non-affecting module?"
• You could also try: "before the first non-affecting module?"

I suspect you could collect some data to guide this, such as, for loaded locations (you could ignore "local" locations since they already have a peephole):

• Histogram of "loaded" vs. "between" vs. "after" non-affecting modules.
• Histogram of "same as last" vs. "same as last-1" vs. "different from last 2".
• [...]

Other things that might be useful to know:

• What effect is the merging having (or would it have)? (i.e., what's the histogram of "adjacent" non-affecting files? (e.g.: 9 ranges of non-affecting files, with two blocks of 5 files and seven blocks of 1 (which aren't adjacent to any others)))
• Is there a change in cycles/instructions when the module cache is hot? (presumably the common case)
• Are the PCM artifacts smaller?
• Are the PCMs bit-for-bit identical now when a non-affecting module is added to the input? (If not, why not?)
• What's the data for implicitly-discovered, explicitly-built modules?

[jansvoboda11]

In D136624#3893607, @dexonsmith wrote:

I'm curious; are system modules allowed to be non-affecting yet, or are they still assumed to be affecting? (It's the system modules that I think are most likely to be adjacent.)

Yes, all the measurements are done with system modules allowed to be non-affecting.

My intuition is that there is likely some peephole that would be quite effective, that might not be useful for general getFileID() lookups.

• I already suggested "same as last lookup?"... I'm curious if that'll help. Maybe that's already in getFileID(), but now that you've factored out that call, it could be useful to replicate.

I tried that and we seemingly never hit that case. I think that's because these two optimizations take precedence:

• You could also try: "past the the last non-affecting module?"
• You could also try: "before the first non-affecting module?"

These avoid the binary search for the vast majority of calls to getAdjustment() with local offsets.

I suspect you could collect some data to guide this, such as, for loaded locations (you could ignore "local" locations since they already have a peephole):

• Histogram of "loaded" vs. "between" vs. "after" non-affecting modules.
• Histogram of "same as last" vs. "same as last-1" vs. "different from last 2".
• [...]

Loaded locations make up 68% of all calls to getAdjustment(), so it's good it's the first thing we check for. Around 4% of all calls are for locations before the first non-affecting module. That's the second special case. Around 28% are pointing after the last non-affecting module, and the current revision has a special case for that as well. I think it would make sense to prioritize this check. Only the remaining 0.3% of calls do the actual binary search.

Other things that might be useful to know:

• What effect is the merging having (or would it have)? (i.e., what's the histogram of "adjacent" non-affecting files? (e.g.: 9 ranges of non-affecting files, with two blocks of 5 files and seven blocks of 1 (which aren't adjacent to any others)))

Big one. We usually get between 70-100 non-affecting module maps, but they merge into 4-6 consecutive regions.

• Is there a change in cycles/instructions when the module cache is hot? (presumably the common case)

I didn't notice this (but didn't look for it specifically). How could that affect performance for PCM writes?

• Are the PCM artifacts smaller?

Yes, we leave out the 70-100 SLocEntries. Nothing much changes otherwise.

• Are the PCMs bit-for-bit identical now when a non-affecting module is added to the input? (If not, why not?)

They are not in the current revision, but I'll create a patch that makes them so (we also need to adjust the NumCreatedFileIDs).

• What's the data for implicitly-discovered, explicitly-built modules?

I didn't measure that. I don't expect much of a difference, though. The scanner has an empty AST, so the number of SourceLocations will be pretty low. (Other parts of the PCM file, like the metadata, don't write much SourceManager offsets.) Even if there was a small regression, I'd be fine with it, since we'll be moving off of implicit build in the scanner. For the explicit builds themselves, this shouldn't affect performance at all. Implicit module map search is disabled, all necessary/affecting module maps are deserialized from the PCM files or provided on the command-line.

[jansvoboda11]

clang/include/clang/Serialization/ASTWriter.h

449–452

Usually 4-6 elements. Making them a SmallVector<T, 8> didn't affect performance, though.

clang/lib/Serialization/ASTWriter.cpp

5282–5283

Not that often, see my top-level comment.

5289–5290

This ended up being faster due to merging of non-affecting files. Thanks for the suggestion!

clang/test/Modules/non-affecting-module-maps-source-locations.m

32

Yes, I'll probably drop this test entirely and just check the PCM files are bit-for-bit identical when a non-affecting file is not loaded at all.

[dexonsmith]

In D136624#3900051, @jansvoboda11 wrote:

• Is there a change in cycles/instructions when the module cache is hot? (presumably the common case)

I didn't notice this (but didn't look for it specifically). How could that affect performance for PCM writes?

It wouldn't check write perf, but it's part of the overall build perf impact. This being neutral (no regression) could help to justify landing the change even when there's a penalty for a cold modules cache. My interest in (implicitly-discovered) explicitly-built modules was similar (if that's neutral, then a regression here is less critical).

Partly, trying to dig into why read speeds got slower. But maybe that was noise that went away though when you switched to cycles/instructions?

Loaded locations make up 68% of all calls to getAdjustment(), so it's good it's the first thing we check for. Around 4% of all calls are for locations before the first non-affecting module. That's the second special case. Around 28% are pointing after the last non-affecting module, and the current revision has a special case for that as well. I think it would make sense to prioritize this check. Only the remaining 0.3% of calls do the actual binary search.

Great; looking forward to seeing new numbers.

(BTW, if you check before/after last non-affecting module, does one of those subsume the is-loaded check entirely? Looking at isLoadedOffset() makes me think it might. If so, maybe you can replace that check with an assertion.)

In D136624#3900051, @jansvoboda11 wrote:

• Are the PCMs bit-for-bit identical now when a non-affecting module is added to the input? (If not, why not?)

They are not in the current revision, but I'll create a patch that makes them so (we also need to adjust the NumCreatedFileIDs).

Great. Seems like another motivation to land this change (despite a regression), as it can impact meta-build perf if/when artifacts are being archived/shared in a larger context. Specifically, if the PCM artifacts are more stable, then:

• They take up less aggregate space in CAS storage.
• Later build steps get more cache hits.

Overall this patch seems like a good step to me; I don't think 1-2% for a single compilation on a cold cache is that bad, assuming a hot cache doesn't regress:

• Across a single build, the cache gets hot pretty quickly as most compilations reuse the same modules.
• Across incremental builds, most of the cache (at least system modules) will (usually) be hot.
• For explicit builds, it sounds like there's no regression anyway.

clang/include/clang/Basic/SourceManager.h

1831–1833

The logic for isLoadedOffset() suggests that it could maybe be subsumed with "location past the end"?

clang/test/Modules/non-affecting-module-maps-source-locations.m

32

That sounds great.

[jansvoboda11]

In D136624#3900183, @dexonsmith wrote:

Partly, trying to dig into why read speeds got slower. But maybe that was noise that went away though when you switched to cycles/instructions?
Great; looking forward to seeing new numbers.

Ah, I forgot to mention this. Building the modules is now only 0.2% slower and importing them 1.2% faster (compared to PCMs with all input files serialized).

clang/include/clang/Basic/SourceManager.h

1831–1833

I don't think so - we don't want to adjust loaded offsets. Their invariant is that they grow from 2^31 downwards.

We do want to adjust local offsets past the last non-affecting file though.

[dexonsmith]

In D136624#3900593, @jansvoboda11 wrote:

Ah, I forgot to mention this. Building the modules is now only 0.2% slower and importing them 1.2% faster (compared to PCMs with all input files serialized).

Awesome. All upside then :).

I added a few nitpick-y comments inline. I can take another look once you've made the PCMs bit-for-bit identical and updated the test (is that happening here, or in a separate review)?

clang/include/clang/Basic/SourceManager.h

1831–1833

Oops, right!

clang/lib/Serialization/ASTWriter.cpp

4541

Not sure this comment adds much on top of the code on the next line. A sentence before the for loop describing the overall approach might be useful though.

4553–4557

You can reduce nesting by inverting this condition and continue.

4565–4566

I'd slightly prefer the comment *before* the if, due to how folding tends to work in editors (see the comment even when the code is folded). Probably relies on dropping the else (see below).

4571

I suggest continue before the else to avoid adding nesting for insertion.

[jansvoboda11]

Rebase, decrease nesting, test using diff

[dexonsmith]

LGTM, with one suggestion for the test inline.

clang/test/Modules/add-remove-irrelevant-module-map.m

22 ↗

(On Diff #472466)

Maybe the CC1s should add -verify and test-simple.m should have // expected-no-diagnostics to help protect against bitrot?

clang/include/clang/Basic/SourceManager.h

public:
 
   /// Return the FileID for a SourceLocation.
   ///
   /// This is a very hot method that is used for all SourceManager queries
   /// that start with a SourceLocation object.  It is responsible for finding
   /// the entry in SLocEntryTable which contains the specified location.
   ///
   FileID getFileID(SourceLocation SpellingLoc) const {
-    SourceLocation::UIntTy SLocOffset = SpellingLoc.getOffset();
-
-    // If our one-entry cache covers this offset, just return it.
-    if (isOffsetInFileID(LastFileIDLookup, SLocOffset))
-      return LastFileIDLookup;
-
-    return getFileIDSlow(SLocOffset);
+    return getFileID(SpellingLoc.getOffset());
   }
 
   /// Return the filename of the file containing a SourceLocation.
   StringRef getFilename(SourceLocation SpellingLoc) const;
 
   /// Return the source location corresponding to the first byte of
   /// the specified file.
   SourceLocation getLocForStartOfFile(FileID FID) const {
public:
   /// in the global source view. The lowest ID and the base offset of the
   /// entries will be returned.
   std::pair<int, SourceLocation::UIntTy>
   AllocateLoadedSLocEntries(unsigned NumSLocEntries,
                             SourceLocation::UIntTy TotalSize);
 
   /// Returns true if \p Loc came from a PCH/Module.
   bool isLoadedSourceLocation(SourceLocation Loc) const {
-    return Loc.getOffset() >= CurrentLoadedOffset;
+    return isLoadedOffset(Loc.getOffset());
   }
 
   /// Returns true if \p Loc did not come from a PCH/Module.
   bool isLocalSourceLocation(SourceLocation Loc) const {
-    return Loc.getOffset() < NextLocalOffset;
+    return isLocalOffset(Loc.getOffset());
   }
 
   /// Returns true if \p FID came from a PCH/Module.
   bool isLoadedFileID(FileID FID) const {
     assert(FID.ID != -1 && "Using FileID sentinel value");
     return FID.ID < 0;
   }
 
const SrcMgr::SLocEntry &getSLocEntryByID(int ID,
     return getLocalSLocEntry(static_cast<unsigned>(ID));
   }
 
   const SrcMgr::SLocEntry &
   getLoadedSLocEntryByID(int ID, bool *Invalid = nullptr) const {
     return getLoadedSLocEntry(static_cast<unsigned>(-ID - 2), Invalid);
   }
 
+  FileID getFileID(SourceLocation::UIntTy SLocOffset) const {
+    // If our one-entry cache covers this offset, just return it.
+    if (isOffsetInFileID(LastFileIDLookup, SLocOffset))
+      return LastFileIDLookup;
+
+    return getFileIDSlow(SLocOffset);
+  }
+
+  bool isLocalOffset(SourceLocation::UIntTy SLocOffset) const {
+    return SLocOffset < CurrentLoadedOffset;
+  }
+
+  bool isLoadedOffset(SourceLocation::UIntTy SLocOffset) const {
+    return SLocOffset >= CurrentLoadedOffset;
+  }
+
   /// Implements the common elements of storing an expansion info struct into
   /// the SLocEntry table and producing a source location that refers to it.
   SourceLocation
   createExpansionLocImpl(const SrcMgr::ExpansionInfo &Expansion,
                          unsigned Length, int LoadedID = 0,
                          SourceLocation::UIntTy LoadedOffset = 0);
 
   /// Return true if the specified FileID contains the

clang/include/clang/Serialization/ASTWriter.h

private:
   /// A mapping from each known submodule to its ID number, which will
   /// be a positive integer.
   llvm::DenseMap<const Module *, unsigned> SubmoduleIDs;
 
   /// A list of the module file extension writers.
   std::vector<std::unique_ptr<ModuleFileExtensionWriter>>
       ModuleFileExtensionWriters;
 
-  /// User ModuleMaps skipped when writing control block.
-  std::set<const FileEntry *> SkippedModuleMaps;
+  /// Sorted input files that did not affect the compilation.
+  std::vector<FileID> NonAffectingInputs;
+  /// Exclusive prefix sum of the lengths of preceding non-affecting inputs.
+  std::vector<SourceLocation::UIntTy> NonAffectingInputOffsetAdjustments;
+  /// Exclusive prefix sum of the count of preceding non-affecting inputs.
+  std::vector<unsigned> NonAffectingInputFileIDAdjustments;
+
+  /// Collects input files that didn't affect compilation of the current module,
+  /// and initializes data structures necessary for leaving those files out
+  /// during \c SourceManager serialization.
+  void collectNonAffectingInputFiles();
+
+  /// Prepares the specified \c FileID for serialization, accounting for any
+  /// non-affecting files.
+  FileID getAdjustedFileID(FileID FID) const;
+  /// Prepares the specified \c SourceLocation for serialization, accounting for
+  /// any non-affecting files.
+  SourceLocation getAdjustedLocation(SourceLocation Loc) const;
+  /// Prepares the specified \c SourceRange for serialization, accounting for
+  /// any non-affecting files.
+  SourceRange getAdjustedRange(SourceRange Range) const;
+  SourceLocation::UIntTy getAdjustedOffset(SourceLocation::UIntTy Offset) const;
+  SourceLocation::UIntTy getAdjustment(SourceLocation::UIntTy Offset) const;
 
   /// Retrieve or create a submodule ID for this module.
   unsigned getSubmoduleID(Module *Mod);
 
   /// Write the given subexpression to the bitstream.
   void WriteSubStmt(Stmt *S);
 
   void WriteBlockInfoBlock();
   void WriteControlBlock(Preprocessor &PP, ASTContext &Context,
                          StringRef isysroot);
 
   /// Write out the signature and diagnostic options, and return the signature.
   ASTFileSignature writeUnhashedControlBlock(Preprocessor &PP,
                                              ASTContext &Context);
 
   /// Calculate hash of the pcm content.
   static std::pair<ASTFileSignature, ASTFileSignature>
   createSignature(StringRef AllBytes, StringRef ASTBlockBytes);
 
-  void WriteInputFiles(SourceManager &SourceMgr, HeaderSearchOptions &HSOpts,
-                       std::set<const FileEntry *> &AffectingModuleMaps);
+  void WriteInputFiles(SourceManager &SourceMgr, HeaderSearchOptions &HSOpts);
   void WriteSourceManagerBlock(SourceManager &SourceMgr,
                                const Preprocessor &PP);
   void writeIncludedFiles(raw_ostream &Out, const Preprocessor &PP);
   void WritePreprocessor(const Preprocessor &PP, bool IsModule);
   void WriteHeaderSearch(const HeaderSearch &HS);
   void WritePreprocessorDetail(PreprocessingRecord &PPRec,
                                uint64_t MacroOffsetsBase);
   void WriteSubmodules(Module *WritingModule);
public:
 
   /// Emit a token.
   void AddToken(const Token &Tok, RecordDataImpl &Record);
 
   /// Emit a AlignPackInfo.
   void AddAlignPackInfo(const Sema::AlignPackInfo &Info,
                         RecordDataImpl &Record);
 
+  /// Emit a FileID.
+  void AddFileID(FileID, RecordDataImpl &Record);
+
   /// Emit a source location.
   void AddSourceLocation(SourceLocation Loc, RecordDataImpl &Record,
                          LocSeq *Seq = nullptr);
 
   /// Emit a source range.
   void AddSourceRange(SourceRange Range, RecordDataImpl &Record,
                       LocSeq *Seq = nullptr);
 

clang/lib/Serialization/ASTWriter.cpp

#include "clang/Serialization/TypeBitCodes.def"
   case Type::Builtin:
     llvm_unreachable("shouldn't be serializing a builtin type this way");
   }
   llvm_unreachable("bad type kind");
 }
 
 namespace {
 
-std::set<const FileEntry *> GetAllModuleMaps(const HeaderSearch &HS,
+std::set<const FileEntry *> GetAffectingModuleMaps(const HeaderSearch &HS,
                                                    Module *RootModule) {
   std::set<const FileEntry *> ModuleMaps{};
   std::set<const Module *> ProcessedModules;
   SmallVector<const Module *> ModulesToProcess{RootModule};
 
   SmallVector<const FileEntry *, 16> FilesByUID;
   HS.getFileMgr().GetUniqueIDMapping(FilesByUID);
 
   if (FilesByUID.size() > HS.header_file_size())
if (const FileEntry *MainFile = SM.getFileEntryForID(SM.getMainFileID())) {
     auto FileAbbrev = std::make_shared<BitCodeAbbrev>();
     FileAbbrev->Add(BitCodeAbbrevOp(ORIGINAL_FILE));
     FileAbbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 6)); // File ID
     FileAbbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob)); // File name
     unsigned FileAbbrevCode = Stream.EmitAbbrev(std::move(FileAbbrev));
 
     Record.clear();
     Record.push_back(ORIGINAL_FILE);
-    Record.push_back(SM.getMainFileID().getOpaqueValue());
+    AddFileID(SM.getMainFileID(), Record);
     EmitRecordWithPath(FileAbbrevCode, Record, MainFile->getName());
   }
 
   Record.clear();
-  Record.push_back(SM.getMainFileID().getOpaqueValue());
+  AddFileID(SM.getMainFileID(), Record);
   Stream.EmitRecord(ORIGINAL_FILE_ID, Record);
 
-  std::set<const FileEntry *> AffectingModuleMaps;
-  if (WritingModule) {
-    AffectingModuleMaps =
-        GetAllModuleMaps(PP.getHeaderSearchInfo(), WritingModule);
-  }
-
   WriteInputFiles(Context.SourceMgr,
-                  PP.getHeaderSearchInfo().getHeaderSearchOpts(),
-                  AffectingModuleMaps);
+                  PP.getHeaderSearchInfo().getHeaderSearchOpts());
   Stream.ExitBlock();
 }
 
 namespace  {
 
 /// An input file.
 struct InputFileEntry {
   const FileEntry *File;
   bool IsSystemFile;
   bool IsTransient;
   bool BufferOverridden;
   bool IsTopLevelModuleMap;
   uint32_t ContentHash[2];
 };
 
 } // namespace
 
-void ASTWriter::WriteInputFiles(
-    SourceManager &SourceMgr, HeaderSearchOptions &HSOpts,
-    std::set<const FileEntry *> &AffectingModuleMaps) {
+void ASTWriter::WriteInputFiles(SourceManager &SourceMgr,
+                                HeaderSearchOptions &HSOpts) {
   using namespace llvm;
 
   Stream.EnterSubblock(INPUT_FILES_BLOCK_ID, 4);
 
   // Create input-file abbreviation.
   auto IFAbbrev = std::make_shared<BitCodeAbbrev>();
   IFAbbrev->Add(BitCodeAbbrevOp(INPUT_FILE));
   IFAbbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 6)); // ID
for (unsigned I = 1, N = SourceMgr.local_sloc_entry_size(); I != N; ++I) {
     // We only care about file entries that were not overridden.
     if (!SLoc->isFile())
       continue;
     const SrcMgr::FileInfo &File = SLoc->getFile();
     const SrcMgr::ContentCache *Cache = &File.getContentCache();
     if (!Cache->OrigEntry)
       continue;
 
+    // Do not emit inputs that do not affect current module.
     if (isModuleMap(File.getFileCharacteristic()) &&
         !isSystem(File.getFileCharacteristic()) &&
-        !AffectingModuleMaps.empty() &&
-        AffectingModuleMaps.find(Cache->OrigEntry) ==
-            AffectingModuleMaps.end()) {
-      SkippedModuleMaps.insert(Cache->OrigEntry);
-      // Do not emit modulemaps that do not affect current module.
+        std::binary_search(NonAffectingInputs.begin(), NonAffectingInputs.end(),
+                           FileID::get(I)))
       continue;
-    }
 
     InputFileEntry Entry;
     Entry.File = Cache->OrigEntry;
     Entry.IsSystemFile = isSystem(File.getFileCharacteristic());
     Entry.IsTransient = Cache->IsTransient;
     Entry.BufferOverridden = Cache->BufferOverridden;
     Entry.IsTopLevelModuleMap = isModuleMap(File.getFileCharacteristic()) &&
                                 File.getIncludeLoc().isInvalid();
if (SLoc->isFile()) {
       } else
         Code = SM_SLOC_BUFFER_ENTRY;
     } else
       Code = SM_SLOC_EXPANSION_ENTRY;
     Record.clear();
     Record.push_back(Code);
 
     // Starting offset of this entry within this module, so skip the dummy.
-    Record.push_back(SLoc->getOffset() - 2);
+    Record.push_back(getAdjustedOffset(SLoc->getOffset()) - 2);
     if (SLoc->isFile()) {
       const SrcMgr::FileInfo &File = SLoc->getFile();
       const SrcMgr::ContentCache *Content = &File.getContentCache();
-      if (Content->OrigEntry && !SkippedModuleMaps.empty() &&
-          SkippedModuleMaps.find(Content->OrigEntry) !=
-              SkippedModuleMaps.end()) {
+      if (Content->OrigEntry &&
+          llvm::is_contained(NonAffectingInputs, FID)) {
         // Do not emit files that were not listed as inputs.
+        SLocEntryOffsets.pop_back();
         continue;
       }
       AddSourceLocation(File.getIncludeLoc(), Record);
       Record.push_back(File.getFileCharacteristic()); // FIXME: stable encoding
       Record.push_back(File.hasLineDirectives());
 
       bool EmitBlob = false;
       if (Content->OrigEntry) {
if (SLoc->isFile()) {
                             : Expansion.getExpansionLocEnd(),
                         Record, Seq);
       Record.push_back(Expansion.isExpansionTokenRange());
 
       // Compute the token length for this macro expansion.
       SourceLocation::UIntTy NextOffset = SourceMgr.getNextLocalOffset();
       if (I + 1 != N)
         NextOffset = SourceMgr.getLocalSLocEntry(I + 1).getOffset();
-      Record.push_back(NextOffset - SLoc->getOffset() - 1);
+      Record.push_back(getAdjustedOffset(NextOffset - SLoc->getOffset()) - 1);
       Stream.EmitRecordWithAbbrev(SLocExpansionAbbrv, Record);
     }
   }
 
   Stream.ExitBlock();
 
   if (SLocEntryOffsets.empty())
     return;
 
   // Write the source-location offsets table into the AST block. This
   // table is used for lazily loading source-location information.
   using namespace llvm;
 
   auto Abbrev = std::make_shared<BitCodeAbbrev>();
   Abbrev->Add(BitCodeAbbrevOp(SOURCE_LOCATION_OFFSETS));
   Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 16)); // # of slocs
   Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 16)); // total size
   Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 32)); // base offset
   Abbrev->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob)); // offsets
   unsigned SLocOffsetsAbbrev = Stream.EmitAbbrev(std::move(Abbrev));
   {
     RecordData::value_type Record[] = {
         SOURCE_LOCATION_OFFSETS, SLocEntryOffsets.size(),
-        SourceMgr.getNextLocalOffset() - 1 /* skip dummy */,
+        getAdjustedOffset(SourceMgr.getNextLocalOffset()) - 1 /* skip dummy */,
         SLocEntryOffsetsBase - SourceManagerBlockOffset};
     Stream.EmitRecordWithBlob(SLocOffsetsAbbrev, Record,
                               bytes(SLocEntryOffsets));
   }
   // Write the source location entry preloads array, telling the AST
   // reader which source locations entries it should load eagerly.
   Stream.EmitRecord(SOURCE_LOCATION_PRELOADS, PreloadSLocs);
 
for (PreprocessingRecord::iterator E = PPRec.local_begin(),
                                   EEnd = PPRec.local_end();
        E != EEnd;
        (void)++E, ++NumPreprocessingRecords, ++NextPreprocessorEntityID) {
     Record.clear();
 
     uint64_t Offset = Stream.GetCurrentBitNo() - MacroOffsetsBase;
     assert((Offset >> 32) == 0 && "Preprocessed entity offset too large");
     PreprocessedEntityOffsets.push_back(
-        PPEntityOffset((*E)->getSourceRange(), Offset));
+        PPEntityOffset(getAdjustedRange((*E)->getSourceRange()), Offset));
 
     if (auto *MD = dyn_cast<MacroDefinitionRecord>(*E)) {
       // Record this macro definition's ID.
       MacroDefinitions[MD] = NextPreprocessorEntityID;
 
       AddIdentifierRef(MD->getName(), Record);
       Stream.EmitRecord(PPD_MACRO_DEFINITION, Record);
       continue;
void ASTWriter::WritePragmaDiagnosticMappings(const DiagnosticsEngine &Diag,
   // Emit the state transitions.
   unsigned NumLocations = 0;
   for (auto &FileIDAndFile : Diag.DiagStatesByLoc.Files) {
     if (!FileIDAndFile.first.isValid() ||
         !FileIDAndFile.second.HasLocalTransitions)
       continue;
     ++NumLocations;
 
+    // TODO: Only write the FileID.
     SourceLocation Loc = Diag.SourceMgr->getComposedLoc(FileIDAndFile.first, 0);
     assert(!Loc.isInvalid() && "start loc for valid FileID is invalid");
     AddSourceLocation(Loc, Record);
 
     Record.push_back(FileIDAndFile.second.StateTransitions.size());
     for (auto &StatePoint : FileIDAndFile.second.StateTransitions) {
-      Record.push_back(StatePoint.Offset);
+      Record.push_back(getAdjustedOffset(StatePoint.Offset));
       AddDiagState(StatePoint.State, false);
     }
   }
 
   // Backpatch the number of locations.
   Record[NumLocationsIdx] = NumLocations;
 
   // Emit CurDiagStateLoc.  Do it last in order to match source order.

 static void AddLazyVectorDecls(ASTWriter &Writer, Vector &Vec,
                                ASTWriter::RecordData &Record) {
   for (typename Vector::iterator I = Vec.begin(nullptr, true), E = Vec.end();
        I != E; ++I) {
     Writer.AddDeclRef(*I, Record);
   }
 }
 
+void ASTWriter::collectNonAffectingInputFiles() {
+  if (!WritingModule)
+    return;
+
+  std::set<const FileEntry *> AffectingModuleMaps =
+      GetAffectingModuleMaps(PP->getHeaderSearchInfo(), WritingModule);
+
+  SourceManager &SrcMgr = PP->getSourceManager();
+
+  for (unsigned I = 1, N = SrcMgr.local_sloc_entry_size(); I != N; ++I) {
+    // Get this source location entry.
+    const SrcMgr::SLocEntry *SLoc = &SrcMgr.getLocalSLocEntry(I);
+    assert(&SrcMgr.getSLocEntry(FileID::get(I)) == SLoc);
+
+    if (!SLoc->isFile())
+      continue;
+    const SrcMgr::FileInfo &File = SLoc->getFile();
+    const SrcMgr::ContentCache *Cache = &File.getContentCache();
+    if (!Cache->OrigEntry)
+      continue;
+
+    if (isModuleMap(File.getFileCharacteristic()) &&
+        !isSystem(File.getFileCharacteristic()) &&
+        !AffectingModuleMaps.empty() &&
+        AffectingModuleMaps.find(Cache->OrigEntry) ==
+            AffectingModuleMaps.end()) {
+      NonAffectingInputs.push_back(FileID::get(I));
+    }
+  }
+
+  llvm::sort(NonAffectingInputs);
+
+  unsigned Size = 0;
+  unsigned Count = 0;
+  NonAffectingInputOffsetAdjustments.reserve(NonAffectingInputs.size() + 1);
+  NonAffectingInputFileIDAdjustments.reserve(NonAffectingInputs.size() + 1);
+  NonAffectingInputOffsetAdjustments.push_back(Size);
+  NonAffectingInputFileIDAdjustments.push_back(Count);
+  for (FileID FID : NonAffectingInputs) {
+    Size += SrcMgr.getFileIDSize(FID);
+    Count += 1;
+    NonAffectingInputOffsetAdjustments.push_back(Size);
+    NonAffectingInputFileIDAdjustments.push_back(Count);
+  }
+}
+
 ASTFileSignature ASTWriter::WriteASTCore(Sema &SemaRef, StringRef isysroot,
                                          Module *WritingModule) {
   using namespace llvm;
 
   bool isModule = WritingModule != nullptr;
 
   // Make sure that the AST reader knows to finalize itself.
   if (Chain)
     Chain->finalizeForWriting();
 
   ASTContext &Context = SemaRef.Context;
   Preprocessor &PP = SemaRef.PP;
 
+  collectNonAffectingInputFiles();
+
   // Set up predefined declaration IDs.
   auto RegisterPredefDecl = [&] (Decl *D, PredefinedDeclIDs ID) {
     if (D) {
       assert(D->isCanonicalDecl() && "predefined decl is not canonical");
       DeclIDs[D] = ID;
     }
   };
   RegisterPredefDecl(Context.getTranslationUnitDecl(),

 }
 
 void ASTWriter::AddAlignPackInfo(const Sema::AlignPackInfo &Info,
                                  RecordDataImpl &Record) {
   uint32_t Raw = Sema::AlignPackInfo::getRawEncoding(Info);
   Record.push_back(Raw);
 }
 
+void ASTWriter::AddFileID(FileID FID, RecordDataImpl &Record) {
+  Record.push_back(getAdjustedFileID(FID).getOpaqueValue());
+}
+
+FileID ASTWriter::getAdjustedFileID(FileID FID) const {
+  if (FID.isInvalid() || PP->getSourceManager().isLoadedFileID(FID) ||
+      NonAffectingInputs.empty())
+    return FID;
+  auto It = llvm::lower_bound(NonAffectingInputs, FID);
+  unsigned Idx = std::distance(NonAffectingInputs.begin(), It);
+  unsigned Offset = NonAffectingInputFileIDAdjustments[Idx];
+  return FileID::get(FID.getOpaqueValue() - Offset);
+}
+
+SourceLocation::UIntTy
+ASTWriter::getAdjustment(SourceLocation::UIntTy Offset) const {
+  if (PP->getSourceManager().isLoadedOffset(Offset) ||
+      NonAffectingInputs.empty())
+    return 0;
+  auto It = (Offset == PP->getSourceManager().getNextLocalOffset())
+                ? NonAffectingInputs.end()
+                : llvm::lower_bound(NonAffectingInputs,
+                                    PP->getSourceManager().getFileID(Offset));
+  unsigned Idx = std::distance(NonAffectingInputs.begin(), It);
+  return NonAffectingInputOffsetAdjustments[Idx];
+}
+
+SourceLocation::UIntTy
+ASTWriter::getAdjustedOffset(SourceLocation::UIntTy Offset) const {
+  return Offset - getAdjustment(Offset);
+}
+
+SourceLocation ASTWriter::getAdjustedLocation(SourceLocation Loc) const {
+  if (Loc.isInvalid())
+    return Loc;
+  return Loc.getLocWithOffset(-getAdjustment(Loc.getOffset()));
+}
+
+SourceRange ASTWriter::getAdjustedRange(SourceRange Range) const {
+  return SourceRange(getAdjustedLocation(Range.getBegin()),
+                     getAdjustedLocation(Range.getEnd()));
+}
+
 void ASTWriter::AddSourceLocation(SourceLocation Loc, RecordDataImpl &Record,
                                   SourceLocationSequence *Seq) {
+  Loc = getAdjustedLocation(Loc);
   Record.push_back(SourceLocationEncoding::encode(Loc, Seq));
 }
 
 void ASTWriter::AddSourceRange(SourceRange Range, RecordDataImpl &Record,
                                SourceLocationSequence *Seq) {
   AddSourceLocation(Range.getBegin(), Record, Seq);
   AddSourceLocation(Range.getEnd(), Record, Seq);
 }

clang/lib/Serialization/ASTWriterDecl.cpp

void ASTWriter::WriteDecl(ASTContext &Context, Decl *D) {
 
   // Emit this declaration to the bitstream.
   uint64_t Offset = W.Emit(D);
 
   // Record the offset for this declaration
   SourceLocation Loc = D->getLocation();
   unsigned Index = ID - FirstDeclID;
   if (DeclOffsets.size() == Index)
-    DeclOffsets.emplace_back(Loc, Offset, DeclTypesBlockStartOffset);
+    DeclOffsets.emplace_back(getAdjustedLocation(Loc), Offset,
+                             DeclTypesBlockStartOffset);
   else if (DeclOffsets.size() < Index) {
     // FIXME: Can/should this happen?
     DeclOffsets.resize(Index+1);
-    DeclOffsets[Index].setLocation(Loc);
+    DeclOffsets[Index].setLocation(getAdjustedLocation(Loc));
     DeclOffsets[Index].setBitOffset(Offset, DeclTypesBlockStartOffset);
   } else {
     llvm_unreachable("declarations should be emitted in ID order");
   }
 
   SourceManager &SM = Context.getSourceManager();
   if (Loc.isValid() && SM.isLocalSourceLocation(Loc))
     associateDeclWithFile(D, ID);

clang/test/Modules/non-affecting-module-maps-source-locations.m

+// This patch tests that non-affecting files are serialized in such way that
+// does not break subsequent source locations (e.g. in serialized pragma
+// diagnostic mappings).
+
+// RUN: rm -rf %t
+// RUN: split-file %s %t
+
+//--- tu.m
+#include "first.h"
+
+//--- first/module.modulemap
+module first { header "first.h" }
+//--- first/first.h
+@import second;
+#pragma clang diagnostic ignored "-Wunreachable-code"
+
+//--- second/extra/module.modulemap
+module second_extra {}
+//--- second/module.modulemap
+module second { module sub { header "second_sub.h" } }
+//--- second/second_sub.h
+@import third;
+
+//--- third/module.modulemap
+module third { header "third.h" }
+//--- third/third.h
+#if __has_feature(nullability)
+#endif
+#if __has_feature(nullability)
+#endif
+
+// RUN: %clang_cc1 -I %t/first -I %t/second -I %t/third -fmodules -fimplicit-module-maps -fmodules-cache-path=%t/cache %t/tu.m -o %t/tu.o