Split BuiltinOps out of StandardOps.

* Move Return, Constant and AffineApply out into BuiltinOps;
* BuiltinOps are always registered, while StandardOps follow the same dynamic registration;
* Kept isValidX in MLValue as we don't have a verify on AffineMap so need to keep it callable from Parser (I wanted to move it to be called in verify instead);

PiperOrigin-RevId: 216592527

20 files changed, 642 insertions, 518 deletions

diff --git a/mlir/include/mlir/IR/BuiltinOps.h b/mlir/include/mlir/IR/BuiltinOps.h
new file mode 100644
index 00000000000..68bf797a73d
--- /dev/null
+++ b/mlir/include/mlir/IR/BuiltinOps.h
@@ -0,0 +1,217 @@
+//===- BuiltinOps.h - Builtin MLIR Operations -----------------*- C++ -*-===//
+//
+// 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.
+// =============================================================================
+//
+// This file defines convenience types for working with builtin operations
+// in the MLIR instruction set.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_BUILTINOPS_H
+#define MLIR_IR_BUILTINOPS_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+class Builder;
+class MLValue;
+class OperationSet;
+
+/// The "affine_apply" operation applies an affine map to a list of operands,
+/// yielding a list of results. The operand and result list sizes must be the
+/// same. All operands and results are of type 'Index'. This operation
+/// requires a single affine map attribute named "map".
+/// For example:
+///
+///   %y = "affine_apply" (%x) { map: (d0) -> (d0 + 1) } :
+///          (index) -> (index)
+///
+/// equivalently:
+///
+///   #map42 = (d0)->(d0+1)
+///   %y = affine_apply #map42(%x)
+///
+class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands,
+                                OpTrait::VariadicResults> {
+public:
+  /// Builds an affine apply op with the specified map and operands.
+  static void build(Builder *builder, OperationState *result, AffineMap map,
+                    ArrayRef<SSAValue *> operands);
+
+  /// Returns the affine map to be applied by this operation.
+  AffineMap getAffineMap() const {
+    return getAttrOfType<AffineMapAttr>("map")->getValue();
+  }
+
+  /// Returns true if the result of this operation can be used as dimension id.
+  bool isValidDim() const;
+
+  /// Returns true if the result of this operation is a symbol.
+  bool isValidSymbol() const;
+
+  static StringRef getOperationName() { return "affine_apply"; }
+
+  // Hooks to customize behavior of this op.
+  static bool parse(OpAsmParser *parser, OperationState *result);
+  void print(OpAsmPrinter *p) const;
+  bool verify() const;
+  bool constantFold(ArrayRef<Attribute *> operands,
+                    SmallVectorImpl<Attribute *> &results,
+                    MLIRContext *context) const;
+
+private:
+  friend class Operation;
+  explicit AffineApplyOp(const Operation *state) : Op(state) {}
+};
+
+/// The "constant" operation requires a single attribute named "value".
+/// It returns its value as an SSA value.  For example:
+///
+///   %1 = "constant"(){value: 42} : i32
+///   %2 = "constant"(){value: @foo} : (f32)->f32
+///
+class ConstantOp
+    : public Op<ConstantOp, OpTrait::ZeroOperands, OpTrait::OneResult> {
+public:
+  /// Builds a constant op with the specified attribute value and result type.
+  static void build(Builder *builder, OperationState *result, Attribute *value,
+                    Type *type);
+
+  Attribute *getValue() const { return getAttr("value"); }
+
+  static StringRef getOperationName() { return "constant"; }
+
+  // Hooks to customize behavior of this op.
+  static bool parse(OpAsmParser *parser, OperationState *result);
+  void print(OpAsmPrinter *p) const;
+  bool verify() const;
+  Attribute *constantFold(ArrayRef<Attribute *> operands,
+                          MLIRContext *context) const;
+
+protected:
+  friend class Operation;
+  explicit ConstantOp(const Operation *state) : Op(state) {}
+};
+
+/// This is a refinement of the "constant" op for the case where it is
+/// returning a float value of FloatType.
+///
+///   %1 = "constant"(){value: 42.0} : bf16
+///
+class ConstantFloatOp : public ConstantOp {
+public:
+  /// Builds a constant float op producing a float of the specified type.
+  static void build(Builder *builder, OperationState *result, double value,
+                    FloatType *type);
+
+  double getValue() const {
+    return getAttrOfType<FloatAttr>("value")->getValue();
+  }
+
+  static bool isClassFor(const Operation *op);
+
+private:
+  friend class Operation;
+  explicit ConstantFloatOp(const Operation *state) : ConstantOp(state) {}
+};
+
+/// This is a refinement of the "constant" op for the case where it is
+/// returning an integer value of IntegerType.
+///
+///   %1 = "constant"(){value: 42} : i32
+///
+class ConstantIntOp : public ConstantOp {
+public:
+  /// Build a constant int op producing an integer of the specified width.
+  static void build(Builder *builder, OperationState *result, int64_t value,
+                    unsigned width);
+
+  int64_t getValue() const {
+    return getAttrOfType<IntegerAttr>("value")->getValue();
+  }
+
+  static bool isClassFor(const Operation *op);
+
+private:
+  friend class Operation;
+  explicit ConstantIntOp(const Operation *state) : ConstantOp(state) {}
+};
+
+/// This is a refinement of the "constant" op for the case where it is
+/// returning an integer value of Index type.
+///
+///   %1 = "constant"(){value: 99} : () -> index
+///
+class ConstantIndexOp : public ConstantOp {
+public:
+  /// Build a constant int op producing an index.
+  static void build(Builder *builder, OperationState *result, int64_t value);
+
+  int64_t getValue() const {
+    return getAttrOfType<IntegerAttr>("value")->getValue();
+  }
+
+  static bool isClassFor(const Operation *op);
+
+private:
+  friend class Operation;
+  explicit ConstantIndexOp(const Operation *state) : ConstantOp(state) {}
+};
+
+/// The "return" operation represents a return statement of an ML function.
+/// The operation takes variable number of operands and produces no results.
+/// The operand number and types must match the signature of the ML function
+/// that contains the operation. For example:
+///
+///   mlfunc @foo() : (i32, f8) {
+///   ...
+///   return %0, %1 : i32, f8
+///
+class ReturnOp
+    : public Op<ReturnOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
+public:
+  static StringRef getOperationName() { return "return"; }
+
+  static void build(Builder *builder, OperationState *result,
+                    ArrayRef<SSAValue *> results);
+
+  // Hooks to customize behavior of this op.
+  static bool parse(OpAsmParser *parser, OperationState *result);
+  void print(OpAsmPrinter *p) const;
+  bool verify() const;
+
+private:
+  friend class Operation;
+  explicit ReturnOp(const Operation *state) : Op(state) {}
+};
+
+// Prints dimension and symbol list.
+void printDimAndSymbolList(Operation::const_operand_iterator begin,
+                           Operation::const_operand_iterator end,
+                           unsigned numDims, OpAsmPrinter *p);
+
+// Parses dimension and symbol list and returns true if parsing failed.
+bool parseDimAndSymbolList(OpAsmParser *parser,
+                           SmallVector<SSAValue *, 4> &operands,
+                           unsigned &numDims);
+
+/// Install the builtin operations in the specified MLIRContext.
+void registerBuiltinOperations(MLIRContext *ctx);
+
+} // end namespace mlir
+
+#endif
diff --git a/mlir/include/mlir/IR/MLValue.h b/mlir/include/mlir/IR/MLValue.h
index a6c8c759ef0..1961da13d5e 100644
--- a/mlir/include/mlir/IR/MLValue.h
+++ b/mlir/include/mlir/IR/MLValue.h
@@ -45,9 +45,10 @@ using StmtOperand = IROperandImpl<MLValue, Statement>;
 /// MLValue is the base class for SSA values in ML functions.
 class MLValue : public SSAValueImpl<StmtOperand, Statement, MLValueKind> {
 public:
-  /// Returns true if this MLValue can be used as a dimension id.
+  /// Returns true if the given MLValue can be used as a dimension id.
   bool isValidDim() const;
-  /// Returns true if this MLValue can be used as a symbol.
+
+  /// Returns true if the given MLValue can be used as a symbol.
   bool isValidSymbol() const;
 
   static bool classof(const SSAValue *value) {
diff --git a/mlir/include/mlir/IR/StandardOps.h b/mlir/include/mlir/StandardOps/StandardOps.h
index 63830712622..d559703dab4 100644
--- a/mlir/include/mlir/IR/StandardOps.h
+++ b/mlir/include/mlir/StandardOps/StandardOps.h
@@ -20,8 +20,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_IR_STANDARDOPS_H
-#define MLIR_IR_STANDARDOPS_H
+#ifndef MLIR_STANDARDOPS_STANDARDOPS_H
+#define MLIR_STANDARDOPS_STANDARDOPS_H
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/OpDefinition.h"
@@ -29,6 +29,7 @@
 namespace mlir {
 class OperationSet;
 class Builder;
+class MLValue;
 
 /// The "addf" operation takes two operands and returns one result, each of
 /// these is required to be of the same type.  This type may be a floating point
@@ -68,53 +69,6 @@ private:
   explicit AddIOp(const Operation *state) : BinaryOp(state) {}
 };
 
-/// The "affine_apply" operation applies an affine map to a list of operands,
-/// yielding a list of results. The operand and result list sizes must be the
-/// same. All operands and results are of type 'Index'. This operation
-/// requires a single affine map attribute named "map".
-/// For example:
-///
-///   %y = "affine_apply" (%x) { map: (d0) -> (d0 + 1) } :
-///          (index) -> (index)
-///
-/// equivalently:
-///
-///   #map42 = (d0)->(d0+1)
-///   %y = affine_apply #map42(%x)
-///
-class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands,
-                                OpTrait::VariadicResults> {
-public:
-  /// Builds an affine apply op with the specified map and operands.
-  static void build(Builder *builder, OperationState *result, AffineMap map,
-                    ArrayRef<SSAValue *> operands);
-
-  /// Returns the affine map to be applied by this operation.
-  AffineMap getAffineMap() const {
-    return getAttrOfType<AffineMapAttr>("map")->getValue();
-  }
-
-  /// Returns true if the result of this operation can be used as dimension id.
-  bool isValidDim() const;
-
-  /// Returns true if the result of this operation is a symbol.
-  bool isValidSymbol() const;
-
-  static StringRef getOperationName() { return "affine_apply"; }
-
-  // Hooks to customize behavior of this op.
-  static bool parse(OpAsmParser *parser, OperationState *result);
-  void print(OpAsmPrinter *p) const;
-  bool verify() const;
-  bool constantFold(ArrayRef<Attribute *> operands,
-                    SmallVectorImpl<Attribute *> &results,
-                    MLIRContext *context) const;
-
-private:
-  friend class Operation;
-  explicit AffineApplyOp(const Operation *state) : Op(state) {}
-};
-
 /// The "alloc" operation allocates a region of memory, as specified by its
 /// memref type. For example:
 ///
@@ -211,100 +165,6 @@ protected:
   explicit CallIndirectOp(const Operation *state) : Op(state) {}
 };
 
-/// The "constant" operation requires a single attribute named "value".
-/// It returns its value as an SSA value.  For example:
-///
-///   %1 = "constant"(){value: 42} : i32
-///   %2 = "constant"(){value: @foo} : (f32)->f32
-///
-class ConstantOp
-    : public Op<ConstantOp, OpTrait::ZeroOperands, OpTrait::OneResult> {
-public:
-  /// Builds a constant op with the specified attribute value and result type.
-  static void build(Builder *builder, OperationState *result, Attribute *value,
-                    Type *type);
-
-  Attribute *getValue() const { return getAttr("value"); }
-
-  static StringRef getOperationName() { return "constant"; }
-
-  // Hooks to customize behavior of this op.
-  static bool parse(OpAsmParser *parser, OperationState *result);
-  void print(OpAsmPrinter *p) const;
-  bool verify() const;
-  Attribute *constantFold(ArrayRef<Attribute *> operands,
-                          MLIRContext *context) const;
-
-protected:
-  friend class Operation;
-  explicit ConstantOp(const Operation *state) : Op(state) {}
-};
-
-/// This is a refinement of the "constant" op for the case where it is
-/// returning a float value of FloatType.
-///
-///   %1 = "constant"(){value: 42.0} : bf16
-///
-class ConstantFloatOp : public ConstantOp {
-public:
-  /// Builds a constant float op producing a float of the specified type.
-  static void build(Builder *builder, OperationState *result, double value,
-                    FloatType *type);
-
-  double getValue() const {
-    return getAttrOfType<FloatAttr>("value")->getValue();
-  }
-
-  static bool isClassFor(const Operation *op);
-
-private:
-  friend class Operation;
-  explicit ConstantFloatOp(const Operation *state) : ConstantOp(state) {}
-};
-
-/// This is a refinement of the "constant" op for the case where it is
-/// returning an integer value of IntegerType.
-///
-///   %1 = "constant"(){value: 42} : i32
-///
-class ConstantIntOp : public ConstantOp {
-public:
-  /// Build a constant int op producing an integer of the specified width.
-  static void build(Builder *builder, OperationState *result, int64_t value,
-                    unsigned width);
-
-  int64_t getValue() const {
-    return getAttrOfType<IntegerAttr>("value")->getValue();
-  }
-
-  static bool isClassFor(const Operation *op);
-
-private:
-  friend class Operation;
-  explicit ConstantIntOp(const Operation *state) : ConstantOp(state) {}
-};
-
-/// This is a refinement of the "constant" op for the case where it is
-/// returning an integer value of Index type.
-///
-///   %1 = "constant"(){value: 99} : () -> index
-///
-class ConstantIndexOp : public ConstantOp {
-public:
-  /// Build a constant int op producing an index.
-  static void build(Builder *builder, OperationState *result, int64_t value);
-
-  int64_t getValue() const {
-    return getAttrOfType<IntegerAttr>("value")->getValue();
-  }
-
-  static bool isClassFor(const Operation *op);
-
-private:
-  friend class Operation;
-  explicit ConstantIndexOp(const Operation *state) : ConstantOp(state) {}
-};
-
 /// The "dealloc" operation frees the region of memory referenced by a memref
 /// which was originally created by the "alloc" operation.
 /// The "dealloc" operation should not be called on memrefs which alias an
@@ -601,33 +461,6 @@ private:
   explicit MulIOp(const Operation *state) : BinaryOp(state) {}
 };
 
-/// The "return" operation represents a return statement of an ML function.
-/// The operation takes variable number of operands and produces no results.
-/// The operand number and types must match the signature of the ML function
-/// that contains the operation. For example:
-///
-///   mlfunc @foo() : (i32, f8) {
-///   ...
-///   return %0, %1 : i32, f8
-///
-class ReturnOp
-    : public Op<ReturnOp, OpTrait::VariadicOperands, OpTrait::ZeroResult> {
-public:
-  static StringRef getOperationName() { return "return"; }
-
-  static void build(Builder *builder, OperationState *result,
-                    ArrayRef<SSAValue *> results);
-
-  // Hooks to customize behavior of this op.
-  static bool parse(OpAsmParser *parser, OperationState *result);
-  void print(OpAsmPrinter *p) const;
-  bool verify() const;
-
-private:
-  friend class Operation;
-  explicit ReturnOp(const Operation *state) : Op(state) {}
-};
-
 /// The "shape_cast" operation converts a tensor from one type to an equivalent
 /// type without changing any data elements.  The source and destination types
 /// must both be tensor types with the same element type, and the source and
@@ -742,7 +575,7 @@ private:
 };
 
 /// Install the standard operations in the specified operation set.
-void registerStandardOperations(OperationSet &opSet);
+void registerStandardOperations(MLIRContext *ctx);
 
 } // end namespace mlir
 
diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp
index 19b6638a28a..f332836180a 100644
--- a/mlir/lib/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Analysis/AffineAnalysis.cpp
@@ -23,7 +23,7 @@
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/IR/AffineExprVisitor.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Statements.h"
 
 using namespace mlir;
diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 5460678e8c1..af561ff3bab 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -23,9 +23,9 @@
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLValue.h"
-#include "mlir/IR/StandardOps.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/raw_ostream.h"
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 5263c80e8e2..5fb69e08e2b 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -23,13 +23,13 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSet.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/Statements.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "mlir/IR/Types.h"
@@ -39,6 +39,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
+
 using namespace mlir;
 
 void Identifier::print(raw_ostream &os) const { os << str(); }
diff --git a/mlir/lib/IR/BuiltinOps.cpp b/mlir/lib/IR/BuiltinOps.cpp
new file mode 100644
index 00000000000..ae28bf6a755
--- /dev/null
+++ b/mlir/lib/IR/BuiltinOps.cpp
@@ -0,0 +1,326 @@
+//===- BuiltinOps.cpp - Builtin MLIR Operations -------------------------===//
+//
+// 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/BuiltinOps.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/OperationSet.h"
+#include "mlir/IR/SSAValue.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/MathExtras.h"
+#include "mlir/Support/STLExtras.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+void mlir::printDimAndSymbolList(Operation::const_operand_iterator begin,
+                                 Operation::const_operand_iterator end,
+                                 unsigned numDims, OpAsmPrinter *p) {
+  *p << '(';
+  p->printOperands(begin, begin + numDims);
+  *p << ')';
+
+  if (begin + numDims != end) {
+    *p << '[';
+    p->printOperands(begin + numDims, end);
+    *p << ']';
+  }
+}
+
+// Parses dimension and symbol list, and sets 'numDims' to the number of
+// dimension operands parsed.
+// Returns 'false' on success and 'true' on error.
+bool mlir::parseDimAndSymbolList(OpAsmParser *parser,
+                                 SmallVector<SSAValue *, 4> &operands,
+                                 unsigned &numDims) {
+  SmallVector<OpAsmParser::OperandType, 8> opInfos;
+  if (parser->parseOperandList(opInfos, -1, OpAsmParser::Delimiter::Paren))
+    return true;
+  // Store number of dimensions for validation by caller.
+  numDims = opInfos.size();
+
+  // Parse the optional symbol operands.
+  auto *affineIntTy = parser->getBuilder().getIndexType();
+  if (parser->parseOperandList(opInfos, -1,
+                               OpAsmParser::Delimiter::OptionalSquare) ||
+      parser->resolveOperands(opInfos, affineIntTy, operands))
+    return true;
+  return false;
+}
+
+//===----------------------------------------------------------------------===//
+// AffineApplyOp
+//===----------------------------------------------------------------------===//
+
+void AffineApplyOp::build(Builder *builder, OperationState *result,
+                          AffineMap map, ArrayRef<SSAValue *> operands) {
+  result->addOperands(operands);
+  result->types.append(map.getNumResults(), builder->getIndexType());
+  result->addAttribute("map", builder->getAffineMapAttr(map));
+}
+
+bool AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) {
+  auto &builder = parser->getBuilder();
+  auto *affineIntTy = builder.getIndexType();
+
+  AffineMapAttr *mapAttr;
+  unsigned numDims;
+  if (parser->parseAttribute(mapAttr, "map", result->attributes) ||
+      parseDimAndSymbolList(parser, result->operands, numDims) ||
+      parser->parseOptionalAttributeDict(result->attributes))
+    return true;
+  auto map = mapAttr->getValue();
+
+  if (map.getNumDims() != numDims ||
+      numDims + map.getNumSymbols() != result->operands.size()) {
+    return parser->emitError(parser->getNameLoc(),
+                             "dimension or symbol index mismatch");
+  }
+
+  result->types.append(map.getNumResults(), affineIntTy);
+  return false;
+}
+
+void AffineApplyOp::print(OpAsmPrinter *p) const {
+  auto map = getAffineMap();
+  *p << "affine_apply " << map;
+  printDimAndSymbolList(operand_begin(), operand_end(), map.getNumDims(), p);
+  p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"map");
+}
+
+bool AffineApplyOp::verify() const {
+  // Check that affine map attribute was specified.
+  auto *affineMapAttr = getAttrOfType<AffineMapAttr>("map");
+  if (!affineMapAttr)
+    return emitOpError("requires an affine map");
+
+  // Check input and output dimensions match.
+  auto map = affineMapAttr->getValue();
+
+  // Verify that operand count matches affine map dimension and symbol count.
+  if (getNumOperands() != map.getNumDims() + map.getNumSymbols())
+    return emitOpError(
+        "operand count and affine map dimension and symbol count must match");
+
+  // Verify that result count matches affine map result count.
+  if (getNumResults() != map.getNumResults())
+    return emitOpError("result count and affine map result count must match");
+
+  return false;
+}
+
+// The result of the affine apply operation can be used as a dimension id if it
+// is a CFG value or if it is an MLValue, and all the operands are valid
+// dimension ids.
+bool AffineApplyOp::isValidDim() const {
+  for (auto *op : getOperands()) {
+    if (auto *v = dyn_cast<MLValue>(op))
+      if (!v->isValidDim())
+        return false;
+  }
+  return true;
+}
+
+// The result of the affine apply operation can be used as a symbol if it is
+// a CFG value or if it is an MLValue, and all the operands are symbols.
+bool AffineApplyOp::isValidSymbol() const {
+  for (auto *op : getOperands()) {
+    if (auto *v = dyn_cast<MLValue>(op))
+      if (!v->isValidSymbol())
+        return false;
+  }
+  return true;
+}
+
+bool AffineApplyOp::constantFold(ArrayRef<Attribute *> operandConstants,
+                                 SmallVectorImpl<Attribute *> &results,
+                                 MLIRContext *context) const {
+  auto map = getAffineMap();
+  if (map.constantFold(operandConstants, results))
+    return true;
+  // Return false on success.
+  return false;
+}
+
+//===----------------------------------------------------------------------===//
+// Constant*Op
+//===----------------------------------------------------------------------===//
+
+/// Builds a constant op with the specified attribute value and result type.
+void ConstantOp::build(Builder *builder, OperationState *result,
+                       Attribute *value, Type *type) {
+  result->addAttribute("value", value);
+  result->types.push_back(type);
+}
+
+void ConstantOp::print(OpAsmPrinter *p) const {
+  *p << "constant " << *getValue();
+  p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"value");
+
+  if (!isa<FunctionAttr>(getValue()))
+    *p << " : " << *getType();
+}
+
+bool ConstantOp::parse(OpAsmParser *parser, OperationState *result) {
+  Attribute *valueAttr;
+  Type *type;
+
+  if (parser->parseAttribute(valueAttr, "value", result->attributes) ||
+      parser->parseOptionalAttributeDict(result->attributes))
+    return true;
+
+  // 'constant' taking a function reference doesn't get a redundant type
+  // specifier.  The attribute itself carries it.
+  if (auto *fnAttr = dyn_cast<FunctionAttr>(valueAttr))
+    return parser->addTypeToList(fnAttr->getValue()->getType(), result->types);
+
+  return parser->parseColonType(type) ||
+         parser->addTypeToList(type, result->types);
+}
+
+/// The constant op requires an attribute, and furthermore requires that it
+/// matches the return type.
+bool ConstantOp::verify() const {
+  auto *value = getValue();
+  if (!value)
+    return emitOpError("requires a 'value' attribute");
+
+  auto *type = this->getType();
+  if (isa<IntegerType>(type) || type->isIndex()) {
+    if (!isa<IntegerAttr>(value))
+      return emitOpError(
+          "requires 'value' to be an integer for an integer result type");
+    return false;
+  }
+
+  if (isa<FloatType>(type)) {
+    if (!isa<FloatAttr>(value))
+      return emitOpError("requires 'value' to be a floating point constant");
+    return false;
+  }
+
+  if (type->isTFString()) {
+    if (!isa<StringAttr>(value))
+      return emitOpError("requires 'value' to be a string constant");
+    return false;
+  }
+
+  if (isa<FunctionType>(type)) {
+    if (!isa<FunctionAttr>(value))
+      return emitOpError("requires 'value' to be a function reference");
+    return false;
+  }
+
+  return emitOpError(
+      "requires a result type that aligns with the 'value' attribute");
+}
+
+Attribute *ConstantOp::constantFold(ArrayRef<Attribute *> operands,
+                                    MLIRContext *context) const {
+  assert(operands.empty() && "constant has no operands");
+  return getValue();
+}
+
+void ConstantFloatOp::build(Builder *builder, OperationState *result,
+                            double value, FloatType *type) {
+  ConstantOp::build(builder, result, builder->getFloatAttr(value), type);
+}
+
+bool ConstantFloatOp::isClassFor(const Operation *op) {
+  return ConstantOp::isClassFor(op) &&
+         isa<FloatType>(op->getResult(0)->getType());
+}
+
+/// ConstantIntOp only matches values whose result type is an IntegerType.
+bool ConstantIntOp::isClassFor(const Operation *op) {
+  return ConstantOp::isClassFor(op) &&
+         isa<IntegerType>(op->getResult(0)->getType());
+}
+
+void ConstantIntOp::build(Builder *builder, OperationState *result,
+                          int64_t value, unsigned width) {
+  ConstantOp::build(builder, result, builder->getIntegerAttr(value),
+                    builder->getIntegerType(width));
+}
+
+/// ConstantIndexOp only matches values whose result type is Index.
+bool ConstantIndexOp::isClassFor(const Operation *op) {
+  return ConstantOp::isClassFor(op) && op->getResult(0)->getType()->isIndex();
+}
+
+void ConstantIndexOp::build(Builder *builder, OperationState *result,
+                            int64_t value) {
+  ConstantOp::build(builder, result, builder->getIntegerAttr(value),
+                    builder->getIndexType());
+}
+
+//===----------------------------------------------------------------------===//
+// ReturnOp
+//===----------------------------------------------------------------------===//
+
+void ReturnOp::build(Builder *builder, OperationState *result,
+                     ArrayRef<SSAValue *> results) {
+  result->addOperands(results);
+}
+
+bool ReturnOp::parse(OpAsmParser *parser, OperationState *result) {
+  SmallVector<OpAsmParser::OperandType, 2> opInfo;
+  SmallVector<Type *, 2> types;
+  llvm::SMLoc loc;
+  return parser->getCurrentLocation(&loc) || parser->parseOperandList(opInfo) ||
+         (!opInfo.empty() && parser->parseColonTypeList(types)) ||
+         parser->resolveOperands(opInfo, types, loc, result->operands);
+}
+
+void ReturnOp::print(OpAsmPrinter *p) const {
+  *p << "return";
+  if (getNumOperands() > 0) {
+    *p << ' ';
+    p->printOperands(operand_begin(), operand_end());
+    *p << " : ";
+    interleave(operand_begin(), operand_end(),
+               [&](const SSAValue *e) { p->printType(e->getType()); },
+               [&]() { *p << ", "; });
+  }
+}
+
+bool ReturnOp::verify() const {
+  // ReturnOp must be part of an ML function.
+  if (auto *stmt = dyn_cast<OperationStmt>(getOperation())) {
+    StmtBlock *block = stmt->getBlock();
+    if (!block || !isa<MLFunction>(block) || &block->back() != stmt)
+      return emitOpError("must be the last statement in the ML function");
+
+    // Return success. Checking that operand types match those in the function
+    // signature is performed in the ML function verifier.
+    return false;
+  }
+  return emitOpError("cannot occur in a CFG function");
+}
+
+//===----------------------------------------------------------------------===//
+// Register operations.
+//===----------------------------------------------------------------------===//
+
+/// Install the builtin operations in the specified MLIRContext..
+void mlir::registerBuiltinOperations(MLIRContext *ctx) {
+  auto &opSet = OperationSet::get(ctx);
+  opSet.addOperations<AffineApplyOp, ConstantOp, ReturnOp>(
+      /*prefix=*/"");
+}
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 9baf35463f9..f2caa15eb81 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -23,12 +23,12 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OperationSet.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
@@ -279,9 +279,7 @@ public:
       splatElementsAttrs;
 
 public:
-  MLIRContextImpl() : filenames(locationAllocator), identifiers(allocator) {
-    registerStandardOperations(operationSet);
-  }
+  MLIRContextImpl() : filenames(locationAllocator), identifiers(allocator) {}
 
   /// Copy the specified array of elements into memory managed by our bump
   /// pointer allocator.  This assumes the elements are all PODs.
@@ -294,6 +292,7 @@ public:
 } // end namespace mlir
 
 MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
+  registerBuiltinOperations(this);
   initializeAllRegisteredOps(this);
 }
 
diff --git a/mlir/lib/IR/SSAValue.cpp b/mlir/lib/IR/SSAValue.cpp
index 469825fbabc..9b40baffac2 100644
--- a/mlir/lib/IR/SSAValue.cpp
+++ b/mlir/lib/IR/SSAValue.cpp
@@ -19,8 +19,8 @@
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/Instructions.h"
 #include "mlir/IR/MLFunction.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/Statements.h"
+
 using namespace mlir;
 
 /// If this value is the result of an OperationInst, return the instruction
@@ -91,39 +91,3 @@ CFGFunction *BBArgument::getFunction() {
 MLFunction *MLValue::getFunction() {
   return cast<MLFunction>(static_cast<SSAValue *>(this)->getFunction());
 }
-
-// MLValue can be used a a dimension id if it is valid as a symbol, or
-// it is an induction variable, or it is a result of affine apply operation
-// with dimension id arguments.
-bool MLValue::isValidDim() const {
-  if (auto *stmt = getDefiningStmt()) {
-    // Top level statement or constant operation is ok.
-    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
-      return true;
-    // Affine apply operation is ok if all of its operands are ok.
-    if (auto op = stmt->getAs<AffineApplyOp>())
-      return op->isValidDim();
-    return false;
-  }
-  // This value is either a function argument or an induction variable. Both
-  // are ok.
-  return true;
-}
-
-// MLValue can be used as a symbol if it is a constant, or it is defined at
-// the top level, or it is a result of affine apply operation with symbol
-// arguments.
-bool MLValue::isValidSymbol() const {
-  if (auto *stmt = getDefiningStmt()) {
-    // Top level statement or constant operation is ok.
-    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
-      return true;
-    // Affine apply operation is ok if all of its operands are ok.
-    if (auto op = stmt->getAs<AffineApplyOp>())
-      return op->isValidSymbol();
-    return false;
-  }
-  // This value is either a function argument or an induction variable.
-  // Function argument is ok, induction variable is not.
-  return isa<MLFuncArgument>(this);
-}
diff --git a/mlir/lib/IR/Statement.cpp b/mlir/lib/IR/Statement.cpp
index e4cb3e5bd87..9dd22852bac 100644
--- a/mlir/lib/IR/Statement.cpp
+++ b/mlir/lib/IR/Statement.cpp
@@ -17,10 +17,10 @@
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/Statements.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "llvm/ADT/DenseMap.h"
@@ -92,6 +92,42 @@ const MLValue *Statement::getOperand(unsigned idx) const {
   return getStmtOperand(idx).get();
 }
 
+// MLValue can be used as a dimension id if it is valid as a symbol, or
+// it is an induction variable, or it is a result of affine apply operation
+// with dimension id arguments.
+bool MLValue::isValidDim() const {
+  if (auto *stmt = getDefiningStmt()) {
+    // Top level statement or constant operation is ok.
+    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
+      return true;
+    // Affine apply operation is ok if all of its operands are ok.
+    if (auto op = stmt->getAs<AffineApplyOp>())
+      return op->isValidDim();
+    return false;
+  }
+  // This value is either a function argument or an induction variable. Both
+  // are ok.
+  return true;
+}
+
+// MLValue can be used as a symbol if it is a constant, or it is defined at
+// the top level, or it is a result of affine apply operation with symbol
+// arguments.
+bool MLValue::isValidSymbol() const {
+  if (auto *stmt = getDefiningStmt()) {
+    // Top level statement or constant operation is ok.
+    if (stmt->getParentStmt() == nullptr || stmt->is<ConstantOp>())
+      return true;
+    // Affine apply operation is ok if all of its operands are ok.
+    if (auto op = stmt->getAs<AffineApplyOp>())
+      return op->isValidSymbol();
+    return false;
+  }
+  // This value is either a function argument or an induction variable.
+  // Function argument is ok, induction variable is not.
+  return isa<MLFuncArgument>(this);
+}
+
 void Statement::setOperand(unsigned idx, MLValue *value) {
   getStmtOperand(idx).set(value);
 }
diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 7315d019075..2777540a465 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -25,6 +25,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLFunction.h"
diff --git a/mlir/lib/StandardOps/StandardOpRegistration.cpp b/mlir/lib/StandardOps/StandardOpRegistration.cpp
new file mode 100644
index 00000000000..5806f9cdba4
--- /dev/null
+++ b/mlir/lib/StandardOps/StandardOpRegistration.cpp
@@ -0,0 +1,24 @@
+//===- StandardOpsRegistration.cpp - Register standard Op types -*- C++ -*-===//
+//
+// 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/OperationSet.h"
+#include "mlir/StandardOps/StandardOps.h"
+
+using namespace mlir;
+
+// Static initialization for standard op registration.
+static OpInitializeRegistration StandardOps(registerStandardOperations);
diff --git a/mlir/lib/IR/StandardOps.cpp b/mlir/lib/StandardOps/StandardOps.cpp
index 1099dc45ab7..1ae41094a86 100644
--- a/mlir/lib/IR/StandardOps.cpp
+++ b/mlir/lib/StandardOps/StandardOps.cpp
@@ -15,10 +15,11 @@
 // limitations under the License.
 // =============================================================================
 
-#include "mlir/IR/StandardOps.h"
+#include "mlir/StandardOps/StandardOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSet.h"
 #include "mlir/IR/SSAValue.h"
@@ -29,41 +30,6 @@
 
 using namespace mlir;
 
-static void printDimAndSymbolList(Operation::const_operand_iterator begin,
-                                  Operation::const_operand_iterator end,
-                                  unsigned numDims, OpAsmPrinter *p) {
-  *p << '(';
-  p->printOperands(begin, begin + numDims);
-  *p << ')';
-
-  if (begin + numDims != end) {
-    *p << '[';
-    p->printOperands(begin + numDims, end);
-    *p << ']';
-  }
-}
-
-// Parses dimension and symbol list, and sets 'numDims' to the number of
-// dimension operands parsed.
-// Returns 'false' on success and 'true' on error.
-static bool parseDimAndSymbolList(OpAsmParser *parser,
-                                  SmallVector<SSAValue *, 4> &operands,
-                                  unsigned &numDims) {
-  SmallVector<OpAsmParser::OperandType, 8> opInfos;
-  if (parser->parseOperandList(opInfos, -1, OpAsmParser::Delimiter::Paren))
-    return true;
-  // Store number of dimensions for validation by caller.
-  numDims = opInfos.size();
-
-  // Parse the optional symbol operands.
-  auto *affineIntTy = parser->getBuilder().getIndexType();
-  if (parser->parseOperandList(opInfos, -1,
-                               OpAsmParser::Delimiter::OptionalSquare) ||
-      parser->resolveOperands(opInfos, affineIntTy, operands))
-    return true;
-  return false;
-}
-
 //===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
@@ -97,100 +63,6 @@ Attribute *AddIOp::constantFold(ArrayRef<Attribute *> operands,
 }
 
 //===----------------------------------------------------------------------===//
-// AffineApplyOp
-//===----------------------------------------------------------------------===//
-
-void AffineApplyOp::build(Builder *builder, OperationState *result,
-                          AffineMap map, ArrayRef<SSAValue *> operands) {
-  result->addOperands(operands);
-  result->types.append(map.getNumResults(), builder->getIndexType());
-  result->addAttribute("map", builder->getAffineMapAttr(map));
-}
-
-bool AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) {
-  auto &builder = parser->getBuilder();
-  auto *affineIntTy = builder.getIndexType();
-
-  AffineMapAttr *mapAttr;
-  unsigned numDims;
-  if (parser->parseAttribute(mapAttr, "map", result->attributes) ||
-      parseDimAndSymbolList(parser, result->operands, numDims) ||
-      parser->parseOptionalAttributeDict(result->attributes))
-    return true;
-  auto map = mapAttr->getValue();
-
-  if (map.getNumDims() != numDims ||
-      numDims + map.getNumSymbols() != result->operands.size()) {
-    return parser->emitError(parser->getNameLoc(),
-                             "dimension or symbol index mismatch");
-  }
-
-  result->types.append(map.getNumResults(), affineIntTy);
-  return false;
-}
-
-void AffineApplyOp::print(OpAsmPrinter *p) const {
-  auto map = getAffineMap();
-  *p << "affine_apply " << map;
-  printDimAndSymbolList(operand_begin(), operand_end(), map.getNumDims(), p);
-  p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"map");
-}
-
-bool AffineApplyOp::verify() const {
-  // Check that affine map attribute was specified.
-  auto *affineMapAttr = getAttrOfType<AffineMapAttr>("map");
-  if (!affineMapAttr)
-    return emitOpError("requires an affine map");
-
-  // Check input and output dimensions match.
-  auto map = affineMapAttr->getValue();
-
-  // Verify that operand count matches affine map dimension and symbol count.
-  if (getNumOperands() != map.getNumDims() + map.getNumSymbols())
-    return emitOpError(
-        "operand count and affine map dimension and symbol count must match");
-
-  // Verify that result count matches affine map result count.
-  if (getNumResults() != map.getNumResults())
-    return emitOpError("result count and affine map result count must match");
-
-  return false;
-}
-
-// The result of the affine apply operation can be used as a dimension id if it
-// is a CFG value or if it is an MLValue, and all the operands are valid
-// dimension ids.
-bool AffineApplyOp::isValidDim() const {
-  for (auto *op : getOperands()) {
-    if (auto *v = dyn_cast<MLValue>(op))
-      if (!v->isValidDim())
-        return false;
-  }
-  return true;
-}
-
-// The result of the affine apply operation can be used as a symbol if it is
-// a CFG value or if it is an MLValue, and all the operands are symbols.
-bool AffineApplyOp::isValidSymbol() const {
-  for (auto *op : getOperands()) {
-    if (auto *v = dyn_cast<MLValue>(op))
-      if (!v->isValidSymbol())
-        return false;
-  }
-  return true;
-}
-
-bool AffineApplyOp::constantFold(ArrayRef<Attribute *> operandConstants,
-                                 SmallVectorImpl<Attribute *> &results,
-                                 MLIRContext *context) const {
-  auto map = getAffineMap();
-  if (map.constantFold(operandConstants, results))
-    return true;
-  // Return false on success.
-  return false;
-}
-
-//===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
@@ -402,118 +274,6 @@ bool CallIndirectOp::verify() const {
 }
 
 //===----------------------------------------------------------------------===//
-// Constant*Op
-//===----------------------------------------------------------------------===//
-
-/// Builds a constant op with the specified attribute value and result type.
-void ConstantOp::build(Builder *builder, OperationState *result,
-                       Attribute *value, Type *type) {
-  result->addAttribute("value", value);
-  result->types.push_back(type);
-}
-
-void ConstantOp::print(OpAsmPrinter *p) const {
-  *p << "constant " << *getValue();
-  p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"value");
-
-  if (!isa<FunctionAttr>(getValue()))
-    *p << " : " << *getType();
-}
-
-bool ConstantOp::parse(OpAsmParser *parser, OperationState *result) {
-  Attribute *valueAttr;
-  Type *type;
-
-  if (parser->parseAttribute(valueAttr, "value", result->attributes) ||
-      parser->parseOptionalAttributeDict(result->attributes))
-    return true;
-
-  // 'constant' taking a function reference doesn't get a redundant type
-  // specifier.  The attribute itself carries it.
-  if (auto *fnAttr = dyn_cast<FunctionAttr>(valueAttr))
-    return parser->addTypeToList(fnAttr->getValue()->getType(), result->types);
-
-  return parser->parseColonType(type) ||
-         parser->addTypeToList(type, result->types);
-}
-
-/// The constant op requires an attribute, and furthermore requires that it
-/// matches the return type.
-bool ConstantOp::verify() const {
-  auto *value = getValue();
-  if (!value)
-    return emitOpError("requires a 'value' attribute");
-
-  auto *type = this->getType();
-  if (isa<IntegerType>(type) || type->isIndex()) {
-    if (!isa<IntegerAttr>(value))
-      return emitOpError(
-          "requires 'value' to be an integer for an integer result type");
-    return false;
-  }
-
-  if (isa<FloatType>(type)) {
-    if (!isa<FloatAttr>(value))
-      return emitOpError("requires 'value' to be a floating point constant");
-    return false;
-  }
-
-  if (type->isTFString()) {
-    if (!isa<StringAttr>(value))
-      return emitOpError("requires 'value' to be a string constant");
-    return false;
-  }
-
-  if (isa<FunctionType>(type)) {
-    if (!isa<FunctionAttr>(value))
-      return emitOpError("requires 'value' to be a function reference");
-    return false;
-  }
-
-  return emitOpError(
-      "requires a result type that aligns with the 'value' attribute");
-}
-
-Attribute *ConstantOp::constantFold(ArrayRef<Attribute *> operands,
-                                    MLIRContext *context) const {
-  assert(operands.empty() && "constant has no operands");
-  return getValue();
-}
-
-void ConstantFloatOp::build(Builder *builder, OperationState *result,
-                            double value, FloatType *type) {
-  ConstantOp::build(builder, result, builder->getFloatAttr(value), type);
-}
-
-bool ConstantFloatOp::isClassFor(const Operation *op) {
-  return ConstantOp::isClassFor(op) &&
-         isa<FloatType>(op->getResult(0)->getType());
-}
-
-/// ConstantIntOp only matches values whose result type is an IntegerType.
-bool ConstantIntOp::isClassFor(const Operation *op) {
-  return ConstantOp::isClassFor(op) &&
-         isa<IntegerType>(op->getResult(0)->getType());
-}
-
-void ConstantIntOp::build(Builder *builder, OperationState *result,
-                          int64_t value, unsigned width) {
-  ConstantOp::build(builder, result, builder->getIntegerAttr(value),
-                    builder->getIntegerType(width));
-}
-
-/// ConstantIndexOp only matches values whose result type is Index.
-bool ConstantIndexOp::isClassFor(const Operation *op) {
-  return ConstantOp::isClassFor(op) && op->getResult(0)->getType()->isIndex();
-}
-
-void ConstantIndexOp::build(Builder *builder, OperationState *result,
-                            int64_t value) {
-  ConstantOp::build(builder, result, builder->getIntegerAttr(value),
-                    builder->getIndexType());
-}
-
-//===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
@@ -903,50 +663,6 @@ Attribute *MulIOp::constantFold(ArrayRef<Attribute *> operands,
 }
 
 //===----------------------------------------------------------------------===//
-// ReturnOp
-//===----------------------------------------------------------------------===//
-
-void ReturnOp::build(Builder *builder, OperationState *result,
-                     ArrayRef<SSAValue *> results) {
-  result->addOperands(results);
-}
-
-bool ReturnOp::parse(OpAsmParser *parser, OperationState *result) {
-  SmallVector<OpAsmParser::OperandType, 2> opInfo;
-  SmallVector<Type *, 2> types;
-  llvm::SMLoc loc;
-  return parser->getCurrentLocation(&loc) || parser->parseOperandList(opInfo) ||
-         (!opInfo.empty() && parser->parseColonTypeList(types)) ||
-         parser->resolveOperands(opInfo, types, loc, result->operands);
-}
-
-void ReturnOp::print(OpAsmPrinter *p) const {
-  *p << "return";
-  if (getNumOperands() > 0) {
-    *p << ' ';
-    p->printOperands(operand_begin(), operand_end());
-    *p << " : ";
-    interleave(operand_begin(), operand_end(),
-               [&](const SSAValue *e) { p->printType(e->getType()); },
-               [&]() { *p << ", "; });
-  }
-}
-
-bool ReturnOp::verify() const {
-  // ReturnOp must be part of an ML function.
-  if (auto *stmt = dyn_cast<OperationStmt>(getOperation())) {
-    StmtBlock *block = stmt->getBlock();
-    if (!block || !isa<MLFunction>(block) || &block->back() != stmt)
-      return emitOpError("must be the last statement in the ML function");
-
-    // Return success. Checking that operand types match those in the function
-    // signature is performed in the ML function verifier.
-    return false;
-  }
-  return emitOpError("cannot occur in a CFG function");
-}
-
-//===----------------------------------------------------------------------===//
 // ShapeCastOp
 //===----------------------------------------------------------------------===//
 
@@ -1106,11 +822,12 @@ Attribute *SubIOp::constantFold(ArrayRef<Attribute *> operands,
 // Register operations.
 //===----------------------------------------------------------------------===//
 
-/// Install the standard operations in the specified operation set.
-void mlir::registerStandardOperations(OperationSet &opSet) {
-  opSet.addOperations<AddFOp, AddIOp, AffineApplyOp, AllocOp, CallOp,
-                      CallIndirectOp, ConstantOp, DeallocOp, DimOp, DmaStartOp,
-                      DmaWaitOp, ExtractElementOp, LoadOp, MulFOp, MulIOp,
-                      ReturnOp, ShapeCastOp, StoreOp, SubFOp, SubIOp>(
-      /*prefix=*/"");
+/// Install the standard operations in the specified MLIRContext.
+void mlir::registerStandardOperations(MLIRContext *ctx) {
+  auto &opSet = OperationSet::get(ctx);
+  opSet
+      .addOperations<AddFOp, AddIOp, AllocOp, CallOp, CallIndirectOp, DeallocOp,
+                     DimOp, DmaStartOp, DmaWaitOp, ExtractElementOp, LoadOp,
+                     MulFOp, MulIOp, ShapeCastOp, StoreOp, SubFOp, SubIOp>(
+          /*prefix=*/"");
 }
diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp
index d2f24ba8de8..0aa593202e1 100644
--- a/mlir/lib/Transforms/ComposeAffineMaps.cpp
+++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp
@@ -61,8 +61,9 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/Support/CommandLine.h"
diff --git a/mlir/lib/Transforms/ConstantFold.cpp b/mlir/lib/Transforms/ConstantFold.cpp
index e8d7033e85c..8af7148038f 100644
--- a/mlir/lib/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Transforms/ConstantFold.cpp
@@ -16,11 +16,12 @@
 // =============================================================================
 
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/CFGFunction.h"
-#include "mlir/IR/StandardOps.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Passes.h"
+
 using namespace mlir;
 
 namespace {
diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index 8738d74acdd..2211a93cd54 100644
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -25,7 +25,7 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Pass.h"
diff --git a/mlir/lib/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
index 80ea0f55ba7..b8e98379051 100644
--- a/mlir/lib/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
@@ -48,7 +48,7 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/StmtVisitor.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Pass.h"
diff --git a/mlir/lib/Transforms/LoopUtils.cpp b/mlir/lib/Transforms/LoopUtils.cpp
index 95726653551..26dbe33c75f 100644
--- a/mlir/lib/Transforms/LoopUtils.cpp
+++ b/mlir/lib/Transforms/LoopUtils.cpp
@@ -25,9 +25,10 @@
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Statements.h"
 #include "mlir/IR/StmtVisitor.h"
+#include "mlir/StandardOps/StandardOps.h"
 #include "llvm/ADT/DenseMap.h"
 
 using namespace mlir;
diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp
index 87899564172..dd8b9a7615c 100644
--- a/mlir/lib/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp
@@ -23,7 +23,8 @@
 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Pass.h"
 #include "mlir/Transforms/Utils.h"
diff --git a/mlir/lib/Transforms/Utils.cpp b/mlir/lib/Transforms/Utils.cpp
index e1601c4f75d..2e8f0d32736 100644
--- a/mlir/lib/Transforms/Utils.cpp
+++ b/mlir/lib/Transforms/Utils.cpp
@@ -24,7 +24,8 @@
 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/StandardOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/StandardOps/StandardOps.h"
 #include "llvm/ADT/DenseMap.h"
 
 using namespace mlir;