[mlir:Bytecode] Add shared_ptr<SourceMgr> overloads to allow safe mmap of data

[rriddle]

Summary

The bytecode reader currently has no mechanism that allows for directly referencing
data from the input buffer safely. This commit adds shared_ptr<SourceMgr> overloads
that provide an explicit and safe way of extending the lifetime of the input. The usage of
these new overloads is adopted in all of our tooling, and is implicitly used in the filename
only parser methods.

[jpienaar]

Yay for timing as I was just talking about mmap'd resources.

Do we need shared_ptr and do we need it with the IR? (I have a strong dislike with shared pointers as you know :)). We did talk about it before, but just to refresh.

Do we expect the SourceMgr to be scoped to a MLIRContext or be shared between many?

[rriddle]

In D139366#3975092, @jpienaar wrote:

Yay for timing as I was just talking about mmap'd resources.

Do we need shared_ptr and do we need it with the IR? (I have a strong dislike with shared pointers as you know :)). We did talk about it before, but just to refresh.

We need something that allows for extending a reference, otherwise there is no way to correctly and properly extend the lifetime. shared_ptr was the most effective that required least amount of mucking about. Do you have an alternate suggestion?

Do we expect the SourceMgr to be scoped to a MLIRContext or be shared between many?

I would expect in most cases the SourceMgr is for that specific context (given that most cases only use a single source manager anyways).

[jpienaar]

Makes sense, thanks for explaining / clarifying.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

567

Is this intentionally non-mutable only here?

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

1271 ↗

(On Diff #480225)

Could we have an overload for original case? (to reduce number of changes needed for cases where this optimization isn't needed and folks opt-in to higher performance one). [not blocking]

[mehdi_amini]

In D139366#3975096, @rriddle wrote:

In D139366#3975092, @jpienaar wrote:

Yay for timing as I was just talking about mmap'd resources.

Do we need shared_ptr and do we need it with the IR? (I have a strong dislike with shared pointers as you know :)). We did talk about it before, but just to refresh.

We need something that allows for extending a reference, otherwise there is no way to correctly and properly extend the lifetime. shared_ptr was the most effective that required least amount of mucking about. Do you have an alternate suggestion?

An alternative way is to use parsing flags, so that instead of implicit extension of lifetime it could be just explicit on the client side: "here is an input buffer that I promise to keep alive beyond the IR".

Do we expect the SourceMgr to be scoped to a MLIRContext or be shared between many?

I would expect in most cases the SourceMgr is for that specific context (given that most cases only use a single source manager anyways).

But we should be allowed to share it right?

mlir/include/mlir/Parser/Parser.h

96

Missing an empty line before this comment block? Unless there is some implicit grouping I'm not getting? (same elsewhere)

219

Can you expand these comments with some context? That is include the "why" this is useful/important:

This is useful for example to avoid copying some large resources
into the MLIRContext and instead referencing the data directly from
the input buffers.

[rriddle]

In D139366#3980424, @mehdi_amini wrote:

In D139366#3975096, @rriddle wrote:

In D139366#3975092, @jpienaar wrote:

Yay for timing as I was just talking about mmap'd resources.

Do we need shared_ptr and do we need it with the IR? (I have a strong dislike with shared pointers as you know :)). We did talk about it before, but just to refresh.

We need something that allows for extending a reference, otherwise there is no way to correctly and properly extend the lifetime. shared_ptr was the most effective that required least amount of mucking about. Do you have an alternate suggestion?

An alternative way is to use parsing flags, so that instead of implicit extension of lifetime it could be just explicit on the client side: "here is an input buffer that I promise to keep alive beyond the IR".

Yeah I considered that, but the implicitness of it didn't feel great (well worse to me than using shared_ptr). I wanted it to be explicit in the API and difficult to mess up.

Do we expect the SourceMgr to be scoped to a MLIRContext or be shared between many?

I would expect in most cases the SourceMgr is for that specific context (given that most cases only use a single source manager anyways).

But we should be allowed to share it right?

Yeah, if desired it should be shareable. Just in most cases it is specific.

mlir/include/mlir/Bytecode/BytecodeReader.h

 #ifndef MLIR_BYTECODE_BYTECODEREADER_H
 #define MLIR_BYTECODE_BYTECODEREADER_H
 
 #include "mlir/IR/AsmState.h"
 #include "mlir/Support/LLVM.h"
 
 namespace llvm {
 class MemoryBufferRef;
+class SourceMgr;
 } // namespace llvm
 
 namespace mlir {
 /// Returns true if the given buffer starts with the magic bytes that signal
 /// MLIR bytecode.
 bool isBytecode(llvm::MemoryBufferRef buffer);
 
 /// Read the operations defined within the given memory buffer, containing MLIR
 /// bytecode, into the provided block.
 LogicalResult readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
                                const ParserConfig &config);
+/// An overload with a source manager whose main file buffer is used for
+/// parsing. The lifetime of the source manager may be freely extended during
+/// parsing such that the source manager is not destroyed before the parsed IR.
+LogicalResult
+readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                 Block *block, const ParserConfig &config);
 } // namespace mlir
 
 #endif // MLIR_BYTECODE_BYTECODEREADER_H

mlir/include/mlir/IR/AsmState.h

 /// This class provides a simple utility wrapper for creating "unmanaged"
 /// AsmResourceBlobs. The lifetime of the data provided to these blobs is
 /// guaranteed to persist beyond the lifetime of this reference.
 class UnmanagedAsmResourceBlob {
 public:
   /// Create a new unmanaged resource directly referencing the provided data.
   /// `dataIsMutable` indicates if the allocated data can be mutated. By
   /// default, we treat unmanaged blobs as immutable.
-  static AsmResourceBlob allocateWithAlign(ArrayRef<char> data, size_t align,
-                                           bool dataIsMutable = false) {
-    return AsmResourceBlob(data, align, /*deleter=*/{},
-                           /*dataIsMutable=*/false);
+  static AsmResourceBlob
+  allocateWithAlign(ArrayRef<char> data, size_t align,
+                    AsmResourceBlob::DeleterFn deleter = {},
+    return AsmResourceBlob(data, align, std::move(deleter), dataIsMutable);
   }
   template <typename T>
-  static AsmResourceBlob allocateInferAlign(ArrayRef<T> data,
-                                            bool dataIsMutable = false) {
+  static AsmResourceBlob
+  allocateInferAlign(ArrayRef<T> data, AsmResourceBlob::DeleterFn deleter = {},
     return allocateWithAlign(
         ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
-        alignof(T));
+        alignof(T), std::move(deleter), dataIsMutable);
   }
 };
 
 /// This class is used to build resource entries for use by the printer. Each
 /// resource entry is represented using a key/value pair. The provided key must
 /// be unique within the current context, which allows for a client to provide
 /// resource entries without worrying about overlap with other clients.
 class AsmResourceBuilder {

mlir/include/mlir/Parser/Parser.h

 /// placed before the current terminator. If parsing is successful, success is
 /// returned. Otherwise, an error message is emitted through the error handler
 /// registered in the context, and failure is returned. If `sourceFileLoc` is
 /// non-null, it is populated with a file location representing the start of the
 /// source file that is being parsed.
 LogicalResult parseSourceFile(const llvm::SourceMgr &sourceMgr, Block *block,
                               const ParserConfig &config,
                               LocationAttr *sourceFileLoc = nullptr);
+/// An overload with a source manager that may have references taken during the
+/// parsing process, and whose lifetime can be freely extended (such that the
+/// source manager is not destroyed before the parsed IR). This is useful, for
+/// example, to avoid copying some large resources into the MLIRContext and
+/// instead referencing the data directly from the input buffers.
+LogicalResult parseSourceFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                              Block *block, const ParserConfig &config,
+                              LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the file specified by the indicated filename and appends parsed
 /// operations to the given block. If the block is non-empty, the operations are
 /// placed before the current terminator. If parsing is successful, success is
 /// returned. Otherwise, an error message is emitted through the error handler
 /// registered in the context, and failure is returned. If `sourceFileLoc` is
 /// non-null, it is populated with a file location representing the start of the
 /// source file that is being parsed.
 LogicalResult parseSourceFile(llvm::StringRef filename, Block *block,
                               const ParserConfig &config,
                               LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the file specified by the indicated filename using the provided
 /// SourceMgr and appends parsed operations to the given block. If the block is
 /// non-empty, the operations are placed before the current terminator. If
 /// parsing is successful, success is returned. Otherwise, an error message is
 /// emitted through the error handler registered in the context, and failure is
 /// returned. If `sourceFileLoc` is non-null, it is populated with a file
 /// location representing the start of the source file that is being parsed.
 LogicalResult parseSourceFile(llvm::StringRef filename,
                               llvm::SourceMgr &sourceMgr, Block *block,
                               const ParserConfig &config,
                               LocationAttr *sourceFileLoc = nullptr);
+/// An overload with a source manager that may have references taken during the
+/// parsing process, and whose lifetime can be freely extended (such that the
+/// source manager is not destroyed before the parsed IR). This is useful, for
+/// example, to avoid copying some large resources into the MLIRContext and
+/// instead referencing the data directly from the input buffers.
+LogicalResult parseSourceFile(llvm::StringRef filename,
+                              const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                              Block *block, const ParserConfig &config,
+                              LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the IR string and appends parsed operations to the given block.
 /// If the block is non-empty, the operations are placed before the current
 /// terminator. If parsing is successful, success is returned. Otherwise, an
 /// error message is emitted through the error handler registered in the
 /// context, and failure is returned. If `sourceFileLoc` is non-null, it is
 /// populated with a file location representing the start of the source file
 /// that is being parsed.

 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the
 /// `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT = Operation *>
 inline OwningOpRef<ContainerOpT>
 parseSourceFile(const llvm::SourceMgr &sourceMgr, const ParserConfig &config) {
   return detail::parseSourceFile<ContainerOpT>(config, sourceMgr);
 }
+/// An overload with a source manager that may have references taken during the
+/// parsing process, and whose lifetime can be freely extended (such that the
+/// source manager is not destroyed before the parsed IR). This is useful, for
+/// example, to avoid copying some large resources into the MLIRContext and
+/// instead referencing the data directly from the input buffers.
+template <typename ContainerOpT = Operation *>
+inline OwningOpRef<ContainerOpT>
+parseSourceFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                const ParserConfig &config) {
+  return detail::parseSourceFile<ContainerOpT>(config, sourceMgr);
+}
 
 /// This parses the file specified by the indicated filename. If the source IR
 /// contained a single instance of `ContainerOpT`, it is returned. Otherwise, a
 /// new instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the

 /// required to have a single region containing a single block, and must
 /// implement the `SingleBlockImplicitTerminator` trait.
 template <typename ContainerOpT = Operation *>
 inline OwningOpRef<ContainerOpT> parseSourceFile(llvm::StringRef filename,
                                                  llvm::SourceMgr &sourceMgr,
                                                  const ParserConfig &config) {
   return detail::parseSourceFile<ContainerOpT>(config, filename, sourceMgr);
 }
+/// An overload with a source manager that may have references taken during the
+/// parsing process, and whose lifetime can be freely extended (such that the
+/// source manager is not destroyed before the parsed IR). This is useful, for
+/// example, to avoid copying some large resources into the MLIRContext and
+/// instead referencing the data directly from the input buffers.
+template <typename ContainerOpT = Operation *>
+inline OwningOpRef<ContainerOpT>
+parseSourceFile(llvm::StringRef filename,
+                const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                const ParserConfig &config) {
+  return detail::parseSourceFile<ContainerOpT>(config, filename, sourceMgr);
+}
 
 /// This parses the provided string containing MLIR. If the source IR contained
 /// a single instance of `ContainerOpT`, it is returned. Otherwise, a new
 /// instance of `ContainerOpT` is constructed containing all of the parsed
 /// operations. If parsing was not successful, null is returned and an error
 /// message is emitted through the error handler registered in the context, and
 /// failure is returned. `ContainerOpT` is required to have a single region
 /// containing a single block, and must implement the

mlir/include/mlir/Tools/ParseUtilities.h

 
 namespace mlir {
 /// This parses the file specified by the indicated SourceMgr. If parsing was
 /// not successful, null is returned and an error message is emitted through the
 /// error handler registered in the context.
 /// If 'insertImplicitModule' is true a top-level 'builtin.module' op will be
 /// inserted that contains the parsed IR, unless one exists already.
 inline OwningOpRef<Operation *>
-parseSourceFileForTool(llvm::SourceMgr &sourceMgr, const ParserConfig &config,
-                       bool insertImplicitModule) {
+parseSourceFileForTool(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                       const ParserConfig &config, bool insertImplicitModule) {
   if (insertImplicitModule) {
     // TODO: Move implicit module logic out of 'parseSourceFile' and into here.
     return parseSourceFile<ModuleOp>(sourceMgr, config);
   }
   return parseSourceFile(sourceMgr, config);
 }
 } // namespace mlir
 
 #endif // MLIR_TOOLS_PARSEUTILITIES_H

mlir/include/mlir/Tools/mlir-translate/Translation.h

 template <typename OpTy>
 class OwningOpRef;
 
 /// Interface of the function that translates the sources managed by `sourceMgr`
 /// to MLIR. The source manager has at least one buffer. The implementation
 /// should create a new MLIR Operation in the given context and return a
 /// pointer to it, or a nullptr in case of any error.
 using TranslateSourceMgrToMLIRFunction = std::function<OwningOpRef<Operation *>(
-    llvm::SourceMgr &sourceMgr, MLIRContext *)>;
+    const std::shared_ptr<llvm::SourceMgr> &sourceMgr, MLIRContext *)>;
+using TranslateRawSourceMgrToMLIRFunction =
+    std::function<OwningOpRef<Operation *>(llvm::SourceMgr &sourceMgr,
+                                           MLIRContext *)>;
 
 /// Interface of the function that translates the given string to MLIR. The
 /// implementation should create a new MLIR Operation in the given context. If
 /// source-related error reporting is required from within the function, use
 /// TranslateSourceMgrToMLIRFunction instead.
 using TranslateStringRefToMLIRFunction =
     std::function<OwningOpRef<Operation *>(llvm::StringRef, MLIRContext *)>;
 
 /// Interface of the function that translates MLIR to a different format and
 /// outputs the result to a stream. It is allowed to modify the operation.
 using TranslateFromMLIRFunction =
     std::function<LogicalResult(Operation *, llvm::raw_ostream &output)>;
 
 /// Interface of the function that performs file-to-file translation involving
 /// MLIR. The input file is held in the given MemoryBuffer; the output file
 /// should be written to the given raw_ostream. The implementation should create
 /// all MLIR constructs needed during the process inside the given context. This
 /// can be used for round-tripping external formats through the MLIR system.
 using TranslateFunction = std::function<LogicalResult(
-    llvm::SourceMgr &sourceMgr, llvm::raw_ostream &output, MLIRContext *)>;
+    const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+    llvm::raw_ostream &output, MLIRContext *)>;
 
 /// This class contains all of the components necessary for performing a
 /// translation.
 class Translation {
 public:
   Translation() = default;
   Translation(TranslateFunction function, StringRef description,
               Optional<llvm::Align> inputAlignment)
       : function(std::move(function)), description(description),
         inputAlignment(inputAlignment) {}
 
   /// Return the description of this translation.
   StringRef getDescription() const { return description; }
 
   /// Return the optional alignment desired for the input of the translation.
   Optional<llvm::Align> getInputAlignment() const { return inputAlignment; }
 
   /// Invoke the translation function with the given input and output streams.
-  LogicalResult operator()(llvm::SourceMgr &sourceMgr,
+  LogicalResult operator()(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
                            llvm::raw_ostream &output,
                            MLIRContext *context) const {
     return function(sourceMgr, output, context);
   }
 
 private:
   /// The underlying translation function.
   TranslateFunction function;

 /// \{
 struct TranslateToMLIRRegistration {
   TranslateToMLIRRegistration(
       llvm::StringRef name, llvm::StringRef description,
       const TranslateSourceMgrToMLIRFunction &function,
       Optional<llvm::Align> inputAlignment = std::nullopt);
   TranslateToMLIRRegistration(
       llvm::StringRef name, llvm::StringRef description,
+      const TranslateRawSourceMgrToMLIRFunction &function,
+      Optional<llvm::Align> inputAlignment = std::nullopt);
+  TranslateToMLIRRegistration(
+      llvm::StringRef name, llvm::StringRef description,
       const TranslateStringRefToMLIRFunction &function,
       Optional<llvm::Align> inputAlignment = std::nullopt);
 };
 
 struct TranslateFromMLIRRegistration {
   TranslateFromMLIRRegistration(
       llvm::StringRef name, llvm::StringRef description,
       const TranslateFromMLIRFunction &function,

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Verifier.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/SaveAndRestore.h"
+#include "llvm/Support/SourceMgr.h"
 #include <optional>
 
 #define DEBUG_TYPE "mlir-bytecode-reader"
 
 using namespace mlir;
 
 /// Stringify the given section ID.
 static std::string toString(bytecode::Section::ID sectionID) {

 // ResourceSectionReader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read the resource section from the bytecode.
 class ResourceSectionReader {
 public:
   /// Initialize the resource section reader with the given section data.
-  LogicalResult initialize(Location fileLoc, const ParserConfig &config,
-                           MutableArrayRef<BytecodeDialect> dialects,
-                           StringSectionReader &stringReader,
-                           ArrayRef<uint8_t> sectionData,
-                           ArrayRef<uint8_t> offsetSectionData);
+  LogicalResult
+  initialize(Location fileLoc, const ParserConfig &config,
+             StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
+             ArrayRef<uint8_t> offsetSectionData,
+             const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef);
 
   /// Parse a dialect resource handle from the resource section.
   LogicalResult parseResourceHandle(EncodingReader &reader,
                                     AsmDialectResourceHandle &result) {
     return parseEntry(reader, dialectResources, result, "resource handle");
   }
 
 private:
   /// The table of dialect resources within the bytecode file.
   SmallVector<AsmDialectResourceHandle> dialectResources;
 };
 
 class ParsedResourceEntry : public AsmParsedResourceEntry {
 public:
   ParsedResourceEntry(StringRef key, AsmResourceEntryKind kind,
-                      EncodingReader &reader, StringSectionReader &stringReader)
-      : key(key), kind(kind), reader(reader), stringReader(stringReader) {}
+                      EncodingReader &reader, StringSectionReader &stringReader,
+                      const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
+      : key(key), kind(kind), reader(reader), stringReader(stringReader),
+        bufferOwnerRef(bufferOwnerRef) {}
   ~ParsedResourceEntry() override = default;
 
   StringRef getKey() const final { return key; }
 
   InFlightDiagnostic emitError() const final { return reader.emitError(); }
 
   AsmResourceEntryKind getKind() const final { return kind; }
 
if (kind != AsmResourceEntryKind::Blob)
       return emitError() << "expected a blob resource entry, but found a "
                          << toString(kind) << " entry instead";
 
     ArrayRef<uint8_t> data;
     uint64_t alignment;
     if (failed(reader.parseBlobAndAlignment(data, alignment)))
       return failure();
 
+    // If we have an extendable reference to the buffer owner, we don't need to
+    // allocate a new buffer for the data, and can use the data directly.
+    if (bufferOwnerRef) {
+      ArrayRef<char> charData(reinterpret_cast<const char *>(data.data()),
+                              data.size());
+
+      // Allocate an unmanager buffer which captures a reference to the owner.
+      // For now we just mark this as immutable, but in the future we should
+      // explore marking this as mutable when desired.
+      return UnmanagedAsmResourceBlob::allocateWithAlign(
+          charData, alignment,
+          [bufferOwnerRef = bufferOwnerRef](void *, size_t, size_t) {});
+    }
+
     // Allocate memory for the blob using the provided allocator and copy the
     // data into it.
-    // FIXME: If the current holder of the bytecode can ensure its lifetime
-    // (e.g. when mmap'd), we should not copy the data. We should use the data
-    // from the bytecode directly.
     AsmResourceBlob blob = allocator(data.size(), alignment);
     assert(llvm::isAddrAligned(llvm::Align(alignment), blob.getData().data()) &&
            blob.isMutable() &&
            "blob allocator did not return a properly aligned address");
     memcpy(blob.getMutableData().data(), data.data(), data.size());
     return blob;
   }
 
 private:
   StringRef key;
   AsmResourceEntryKind kind;
   EncodingReader &reader;
   StringSectionReader &stringReader;
+  const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
 };
 } // namespace
 
 template <typename T>
 static LogicalResult
 parseResourceGroup(Location fileLoc, bool allowEmpty,
                    EncodingReader &offsetReader, EncodingReader &resourceReader,
                    StringSectionReader &stringReader, T *handler,
+                   const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef,
                    function_ref<LogicalResult(StringRef)> processKeyFn = {}) {
   uint64_t numResources;
   if (failed(offsetReader.parseVarInt(numResources)))
     return failure();
 
   for (uint64_t i = 0; i < numResources; ++i) {
     StringRef key;
     AsmResourceEntryKind kind;
if (allowEmpty && data.empty())
       continue;
 
     // Ignore the entry if we don't have a valid handler.
     if (!handler)
       continue;
 
     // Otherwise, parse the resource value.
     EncodingReader entryReader(data, fileLoc);
-    ParsedResourceEntry entry(key, kind, entryReader, stringReader);
+    ParsedResourceEntry entry(key, kind, entryReader, stringReader,
+                              bufferOwnerRef);
     if (failed(handler->parseResource(entry)))
       return failure();
     if (!entryReader.empty()) {
       return entryReader.emitError(
           "unexpected trailing bytes in resource entry '", key, "'");
     }
   }
   return success();
 }
 
-LogicalResult
-ResourceSectionReader::initialize(Location fileLoc, const ParserConfig &config,
-                                  MutableArrayRef<BytecodeDialect> dialects,
-                                  StringSectionReader &stringReader,
-                                  ArrayRef<uint8_t> sectionData,
-                                  ArrayRef<uint8_t> offsetSectionData) {
+LogicalResult ResourceSectionReader::initialize(
+    Location fileLoc, const ParserConfig &config,
+    StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
+    ArrayRef<uint8_t> offsetSectionData,
+    const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
   EncodingReader resourceReader(sectionData, fileLoc);
   EncodingReader offsetReader(offsetSectionData, fileLoc);
 
   // Read the number of external resource providers.
   uint64_t numExternalResourceGroups;
   if (failed(offsetReader.parseVarInt(numExternalResourceGroups)))
     return failure();
 
   // Utility functor that dispatches to `parseResourceGroup`, but implicitly
   // provides most of the arguments.
   auto parseGroup = [&](auto *handler, bool allowEmpty = false,
                         function_ref<LogicalResult(StringRef)> keyFn = {}) {
     return parseResourceGroup(fileLoc, allowEmpty, offsetReader, resourceReader,
-                              stringReader, handler, keyFn);
+                              stringReader, handler, bufferOwnerRef, keyFn);
   };
 
   // Read the external resources from the bytecode.
   for (uint64_t i = 0; i < numExternalResourceGroups; ++i) {
     StringRef key;
     if (failed(stringReader.parseString(offsetReader, key)))
       return failure();
 

 //===----------------------------------------------------------------------===//
 // Bytecode Reader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read a bytecode buffer and translate it into MLIR.
 class BytecodeReader {
 public:
-  BytecodeReader(Location fileLoc, const ParserConfig &config)
+  BytecodeReader(Location fileLoc, const ParserConfig &config,
+                 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
       : config(config), fileLoc(fileLoc),
         attrTypeReader(stringReader, resourceReader, fileLoc),
         // Use the builtin unrealized conversion cast operation to represent
         // forward references to values that aren't yet defined.
         forwardRefOpState(UnknownLoc::get(config.getContext()),
                           "builtin.unrealized_conversion_cast", ValueRange(),
-                          NoneType::get(config.getContext())) {}
+                          NoneType::get(config.getContext())),
+        bufferOwnerRef(bufferOwnerRef) {}
 
   /// Read the bytecode defined within `buffer` into the given block.
   LogicalResult read(llvm::MemoryBufferRef buffer, Block *block);
 
 private:
   /// Return the context for this config.
   MLIRContext *getContext() const { return config.getContext(); }
 
private:
   /// A block containing the set of operations defined to create forward
   /// references.
   Block forwardRefOps;
   /// A block containing previously created, and no longer used, forward
   /// reference operations.
   Block openForwardRefOps;
   /// An operation state used when instantiating forward references.
   OperationState forwardRefOpState;
+
+  /// The optional owning source manager, which when present may be used to
+  /// extend the lifetime of the input buffer.
+  const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
 };
 } // namespace
 
 LogicalResult BytecodeReader::read(llvm::MemoryBufferRef buffer, Block *block) {
   EncodingReader reader(buffer.getBuffer(), fileLoc);
 
   // Skip over the bytecode header, this should have already been checked.
   if (failed(reader.skipBytes(StringRef("ML\xefR").size())))
LogicalResult BytecodeReader::parseResourceSection(
   }
 
   // If the resource sections are absent, there is nothing to do.
   if (!resourceData)
     return success();
 
   // Initialize the resource reader with the resource sections.
   return resourceReader.initialize(fileLoc, config, dialects, stringReader,
-                                   *resourceData, *resourceOffsetData);
+                                   *resourceData, *resourceOffsetData,
+                                   bufferOwnerRef);
 }
 
 //===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
                                              Block *block) {
   EncodingReader reader(sectionData, fileLoc);

 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
 bool mlir::isBytecode(llvm::MemoryBufferRef buffer) {
   return buffer.getBuffer().startswith("ML\xefR");
 }
 
-LogicalResult mlir::readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
-                                     const ParserConfig &config) {
+/// Read the bytecode from the provided memory buffer reference.
+/// `bufferOwnerRef` if provided is the owning source manager for the buffer,
+/// and may be used to extend the lifetime of the buffer.
+static LogicalResult
+readBytecodeFileImpl(llvm::MemoryBufferRef buffer, Block *block,
+                     const ParserConfig &config,
+                     const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
   Location sourceFileLoc =
       FileLineColLoc::get(config.getContext(), buffer.getBufferIdentifier(),
                           /*line=*/0, /*column=*/0);
   if (!isBytecode(buffer)) {
     return emitError(sourceFileLoc,
                      "input buffer is not an MLIR bytecode file");
   }
 
-  BytecodeReader reader(sourceFileLoc, config);
+  BytecodeReader reader(sourceFileLoc, config, bufferOwnerRef);
   return reader.read(buffer, block);
 }
+
+LogicalResult mlir::readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
+                                     const ParserConfig &config) {
+  return readBytecodeFileImpl(buffer, block, config, /*bufferOwnerRef=*/{});
+}
+LogicalResult
+mlir::readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                       Block *block, const ParserConfig &config) {
+  return readBytecodeFileImpl(
+      *sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID()), block, config,
+      sourceMgr);
+}

mlir/lib/ExecutionEngine/JitRunner.cpp

static OwningOpRef<Operation *> parseMLIRInput(StringRef inputFilename,
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
     return nullptr;
   }
 
-  llvm::SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
+  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+  sourceMgr->AddNewSourceBuffer(std::move(file), SMLoc());
   OwningOpRef<Operation *> module =
       parseSourceFileForTool(sourceMgr, context, insertImplicitModule);
   if (!module)
     return nullptr;
   if (!module.get()->hasTrait<OpTrait::SymbolTable>()) {
     llvm::errs() << "Error: top-level op must be a symbol table.\n";
     return nullptr;
   }

mlir/lib/Parser/Parser.cpp

if (sourceFileLoc) {
     *sourceFileLoc = FileLineColLoc::get(config.getContext(),
                                          sourceBuf->getBufferIdentifier(),
                                          /*line=*/0, /*column=*/0);
   }
   if (isBytecode(*sourceBuf))
     return readBytecodeFile(*sourceBuf, block, config);
   return parseAsmSourceFile(sourceMgr, block, config);
 }
+LogicalResult
+mlir::parseSourceFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                      Block *block, const ParserConfig &config,
+                      LocationAttr *sourceFileLoc) {
+  const auto *sourceBuf =
+      sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID());
+  if (sourceFileLoc) {
+    *sourceFileLoc = FileLineColLoc::get(config.getContext(),
+                                         sourceBuf->getBufferIdentifier(),
+                                         /*line=*/0, /*column=*/0);
+  }
+  if (isBytecode(*sourceBuf))
+    return readBytecodeFile(sourceMgr, block, config);
+  return parseAsmSourceFile(*sourceMgr, block, config);
+}
 
 LogicalResult mlir::parseSourceFile(llvm::StringRef filename, Block *block,
                                     const ParserConfig &config,
                                     LocationAttr *sourceFileLoc) {
-  llvm::SourceMgr sourceMgr;
+  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
   return parseSourceFile(filename, sourceMgr, block, config, sourceFileLoc);
 }
 
-LogicalResult mlir::parseSourceFile(llvm::StringRef filename,
-                                    llvm::SourceMgr &sourceMgr, Block *block,
-                                    const ParserConfig &config,
-                                    LocationAttr *sourceFileLoc) {
+static LogicalResult loadSourceFileBuffer(llvm::StringRef filename,
+                                          llvm::SourceMgr &sourceMgr,
+                                          MLIRContext *ctx) {
   if (sourceMgr.getNumBuffers() != 0) {
     // TODO: Extend to support multiple buffers.
-    return emitError(mlir::UnknownLoc::get(config.getContext()),
+    return emitError(mlir::UnknownLoc::get(ctx),
                      "only main buffer parsed at the moment");
   }
   auto fileOrErr = llvm::MemoryBuffer::getFileOrSTDIN(filename);
   if (std::error_code error = fileOrErr.getError())
-    return emitError(mlir::UnknownLoc::get(config.getContext()),
+    return emitError(mlir::UnknownLoc::get(ctx),
                      "could not open input file " + filename);
 
   // Load the MLIR source file.
   sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), SMLoc());
+  return success();
+}
+
+LogicalResult mlir::parseSourceFile(llvm::StringRef filename,
+                                    llvm::SourceMgr &sourceMgr, Block *block,
+                                    const ParserConfig &config,
+                                    LocationAttr *sourceFileLoc) {
+  if (failed(loadSourceFileBuffer(filename, sourceMgr, config.getContext())))
+    return failure();
+  return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
+}
+LogicalResult mlir::parseSourceFile(
+    llvm::StringRef filename, const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+    Block *block, const ParserConfig &config, LocationAttr *sourceFileLoc) {
+  if (failed(loadSourceFileBuffer(filename, *sourceMgr, config.getContext())))
+    return failure();
   return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
 }
 
 LogicalResult mlir::parseSourceString(llvm::StringRef sourceStr, Block *block,
                                       const ParserConfig &config,
                                       LocationAttr *sourceFileLoc) {
   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(sourceStr);
   if (!memBuffer)
     return failure();
 
   llvm::SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   return parseSourceFile(sourceMgr, block, config, sourceFileLoc);
 }

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

 using namespace llvm;
 
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
 /// This typically parses the main source file, runs zero or more optimization
 /// passes, then prints the output.
 ///
-static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
-                                    bool verifyPasses, SourceMgr &sourceMgr,
-                                    MLIRContext *context,
-                                    PassPipelineFn passManagerSetupFn,
-                                    bool emitBytecode, bool implicitModule) {
+static LogicalResult
+performActions(raw_ostream &os, bool verifyDiagnostics, bool verifyPasses,
+               const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+               MLIRContext *context, PassPipelineFn passManagerSetupFn,
   DefaultTimingManager tm;
   applyDefaultTimingManagerCLOptions(tm);
   TimingScope timing = tm.getRootScope();
 
   // Disable multi-threading when parsing the input file. This removes the
   // unnecessary/costly context synchronization when parsing.
   bool wasThreadingEnabled = context->isMultithreadingEnabled();
   context->disableMultithreading();

 static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
               bool verifyDiagnostics, bool verifyPasses,
               bool allowUnregisteredDialects, bool preloadDialectsInContext,
               bool emitBytecode, bool implicitModule,
               PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
               llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
-  SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
+  auto sourceMgr = std::make_shared<SourceMgr>();
+  sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
 
   // Parse the input file.
   if (preloadDialectsInContext)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   if (verifyDiagnostics)
     context.printOpOnDiagnostic(false);
   context.getDebugActionManager().registerActionHandler<DebugCounter>();
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
   if (!verifyDiagnostics) {
-    SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
+    SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
                           &context, passManagerSetupFn, emitBytecode,
                           implicitModule);
   }
 
-  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
+  SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr, &context);
 
   // Do any processing requested by command line flags.  We don't care whether
   // these actions succeed or fail, we only care what diagnostics they produce
   // and whether they match our expectations.
   (void)performActions(os, verifyDiagnostics, verifyPasses, sourceMgr, &context,
                        passManagerSetupFn, emitBytecode, implicitModule);
 
   // Verify the diagnostic handler to make sure that each of the diagnostics

mlir/lib/Tools/mlir-reduce/MlirReduceMain.cpp

OwningOpRef<Operation *> loadModule(MLIRContext &context,
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
     return nullptr;
   }
 
-  llvm::SourceMgr sourceMgr;
-  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
+  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+  sourceMgr->AddNewSourceBuffer(std::move(file), SMLoc());
   return parseSourceFileForTool(sourceMgr, &context, insertImplictModule);
 }
 
 LogicalResult mlir::mlirReduceMain(int argc, char **argv,
                                    MLIRContext &context) {
   // Override the default '-h' and use the default PrintHelpMessage() which
   // won't print options in categories.
   static llvm::cl::opt<bool> help("h", llvm::cl::desc("Alias for -help"),

mlir/lib/Tools/mlir-translate/MlirTranslateMain.cpp

LogicalResult mlir::mlirTranslateMain(int argc, char **argv,
   }
 
   // Processes the memory buffer with a new MLIRContext.
   auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
                            raw_ostream &os) {
     MLIRContext context;
     context.allowUnregisteredDialects(allowUnregisteredDialects);
     context.printOpOnDiagnostic(!verifyDiagnostics);
-    llvm::SourceMgr sourceMgr;
-    sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
+    auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+    sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
     if (!verifyDiagnostics) {
-      SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
+      SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
       return (*translationRequested)(sourceMgr, os, &context);
     }
 
     // In the diagnostic verification flow, we ignore whether the translation
     // failed (in most cases, it is expected to fail). Instead, we check if the
     // diagnostics were produced as expected.
-    SourceMgrDiagnosticVerifierHandler sourceMgrHandler(sourceMgr, &context);
+    SourceMgrDiagnosticVerifierHandler sourceMgrHandler(*sourceMgr, &context);
     (void)(*translationRequested)(sourceMgr, os, &context);
     return sourceMgrHandler.verify();
   };
 
   if (failed(splitAndProcessBuffer(std::move(input), processBuffer,
                                    output->os(), splitInputFile)))
     return failure();
 
   output->keep();
   return success();
 }

mlir/lib/Tools/mlir-translate/Translation.cpp

 // Translation to MLIR
 //===----------------------------------------------------------------------===//
 
 // Puts `function` into the to-MLIR translation registry unless there is already
 // a function registered for the same name.
 static void registerTranslateToMLIRFunction(
     StringRef name, StringRef description, Optional<llvm::Align> inputAlignment,
     const TranslateSourceMgrToMLIRFunction &function) {
-  auto wrappedFn = [function](llvm::SourceMgr &sourceMgr, raw_ostream &output,
-                              MLIRContext *context) {
+  auto wrappedFn = [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                              raw_ostream &output, MLIRContext *context) {
     OwningOpRef<Operation *> op = function(sourceMgr, context);
     if (!op || failed(verify(*op)))
       return failure();
     op.get()->print(output);
     return success();
   };
   registerTranslation(name, description, inputAlignment, wrappedFn);
 }
 
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, StringRef description,
     const TranslateSourceMgrToMLIRFunction &function,
     Optional<llvm::Align> inputAlignment) {
   registerTranslateToMLIRFunction(name, description, inputAlignment, function);
 }
+TranslateToMLIRRegistration::TranslateToMLIRRegistration(
+    StringRef name, StringRef description,
+    const TranslateRawSourceMgrToMLIRFunction &function,
+    Optional<llvm::Align> inputAlignment) {
+  registerTranslateToMLIRFunction(
+      name, description, inputAlignment,
+      [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                 MLIRContext *ctx) { return function(*sourceMgr, ctx); });
+}
 /// Wraps `function` with a lambda that extracts a StringRef from a source
 /// manager and registers the wrapper lambda as a to-MLIR conversion.
 TranslateToMLIRRegistration::TranslateToMLIRRegistration(
     StringRef name, StringRef description,
     const TranslateStringRefToMLIRFunction &function,
     Optional<llvm::Align> inputAlignment) {
   registerTranslateToMLIRFunction(
       name, description, inputAlignment,
-      [function](llvm::SourceMgr &sourceMgr, MLIRContext *ctx) {
+      [function](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                 MLIRContext *ctx) {
         const llvm::MemoryBuffer *buffer =
-            sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
+            sourceMgr->getMemoryBuffer(sourceMgr->getMainFileID());
         return function(buffer->getBuffer(), ctx);
       });
 }
 
 //===----------------------------------------------------------------------===//
 // Translation from MLIR
 //===----------------------------------------------------------------------===//
 
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, StringRef description,
     const TranslateFromMLIRFunction &function,
     const std::function<void(DialectRegistry &)> &dialectRegistration) {
   registerTranslation(
       name, description, /*inputAlignment=*/std::nullopt,
-      [function, dialectRegistration](llvm::SourceMgr &sourceMgr,
-                                      raw_ostream &output,
-                                      MLIRContext *context) {
+      [function,
+       dialectRegistration](const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
+                            raw_ostream &output, MLIRContext *context) {
         DialectRegistry registry;
         dialectRegistration(registry);
         context->appendDialectRegistry(registry);
         bool implicitModule =
             (!clOptions.isConstructed() || !clOptions->noImplicitModule);
         OwningOpRef<Operation *> op =
             parseSourceFileForTool(sourceMgr, context, implicitModule);
         if (!op || failed(verify(*op)))