[IR] Support importing modules with invalid data layouts.

[jsilvanus]

Summary

Use the existing mechanism to change the data layout using callbacks.

Before this patch, we had a callback type DataLayoutCallbackTy that receives
a single StringRef specifying the target triple, and optionally returns
the data layout string to be used. Module loaders (both IR and BC) then
apply the callback to potentially override the module's data layout,
after first having imported and parsed the data layout from the file.

We can't do the same to fix invalid data layouts, because the import will already
fail, before the callback has a chance to fix it.
Instead, module loaders now tentatively parse the data layout into a string,
wait until the target triple has been parsed, apply the override callback
to the imported string and only then parse the tentative string as a data layout.

Moreover, add the old data layout string S as second argument to the callback,
in addition to the already existing target triple argument.
S is either the default data layout string in case none is specified, or the data
layout string specified in the module, possibly after auto-upgrades (for the BitcodeReader).
This allows callbacks to inspect the old data layout string,
and fix it instead of setting a fixed data layout.

Also allow to pass data layout override callbacks to lazy bitcode module
loader functions.

[jsilvanus]

Add lit test case.

[jsilvanus]

For context, this prepares requiring naturally aligned i8 as discussed in https://discourse.llvm.org/t/status-of-overaligned-i8/66913/

This will render some existing modules (e.g. DXIL) invalid, but some users still need to import such modules.

[sebastian-ne]

The changes look good from my side. Can you add a unittest that converts an invalid DataLayout string to something legal?

llvm/include/llvm/Bitcode/BitcodeReader.h

38–39

Instead of “old” data layout string, maybe something like “data layout string from the input”?

[arsenm]

Is DataLayout compatibility still determined with string equality and not semantic equality?

[jsilvanus]

In D140985#4028591, @arsenm wrote:

Is DataLayout compatibility still determined with string equality and not semantic equality?

What exactly do you mean?
The callback mechanism changed in this patch was and is string-based, and the parsing of the string into a DL happens after the callback.

[jsilvanus]

Address review feedback

 - improve comment on callback args
 - add callback unit test

[nikic]

llvm/lib/AsmParser/LLParser.cpp

384

It would be better to go through DataLayout::parse here as well, and report without location. setDataLayout will cause a fatal error for invalid data layouts, I believe.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

4210

High level question: Why is the existing UpgradeDataLayoutString mechanigm not sufficient?

[jsilvanus]

Address review feedback: Handle invalid data layouts gracefully

[jsilvanus]

Thanks for reviewing @nikic. Comments inline.

llvm/lib/AsmParser/LLParser.cpp

384

That's correct. I've changed it to always use DataLayout::parse and return an error instead of running into a fatal error.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

4210

In the case this is intended for (requiring natural alignment for i8), it is true that there is a straightforward fix of a broken data layout string: Just replace i8s alignment with natural alignment.

But depending on how i8 is used in the module, this can break semantics, requiring nontrivial fixup (e.g. adding manual padding). This does not seem to be necessary for DXIL though.

I wanted to avoid auto-upgrade (silently?) breaking a module, and adding the manual padding as part of auto-upgrade seems excessive, at least given that we're not aware of any cases actually requiring it.

On the other hand, letting the user fix the layout string avoids that issue of auto-upgrade breaking a module. If the user is sure that i8 alignment does not matter, the callback provides the option to fix the DL string.

[arsenm]

In D140985#4028694, @jsilvanus wrote:

In D140985#4028591, @arsenm wrote:

Is DataLayout compatibility still determined with string equality and not semantic equality?

What exactly do you mean?
The callback mechanism changed in this patch was and is string-based, and the parsing of the string into a DL happens after the callback.

I mean llvm-link complains if you link modules with different datalayouts, but it does dumb string comparison so if you reorder fields (or rely on defaulted fields) it prints a warning

[nikic]

LGTM

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

4210

I see. If this is intended for manual user intervention, then that makes sense.

[cmtice]

Just a heads up: It looks like this change is causing many of our internal tests to fail. I'm working on trying to create a reproducer for you.

llvm/include/llvm/AsmParser/LLParser.h

namespace llvm {
   public:
     LLParser(StringRef F, SourceMgr &SM, SMDiagnostic &Err, Module *M,
              ModuleSummaryIndex *Index, LLVMContext &Context,
              SlotMapping *Slots = nullptr)
         : Context(Context), OPLex(F, SM, Err, Context),
           Lex(F, SM, Err, Context), M(M), Index(Index), Slots(Slots),
           BlockAddressPFS(nullptr) {}
     bool Run(
-        bool UpgradeDebugInfo, DataLayoutCallbackTy DataLayoutCallback =
-                                   [](StringRef) { return std::nullopt; });
+        bool UpgradeDebugInfo,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     bool parseStandaloneConstantValue(Constant *&C, const SlotMapping *Slots);
 
     bool parseTypeAtBeginning(Type *&Ty, unsigned &Read,
                               const SlotMapping *Slots);
 
     LLVMContext &getContext() { return Context; }
 
bool parseIndexList(SmallVectorImpl<unsigned> &Indices) {
         return tokError("expected index");
       return false;
     }
 
     // Top-Level Entities
     bool parseTopLevelEntities();
     bool validateEndOfModule(bool UpgradeDebugInfo);
     bool validateEndOfIndex();
-    bool parseTargetDefinitions();
-    bool parseTargetDefinition();
+    bool parseTargetDefinitions(DataLayoutCallbackTy DataLayoutCallback);
+    bool parseTargetDefinition(std::string &TentativeDLStr, LocTy &DLStrLoc);
     bool parseModuleAsm();
     bool parseSourceFileName();
     bool parseUnnamedType();
     bool parseNamedType();
     bool parseDeclare();
     bool parseDefine();
 
     bool parseGlobalType(bool &IsConstant);

llvm/include/llvm/AsmParser/Parser.h

 class LLVMContext;
 class MemoryBufferRef;
 class Module;
 class ModuleSummaryIndex;
 struct SlotMapping;
 class SMDiagnostic;
 class Type;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// This function is a main interface to the LLVM Assembly Parser. It parses
 /// an ASCII file that (presumably) contains LLVM Assembly code. It returns a
 /// Module (intermediate representation) with the corresponding features. Note
 /// that this does not verify that the generated Module is valid, so you should
 /// run the verifier after parsing the file to check that it is okay.
 /// Parse LLVM Assembly from a file

 /// \param Err Error result info.
 /// \param Context Context in which to allocate globals info.
 /// \param Slots The optional slot mapping that will be initialized during
 ///              parsing.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 ParsedModuleAndIndex parseAssemblyFileWithIndex(
     StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
 /// Only for use in llvm-as for testing; this does not produce a valid module.
 ParsedModuleAndIndex parseAssemblyFileWithIndexNoUpgradeDebugInfo(
     StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots, DataLayoutCallbackTy DataLayoutCallback);
 

 /// \param F The MemoryBuffer containing assembly
 /// \param Err Error result info.
 /// \param Slots The optional slot mapping that will be initialized during
 ///              parsing.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 std::unique_ptr<Module> parseAssembly(
     MemoryBufferRef F, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
 /// Parse LLVM Assembly including the summary index from a MemoryBuffer.
 ///
 /// \param F The MemoryBuffer containing assembly with summary
 /// \param Err Error result info.
 /// \param Slots The optional slot mapping that will be initialized during

 /// \param Err Error result info.
 /// \param Slots The optional slot mapping that will be initialized during
 ///              parsing.
 /// \return true on error.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 bool parseAssemblyInto(
     MemoryBufferRef F, Module *M, ModuleSummaryIndex *Index, SMDiagnostic &Err,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
 /// Parse a type and a constant value in the given string.
 ///
 /// The constant value can be any LLVM constant, including a constant
 /// expression.
 ///

llvm/include/llvm/Bitcode/BitcodeReader.h

 #include <vector>
 namespace llvm {
 
 class LLVMContext;
 class Module;
 class MemoryBuffer;
 class ModuleSummaryIndex;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+// Callback to override the data layout string of an imported bitcode module.
+// The first argument is the target triple, the second argument the data layout
+// string from the input, or a default string. It will be used if the callback
+// returns std::nullopt.
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
   // These functions are for converting Expected/Error values to
   // ErrorOr/std::error_code for compatibility with legacy clients. FIXME:
   // Remove these functions once no longer needed by the C and libLTO APIs.
 
   std::error_code errorToErrorCodeAndEmitErrors(LLVMContext &Ctx, Error Err);
 
public:
     StringRef getStrtab() const { return Strtab; }
 
     StringRef getModuleIdentifier() const { return ModuleIdentifier; }
 
     /// Read the bitcode module and prepare for lazy deserialization of function
     /// bodies. If ShouldLazyLoadMetadata is true, lazily load metadata as well.
     /// If IsImporting is true, this module is being parsed for ThinLTO
     /// importing into another module.
-    Expected<std::unique_ptr<Module>> getLazyModule(LLVMContext &Context,
-                                                    bool ShouldLazyLoadMetadata,
-                                                    bool IsImporting);
+    Expected<std::unique_ptr<Module>> getLazyModule(
+        LLVMContext &Context, bool ShouldLazyLoadMetadata, bool IsImporting,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     /// Read the entire bitcode module and return it.
     Expected<std::unique_ptr<Module>> parseModule(
-        LLVMContext &Context, DataLayoutCallbackTy DataLayoutCallback =
-                                  [](StringRef) { return std::nullopt; });
+        LLVMContext &Context,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     /// Returns information about the module to be used for LTO: whether to
     /// compile with ThinLTO, and whether it has a summary.
     Expected<BitcodeLTOInfo> getLTOInfo();
 
     /// Parse the specified bitcode buffer, returning the module summary index.
     Expected<std::unique_ptr<ModuleSummaryIndex>> getSummary();
 
typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
   /// Returns a list of modules in the specified bitcode buffer.
   Expected<std::vector<BitcodeModule>>
   getBitcodeModuleList(MemoryBufferRef Buffer);
 
   /// Read the header of the specified bitcode buffer and prepare for lazy
   /// deserialization of function bodies. If ShouldLazyLoadMetadata is true,
   /// lazily load metadata as well. If IsImporting is true, this module is
   /// being parsed for ThinLTO importing into another module.
-  Expected<std::unique_ptr<Module>>
-  getLazyBitcodeModule(MemoryBufferRef Buffer, LLVMContext &Context,
-                       bool ShouldLazyLoadMetadata = false,
-                       bool IsImporting = false);
+  Expected<std::unique_ptr<Module>> getLazyBitcodeModule(
+      MemoryBufferRef Buffer, LLVMContext &Context,
+      bool ShouldLazyLoadMetadata = false, bool IsImporting = false,
+      DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+        return std::nullopt;
+      });
 
   /// Like getLazyBitcodeModule, except that the module takes ownership of
   /// the memory buffer if successful. If successful, this moves Buffer. On
   /// error, this *does not* move Buffer. If IsImporting is true, this module is
   /// being parsed for ThinLTO importing into another module.
   Expected<std::unique_ptr<Module>> getOwningLazyBitcodeModule(
       std::unique_ptr<MemoryBuffer> &&Buffer, LLVMContext &Context,
       bool ShouldLazyLoadMetadata = false, bool IsImporting = false);
typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
   /// Read the header of the specified bitcode buffer and extract just the
   /// producer string information. If successful, this returns a string. On
   /// error, this returns "".
   Expected<std::string> getBitcodeProducerString(MemoryBufferRef Buffer);
 
   /// Read the specified bitcode file, returning the module.
   Expected<std::unique_ptr<Module>> parseBitcodeFile(
       MemoryBufferRef Buffer, LLVMContext &Context,
-      DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+      DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
         return std::nullopt;
       });
 
   /// Returns LTO information for the specified bitcode file.
   Expected<BitcodeLTOInfo> getBitcodeLTOInfo(MemoryBufferRef Buffer);
 
   /// Parse the specified bitcode buffer, returning the module summary index.
   Expected<std::unique_ptr<ModuleSummaryIndex>>

llvm/include/llvm/CodeGen/MIRParser/MIRParser.h

 class LLVMContext;
 class MemoryBuffer;
 class Module;
 class MIRParserImpl;
 class MachineModuleInfo;
 class SMDiagnostic;
 class StringRef;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// This class initializes machine functions by applying the state loaded from
 /// a MIR file.
 class MIRParser {
   std::unique_ptr<MIRParserImpl> Impl;
 
 public:
   MIRParser(std::unique_ptr<MIRParserImpl> Impl);
   MIRParser(const MIRParser &) = delete;
   ~MIRParser();
 
   /// Parses the optional LLVM IR module in the MIR file.
   ///
   /// A new, empty module is created if the LLVM IR isn't present.
   /// \returns nullptr if a parsing error occurred.
   std::unique_ptr<Module>
-  parseIRModule(DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
-    return std::nullopt;
-  });
+  parseIRModule(DataLayoutCallbackTy DataLayoutCallback =
+                    [](StringRef, StringRef) { return std::nullopt; });
 
   /// Parses MachineFunctions in the MIR file and add them to the given
   /// MachineModuleInfo \p MMI.
   ///
   /// \returns true if an error occurred.
   bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
 };
 

llvm/include/llvm/IRReader/IRReader.h

 namespace llvm {
 
 class MemoryBuffer;
 class MemoryBufferRef;
 class Module;
 class SMDiagnostic;
 class LLVMContext;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// If the given MemoryBuffer holds a bitcode image, return a Module
 /// for it which does lazy deserialization of function bodies.  Otherwise,
 /// attempt to parse it as LLVM Assembly and return a fully populated
 /// Module. The ShouldLazyLoadMetadata flag is passed down to the bitcode
 /// reader to optionally enable lazy metadata loading. This takes ownership
 /// of \p Buffer.
getLazyIRFileModule(StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
                     bool ShouldLazyLoadMetadata = false);
 
 /// If the given MemoryBuffer holds a bitcode image, return a Module
 /// for it.  Otherwise, attempt to parse it as LLVM Assembly and return
 /// a Module for it.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 std::unique_ptr<Module> parseIR(
     MemoryBufferRef Buffer, SMDiagnostic &Err, LLVMContext &Context,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
 /// If the given file holds a bitcode image, return a Module for it.
 /// Otherwise, attempt to parse it as LLVM Assembly and return a Module
 /// for it.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 std::unique_ptr<Module> parseIRFile(
     StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 }
 
 #endif

llvm/lib/AsmParser/LLParser.cpp

bool LLParser::Run(bool UpgradeDebugInfo,
   Lex.Lex();
 
   if (Context.shouldDiscardValueNames())
     return error(
         Lex.getLoc(),
         "Can't read textual IR with a Context that discards named Values");
 
   if (M) {
-    if (parseTargetDefinitions())
+    if (parseTargetDefinitions(DataLayoutCallback))
       return true;
-
-    if (auto LayoutOverride = DataLayoutCallback(M->getTargetTriple()))
-      M->setDataLayout(*LayoutOverride);
   }
 
   return parseTopLevelEntities() || validateEndOfModule(UpgradeDebugInfo) ||
          validateEndOfIndex();
 }
 
 bool LLParser::parseStandaloneConstantValue(Constant *&C,
                                             const SlotMapping *Slots) {
bool LLParser::validateEndOfIndex() {
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Top-Level Entities
 //===----------------------------------------------------------------------===//
 
-bool LLParser::parseTargetDefinitions() {
-  while (true) {
+bool LLParser::parseTargetDefinitions(DataLayoutCallbackTy DataLayoutCallback) {
+  // Delay parsing of the data layout string until the target triple is known.
+  // Then, pass both the the target triple and the tentative data layout string
+  // to DataLayoutCallback, allowing to override the DL string.
+  // This enables importing modules with invalid DL strings.
+  std::string TentativeDLStr = M->getDataLayoutStr();
+  LocTy DLStrLoc;
+
+  bool Done = false;
+  while (!Done) {
     switch (Lex.getKind()) {
     case lltok::kw_target:
-      if (parseTargetDefinition())
+      if (parseTargetDefinition(TentativeDLStr, DLStrLoc))
         return true;
       break;
     case lltok::kw_source_filename:
       if (parseSourceFileName())
         return true;
       break;
     default:
-      return false;
-    }
+      Done = true;
+    }
+  // Run the override callback to potentially change the data layout string, and
+  // parse the data layout string.
+  if (auto LayoutOverride =
+          DataLayoutCallback(M->getTargetTriple(), TentativeDLStr)) {
+    TentativeDLStr = *LayoutOverride;
+    DLStrLoc = {};
   }
+  Expected<DataLayout> MaybeDL = DataLayout::parse(TentativeDLStr);
+  if (!MaybeDL)
+    return error(DLStrLoc, toString(MaybeDL.takeError()));
+  M->setDataLayout(MaybeDL.get());
+  return false;
 }
 
 bool LLParser::parseTopLevelEntities() {
   // If there is no Module, then parse just the summary index entries.
   if (!M) {
     while (true) {
       switch (Lex.getKind()) {
       case lltok::Eof:
bool LLParser::parseModuleAsm() {
 
   M->appendModuleInlineAsm(AsmStr);
   return false;
 }
 
 /// toplevelentity
 ///   ::= 'target' 'triple' '=' STRINGCONSTANT
 ///   ::= 'target' 'datalayout' '=' STRINGCONSTANT
-bool LLParser::parseTargetDefinition() {
+bool LLParser::parseTargetDefinition(std::string &TentativeDLStr,
+                                     LocTy &DLStrLoc) {
   assert(Lex.getKind() == lltok::kw_target);
   std::string Str;
   switch (Lex.Lex()) {
   default:
     return tokError("unknown target property");
   case lltok::kw_triple:
     Lex.Lex();
     if (parseToken(lltok::equal, "expected '=' after target triple") ||
         parseStringConstant(Str))
       return true;
     M->setTargetTriple(Str);
     return false;
   case lltok::kw_datalayout:
     Lex.Lex();
     if (parseToken(lltok::equal, "expected '=' after target datalayout"))
       return true;
-    LocTy Loc = Lex.getLoc();
-    if (parseStringConstant(Str))
+    DLStrLoc = Lex.getLoc();
+    if (parseStringConstant(TentativeDLStr))
       return true;
-    Expected<DataLayout> MaybeDL = DataLayout::parse(Str);
-    if (!MaybeDL)
-      return error(Loc, toString(MaybeDL.takeError()));
-    M->setDataLayout(MaybeDL.get());
     return false;
   }
 }
 
 /// toplevelentity
 ///   ::= 'source_filename' '=' STRINGCONSTANT
 bool LLParser::parseSourceFileName() {
   assert(Lex.getKind() == lltok::kw_source_filename);

llvm/lib/AsmParser/Parser.cpp

parseAssemblyWithIndex(MemoryBufferRef F, SMDiagnostic &Err,
 
   return {std::move(M), std::move(Index)};
 }
 
 ParsedModuleAndIndex llvm::parseAssemblyWithIndex(MemoryBufferRef F,
                                                   SMDiagnostic &Err,
                                                   LLVMContext &Context,
                                                   SlotMapping *Slots) {
-  return ::parseAssemblyWithIndex(F, Err, Context, Slots,
-                                  /*UpgradeDebugInfo*/ true,
-                                  [](StringRef) { return std::nullopt; });
+  return ::parseAssemblyWithIndex(
+      F, Err, Context, Slots,
+      [](StringRef, StringRef) { return std::nullopt; });
 }
 
 static ParsedModuleAndIndex
 parseAssemblyFileWithIndex(StringRef Filename, SMDiagnostic &Err,
                            LLVMContext &Context, SlotMapping *Slots,
                            bool UpgradeDebugInfo,
                            DataLayoutCallbackTy DataLayoutCallback) {
   ErrorOr<std::unique_ptr<MemoryBuffer>> FileOrErr =
static bool parseSummaryIndexAssemblyInto(MemoryBufferRef F,
   SourceMgr SM;
   std::unique_ptr<MemoryBuffer> Buf = MemoryBuffer::getMemBuffer(F);
   SM.AddNewSourceBuffer(std::move(Buf), SMLoc());
 
   // The parser holds a reference to a context that is unused when parsing the
   // index, but we need to initialize it.
   LLVMContext unusedContext;
   return LLParser(F.getBuffer(), SM, Err, nullptr, &Index, unusedContext)
-      .Run(true, [](StringRef) { return std::nullopt; });
+      .Run(true, [](StringRef, StringRef) { return std::nullopt; });
 }
 
 std::unique_ptr<ModuleSummaryIndex>
 llvm::parseSummaryIndexAssembly(MemoryBufferRef F, SMDiagnostic &Err) {
   std::unique_ptr<ModuleSummaryIndex> Index =
       std::make_unique<ModuleSummaryIndex>(/*HaveGVs=*/false);
 
   if (parseSummaryIndexAssemblyInto(F, *Index, Err))

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

private:
 
   /// Converts alignment exponent (i.e. power of two (or zero)) to the
   /// corresponding alignment to use. If alignment is too large, returns
   /// a corresponding error code.
   Error parseAlignmentValue(uint64_t Exponent, MaybeAlign &Alignment);
   Error parseAttrKind(uint64_t Code, Attribute::AttrKind *Kind);
   Error parseModule(
       uint64_t ResumeBit, bool ShouldLazyLoadMetadata = false,
-      DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+      DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
         return std::nullopt;
       });
 
   Error parseComdatRecord(ArrayRef<uint64_t> Record);
   Error parseGlobalVarRecord(ArrayRef<uint64_t> Record);
   Error parseFunctionRecord(ArrayRef<uint64_t> Record);
   Error parseGlobalIndirectSymbolRecord(unsigned BitCode,
                                         ArrayRef<uint64_t> Record);
Error BitcodeReader::parseModule(uint64_t ResumeBit,
   } else if (Error Err = Stream.EnterSubBlock(bitc::MODULE_BLOCK_ID))
     return Err;
 
   SmallVector<uint64_t, 64> Record;
 
   // Parts of bitcode parsing depend on the datalayout.  Make sure we
   // finalize the datalayout before we run any of that code.
   bool ResolvedDataLayout = false;
-  auto ResolveDataLayout = [&] {
+  // In order to support importing modules with illegal data layout strings,
+  // delay parsing the data layout string until after upgrades and overrides
+  // have been applied, allowing to fix illegal data layout strings.
+  // Initialize to the current module's layout string in case none is specified.
+  std::string TentativeDataLayoutStr = TheModule->getDataLayoutStr();
+
+  auto ResolveDataLayout = [&]() -> Error {
     if (ResolvedDataLayout)
-      return;
+      return Error::success();
 
-    // datalayout and triple can't be parsed after this point.
+    // Datalayout and triple can't be parsed after this point.
     ResolvedDataLayout = true;
 
-    // Upgrade data layout string.
-    std::string DL = llvm::UpgradeDataLayoutString(
-        TheModule->getDataLayoutStr(), TheModule->getTargetTriple());
-    TheModule->setDataLayout(DL);
-
-    if (auto LayoutOverride =
-            DataLayoutCallback(TheModule->getTargetTriple()))
-      TheModule->setDataLayout(*LayoutOverride);
+    // Auto-upgrade the layout string
+    TentativeDataLayoutStr = llvm::UpgradeDataLayoutString(
+        TentativeDataLayoutStr, TheModule->getTargetTriple());
+
+    // Apply override
+    if (auto LayoutOverride = DataLayoutCallback(TheModule->getTargetTriple(),
+                                                 TentativeDataLayoutStr))
+      TentativeDataLayoutStr = *LayoutOverride;
+
+    // Now the layout string is finalized in TentativeDataLayoutStr. Parse it.
+    Expected<DataLayout> MaybeDL = DataLayout::parse(TentativeDataLayoutStr);
+    if (!MaybeDL)
+      return MaybeDL.takeError();
+
+    TheModule->setDataLayout(MaybeDL.get());
+    return Error::success();
   };
 
   // Read all the records for this module.
   while (true) {
     Expected<llvm::BitstreamEntry> MaybeEntry = Stream.advance();
     if (!MaybeEntry)
       return MaybeEntry.takeError();
     llvm::BitstreamEntry Entry = MaybeEntry.get();
 
     switch (Entry.Kind) {
     case BitstreamEntry::Error:
       return error("Malformed block");
     case BitstreamEntry::EndBlock:
-      ResolveDataLayout();
+      if (Error Err = ResolveDataLayout())
+        return Err;
       return globalCleanup();
 
     case BitstreamEntry::SubBlock:
       switch (Entry.ID) {
       default:  // Skip unknown content.
         if (Error Err = Stream.SkipBlock())
           return Err;
         break;
case BitstreamEntry::SubBlock:
         if (Error Err = MDLoader->parseModuleMetadata())
           return Err;
         break;
       case bitc::METADATA_KIND_BLOCK_ID:
         if (Error Err = MDLoader->parseMetadataKinds())
           return Err;
         break;
       case bitc::FUNCTION_BLOCK_ID:
-        ResolveDataLayout();
+        if (Error Err = ResolveDataLayout())
+          return Err;
 
         // If this is the first function body we've seen, reverse the
         // FunctionsWithBodies list.
         if (!SeenFirstFunctionBody) {
           std::reverse(FunctionsWithBodies.begin(), FunctionsWithBodies.end());
           if (Error Err = globalCleanup())
             return Err;
           SeenFirstFunctionBody = true;
case bitc::MODULE_CODE_TRIPLE: {  // TRIPLE: [strchr x N]
       if (convertToString(Record, 0, S))
         return error("Invalid record");
       TheModule->setTargetTriple(S);
       break;
     }
     case bitc::MODULE_CODE_DATALAYOUT: {  // DATALAYOUT: [strchr x N]
       if (ResolvedDataLayout)
         return error("datalayout too late in module");
-      std::string S;
-      if (convertToString(Record, 0, S))
+      if (convertToString(Record, 0, TentativeDataLayoutStr))
         return error("Invalid record");
-      Expected<DataLayout> MaybeDL = DataLayout::parse(S);
-      if (!MaybeDL)
-        return MaybeDL.takeError();
-      TheModule->setDataLayout(MaybeDL.get());
       break;
     }
     case bitc::MODULE_CODE_ASM: {  // ASM: [strchr x N]
       std::string S;
       if (convertToString(Record, 0, S))
         return error("Invalid record");
       TheModule->setModuleInlineAsm(S);
       break;
while (true) {
     // `ThinLinkBitcodeWriter::writeSimplifiedModuleInfo` for the format of each
     // record
     // (https://github.com/llvm/llvm-project/blob/b6a93967d9c11e79802b5e75cec1584d6c8aa472/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp#L4714)
     case bitc::MODULE_CODE_GLOBALVAR:
       if (Error Err = parseGlobalVarRecord(Record))
         return Err;
       break;
     case bitc::MODULE_CODE_FUNCTION:
-      ResolveDataLayout();
+      if (Error Err = ResolveDataLayout())
+        return Err;
       if (Error Err = parseFunctionRecord(Record))
         return Err;
       break;
     case bitc::MODULE_CODE_IFUNC:
     case bitc::MODULE_CODE_ALIAS:
     case bitc::MODULE_CODE_ALIAS_OLD:
       if (Error Err = parseGlobalIndirectSymbolRecord(BitCode, Record))
         return Err;
case bitc::MODULE_CODE_SOURCE_FILENAME:
       SmallString<128> ValueName;
       if (convertToString(Record, 0, ValueName))
         return error("Invalid record");
       TheModule->setSourceFileName(ValueName);
       break;
     }
     Record.clear();
   }
+  return Error::success();
 }
 
 Error BitcodeReader::parseBitcodeInto(Module *M, bool ShouldLazyLoadMetadata,
                                       bool IsImporting,
                                       DataLayoutCallbackTy DataLayoutCallback) {
   TheModule = M;
   MDLoader = MetadataLoader(Stream, *M, ValueList, IsImporting,
                             [&](unsigned ID) { return getTypeByID(ID); });
if (MaterializeAll) {
     if (Error Err = R->materializeForwardReferencedFunctions())
       return std::move(Err);
   }
   return std::move(M);
 }
 
 Expected<std::unique_ptr<Module>>
 BitcodeModule::getLazyModule(LLVMContext &Context, bool ShouldLazyLoadMetadata,
-                             bool IsImporting) {
+                             bool IsImporting,
+                             DataLayoutCallbackTy DataLayoutCallback) {
   return getModuleImpl(Context, false, ShouldLazyLoadMetadata, IsImporting,
-                       [](StringRef) { return std::nullopt; });
+                       DataLayoutCallback);
 }
 
 // Parse the specified bitcode buffer and merge the index into CombinedIndex.
 // We don't use ModuleIdentifier here because the client may need to control the
 // module path used in the combined summary (e.g. when reading summaries for
 // regular LTO modules).
 Error BitcodeModule::readSummary(
     ModuleSummaryIndex &CombinedIndex, StringRef ModulePath, uint64_t ModuleId,
static Expected<BitcodeModule> getSingleModule(MemoryBufferRef Buffer) {
   if (MsOrErr->size() != 1)
     return error("Expected a single module");
 
   return (*MsOrErr)[0];
 }
 
 Expected<std::unique_ptr<Module>>
 llvm::getLazyBitcodeModule(MemoryBufferRef Buffer, LLVMContext &Context,
-                           bool ShouldLazyLoadMetadata, bool IsImporting) {
+                           bool ShouldLazyLoadMetadata, bool IsImporting,
+                           DataLayoutCallbackTy DataLayoutCallback) {
   Expected<BitcodeModule> BM = getSingleModule(Buffer);
   if (!BM)
     return BM.takeError();
 
-  return BM->getLazyModule(Context, ShouldLazyLoadMetadata, IsImporting);
+  return BM->getLazyModule(Context, ShouldLazyLoadMetadata, IsImporting,
+                           DataLayoutCallback);
 }
 
 Expected<std::unique_ptr<Module>> llvm::getOwningLazyBitcodeModule(
     std::unique_ptr<MemoryBuffer> &&Buffer, LLVMContext &Context,
     bool ShouldLazyLoadMetadata, bool IsImporting) {
   auto MOrErr = getLazyBitcodeModule(*Buffer, Context, ShouldLazyLoadMetadata,
                                      IsImporting);
   if (MOrErr)

llvm/lib/CodeGen/MIRParser/MIRParser.cpp

 std::unique_ptr<Module>
 MIRParserImpl::parseIRModule(DataLayoutCallbackTy DataLayoutCallback) {
   if (!In.setCurrentDocument()) {
     if (In.error())
       return nullptr;
     // Create an empty module when the MIR file is empty.
     NoMIRDocuments = true;
     auto M = std::make_unique<Module>(Filename, Context);
-    if (auto LayoutOverride = DataLayoutCallback(M->getTargetTriple()))
+    if (auto LayoutOverride =
+            DataLayoutCallback(M->getTargetTriple(), M->getDataLayoutStr()))
       M->setDataLayout(*LayoutOverride);
     return M;
   }
 
   std::unique_ptr<Module> M;
   // Parse the block scalar manually so that we can return unique pointer
   // without having to go trough YAML traits.
   if (const auto *BSN =
           dyn_cast_or_null<yaml::BlockScalarNode>(In.getCurrentNode())) {
     SMDiagnostic Error;
     M = parseAssembly(MemoryBufferRef(BSN->getValue(), Filename), Error,
                       Context, &IRSlots, DataLayoutCallback);
     if (!M) {
       reportDiagnostic(diagFromBlockStringDiag(Error, BSN->getSourceRange()));
       return nullptr;
     }
     In.nextDocument();
     if (!In.setCurrentDocument())
       NoMIRDocuments = true;
   } else {
     // Create an new, empty module.
     M = std::make_unique<Module>(Filename, Context);
-    if (auto LayoutOverride = DataLayoutCallback(M->getTargetTriple()))
+    if (auto LayoutOverride =
+            DataLayoutCallback(M->getTargetTriple(), M->getDataLayoutStr()))
       M->setDataLayout(*LayoutOverride);
     NoLLVMIR = true;
   }
   return M;
 }
 
 bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
   if (NoMIRDocuments)

llvm/test/Assembler/invalid-datalayout-override.ll

+; Check that importing this file gives an error due to the broken data layout:
+; RUN: not llvm-as < %s 2>&1 | FileCheck %s
+; Check that specifying a valid data layout allows to import this file:
+; RUN: llvm-as -data-layout "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128" < %s
+
+target datalayout = "A16777216"
+; CHECK: Invalid address space, must be a 24-bit integer

llvm/test/Bitcode/function-default-address-spaces.ll

 ; RUN: llvm-as %s  -o - | llvm-dis - | FileCheck %s -check-prefixes CHECK,PROG-AS0
 ; RUN: llvm-as -data-layout "P200" %s  -o - | llvm-dis | FileCheck %s -check-prefixes CHECK,PROG-AS200
-; RUN: not --crash llvm-as -data-layout "P123456789" %s -o /dev/null 2>&1 | FileCheck %s -check-prefix BAD-DATALAYOUT
-; BAD-DATALAYOUT: LLVM ERROR: Invalid address space, must be a 24-bit integer
+; RUN: not llvm-as -data-layout "P123456789" %s -o /dev/null 2>&1 | FileCheck %s -check-prefix BAD-DATALAYOUT
+; BAD-DATALAYOUT: error: Invalid address space, must be a 24-bit integer
 
 ; PROG-AS0-NOT: target datalayout
 ; PROG-AS200: target datalayout = "P200"
 
 ; Check that a function declaration without an address space (i.e. AS0) does not
 ; have the addrspace() attribute printed if it is address space zero and it is
 ; equal to the program address space.
 

llvm/tools/llc/llc.cpp

static int compileModule(char **argv, LLVMContext &Context) {
   std::optional<Reloc::Model> RM = codegen::getExplicitRelocModel();
   std::optional<CodeModel::Model> CM = codegen::getExplicitCodeModel();
 
   const Target *TheTarget = nullptr;
   std::unique_ptr<TargetMachine> Target;
 
   // If user just wants to list available options, skip module loading
   if (!SkipModule) {
-    auto SetDataLayout =
-        [&](StringRef DataLayoutTargetTriple) -> std::optional<std::string> {
+    auto SetDataLayout = [&](StringRef DataLayoutTargetTriple,
+                             StringRef OldDLStr) -> std::optional<std::string> {
       // If we are supposed to override the target triple, do so now.
       std::string IRTargetTriple = DataLayoutTargetTriple.str();
       if (!TargetTriple.empty())
         IRTargetTriple = Triple::normalize(TargetTriple);
       TheTriple = Triple(IRTargetTriple);
       if (TheTriple.getTriple().empty())
         TheTriple.setTriple(sys::getDefaultTargetTriple());
 

llvm/tools/llvm-as/llvm-as.cpp

 int main(int argc, char **argv) {
   InitLLVM X(argc, argv);
   cl::HideUnrelatedOptions(AsCat);
   cl::ParseCommandLineOptions(argc, argv, "llvm .ll -> .bc assembler\n");
   LLVMContext Context;
 
   // Parse the file now...
   SMDiagnostic Err;
-  auto SetDataLayout = [](StringRef) -> std::optional<std::string> {
+  auto SetDataLayout = [](StringRef, StringRef) -> std::optional<std::string> {
     if (ClDataLayout.empty())
       return std::nullopt;
     return ClDataLayout;
   };
   ParsedModuleAndIndex ModuleAndIndex;
   if (DisableVerify) {
     ModuleAndIndex = parseAssemblyFileWithIndexNoUpgradeDebugInfo(
         InputFilename, Err, Context, nullptr, SetDataLayout);

llvm/tools/llvm-reduce/ReducerWorkItem.cpp

if (IsMIR) {
     if (std::error_code EC = FileOrErr.getError()) {
       WithColor::error(errs(), ToolName) << EC.message() << '\n';
       return {nullptr, false};
     }
 
     std::unique_ptr<MIRParser> MParser =
         createMIRParser(std::move(FileOrErr.get()), Ctxt);
 
-    auto SetDataLayout =
-        [&](StringRef DataLayoutTargetTriple) -> std::optional<std::string> {
+    auto SetDataLayout = [&](StringRef DataLayoutTargetTriple,
+                             StringRef OldDLStr) -> std::optional<std::string> {
       // If we are supposed to override the target triple, do so now.
       std::string IRTargetTriple = DataLayoutTargetTriple.str();
       if (!TargetTriple.empty())
         IRTargetTriple = Triple::normalize(TargetTriple);
       TheTriple = Triple(IRTargetTriple);
       if (TheTriple.getTriple().empty())
         TheTriple.setTriple(sys::getDefaultTargetTriple());
 

llvm/tools/opt/opt.cpp

Expected<std::unique_ptr<ToolOutputFile>> RemarksFileOrErr =
                                    RemarksHotnessThreshold);
   if (Error E = RemarksFileOrErr.takeError()) {
     errs() << toString(std::move(E)) << '\n';
     return 1;
   }
   std::unique_ptr<ToolOutputFile> RemarksFile = std::move(*RemarksFileOrErr);
 
   // Load the input module...
-  auto SetDataLayout = [](StringRef) -> std::optional<std::string> {
+  auto SetDataLayout = [](StringRef, StringRef) -> std::optional<std::string> {
     if (ClDataLayout.empty())
       return std::nullopt;
     return ClDataLayout;
   };
   std::unique_ptr<Module> M;
   if (NoUpgradeDebugInfo)
     M = parseAssemblyFileWithIndexNoUpgradeDebugInfo(
             InputFilename, Err, Context, nullptr, SetDataLayout)

llvm/unittests/AsmParser/AsmParserTest.cpp

 //===- llvm/unittest/AsmParser/AsmParserTest.cpp - asm parser unittests ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/AsmParser/SlotMapping.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 
 TEST(AsmParserTest, NullTerminatedInput) {
TEST(AsmParserTest, TypeAtBeginningWithSlotMappingParsing) {
   Ty = parseTypeAtBeginning("i32 garbage", Read, Error, M, &Mapping);
   ASSERT_TRUE(Ty);
   ASSERT_TRUE(Ty->isIntegerTy());
   ASSERT_TRUE(Ty->getPrimitiveSizeInBits() == 32);
   // We go to the next token, i.e., we read "i32" + ' '.
   ASSERT_TRUE(Read == 4);
 }
 
+TEST(AsmParserTest, InvalidDataLayoutStringCallback) {
+  LLVMContext Ctx;
+  SMDiagnostic Error;
+  // Note the invalid i8:7 part
+  // Overalign i32 as marker so we can check that indeed this DL was used,
+  // and not some default.
+  StringRef InvalidDLStr =
+      "e-m:e-p:64:64-i8:7-i16:16-i32:64-i64:64-f80:128-n8:16:32:64";
+  StringRef FixedDLStr =
+      "e-m:e-p:64:64-i8:8-i16:16-i32:64-i64:64-f80:128-n8:16:32:64";
+  Expected<DataLayout> ExpectedFixedDL = DataLayout::parse(FixedDLStr);
+  ASSERT_TRUE(!ExpectedFixedDL.takeError());
+  DataLayout FixedDL = ExpectedFixedDL.get();
+  std::string Source = ("target datalayout = \"" + InvalidDLStr + "\"\n").str();
+  MemoryBufferRef SourceBuffer(Source, "<string>");
+
+  // Check that we reject the source without a DL override.
+  SlotMapping Mapping1;
+  auto Mod1 = parseAssembly(SourceBuffer, Error, Ctx, &Mapping1);
+  EXPECT_TRUE(Mod1 == nullptr);
+
+  // Check that we pass the correct DL str to the callback,
+  // that fixing the DL str from the callback works,
+  // and that the resulting module has the correct DL.
+  SlotMapping Mapping2;
+  auto Mod2 = parseAssembly(
+      SourceBuffer, Error, Ctx, &Mapping2,
+      [&](StringRef Triple, StringRef DLStr) -> std::optional<std::string> {
+        EXPECT_EQ(DLStr, InvalidDLStr);
+        return std::string{FixedDLStr};
+      });
+  ASSERT_TRUE(Mod2 != nullptr);
+  EXPECT_EQ(Mod2->getDataLayout(), FixedDL);
+}
+
 } // end anonymous namespace