[mlir][Parser] Don't hardcode the use of ModuleOp in the parser

[rriddle]

Summary

This was important when ModuleOp was the only top level operation, but that isn't necessarily the case anymore. This is one of the last remaining aspects of the infrastructure that is hardcoded to ModuleOp.

Depends On D92450

[mehdi_amini]

Sweet!

mlir/include/mlir/Dialect/SPIRV/SPIRVModule.h

25

Could this be using OwningSPIRVModuleRef = OwningOpRef<spirv::ModuleOp>; ? (same for OwningModuleRef)

mlir/lib/Parser/Parser.cpp

2035

Does this assume that topLevelBlock isn't empty?

[rriddle]

mlir/include/mlir/Dialect/SPIRV/SPIRVModule.h

25

It can and should, but it involves updating each of the usages that forward declare the class. I'd rather keep that as a simple NFC followup to keep this revision as focused as possible.

mlir/lib/Parser/Parser.cpp

2035

Technically it's fine given that iplist is cyclic(previous node of the sentinel is the sentinel if empty), but switched to explicit check for empty anyways.

mlir/include/mlir/Dialect/SPIRV/SPIRVModule.h

 //===- SPIRVModule.h - SPIR-V Module Utilities ------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_SPIRVMODULE_H
 #define MLIR_DIALECT_SPIRV_SPIRVMODULE_H
 
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
-#include "mlir/IR/OwningOpRefBase.h"
+#include "mlir/IR/OwningOpRef.h"
 
 namespace mlir {
 namespace spirv {
 
 /// This class acts as an owning reference to a SPIR-V module, and will
 /// automatically destroy the held module on destruction if the held module
 /// is valid.
-class OwningSPIRVModuleRef : public OwningOpRefBase<spirv::ModuleOp> {
+// TODO: Remove this class in favor of using OwningOpRef directly.
+class OwningSPIRVModuleRef : public OwningOpRef<spirv::ModuleOp> {
 public:
-  using OwningOpRefBase<spirv::ModuleOp>::OwningOpRefBase;
+  using OwningOpRef<spirv::ModuleOp>::OwningOpRef;
 };
 
 } // end namespace spirv
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVMODULE_H

mlir/include/mlir/IR/BuiltinOps.h

 //===- BuiltinOps.h - MLIR Builtin Operations -------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the Builtin dialect's operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINOPS_H_
 #define MLIR_IR_BUILTINOPS_H_
 
 #include "mlir/IR/FunctionSupport.h"
-#include "mlir/IR/OwningOpRefBase.h"
+#include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 //===----------------------------------------------------------------------===//
 // Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/IR/BuiltinOps.h.inc"
 
 //===----------------------------------------------------------------------===//
 // Dialect Utilities
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 /// This class acts as an owning reference to a module, and will automatically
 /// destroy the held module on destruction if the held module is valid.
-class OwningModuleRef : public OwningOpRefBase<ModuleOp> {
+// TODO: Remove this class in favor of using OwningOpRef directly.
+class OwningModuleRef : public OwningOpRef<ModuleOp> {
 public:
-  using OwningOpRefBase<ModuleOp>::OwningOpRefBase;
+  using OwningOpRef<ModuleOp>::OwningOpRef;
+  OwningModuleRef(OwningOpRef<ModuleOp> &&other)
+      : OwningOpRef<ModuleOp>(std::move(other)) {}
 };
 } // end namespace mlir
 
 namespace llvm {
 // Functions hash just like pointers.
 template <>
 struct DenseMapInfo<mlir::FuncOp> {
   static mlir::FuncOp getEmptyKey() {

mlir/include/mlir/IR/OpDefinition.h

private:
     /// cyclic header inclusion.
     static Operation *buildTerminator(OpBuilder &builder, Location loc) {
       OperationState state(loc, TerminatorOpType::getOperationName());
       TerminatorOpType::build(builder, state);
       return Operation::create(state);
     }
 
   public:
+    /// The type of the operation used as the implicit terminator type.
+    using ImplicitTerminatorOpT = TerminatorOpType;
+
     static LogicalResult verifyTrait(Operation *op) {
       for (unsigned i = 0, e = op->getNumRegions(); i < e; ++i) {
         Region &region = op->getRegion(i);
 
         // Empty regions are fine.
         if (region.empty())
           continue;
 

mlir/include/mlir/IR/OwningOpRef.h

-//===- OwningOpRefBase.h - MLIR OwningOpRefBase -----------------*- C++ -*-===//
+//===- OwningOpRef.h - MLIR OwningOpRef -------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file provides a base class for owning op refs.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_IR_OWNINGOPREFBASE_H
-#define MLIR_IR_OWNINGOPREFBASE_H
+#ifndef MLIR_IR_OWNINGOPREF_H
+#define MLIR_IR_OWNINGOPREF_H
 
 #include <utility>
 
 namespace mlir {
 
 /// This class acts as an owning reference to an op, and will automatically
 /// destroy the held op on destruction if the held op is valid.
 ///
 /// Note that OpBuilder and related functionality should be highly preferred
 /// instead, and this should only be used in situations where existing solutions
 /// are not viable.
 template <typename OpTy>
-class OwningOpRefBase {
+class OwningOpRef {
 public:
-  OwningOpRefBase(std::nullptr_t = nullptr) {}
-  OwningOpRefBase(OpTy op) : op(op) {}
-  OwningOpRefBase(OwningOpRefBase &&other) : op(other.release()) {}
-  ~OwningOpRefBase() {
+  /// The underlying operation type stored in this reference.
+  using OperationT = OpTy;
+
+  OwningOpRef(std::nullptr_t = nullptr) {}
+  OwningOpRef(OpTy op) : op(op) {}
+  OwningOpRef(OwningOpRef &&other) : op(other.release()) {}
+  ~OwningOpRef() {
     if (op)
-      op.erase();
+      op->erase();
   }
 
-  // Assign from another op reference.
-  OwningOpRefBase &operator=(OwningOpRefBase &&other) {
+  /// Assign from another op reference.
+  OwningOpRef &operator=(OwningOpRef &&other) {
     if (op)
-      op.erase();
+      op->erase();
     op = other.release();
     return *this;
   }
 
   /// Allow accessing the internal op.
   OpTy get() const { return op; }
   OpTy operator*() const { return op; }
   OpTy *operator->() { return &op; }
   explicit operator bool() const { return op; }
 
   /// Release the referenced op.
   OpTy release() {
     OpTy released;
     std::swap(released, op);
     return released;
   }
 
 private:
   OpTy op;
 };
 
 } // end namespace mlir
 
-#endif // MLIR_IR_OWNINGOPREFBASE_H
+#endif // MLIR_IR_OWNINGOPREF_H

mlir/include/mlir/Parser.h

 //===- Parser.h - MLIR Parser Library Interface -----------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is contains the interface to the MLIR parser library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PARSER_H
 #define MLIR_PARSER_H
 
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include <cstddef>
 
 namespace llvm {
 class SourceMgr;
 class SMDiagnostic;
 class StringRef;
 } // end namespace llvm
 
 namespace mlir {
-class Attribute;
-class Location;
-class MLIRContext;
-class OwningModuleRef;
-class Type;
-
-/// This parses the file specified by the indicated SourceMgr and returns an
-/// MLIR module if it was valid.  If not, the error message is emitted through
-/// the error handler registered in the context, and a null pointer is returned.
-OwningModuleRef parseSourceFile(const llvm::SourceMgr &sourceMgr,
-                                MLIRContext *context);
-
-/// This parses the file specified by the indicated filename and returns an
-/// MLIR module if it was valid.  If not, the error message is emitted through
-/// the error handler registered in the context, and a null pointer is returned.
-OwningModuleRef parseSourceFile(llvm::StringRef filename, MLIRContext *context);
+namespace detail {
+/// Given a block containing operations that have just been parsed, if the block
+/// contains a single operation of `ContainerOpT` type then remove it from the
+/// block and return it. If the block does not contain just that operation,
+/// create a new operation instance of `ContainerOpT` and move all of the
+/// operations within `parsedBlock` into the first block of the first region.
+/// `ContainerOpT` is required to have a single region containing a single
+/// block, and must implement the `SingleBlockImplicitTerminator` trait.
+template <typename ContainerOpT>
+inline OwningOpRef<ContainerOpT> constructContainerOpForParserIfNecessary(
+    Block *parsedBlock, MLIRContext *context, Location sourceFileLoc) {
+  static_assert(
+      ContainerOpT::template hasTrait<OpTrait::OneRegion>() &&
+          std::is_base_of<typename OpTrait::SingleBlockImplicitTerminator<
+                              typename ContainerOpT::ImplicitTerminatorOpT>::
+                              template Impl<ContainerOpT>,
+                          ContainerOpT>::value,
+      "Expected `ContainerOpT` to have a single region with a single "
+      "block that has an implicit terminator");
+
+  // Check to see if we parsed a single instance of this operation.
+  if (llvm::hasSingleElement(*parsedBlock)) {
+    if (ContainerOpT op = dyn_cast<ContainerOpT>(parsedBlock->front())) {
+      op->remove();
+      return op;
+    }
+  }
+
+  // If not, then build a new one to contain the parsed operations.
+  OpBuilder builder(context);
+  ContainerOpT op = builder.create<ContainerOpT>(sourceFileLoc);
+  OwningOpRef<ContainerOpT> opRef(op);
+  assert(op->getNumRegions() == 1 && llvm::hasSingleElement(op->getRegion(0)) &&
+         "expected generated operation to have a single region with a single "
+         "block");
+  Block *opBlock = &op->getRegion(0).front();
+  opBlock->getOperations().splice(opBlock->begin(),
+                                  parsedBlock->getOperations());
+
+  // After splicing, verify just this operation to ensure it can properly
+  // contain the operations inside of it.
+  if (failed(op.verify()))
+    return OwningOpRef<ContainerOpT>();
+  return std::move(opRef);
+}
+} // end namespace detail
+
+/// This parses the file specified by the indicated 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(const llvm::SourceMgr &sourceMgr, Block *block,
+                              MLIRContext *context,
+                              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,
+                              MLIRContext *context,
+                              LocationAttr *sourceFileLoc = nullptr);
 
 /// This parses the file specified by the indicated filename using the provided
-/// SourceMgr and returns an MLIR module if it was valid.  If not, the error
+/// 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,
+                              MLIRContext *context,
+                              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.
+LogicalResult parseSourceString(llvm::StringRef sourceStr, Block *block,
+                                MLIRContext *context,
+                                LocationAttr *sourceFileLoc = nullptr);
+
+/// This parses the file specified by the indicated SourceMgr. 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
-/// a null pointer is returned.
-OwningModuleRef parseSourceFile(llvm::StringRef filename,
-                                llvm::SourceMgr &sourceMgr,
-                                MLIRContext *context);
+/// failure is returned. `ContainerOpT` is required to have a single region
+/// containing a single block, and must implement the
+/// `SingleBlockImplicitTerminator` trait.
+template <typename ContainerOpT>
+inline OwningOpRef<ContainerOpT>
+parseSourceFile(const llvm::SourceMgr &sourceMgr, MLIRContext *context) {
+  LocationAttr sourceFileLoc;
+  Block block;
+  if (failed(parseSourceFile(sourceMgr, &block, context, &sourceFileLoc)))
+    return OwningOpRef<ContainerOpT>();
+  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
+      &block, context, sourceFileLoc);
+}
+
+/// 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
+/// `SingleBlockImplicitTerminator` trait.
+template <typename ContainerOpT>
+inline OwningOpRef<ContainerOpT> parseSourceFile(llvm::StringRef filename,
+                                                 MLIRContext *context) {
+  LocationAttr sourceFileLoc;
+  Block block;
+  if (failed(parseSourceFile(filename, &block, context, &sourceFileLoc)))
+    return OwningOpRef<ContainerOpT>();
+  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
+      &block, context, sourceFileLoc);
+}
 
-/// This parses the module string to a MLIR module if it was valid.  If not, the
-/// error message is emitted through the error handler registered in the
-/// context, and a null pointer is returned.
-OwningModuleRef parseSourceString(llvm::StringRef moduleStr,
-                                  MLIRContext *context);
+/// This parses the file specified by the indicated filename using the provided
+/// SourceMgr. 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 `SingleBlockImplicitTerminator` trait.
+template <typename ContainerOpT>
+inline OwningOpRef<ContainerOpT> parseSourceFile(llvm::StringRef filename,
+                                                 llvm::SourceMgr &sourceMgr,
+                                                 MLIRContext *context) {
+  LocationAttr sourceFileLoc;
+  Block block;
+  if (failed(parseSourceFile(filename, sourceMgr, &block, context,
+                             &sourceFileLoc)))
+    return OwningOpRef<ContainerOpT>();
+  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
+      &block, context, sourceFileLoc);
+}
+
+/// 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
+/// `SingleBlockImplicitTerminator` trait.
+template <typename ContainerOpT>
+inline OwningOpRef<ContainerOpT> parseSourceString(llvm::StringRef sourceStr,
+                                                   MLIRContext *context) {
+  LocationAttr sourceFileLoc;
+  Block block;
+  if (failed(parseSourceString(sourceStr, &block, context, &sourceFileLoc)))
+    return OwningOpRef<ContainerOpT>();
+  return detail::constructContainerOpForParserIfNecessary<ContainerOpT>(
+      &block, context, sourceFileLoc);
+}
+
+/// TODO: These methods are deprecated in favor of the above template versions.
+/// They should be removed when usages have been updated.
+inline OwningModuleRef parseSourceFile(const llvm::SourceMgr &sourceMgr,
+                                       MLIRContext *context) {
+  return parseSourceFile<ModuleOp>(sourceMgr, context);
+}
+inline OwningModuleRef parseSourceFile(llvm::StringRef filename,
+                                       MLIRContext *context) {
+  return parseSourceFile<ModuleOp>(filename, context);
+}
+inline OwningModuleRef parseSourceFile(llvm::StringRef filename,
+                                       llvm::SourceMgr &sourceMgr,
+                                       MLIRContext *context) {
+  return parseSourceFile<ModuleOp>(filename, sourceMgr, context);
+}
+inline OwningModuleRef parseSourceString(llvm::StringRef moduleStr,
+                                         MLIRContext *context) {
+  return parseSourceString<ModuleOp>(moduleStr, context);
+}
 
 /// This parses a single MLIR attribute to an MLIR context if it was valid.  If
 /// not, an error message is emitted through a new SourceMgrDiagnosticHandler
 /// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
 /// `attrStr`. If the passed `attrStr` has additional tokens that were not part
 /// of the type, an error is emitted.
 // TODO: Improve diagnostic reporting.
 Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context);

mlir/lib/Parser/Parser.cpp

 // OperationParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {
 public:
-  OperationParser(ParserState &state, ModuleOp moduleOp)
-      : Parser(state), opBuilder(moduleOp.getBodyRegion()), moduleOp(moduleOp) {
+  OperationParser(ParserState &state, Operation *topLevelOp)
+      : Parser(state), opBuilder(topLevelOp->getRegion(0)),
+        topLevelOp(topLevelOp) {
+    // The top level operation starts a new name scope.
+    pushSSANameScope(/*isIsolated=*/true);
   }
 
   ~OperationParser();
 
   /// After parsing is finished, this function must be called to see if there
   /// are any remaining issues.
   ParseResult finalize();
 
private:
 
   /// A set of operations whose locations reference aliases that have yet to
   /// be resolved.
   SmallVector<std::pair<Operation *, Token>, 8> opsWithDeferredLocs;
 
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
-  /// The top level module operation.
-  ModuleOp moduleOp;
+  /// The top level operation that holds all of the parsed operations.
+  Operation *topLevelOp;
 };
 } // end anonymous namespace
 
 OperationParser::~OperationParser() {
   for (auto &fwd : forwardRefPlaceholders) {
     // Drop all uses of undefined forward declared reference and destroy
     // defining operation.
     fwd.first.dropAllUses();
for (std::pair<Operation *, Token> &it : opsWithDeferredLocs) {
 
     LocationAttr locAttr = attr.dyn_cast<LocationAttr>();
     if (!locAttr)
       return emitError(tokLoc)
              << "expected location, but found '" << attr << "'";
     it.first->setLoc(locAttr);
   }
 
-  return success();
+  // Pop the top level name scope.
+  return popSSANameScope();
 }
 
 //===----------------------------------------------------------------------===//
 // SSA Value Handling
 //===----------------------------------------------------------------------===//
 
 void OperationParser::pushSSANameScope(bool isIsolated) {
   blocksByName.push_back(DenseMap<StringRef, std::pair<Block *, SMLoc>>());
ParseResult OperationParser::popSSANameScope() {
 
   // Verify that all referenced blocks were defined.
   if (!forwardRefInCurrentScope.empty()) {
     SmallVector<std::pair<const char *, Block *>, 4> errors;
     // Iteration over the map isn't deterministic, so sort by source location.
     for (auto entry : forwardRefInCurrentScope) {
       errors.push_back({entry.second.getPointer(), entry.first});
       // Add this block to the top-level region to allow for automatic cleanup.
-      moduleOp->getRegion(0).push_back(entry.first);
+      topLevelOp->getRegion(0).push_back(entry.first);
     }
     llvm::array_pod_sort(errors.begin(), errors.end());
 
     for (auto entry : errors) {
       auto loc = SMLoc::getFromPointer(entry.first);
       emitError(loc, "reference to an undefined block");
     }
     return failure();
if (getToken().is(Token::l_square)) {
     result.addSuccessors(successors);
   }
 
   // Parse the region list.
   CleanupOpStateRegions guard{result};
   if (consumeIf(Token::l_paren)) {
     do {
       // Create temporary regions with the top level region as parent.
-      result.regions.emplace_back(new Region(moduleOp));
+      result.regions.emplace_back(new Region(topLevelOp));
       if (parseRegion(*result.regions.back(), /*entryArguments=*/{}))
         return nullptr;
     } while (consumeIf(Token::comma));
     if (parseToken(Token::r_paren, "expected ')' to end region list"))
       return nullptr;
   }
 
   if (getToken().is(Token::l_brace)) {

 
 //===----------------------------------------------------------------------===//
 // Top-level entity parsing.
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This parser handles entities that are only valid at the top level of the
 /// file.
-class ModuleParser : public Parser {
+class TopLevelOperationParser : public Parser {
 public:
-  explicit ModuleParser(ParserState &state) : Parser(state) {}
+  explicit TopLevelOperationParser(ParserState &state) : Parser(state) {}
 
-  ParseResult parseModule(ModuleOp module);
+  /// Parse a set of operations into the end of the given Block.
+  ParseResult parse(Block *topLevelBlock, Location parserLoc);
 
 private:
   /// Parse an attribute alias declaration.
   ParseResult parseAttributeAliasDef();
 
   /// Parse an attribute alias declaration.
   ParseResult parseTypeAliasDef();
 };
 } // end anonymous namespace
 
 /// Parses an attribute alias declaration.
 ///
 ///   attribute-alias-def ::= '#' alias-name `=` attribute-value
 ///
-ParseResult ModuleParser::parseAttributeAliasDef() {
+ParseResult TopLevelOperationParser::parseAttributeAliasDef() {
   assert(getToken().is(Token::hash_identifier));
   StringRef aliasName = getTokenSpelling().drop_front();
 
   // Check for redefinitions.
   if (getState().symbols.attributeAliasDefinitions.count(aliasName) > 0)
     return emitError("redefinition of attribute alias id '" + aliasName + "'");
 
   // Make sure this isn't invading the dialect attribute namespace.
ParseResult TopLevelOperationParser::parseAttributeAliasDef() {
   getState().symbols.attributeAliasDefinitions[aliasName] = attr;
   return success();
 }
 
 /// Parse a type alias declaration.
 ///
 ///   type-alias-def ::= '!' alias-name `=` 'type' type
 ///
-ParseResult ModuleParser::parseTypeAliasDef() {
+ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   assert(getToken().is(Token::exclamation_identifier));
   StringRef aliasName = getTokenSpelling().drop_front();
 
   // Check for redefinitions.
   if (getState().symbols.typeAliasDefinitions.count(aliasName) > 0)
     return emitError("redefinition of type alias id '" + aliasName + "'");
 
   // Make sure this isn't invading the dialect type namespace.
ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   if (!aliasedType)
     return failure();
 
   // Register this alias with the parser state.
   getState().symbols.typeAliasDefinitions.try_emplace(aliasName, aliasedType);
   return success();
 }
 
-/// This is the top-level module parser.
-ParseResult ModuleParser::parseModule(ModuleOp module) {
-  OperationParser opParser(getState(), module);
-
-  // Module itself is a name scope.
-  opParser.pushSSANameScope(/*isIsolated=*/true);
-
+ParseResult TopLevelOperationParser::parse(Block *topLevelBlock,
+                                           Location parserLoc) {
+  // Create a top-level operation to contain the parsed state.
+  OwningOpRef<Operation *> topLevelOp(ModuleOp::create(parserLoc));
+  OperationParser opParser(getState(), topLevelOp.get());
   while (true) {
     switch (getToken().getKind()) {
     default:
       // Parse a top-level operation.
       if (opParser.parseOperation())
         return failure();
       break;
 
     // If we got to the end of the file, then we're done.
     case Token::eof: {
       if (opParser.finalize())
         return failure();
 
-      // Handle the case where the top level module was explicitly defined.
-      auto &bodyBlocks = module.getBodyRegion().getBlocks();
-      auto &operations = bodyBlocks.front().getOperations();
-      assert(!operations.empty() && "expected a valid module terminator");
-
-      // Check that the first operation is a module, and it is the only
-      // non-terminator operation.
-      ModuleOp nested = dyn_cast<ModuleOp>(operations.front());
-      if (nested && std::next(operations.begin(), 2) == operations.end()) {
-        // Merge the data of the nested module operation into 'module'.
-        module.setLoc(nested.getLoc());
-        module.setAttrs(nested->getMutableAttrDict());
-        bodyBlocks.splice(bodyBlocks.end(), nested.getBodyRegion().getBlocks());
-
-        // Erase the original module body.
-        bodyBlocks.pop_front();
-      }
+      // Verify that the parsed operations are valid.
+      if (failed(verify(topLevelOp.get())))
+        return failure();
 
-      return opParser.popSSANameScope();
+      // Splice the blocks of the parsed operation over to the provided
+      // top-level block.
+      auto &parsedOps = (*topLevelOp)->getRegion(0).front().getOperations();
+      auto &destOps = topLevelBlock->getOperations();
+      destOps.splice(destOps.empty() ? destOps.end() : std::prev(destOps.end()),
+                     parsedOps, parsedOps.begin(), std::prev(parsedOps.end()));
+      return success();
     }
 
     // If we got an error token, then the lexer already emitted an error, just
     // stop.  Someday we could introduce error recovery if there was demand
     // for it.
     case Token::error:
       return failure();
 
case Token::exclamation_identifier:
         return failure();
       break;
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 
-/// This parses the file specified by the indicated SourceMgr and returns an
-/// MLIR module if it was valid.  If not, it emits diagnostics and returns
-/// null.
-OwningModuleRef mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr,
-                                      MLIRContext *context) {
-  auto sourceBuf = sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
-
-  // This is the result module we are parsing into.
-  OwningModuleRef module(ModuleOp::create(FileLineColLoc::get(
-      sourceBuf->getBufferIdentifier(), /*line=*/0, /*column=*/0, context)));
+LogicalResult mlir::parseSourceFile(const llvm::SourceMgr &sourceMgr,
+                                    Block *block, MLIRContext *context,
+                                    LocationAttr *sourceFileLoc) {
+  const auto *sourceBuf = sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
+
+  Location parserLoc = FileLineColLoc::get(sourceBuf->getBufferIdentifier(),
+                                           /*line=*/0, /*column=*/0, context);
+  if (sourceFileLoc)
+    *sourceFileLoc = parserLoc;
 
   SymbolState aliasState;
   ParserState state(sourceMgr, context, aliasState);
-  if (ModuleParser(state).parseModule(*module))
-    return nullptr;
-
-  // Make sure the parse module has no other structural problems detected by
-  // the verifier.
-  if (failed(verify(*module)))
-    return nullptr;
-
-  return module;
+  return TopLevelOperationParser(state).parse(block, parserLoc);
 }
 
-/// This parses the file specified by the indicated filename and returns an
-/// MLIR module if it was valid.  If not, the error message is emitted through
-/// the error handler registered in the context, and a null pointer is returned.
-OwningModuleRef mlir::parseSourceFile(StringRef filename,
-                                      MLIRContext *context) {
+LogicalResult mlir::parseSourceFile(llvm::StringRef filename, Block *block,
+                                    MLIRContext *context,
+                                    LocationAttr *sourceFileLoc) {
   llvm::SourceMgr sourceMgr;
-  return parseSourceFile(filename, sourceMgr, context);
+  return parseSourceFile(filename, sourceMgr, block, context, sourceFileLoc);
 }
 
-/// This parses the file specified by the indicated filename using the provided
-/// SourceMgr and returns an MLIR module if it was valid.  If not, the error
-/// message is emitted through the error handler registered in the context, and
-/// a null pointer is returned.
-OwningModuleRef mlir::parseSourceFile(StringRef filename,
-                                      llvm::SourceMgr &sourceMgr,
-                                      MLIRContext *context) {
+LogicalResult mlir::parseSourceFile(llvm::StringRef filename,
+                                    llvm::SourceMgr &sourceMgr, Block *block,
+                                    MLIRContext *context,
+                                    LocationAttr *sourceFileLoc) {
   if (sourceMgr.getNumBuffers() != 0) {
     // TODO: Extend to support multiple buffers.
-    emitError(mlir::UnknownLoc::get(context),
+    return emitError(mlir::UnknownLoc::get(context),
                      "only main buffer parsed at the moment");
-    return nullptr;
   }
   auto file_or_err = llvm::MemoryBuffer::getFileOrSTDIN(filename);
-  if (std::error_code error = file_or_err.getError()) {
-    emitError(mlir::UnknownLoc::get(context),
+  if (std::error_code error = file_or_err.getError())
+    return emitError(mlir::UnknownLoc::get(context),
                      "could not open input file " + filename);
-    return nullptr;
-  }
 
-  // Load the MLIR module.
+  // Load the MLIR source file.
   sourceMgr.AddNewSourceBuffer(std::move(*file_or_err), llvm::SMLoc());
-  return parseSourceFile(sourceMgr, context);
+  return parseSourceFile(sourceMgr, block, context, sourceFileLoc);
 }
 
-/// This parses the program string to a MLIR module if it was valid. If not,
-/// it emits diagnostics and returns null.
-OwningModuleRef mlir::parseSourceString(StringRef moduleStr,
-                                        MLIRContext *context) {
-  auto memBuffer = MemoryBuffer::getMemBuffer(moduleStr);
+LogicalResult mlir::parseSourceString(llvm::StringRef sourceStr, Block *block,
+                                      MLIRContext *context,
+                                      LocationAttr *sourceFileLoc) {
+  auto memBuffer = MemoryBuffer::getMemBuffer(sourceStr);
   if (!memBuffer)
-    return nullptr;
+    return failure();
 
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
-  return parseSourceFile(sourceMgr, context);
+  return parseSourceFile(sourceMgr, block, context, sourceFileLoc);
 }