Define a ModuleOp that represents a Module as an Operation.

The ModuleOp contains a single region that must contain a single block. This block must be terminated by a new pseudo operation 'module_terminator'. The syntax for this operations is as follows:

  `module` (`attributes` attr-dict)? region

Example:

  module {
    ...
  }

  module attributes { ... } {
    ...
  }

PiperOrigin-RevId: 254513752

mlir/include/mlir/IR/Module.h

@@ -84,6 +84,50 @@ private:
   /// This is the actual list of functions the module contains.
   FunctionListType functions;
 };
+
+//===--------------------------------------------------------------------===//
+// Module Operation.
+//===--------------------------------------------------------------------===//
+
+/// ModuleOp represents a module, or an operation containing one region with a
+/// single block containing opaque operations. A ModuleOp contains a symbol
+/// table for operations, like FuncOp, held within its region. The region of a
+/// module is not allowed to implicitly capture global values, and all external
+/// references must use attributes.
+class ModuleOp : public Op<ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+                           OpTrait::IsIsolatedFromAbove> {
+public:
+  using Op::Op;
+  static StringRef getOperationName() { return "module"; }
+
+  static void build(Builder *builder, OperationState *result);
+
+  /// Operation hooks.
+  static ParseResult parse(OpAsmParser *parser, OperationState *result);
+  void print(OpAsmPrinter *p);
+  LogicalResult verify();
+
+  /// Return body of this module.
+  Block *getBody();
+};
+
+/// The ModuleTerminatorOp is a special terminator operation for the body of a
+/// ModuleOp, it has no semantic meaning beyond keeping the body of a ModuleOp
+/// well-formed.
+///
+/// This operation does _not_ have a custom syntax. However, ModuleOp will omit
+/// the terminator in their custom syntax for brevity.
+class ModuleTerminatorOp
+    : public Op<ModuleTerminatorOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+                OpTrait::IsTerminator> {
+public:
+  using Op::Op;
+  static StringRef getOperationName() { return "module_terminator"; }
+
+  static void build(Builder *, OperationState *) {}
+  LogicalResult verify();
+};
+
 } // end namespace mlir
 
 #endif // MLIR_IR_MODULE_H

mlir/include/mlir/IR/OpDefinition.h

@@ -881,6 +881,27 @@ ParseResult parseCastOp(OpAsmParser *parser, OperationState *result);
 void printCastOp(Operation *op, OpAsmPrinter *p);
 Value *foldCastOp(Operation *op);
 } // namespace impl
+
+// These functions are out-of-line utilities, which avoids them being template
+// instantiated/duplicated.
+namespace impl {
+/// Insert an operation, generated by `buildTerminatorOp`, at the end of the
+/// region's only block if it does not have a terminator already. If the region
+/// is empty, insert a new block first. `buildTerminatorOp` should return the
+/// terminator operation to insert.
+void ensureRegionTerminator(
+    Region &region, Location loc,
+    llvm::function_ref<Operation *()> buildTerminatorOp);
+/// Templated version that fills the generates the provided operation type.
+template <typename OpTy>
+void ensureRegionTerminator(Region &region, Builder &builder, Location loc) {
+  ensureRegionTerminator(region, loc, [&] {
+    OperationState state(loc->getContext(), loc, OpTy::getOperationName());
+    OpTy::build(&builder, &state);
+    return Operation::create(state);
+  });
+}
+} // namespace impl
 } // end namespace mlir
 
 #endif

mlir/lib/AffineOps/AffineOps.cpp

@@ -717,17 +717,7 @@ static LogicalResult checkHasAffineTerminator(OpState &op, Block &block) {
 // first.
 static void ensureAffineTerminator(Region &region, Builder &builder,
                                    Location loc) {
-  if (region.empty())
-    region.push_back(new Block);
-
-  Block &block = region.back();
-  if (!block.empty() && block.back().isKnownTerminator())
-    return;
-
-  OperationState terminatorState(builder.getContext(), loc,
-                                 AffineTerminatorOp::getOperationName());
-  AffineTerminatorOp::build(&builder, &terminatorState);
-  block.push_back(Operation::create(terminatorState));
+  impl::ensureRegionTerminator<AffineTerminatorOp>(region, builder, loc);
 }
 
 void AffineForOp::build(Builder *builder, OperationState *result,

mlir/lib/IR/MLIRContext.cpp

@@ -27,10 +27,10 @@
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
-#include "mlir/IR/Function.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
+#include "mlir/IR/Module.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/DenseSet.h"
@@ -90,9 +90,9 @@ struct BuiltinDialect : public Dialect {
              MemRefType, NoneType, OpaqueType, RankedTensorType, TupleType,
              UnrankedTensorType, VectorType>();
 
-    // TODO: FuncOp should be moved to a different dialect when it has been
-    // fully decoupled from the core.
-    addOperations<FuncOp>();
+    // TODO: These operations should be moved to a different dialect when they
+    // have been fully decoupled from the core.
+    addOperations<FuncOp, ModuleOp, ModuleTerminatorOp>();
   }
 };
 

mlir/lib/IR/Module.cpp

@@ -0,0 +1,100 @@
+//===- Module.cpp - MLIR Module Operation ---------------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "mlir/IR/Module.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Module Operation.
+//===----------------------------------------------------------------------===//
+
+// Insert `module_terminator` at the end of the region's only block if it does
+// not have a terminator already. If the region is empty, insert a new block
+// first.
+static void ensureModuleTerminator(Region &region, Builder &builder,
+                                   Location loc) {
+  impl::ensureRegionTerminator<ModuleTerminatorOp>(region, builder, loc);
+}
+
+void ModuleOp::build(Builder *builder, OperationState *result) {
+  ensureModuleTerminator(*result->addRegion(), *builder, result->location);
+}
+
+ParseResult ModuleOp::parse(OpAsmParser *parser, OperationState *result) {
+  // If module attributes are present, parse them.
+  if (succeeded(parser->parseOptionalKeyword("attributes")))
+    if (parser->parseOptionalAttributeDict(result->attributes))
+      return failure();
+
+  // Parse the module body.
+  auto *body = result->addRegion();
+  if (parser->parseRegion(*body, llvm::None, llvm::None))
+    return failure();
+
+  // Ensure that this module has a valid terminator.
+  ensureModuleTerminator(*body, parser->getBuilder(), result->location);
+  return success();
+}
+
+void ModuleOp::print(OpAsmPrinter *p) {
+  *p << "module";
+
+  // Print the module attributes.
+  auto attrs = getAttrs();
+  if (!attrs.empty()) {
+    *p << " attributes";
+    p->printOptionalAttrDict(attrs, {});
+  }
+
+  // Print the region.
+  p->printRegion(getOperation()->getRegion(0), /*printEntryBlockArgs=*/false,
+                 /*printBlockTerminators=*/false);
+}
+
+LogicalResult ModuleOp::verify() {
+  auto &bodyRegion = getOperation()->getRegion(0);
+
+  // The body must contain a single basic block.
+  if (bodyRegion.empty() || std::next(bodyRegion.begin()) != bodyRegion.end())
+    return emitOpError("expected body region to have a single block");
+
+  // Check that the body has no block arguments.
+  auto *body = &bodyRegion.front();
+  if (body->getNumArguments() != 0)
+    return emitOpError("expected body to have no arguments");
+
+  if (body->empty() || !isa<ModuleTerminatorOp>(body->back())) {
+    emitOpError("expects region to end with '" +
+                ModuleTerminatorOp::getOperationName() + "'")
+            .attachNote()
+        << "in custom textual format, the absence of terminator implies '"
+        << ModuleTerminatorOp::getOperationName() << "'";
+    return failure();
+  }
+
+  return success();
+}
+
+LogicalResult ModuleTerminatorOp::verify() {
+  if (!isa_and_nonnull<ModuleOp>(getOperation()->getParentOp()))
+    return emitOpError() << "is expected to terminate a '"
+                         << ModuleOp::getOperationName() << "' operation";
+  return success();
+}

mlir/lib/IR/Operation.cpp

@@ -1003,3 +1003,24 @@ Value *impl::foldCastOp(Operation *op) {
     return op->getOperand(0);
   return nullptr;
 }
+
+//===----------------------------------------------------------------------===//
+// CastOp implementation
+//===----------------------------------------------------------------------===//
+
+/// Insert an operation, generated by `buildTerminatorOp`, at the end of the
+/// region's only block if it does not have a terminator already. If the region
+/// is empty, insert a new block first. `buildTerminatorOp` should return the
+/// terminator operation to insert.
+void impl::ensureRegionTerminator(
+    Region &region, Location loc,
+    llvm::function_ref<Operation *()> buildTerminatorOp) {
+  if (region.empty())
+    region.push_back(new Block);
+
+  Block &block = region.back();
+  if (!block.empty() && block.back().isKnownTerminator())
+    return;
+
+  block.push_back(buildTerminatorOp());
+}

mlir/lib/Linalg/IR/LinalgOps.cpp

@@ -128,17 +128,7 @@ static LogicalResult checkHasTerminator(OpState &op, Block &block) {
 // inserted, the location is specified by `loc`. If the region is empty, insert
 // a new block first.
 static void ensureTerminator(Region &region, Builder &builder, Location loc) {
-  if (region.empty())
-    region.push_back(new Block);
-
-  Block &block = region.back();
-  if (!block.empty() && block.back().isKnownTerminator())
-    return;
-
-  OperationState terminatorState(builder.getContext(), loc,
-                                 TerminatorOp::getOperationName());
-  TerminatorOp::build(&builder, &terminatorState);
-  block.push_back(Operation::create(terminatorState));
+  impl::ensureRegionTerminator<TerminatorOp>(region, builder, loc);
 }
 
 void mlir::linalg::ForOp::build(Builder *builder, OperationState *result,

mlir/lib/SPIRV/SPIRVOps.cpp

@@ -126,17 +126,7 @@ static LogicalResult verify(spirv::ConstantOp constOp) {
 //===----------------------------------------------------------------------===//
 
 static void ensureModuleEnd(Region *region, Builder builder, Location loc) {
-  if (region->empty())
-    region->push_back(new Block);
-
-  Block &block = region->back();
-  if (!block.empty() && llvm::isa<spirv::ModuleEndOp>(block.back()))
-    return;
-
-  OperationState state(builder.getContext(), loc,
-                       spirv::ModuleEndOp::getOperationName());
-  spirv::ModuleEndOp::build(&builder, &state);
-  block.push_back(Operation::create(state));
+  impl::ensureRegionTerminator<spirv::ModuleEndOp>(*region, builder, loc);
 }
 
 void spirv::ModuleOp::build(Builder *builder, OperationState *state) {

mlir/test/IR/invalid-module-op.mlir

@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// -----
+
+func @module_op() {
+  // expected-error@+1 {{expected body region to have a single block}}
+  module {
+  ^bb1:
+    "module_terminator"() : () -> ()
+  ^bb2:
+    "module_terminator"() : () -> ()
+  }
+  return
+}
+
+// -----
+
+func @module_op() {
+  // expected-error@+1 {{expected body to have no arguments}}
+  module {
+  ^bb1(%arg: i32):
+    "module_terminator"() : () -> ()
+  }
+  return
+}
+
+// -----
+
+func @module_op() {
+  // expected-error@+2 {{expects region to end with 'module_terminator'}}
+  // expected-note@+1 {{the absence of terminator implies 'module_terminator'}}
+  module {
+    return
+  }
+  return
+}
+
+// -----
+
+func @module_op() {
+  // expected-error@+1 {{is expected to terminate a 'module' operation}}
+  "module_terminator"() : () -> ()
+}

mlir/test/IR/module-op.mlir

@@ -0,0 +1,20 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK-LABEL: func @module_parsing
+func @module_parsing() {
+  // CHECK-NEXT: module {
+  module {
+  }
+
+  // CHECK: module {
+  // CHECK-NEXT: }
+  module {
+    "module_terminator"() : () -> ()
+  }
+
+  // CHECK: module attributes {foo.attr: true} {
+  module attributes {foo.attr: true} {
+  }
+
+  return
+}