Random cleanups:

- Change AllocOp to have a getType() that always returns a MemRefType, since
   that is what it requires.
 - Rename StandardOps/StandardOpRegistration.cpp ->
   StandardOps/OpRegistration.cpp to align with other op sets.
 - Add AffineMap::getContext() helper and use it in the asmprinter.

PiperOrigin-RevId: 218205527

mlir/include/mlir/IR/AffineMap.h

@@ -23,8 +23,6 @@
 #ifndef MLIR_IR_AFFINE_MAP_H
 #define MLIR_IR_AFFINE_MAP_H
 
-#include <vector>
-
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMapInfo.h"
@@ -32,9 +30,7 @@
 namespace mlir {
 
 namespace detail {
-
 class AffineMapStorage;
-
 } // end namespace detail
 
 class AffineExpr;
@@ -60,6 +56,8 @@ public:
   /// Returns a single constant result affine map.
   static AffineMap getConstantMap(int64_t val, MLIRContext *context);
 
+  MLIRContext *getContext() const;
+
   explicit operator bool() { return map; }
   bool operator==(const AffineMap &other) const { return other.map == map; }
 

mlir/include/mlir/StandardOps/StandardOps.h

@@ -100,8 +100,10 @@ private:
 class AllocOp
     : public Op<AllocOp, OpTrait::VariadicOperands, OpTrait::OneResult> {
 public:
-  SSAValue *getMemRef() { return getOperation()->getResult(0); }
-  const SSAValue *getMemRef() const { return getOperation()->getResult(0); }
+  /// The result of an alloc is always a MemRefType.
+  MemRefType *getType() const {
+    return cast<MemRefType>(getResult()->getType());
+  }
 
   static StringRef getOperationName() { return "alloc"; }
 

mlir/lib/IR/AffineMap.cpp

@@ -94,6 +94,8 @@ AffineMap AffineMap::getConstantMap(int64_t val, MLIRContext *context) {
              {getAffineConstantExpr(val, context)}, {});
 }
 
+MLIRContext *AffineMap::getContext() const { return getResult(0).getContext(); }
+
 bool AffineMap::isBounded() const { return !map->rangeSizes.empty(); }
 
 bool AffineMap::isIdentity() const {

mlir/lib/IR/AsmPrinter.cpp

@@ -1595,7 +1595,7 @@ void AffineExpr::dump() const {
 }
 
 void AffineMap::print(raw_ostream &os) const {
-  ModuleState state(/*no context is known*/ nullptr);
+  ModuleState state(getContext());
   ModulePrinter(os, state).printAffineMap(*this);
 }
 

mlir/lib/StandardOps/OpRegistration.cpp

@@ -1,4 +1,4 @@
-//===- StandardOpsRegistration.cpp - Register standard Op types -*- C++ -*-===//
+//===- OpRegistration.cpp - Register standard Op types --------------------===//
 //
 // Copyright 2019 The MLIR Authors.
 //
@@ -16,7 +16,6 @@
 // =============================================================================
 
 #include "mlir/StandardOps/StandardOps.h"
-
 using namespace mlir;
 
 // Static initialization for standard op registration.

mlir/lib/StandardOps/StandardOps.cpp

@@ -83,7 +83,7 @@ void AllocOp::build(Builder *builder, OperationState *result,
 }
 
 void AllocOp::print(OpAsmPrinter *p) const {
-  MemRefType *type = cast<MemRefType>(getMemRef()->getType());
+  MemRefType *type = getType();
   *p << "alloc";
   // Print dynamic dimension operands.
   printDimAndSymbolList(operand_begin(), operand_end(),
@@ -119,7 +119,7 @@ bool AllocOp::parse(OpAsmParser *parser, OperationState *result) {
 }
 
 bool AllocOp::verify() const {
-  auto *memRefType = dyn_cast<MemRefType>(getMemRef()->getType());
+  auto *memRefType = dyn_cast<MemRefType>(getResult()->getType());
   if (!memRefType)
     return emitOpError("result must be a memref");