[mlir][Type] Remove usages of Type::getKind

[rriddle]

Summary

This is in preparation for removing the use of "kinds" within attributes and types in MLIR.

Depends On D85370

[rriddle]

Rebase

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

public:
 
   /// Support for PointerLikeTypeTraits.
   using Type::getAsOpaquePointer;
   static LLVMType getFromOpaquePointer(const void *ptr) {
     return LLVMType(static_cast<ImplType *>(const_cast<void *>(ptr)));
   }
 
   /// Support for isa/cast.
-  static bool classof(Type type) {
-    return type.getKind() >= FIRST_NEW_LLVM_TYPE &&
-           type.getKind() <= LAST_NEW_LLVM_TYPE;
-  }
+  static bool classof(Type type);
 
   LLVMDialect &getDialect();
 
   /// Floating-point type utilities.
   bool isBFloatTy() { return isa<LLVMBFloatType>(); }
   bool isHalfTy() { return isa<LLVMHalfType>(); }
   bool isFloatTy() { return isa<LLVMFloatType>(); }
   bool isDoubleTy() { return isa<LLVMDoubleType>(); }

mlir/include/mlir/Dialect/Quant/QuantTypes.h

public:
   static constexpr unsigned MaxStorageBits = 32;
 
   static LogicalResult
   verifyConstructionInvariants(Location loc, unsigned flags, Type storageType,
                                Type expressedType, int64_t storageTypeMin,
                                int64_t storageTypeMax);
 
   /// Support method to enable LLVM-style type casting.
-  static bool classof(Type type) {
-    return type.getKind() >= Type::FIRST_QUANTIZATION_TYPE &&
-           type.getKind() <= QuantizationTypes::LAST_USED_QUANTIZATION_TYPE;
-  }
+  static bool classof(Type type);
 
   /// Gets the minimum possible stored by a storageType. storageTypeMin must
   /// be greater than or equal to this value.
   static int64_t getDefaultMinimumForInteger(bool isSigned,
                                              unsigned integralWidth) {
     if (isSigned) {
       return llvm::minIntN(integralWidth);
     }

mlir/include/mlir/IR/StandardTypes.h

public:
   /// not take into account any memory layout or widening constraints, e.g. a
   /// vector<3xi57> is reported to occupy 3x57=171 bit, even though in practice
   /// it will likely be stored as in a 4xi64 vector register.  Fail an assertion
   /// if the size cannot be computed statically, i.e. if the type has a dynamic
   /// shape or if its elemental type does not have a known bit width.
   int64_t getSizeInBits() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(Type type) {
-    return type.getKind() == StandardTypes::Vector ||
-           type.getKind() == StandardTypes::RankedTensor ||
-           type.getKind() == StandardTypes::UnrankedTensor ||
-           type.getKind() == StandardTypes::UnrankedMemRef ||
-           type.getKind() == StandardTypes::MemRef;
-  }
+  static bool classof(Type type);
 
   /// Whether the given dimension size indicates a dynamic dimension.
   static constexpr bool isDynamic(int64_t dSize) {
     return dSize == kDynamicSize;
   }
   static constexpr bool isDynamicStrideOrOffset(int64_t dStrideOrOffset) {
     return dStrideOrOffset == kDynamicStrideOrOffset;
   }

 
 /// Tensor types represent multi-dimensional arrays, and have two variants:
 /// RankedTensorType and UnrankedTensorType.
 class TensorType : public ShapedType {
 public:
   using ShapedType::ShapedType;
 
   /// Return true if the specified element type is ok in a tensor.
-  static bool isValidElementType(Type type) {
-    // Note: Non standard/builtin types are allowed to exist within tensor
-    // types. Dialects are expected to verify that tensor types have a valid
-    // element type within that dialect.
-    return type.isa<ComplexType, FloatType, IntegerType, OpaqueType, VectorType,
-                    IndexType>() ||
-           (type.getKind() > Type::Kind::LAST_STANDARD_TYPE);
-  }
+  static bool isValidElementType(Type type);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(Type type) {
-    return type.getKind() == StandardTypes::RankedTensor ||
-           type.getKind() == StandardTypes::UnrankedTensor;
-  }
+  static bool classof(Type type);
 };
 
 //===----------------------------------------------------------------------===//
 // RankedTensorType
 
 /// Ranked tensor types represent multi-dimensional arrays that have a shape
 /// with a fixed number of dimensions. Each shape element can be a non-negative
 /// integer or unknown (represented by -1).

 //===----------------------------------------------------------------------===//
 
 /// Base MemRef for Ranked and Unranked variants
 class BaseMemRefType : public ShapedType {
 public:
   using ShapedType::ShapedType;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(Type type) {
-    return type.getKind() == StandardTypes::MemRef ||
-           type.getKind() == StandardTypes::UnrankedMemRef;
-  }
+  static bool classof(Type type);
 };
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 
 /// MemRef types represent a region of memory that have a shape with a fixed
 /// number of dimensions. Each shape element can be a non-negative integer or
 /// unknown (represented by -1). MemRef types also have an affine map
public:
   /// Return the element type at index 'index'.
   Type getType(size_t index) const {
     assert(index < size() && "invalid index for tuple type");
     return getTypes()[index];
   }
 };
 
 //===----------------------------------------------------------------------===//
+// Deferred Method Definitions
+//===----------------------------------------------------------------------===//
+
+inline bool BaseMemRefType::classof(Type type) {
+  return type.isa<MemRefType, UnrankedMemRefType>();
+}
+
+inline bool ShapedType::classof(Type type) {
+  return type.isa<RankedTensorType, VectorType, UnrankedTensorType,
+                  UnrankedMemRefType, MemRefType>();
+}
+
+inline bool TensorType::classof(Type type) {
+  return type.isa<RankedTensorType, UnrankedTensorType>();
+}
+
+//===----------------------------------------------------------------------===//
 // Type Utilities
 //===----------------------------------------------------------------------===//
 
 /// Returns the strides of the MemRef if the layout map is in strided form.
 /// MemRefs with layout maps in strided form include:
 ///   1. empty or identity layout map, in which case the stride information is
 ///      the canonical form computed from sizes;
 ///   2. single affine map layout of the form `K + k0 * d0 + ... kn * dn`,

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 //===----------------------------------------------------------------------===//
 // LLVMType.
 //===----------------------------------------------------------------------===//
 
+bool LLVMType::classof(Type type) {
+  return llvm::isa<LLVMDialect>(type.getDialect());
+}
+
 LLVMDialect &LLVMType::getDialect() {
   return static_cast<LLVMDialect &>(Type::getDialect());
 }
 
 //----------------------------------------------------------------------------//
 // Integer type utilities.
 
 bool LLVMType::isIntegerTy(unsigned bitwidth) {

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

Type mlir::linalg::LinalgDialect::parseType(DialectAsmParser &parser) const {
   return Type();
 }
 
 /// RangeType prints as just "range".
 static void print(RangeType rt, DialectAsmPrinter &os) { os << "range"; }
 
 void mlir::linalg::LinalgDialect::printType(Type type,
                                             DialectAsmPrinter &os) const {
-  switch (type.getKind()) {
-  default:
-    llvm_unreachable("Unhandled Linalg type");
-  case LinalgTypes::Range:
-    print(type.cast<RangeType>(), os);
-    break;
-  }
 }

mlir/lib/Dialect/Quant/IR/QuantTypes.cpp

 //===- QuantOps.cpp - Quantization Type and Ops Implementation --*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Quant/QuantTypes.h"
 #include "TypeDetail.h"
+#include "mlir/Dialect/Quant/QuantOps.h"
 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/MathExtras.h"
 
 using namespace mlir;
 using namespace mlir::quant;
 using namespace mlir::quant::detail;
 
 unsigned QuantizedType::getFlags() const {
   return static_cast<ImplType *>(impl)->flags;
 }
 
+bool QuantizedType::classof(Type type) {
+  return llvm::isa<QuantizationDialect>(type.getDialect());
+}
+
 LogicalResult QuantizedType::verifyConstructionInvariants(
     Location loc, unsigned flags, Type storageType, Type expressedType,
     int64_t storageTypeMin, int64_t storageTypeMax) {
   // Verify that the storage type is integral.
   // This restriction may be lifted at some point in favor of using bf16
   // or f16 as exact representations on hardware where that is advantageous.
   auto intStorageType = storageType.dyn_cast<IntegerType>();
   if (!intStorageType)

mlir/lib/Dialect/Quant/IR/TypeParser.cpp

llvm::interleave(
         printQuantParams(scales[index], zeroPoints[index], out);
       },
       ",");
   out << "}>";
 }
 
 /// Print a type registered to this dialect.
 void QuantizationDialect::printType(Type type, DialectAsmPrinter &os) const {
-  switch (type.getKind()) {
-  default:
+  if (auto anyType = type.dyn_cast<AnyQuantizedType>())
+    printAnyQuantizedType(anyType, os);
+  else if (auto uniformType = type.dyn_cast<UniformQuantizedType>())
+    printUniformQuantizedType(uniformType, os);
+  else if (auto perAxisType = type.dyn_cast<UniformQuantizedPerAxisType>())
+    printUniformQuantizedPerAxisType(perAxisType, os);
+  else
     llvm_unreachable("Unhandled quantized type");
-  case QuantizationTypes::Any:
-    printAnyQuantizedType(type.cast<AnyQuantizedType>(), os);
-    break;
-  case QuantizationTypes::UniformQuantized:
-    printUniformQuantizedType(type.cast<UniformQuantizedType>(), os);
-    break;
-  case QuantizationTypes::UniformQuantizedPerAxis:
-    printUniformQuantizedPerAxisType(type.cast<UniformQuantizedPerAxisType>(),
-                                     os);
-    break;
-  }
 }

mlir/lib/Dialect/Quant/Utils/UniformSupport.cpp

 using namespace mlir::quant;
 
 static bool isQuantizablePrimitiveType(Type inputType) {
   return inputType.isa<FloatType>();
 }
 
 const ExpressedToQuantizedConverter
 ExpressedToQuantizedConverter::forInputType(Type inputType) {
-  switch (inputType.getKind()) {
-  default:
-    if (isQuantizablePrimitiveType(inputType)) {
-      // Supported primitive type (which just is the expressed type).
-      return ExpressedToQuantizedConverter{inputType, inputType};
-    }
-    // Unsupported.
-    return ExpressedToQuantizedConverter{inputType, nullptr};
-  case StandardTypes::RankedTensor:
-  case StandardTypes::UnrankedTensor:
-  case StandardTypes::Vector: {
-    Type elementType = inputType.cast<ShapedType>().getElementType();
-    if (!isQuantizablePrimitiveType(elementType)) {
-      // Unsupported.
-      return ExpressedToQuantizedConverter{inputType, nullptr};
-    }
-    return ExpressedToQuantizedConverter{
-        inputType, inputType.cast<ShapedType>().getElementType()};
-  }
-  }
-}
+  if (inputType.isa<TensorType, VectorType>()) {
+    Type elementType = inputType.cast<ShapedType>().getElementType();
+    if (!isQuantizablePrimitiveType(elementType))
+      return ExpressedToQuantizedConverter{inputType, nullptr};
+    return ExpressedToQuantizedConverter{inputType, elementType};
+  }
+  if (isQuantizablePrimitiveType(inputType))
 
 Type ExpressedToQuantizedConverter::convert(QuantizedType elementalType) const {
   assert(expressedType && "convert() on unsupported conversion");
+  if (auto tensorType = inputType.dyn_cast<RankedTensorType>())
+    return RankedTensorType::get(tensorType.getShape(), elementalType);
+  if (auto tensorType = inputType.dyn_cast<UnrankedTensorType>())
+    return UnrankedTensorType::get(elementalType);
+  if (auto vectorType = inputType.dyn_cast<VectorType>())
+    return VectorType::get(vectorType.getShape(), elementalType);
 
-  switch (inputType.getKind()) {
-  default:
-    if (elementalType.getExpressedType() == expressedType) {
-      // If the expressed types match, just use the new elemental type.
-      return elementalType;
-    }
-    // Unsupported.
-    return nullptr;
-  case StandardTypes::RankedTensor:
-    return RankedTensorType::get(inputType.cast<RankedTensorType>().getShape(),
-                                 elementalType);
-  case StandardTypes::UnrankedTensor:
-    return UnrankedTensorType::get(elementalType);
-  case StandardTypes::Vector:
-    return VectorType::get(inputType.cast<VectorType>().getShape(),
-                           elementalType);
-  }
+  if (elementalType.getExpressedType() == expressedType)
 }
 
 ElementsAttr
 UniformQuantizedPerAxisValueConverter::convert(Attribute realValue) {
   if (auto attr = realValue.dyn_cast<DenseFPElementsAttr>()) {
     return convert(attr);
   }
   // TODO: handles sparse elements attribute

mlir/lib/Dialect/SPIRV/LayoutUtils.cpp

 Type VulkanLayoutUtils::decorateType(Type type, VulkanLayoutUtils::Size &size,
                                      VulkanLayoutUtils::Size &alignment) {
   if (type.isa<spirv::ScalarType>()) {
     alignment = getScalarTypeAlignment(type);
     // Vulkan spec does not specify any padding for a scalar type.
     size = alignment;
     return type;
   }
-
-  switch (type.getKind()) {
-  case spirv::TypeKind::Struct:
-    return decorateType(type.cast<spirv::StructType>(), size, alignment);
-  case spirv::TypeKind::Array:
-    return decorateType(type.cast<spirv::ArrayType>(), size, alignment);
-  case StandardTypes::Vector:
-    return decorateType(type.cast<VectorType>(), size, alignment);
-  case spirv::TypeKind::RuntimeArray:
+  if (auto structType = type.dyn_cast<spirv::StructType>())
+    return decorateType(structType, size, alignment);
+  if (auto arrayType = type.dyn_cast<spirv::ArrayType>())
+    return decorateType(arrayType, size, alignment);
+  if (auto vectorType = type.dyn_cast<VectorType>())
+    return decorateType(vectorType, size, alignment);
+  if (auto arrayType = type.dyn_cast<spirv::RuntimeArrayType>()) {
     size = std::numeric_limits<Size>().max();
-    return decorateType(type.cast<spirv::RuntimeArrayType>(), alignment);
-  default:
-    llvm_unreachable("unhandled SPIR-V type");
+    return decorateType(arrayType, alignment);
   }
+  llvm_unreachable("unhandled SPIR-V type");
 }
 
 Type VulkanLayoutUtils::decorateType(VectorType vectorType,
                                      VulkanLayoutUtils::Size &size,
                                      VulkanLayoutUtils::Size &alignment) {
   const auto numElements = vectorType.getNumElements();
   auto elementType = vectorType.getElementType();
   Size elementSize = 0;

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Parser.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace spirv {
 #include "mlir/Dialect/SPIRV/SPIRVOpUtils.inc"
 } // namespace spirv
 } // namespace mlir
 

 }
 
 static void print(MatrixType type, DialectAsmPrinter &os) {
   os << "matrix<" << type.getNumColumns() << " x " << type.getColumnType();
   os << ">";
 }
 
 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
-  switch (type.getKind()) {
-  case TypeKind::Array:
-    print(type.cast<ArrayType>(), os);
-    return;
-  case TypeKind::CooperativeMatrix:
-    print(type.cast<CooperativeMatrixNVType>(), os);
-    return;
-  case TypeKind::Pointer:
-    print(type.cast<PointerType>(), os);
-    return;
-  case TypeKind::RuntimeArray:
-    print(type.cast<RuntimeArrayType>(), os);
-    return;
-  case TypeKind::Image:
-    print(type.cast<ImageType>(), os);
-    return;
-  case TypeKind::Struct:
-    print(type.cast<StructType>(), os);
-    return;
-  case TypeKind::Matrix:
-    print(type.cast<MatrixType>(), os);
-    return;
-  default:
-    llvm_unreachable("unhandled SPIR-V type");
-  }
+  TypeSwitch<Type>(type)
+      .Case<ArrayType, CooperativeMatrixNVType, PointerType, RuntimeArrayType,
+            ImageType, StructType, MatrixType>(
+          [&](auto type) { print(type, os); })
+      .Default([](Type) { llvm_unreachable("unhandled SPIR-V type"); });
 }
 
 //===----------------------------------------------------------------------===//
 // Attribute Parsing
 //===----------------------------------------------------------------------===//
 
 /// Parses a comma-separated list of keywords, invokes `processKeyword` on each
 /// of the parsed keyword, and returns failure if any error occurs.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verify(spirv::ConstantOp constOp) {
   return constOp.emitOpError("cannot have value of type ") << valueType;
 }
 
 bool spirv::ConstantOp::isBuildableWith(Type type) {
   // Must be valid SPIR-V type first.
   if (!type.isa<spirv::SPIRVType>())
     return false;
 
-  if (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
-      type.getKind() <= spirv::TypeKind::LAST_SPIRV_TYPE) {
+  if (isa<SPIRVDialect>(type.getDialect())) {
     // TODO: support constant struct
     return type.isa<spirv::ArrayType>();
   }
 
   return true;
 }
 
 spirv::ConstantOp spirv::ConstantOp::getZero(Type type, Location loc,

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::spirv;
 
 // Pull in all enum utility function definitions
 #include "mlir/Dialect/SPIRV/SPIRVEnums.cpp.inc"
 // Pull in all enum type availability query function definitions
 #include "mlir/Dialect/SPIRV/SPIRVEnumAvailability.cpp.inc"
Optional<int64_t> ArrayType::getSizeInBytes() {
   return (*size + getArrayStride()) * getNumElements();
 }
 
 //===----------------------------------------------------------------------===//
 // CompositeType
 //===----------------------------------------------------------------------===//
 
 bool CompositeType::classof(Type type) {
-  switch (type.getKind()) {
-  case TypeKind::Array:
-  case TypeKind::CooperativeMatrix:
-  case TypeKind::Matrix:
-  case TypeKind::RuntimeArray:
-  case TypeKind::Struct:
-    return true;
-  case StandardTypes::Vector:
-    return isValid(type.cast<VectorType>());
-  default:
-    return false;
-  }
+  if (auto vectorType = type.dyn_cast<VectorType>())
+    return isValid(vectorType);
+  return type
+      .isa<spirv::ArrayType, spirv::CooperativeMatrixNVType, spirv::MatrixType,
+           spirv::RuntimeArrayType, spirv::StructType>();
 }
 
 bool CompositeType::isValid(VectorType type) {
   return type.getRank() == 1 && type.getElementType().isa<ScalarType>() &&
          type.getNumElements() >= 2 && type.getNumElements() <= 4;
 }
 
 Type CompositeType::getElementType(unsigned index) const {
-  switch (getKind()) {
-  case spirv::TypeKind::Array:
-    return cast<ArrayType>().getElementType();
-  case spirv::TypeKind::CooperativeMatrix:
-    return cast<CooperativeMatrixNVType>().getElementType();
-  case spirv::TypeKind::Matrix:
-    return cast<MatrixType>().getColumnType();
-  case spirv::TypeKind::RuntimeArray:
-    return cast<RuntimeArrayType>().getElementType();
-  case spirv::TypeKind::Struct:
-    return cast<StructType>().getElementType(index);
-  case StandardTypes::Vector:
-    return cast<VectorType>().getElementType();
-  default:
-    llvm_unreachable("invalid composite type");
-  }
+  return TypeSwitch<Type, Type>(*this)
+      .Case<ArrayType, CooperativeMatrixNVType, RuntimeArrayType, VectorType>(
+          [](auto type) { return type.getElementType(); })
+      .Case<MatrixType>([](MatrixType type) { return type.getColumnType(); })
+      .Case<StructType>(
+          [index](StructType type) { return type.getElementType(index); })
+      .Default(
+          [](Type) -> Type { llvm_unreachable("invalid composite type"); });
 }
 
 unsigned CompositeType::getNumElements() const {
   switch (getKind()) {
   case spirv::TypeKind::Array:
     return cast<ArrayType>().getNumElements();
   case spirv::TypeKind::CooperativeMatrix:
     llvm_unreachable(

mlir/lib/Dialect/Traits.cpp

Type OpTrait::util::getBroadcastedType(Type type1, Type type2,
   if (type1.isa<UnrankedTensorType>() || type2.isa<UnrankedTensorType>()) {
     if (type1.isa<VectorType>() || type2.isa<VectorType>())
       return {};
     return UnrankedTensorType::get(elementType);
   }
 
   // Returns the type kind if the given type is a vector or ranked tensor type.
   // Returns llvm::None otherwise.
-  auto getCompositeTypeKind = [](Type type) -> Optional<StandardTypes::Kind> {
+  auto getCompositeTypeKind = [](Type type) -> Optional<TypeID> {
     if (type.isa<VectorType, RankedTensorType>())
-      return static_cast<StandardTypes::Kind>(type.getKind());
+      return type.getTypeID();
     return llvm::None;
   };
 
   // Make sure the composite type, if has, is consistent.
-  auto compositeKind1 = getCompositeTypeKind(type1);
-  auto compositeKind2 = getCompositeTypeKind(type2);
-  Optional<StandardTypes::Kind> resultCompositeKind;
+  Optional<TypeID> compositeKind1 = getCompositeTypeKind(type1);
+  Optional<TypeID> compositeKind2 = getCompositeTypeKind(type2);
+  Optional<TypeID> resultCompositeKind;
 
   if (compositeKind1 && compositeKind2) {
     // Disallow mixing vector and tensor.
     if (compositeKind1 != compositeKind2)
       return {};
     resultCompositeKind = compositeKind1;
   } else if (compositeKind1) {
     resultCompositeKind = compositeKind1;
   } else if (compositeKind2) {
     resultCompositeKind = compositeKind2;
   }
 
   // Get the shape of each type.
   SmallVector<int64_t, 4> resultShape;
   if (!getBroadcastedShape(getShape(type1), getShape(type2), resultShape))
     return {};
 
   // Compose the final broadcasted type
-  if (resultCompositeKind == StandardTypes::Vector)
+  if (resultCompositeKind == VectorType::getTypeID())
     return VectorType::get(resultShape, elementType);
-  if (resultCompositeKind == StandardTypes::RankedTensor)
+  if (resultCompositeKind == RankedTensorType::getTypeID())
     return RankedTensorType::get(resultShape, elementType);
   return elementType;
 }
 
 /// Returns a tuple corresponding to whether range has tensor or vector type.
 template <typename iterator_range>
 static std::tuple<bool, bool> hasTensorOrVectorType(iterator_range types) {
   return std::make_tuple(

mlir/lib/IR/StandardTypes.cpp

 //===- StandardTypes.cpp - MLIR Standard Type Classes ---------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/StandardTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Dialect.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Type
int64_t ShapedType::getSizeInBits() const {
   // cannot.
   assert(isa<TensorType>() && "unsupported element type");
   assert((elementType.isa<VectorType, TensorType>()) &&
          "unsupported tensor element type");
   return getNumElements() * elementType.cast<ShapedType>().getSizeInBits();
 }
 
 ArrayRef<int64_t> ShapedType::getShape() const {
-  switch (getKind()) {
-  case StandardTypes::Vector:
-    return cast<VectorType>().getShape();
-  case StandardTypes::RankedTensor:
-    return cast<RankedTensorType>().getShape();
-  case StandardTypes::MemRef:
-    return cast<MemRefType>().getShape();
-  default:
-    llvm_unreachable("not a ShapedType or not ranked");
-  }
+  if (auto vectorType = dyn_cast<VectorType>())
+    return vectorType.getShape();
+  if (auto tensorType = dyn_cast<RankedTensorType>())
+    return tensorType.getShape();
 }
 
 int64_t ShapedType::getNumDynamicDims() const {
   return llvm::count_if(getShape(), isDynamic);
 }
 
 bool ShapedType::hasStaticShape() const {
   return hasRank() && llvm::none_of(getShape(), isDynamic);

 ArrayRef<int64_t> VectorType::getShape() const { return getImpl()->getShape(); }
 
 //===----------------------------------------------------------------------===//
 // TensorType
 //===----------------------------------------------------------------------===//
 
 // Check if "elementType" can be an element type of a tensor. Emit errors if
 // location is not nullptr.  Returns failure if check failed.
-static inline LogicalResult checkTensorElementType(Location location,
-                                                   Type elementType) {
+static LogicalResult checkTensorElementType(Location location,
   if (!TensorType::isValidElementType(elementType))
     return emitError(location, "invalid tensor element type");
   return success();
 }
 
+/// Return true if the specified element type is ok in a tensor.
+bool TensorType::isValidElementType(Type type) {
+  // Note: Non standard/builtin types are allowed to exist within tensor
+  // types. Dialects are expected to verify that tensor types have a valid
+  // element type within that dialect.
+  return type.isa<ComplexType, FloatType, IntegerType, OpaqueType, VectorType,
+                  IndexType>() ||
+         !type.getDialect().getNamespace().empty();
+}
+
 //===----------------------------------------------------------------------===//
 // RankedTensorType
 //===----------------------------------------------------------------------===//
 
 RankedTensorType RankedTensorType::get(ArrayRef<int64_t> shape,
                                        Type elementType) {
   return Base::get(elementType.getContext(), StandardTypes::RankedTensor, shape,
                    elementType);