[mlir][spirv] Move type checks from dialect class to type hierarchy

[antiagainst]

Summary

Types should be checked with the type hierarchy. This should result in
better responsibility division and API surface.

Depends On D76242

[mravishankar]

This is so much nicer! Thanks

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

166

You could just make this a dyn_cast<spirv::SPIRVType>()

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

166

I think it's different. If this is not a SPIR-V type, then dyn_cast<spirv::SPIRVType>() gives us a Type() instead of llvm::None. The former will let the framework to try another conversion rule but the later will directly fail. (Although here it is the last conversion rule so probably doesn't matter that much.) Anyway, this is completely changed in the next commit (https://reviews.llvm.org/D76244). :)

[antiagainst]

Rebased

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

def SPIRV_Dialect : Dialect {
   let cppNamespace = "spirv";
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
 
   let extraClassDeclaration = [{
     //===------------------------------------------------------------------===//
-    // Type
-    //===------------------------------------------------------------------===//
-
-    /// Checks if the given `type` is valid in SPIR-V dialect.
-    static bool isValidType(Type type);
-
-    /// Checks if the given `scalar type` is valid in SPIR-V dialect.
-    static bool isValidScalarType(Type type);
-
-    //===------------------------------------------------------------------===//
     // Attribute
     //===------------------------------------------------------------------===//
 
     /// Returns the attribute name to use when specifying decorations on results
     /// of operations.
     static std::string getAttributeName(Decoration decoration);
   }];
 }

mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h

 
 // Base SPIR-V type for providing availability queries.
 class SPIRVType : public Type {
 public:
   using Type::Type;
 
   static bool classof(Type type);
 
+  bool isScalarOrVector();
+
   /// The extension requirements for each type are following the
   /// ((Extension::A OR Extension::B) AND (Extension::C OR Extension::D))
   /// convention.
   using ExtensionArrayRefVector = SmallVectorImpl<ArrayRef<spirv::Extension>>;
 
   /// Appends to `extensions` the extensions needed for this type to appear in
   /// the given `storage` class. This method does not guarantee the uniqueness
   /// of extensions; the same extension may be appended multiple times.

 
 // SPIR-V scalar type: bool type, integer type, floating point type.
 class ScalarType : public SPIRVType {
 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
 
+  /// Returns true if the given integer type is valid for the SPIR-V dialect.
+  static bool isValid(FloatType);
+  /// Returns true if the given float type is valid for the SPIR-V dialect.
+  static bool isValid(IntegerType);
+
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
 };
 
 // SPIR-V composite type: VectorType, SPIR-V ArrayType, or SPIR-V StructType.
 class CompositeType : public SPIRVType {
 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
 
+  /// Returns true if the given vector type is valid for the SPIR-V dialect.
+  static bool isValid(VectorType);
+
   unsigned getNumElements() const;
 
   Type getElementType(unsigned) const;
 
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);

mlir/lib/Dialect/SPIRV/LayoutUtils.cpp

VulkanLayoutUtils::decorateType(spirv::StructType structType,
   size = llvm::alignTo(structMemberOffset, maxMemberAlignment);
   alignment = maxMemberAlignment;
   structType.getMemberDecorations(memberDecorations);
   return spirv::StructType::get(memberTypes, layoutInfo, memberDecorations);
 }
 
 Type VulkanLayoutUtils::decorateType(Type type, VulkanLayoutUtils::Size &size,
                                      VulkanLayoutUtils::Size &alignment) {
-  if (spirv::SPIRVDialect::isValidScalarType(type)) {
+  if (type.isa<spirv::ScalarType>()) {
     alignment = VulkanLayoutUtils::getScalarTypeAlignment(type);
     // Vulkan spec does not specify any padding for a scalar type.
     size = alignment;
     return type;
   }
 
   switch (type.getKind()) {
   case spirv::TypeKind::Struct:

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

 //===- SPIRVCanonicalization.cpp - MLIR SPIR-V canonicalization patterns --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the folders and canonicalization patterns for SPIR-V ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/Functional.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
return std::next(block->begin()) == block->end() &&
            isa<spirv::BranchConditionalOp>(block->front());
   }
 
   bool isSameAttrList(spirv::StoreOp lhs, spirv::StoreOp rhs) const {
     return lhs.getOperation()->getAttrList().getDictionary() ==
            rhs.getOperation()->getAttrList().getDictionary();
   }
 
-  // Checks that given type is valid for `spv.SelectOp`.
-  // According to SPIR-V spec:
-  // "Before version 1.4, Result Type must be a pointer, scalar, or vector.
-  // Starting with version 1.4, Result Type can additionally be a composite type
-  // other than a vector."
-  bool isValidType(Type type) const {
-    return spirv::SPIRVDialect::isValidScalarType(type) ||
-           type.isa<VectorType>();
-  }
 
   // Returns a source value for the given block.
   Value getSrcValue(Block *block) const {
     auto storeOp = cast<spirv::StoreOp>(block->front());
     return storeOp.value();
   }
 
   // Returns a destination value for the given block.
LogicalResult ConvertSelectionOpToSelect::canCanonicalizeSelection(
   auto falseBrBranchOp =
       dyn_cast<spirv::BranchOp>(*std::next(falseBlock->begin()));
 
   if (!trueBrStoreOp || !trueBrBranchOp || !falseBrStoreOp ||
       !falseBrBranchOp) {
     return failure();
   }
 
+  // Checks that given type is valid for `spv.SelectOp`.
+  // According to SPIR-V spec:
+  // "Before version 1.4, Result Type must be a pointer, scalar, or vector.
+  // Starting with version 1.4, Result Type can additionally be a composite type
+  // other than a vector."
+  bool isScalarOrVector = trueBrStoreOp.value()
+                              .getType()
+                              .cast<spirv::SPIRVType>()
+                              .isScalarOrVector();
+
   // Check that each `spv.Store` uses the same pointer, memory access
   // attributes and a valid type of the value.
   if ((trueBrStoreOp.ptr() != falseBrStoreOp.ptr()) ||
-      !isSameAttrList(trueBrStoreOp, falseBrStoreOp) ||
-      !isValidType(trueBrStoreOp.value().getType())) {
+      !isSameAttrList(trueBrStoreOp, falseBrStoreOp) || !isScalarOrVector) {
     return failure();
   }
 
   if ((trueBrBranchOp.getOperation()->getSuccessor(0) != mergeBlock) ||
       (falseBrBranchOp.getOperation()->getSuccessor(0) != mergeBlock)) {
     return failure();
   }
 
   return success();
 }
 } // end anonymous namespace
 
 void spirv::SelectionOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ConvertSelectionOpToSelect>(context);
 }

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

 template <>
 Optional<Type> parseAndVerify<Type>(SPIRVDialect const &dialect,
                                     DialectAsmParser &parser);
 
 template <>
 Optional<uint64_t> parseAndVerify<uint64_t>(SPIRVDialect const &dialect,
                                             DialectAsmParser &parser);
 
-static bool isValidSPIRVIntType(IntegerType type) {
-  return llvm::is_contained(ArrayRef<unsigned>({1, 8, 16, 32, 64}),
-                            type.getWidth());
-}
-
-bool SPIRVDialect::isValidScalarType(Type type) {
-  if (type.isa<FloatType>()) {
-    return !type.isBF16();
-  }
-  if (auto intType = type.dyn_cast<IntegerType>()) {
-    return isValidSPIRVIntType(intType);
-  }
-  return false;
-}
-
-static bool isValidSPIRVVectorType(VectorType type) {
-  return type.getRank() == 1 &&
-         SPIRVDialect::isValidScalarType(type.getElementType()) &&
-         type.getNumElements() >= 2 && type.getNumElements() <= 4;
-}
-
-bool SPIRVDialect::isValidType(Type type) {
-  // Allow SPIR-V dialect types
-  if (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
-      type.getKind() <= TypeKind::LAST_SPIRV_TYPE) {
-    return true;
-  }
-  if (SPIRVDialect::isValidScalarType(type)) {
-    return true;
-  }
-  if (auto vectorType = type.dyn_cast<VectorType>()) {
-    return isValidSPIRVVectorType(vectorType);
-  }
-  return false;
-}
-
 static Type parseAndVerifyType(SPIRVDialect const &dialect,
                                DialectAsmParser &parser) {
   Type type;
   llvm::SMLoc typeLoc = parser.getCurrentLocation();
   if (parser.parseType(type))
     return Type();
 
   // Allow SPIR-V dialect types
   if (&type.getDialect() == &dialect)
     return type;
 
   // Check other allowed types
   if (auto t = type.dyn_cast<FloatType>()) {
     if (type.isBF16()) {
       parser.emitError(typeLoc, "cannot use 'bf16' to compose SPIR-V types");
       return Type();
     }
   } else if (auto t = type.dyn_cast<IntegerType>()) {
-    if (!isValidSPIRVIntType(t)) {
+    if (!ScalarType::isValid(t)) {
       parser.emitError(typeLoc,
                        "only 1/8/16/32/64-bit integer type allowed but found ")
           << type;
       return Type();
     }
   } else if (auto t = type.dyn_cast<VectorType>()) {
     if (t.getRank() != 1) {
       parser.emitError(typeLoc, "only 1-D vector allowed but found ") << t;

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

 #undef STORAGE_SPACE_MAP_FN
 }
 
 #undef STORAGE_SPACE_MAP_LIST
 
 // TODO(ravishankarm): This is a utility function that should probably be
 // exposed by the SPIR-V dialect. Keeping it local till the use case arises.
 static Optional<int64_t> getTypeNumBytes(Type t) {
-  if (spirv::SPIRVDialect::isValidScalarType(t)) {
+  if (t.isa<spirv::ScalarType>()) {
     auto bitWidth = t.getIntOrFloatBitWidth();
     // According to the SPIR-V spec:
     // "There is no physical size or bit pattern defined for values with boolean
     // type. If they are stored (in conjunction with OpVariable), they can only
     // be used with logical addressing operations, not physical, and only with
     // non-externally visible shader Storage Classes: Workgroup, CrossWorkgroup,
     // Private, Function, Input, and Output."
     if (bitWidth == 1) {
static Optional<int64_t> getTypeNumBytes(Type t) {
   // TODO: Add size computation for other types.
   return llvm::None;
 }
 
 SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr)
     : targetEnv(targetAttr) {
   addConversion([](Type type) -> Optional<Type> {
     // If the type is already valid in SPIR-V, directly return.
-    return spirv::SPIRVDialect::isValidType(type) ? type : Optional<Type>();
+    return type.isa<spirv::SPIRVType>() ? type : Optional<Type>();
   });
   addConversion([](IndexType indexType) {
     return SPIRVTypeConverter::getIndexType(indexType.getContext());
   });
   addConversion([this](MemRefType memRefType) -> Type {
     auto elementType = convertType(memRefType.getElementType());
     if (!elementType)
       return Type();

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

default:
     return constOp.emitOpError("cannot have value of type ") << valueType;
   }
 
   return success();
 }
 
 bool spirv::ConstantOp::isBuildableWith(Type type) {
   // Must be valid SPIR-V type first.
-  if (!SPIRVDialect::isValidType(type))
+  if (!type.isa<spirv::SPIRVType>())
     return false;
 
   if (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
       type.getKind() <= spirv::TypeKind::LAST_SPIRV_TYPE) {
     // TODO(antiagainst): support constant struct
     return type.isa<spirv::ArrayType>();
   }
 
static LogicalResult verify(spirv::SpecConstantOp constOp) {
 
   auto value = constOp.default_value();
 
   switch (value.getKind()) {
   case StandardAttributes::Bool:
   case StandardAttributes::Integer:
   case StandardAttributes::Float: {
     // Make sure bitwidth is allowed.
-    if (!spirv::SPIRVDialect::isValidType(value.getType()))
+    if (!value.getType().isa<spirv::SPIRVType>())
       return constOp.emitOpError("default value bitwidth disallowed");
     return success();
   }
   default:
     return constOp.emitOpError(
         "default value can only be a bool, integer, or float scalar");
   }
 }

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 // CompositeType
 //===----------------------------------------------------------------------===//
 
 bool CompositeType::classof(Type type) {
   switch (type.getKind()) {
   case TypeKind::Array:
   case TypeKind::RuntimeArray:
   case TypeKind::Struct:
-  case StandardTypes::Vector:
     return true;
+  case StandardTypes::Vector:
+    return isValid(type.cast<VectorType>());
   default:
     return false;
   }
 }
 
+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::RuntimeArray:
     return cast<RuntimeArrayType>().getElementType();
   case spirv::TypeKind::Struct:
     return cast<StructType>().getElementType(index);
void RuntimeArrayType::getCapabilities(
   }
   getElementType().cast<SPIRVType>().getCapabilities(capabilities, storage);
 }
 
 //===----------------------------------------------------------------------===//
 // ScalarType
 //===----------------------------------------------------------------------===//
 
-bool ScalarType::classof(Type type) { return type.isIntOrFloat(); }
+bool ScalarType::classof(Type type) {
+  if (auto floatType = type.dyn_cast<FloatType>()) {
+    return isValid(floatType);
+  }
+  if (auto intType = type.dyn_cast<IntegerType>()) {
+    return isValid(intType);
+  }
+  return false;
+}
+
+bool ScalarType::isValid(FloatType type) { return !type.isBF16(); }
+
+bool ScalarType::isValid(IntegerType type) {
+  switch (type.getWidth()) {
+  case 1:
+  case 8:
+  case 16:
+  case 32:
+  case 64:
+    return true;
+  default:
+    return false;
+  }
+}
 
 void ScalarType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                                Optional<StorageClass> storage) {
   // 8- or 16-bit integer/floating-point numbers will require extra extensions
   // to appear in interface storage classes. See SPV_KHR_16bit_storage and
   // SPV_KHR_8bit_storage for more details.
   if (!storage)
     return;

 #undef WIDTH_CASE
 }
 
 //===----------------------------------------------------------------------===//
 // SPIRVType
 //===----------------------------------------------------------------------===//
 
 bool SPIRVType::classof(Type type) {
-  return type.isa<ScalarType>() || type.isa<VectorType>() ||
-         (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
-          type.getKind() <= TypeKind::LAST_SPIRV_TYPE);
+  // Allow SPIR-V dialect types
+  if (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
+      type.getKind() <= TypeKind::LAST_SPIRV_TYPE)
+    return true;
+  if (type.isa<ScalarType>())
+    return true;
+  if (auto vectorType = type.dyn_cast<VectorType>())
+    return CompositeType::isValid(vectorType);
+  return false;
+}
+
+bool SPIRVType::isScalarOrVector() {
+  return isIntOrFloat() || isa<VectorType>();
 }
 
 void SPIRVType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                               Optional<StorageClass> storage) {
   if (auto scalarType = dyn_cast<ScalarType>()) {
     scalarType.getExtensions(extensions, storage);
   } else if (auto compositeType = dyn_cast<CompositeType>()) {
     compositeType.getExtensions(extensions, storage);

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVLowering.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
-/// Checks if the `type` is a scalar or vector type. It is assumed that they are
-/// valid for SPIR-V dialect already.
-static bool isScalarOrVectorType(Type type) {
-  return spirv::SPIRVDialect::isValidScalarType(type) || type.isa<VectorType>();
-}
-
 /// Creates a global variable for an argument based on the ABI info.
 static spirv::GlobalVariableOp
 createGlobalVarForEntryPointArgument(OpBuilder &builder, spirv::FuncOp funcOp,
                                      unsigned argIndex,
                                      spirv::InterfaceVarABIAttr abiInfo) {
   auto spirvModule = funcOp.getParentOfType<spirv::ModuleOp>();
   if (!spirvModule)
     return nullptr;
 
   OpBuilder::InsertionGuard moduleInsertionGuard(builder);
   builder.setInsertionPoint(funcOp.getOperation());
   std::string varName =
       funcOp.getName().str() + "_arg_" + std::to_string(argIndex);
 
   // Get the type of variable. If this is a scalar/vector type and has an ABI
   // info create a variable of type !spv.ptr<!spv.struct<elementType>>. If not
   // it must already be a !spv.ptr<!spv.struct<...>>.
   auto varType = funcOp.getType().getInput(argIndex);
-  if (isScalarOrVectorType(varType)) {
+  if (varType.cast<spirv::SPIRVType>().isScalarOrVector()) {
     auto storageClass =
         static_cast<spirv::StorageClass>(abiInfo.storage_class().getInt());
     varType =
         spirv::PointerType::get(spirv::StructType::get(varType), storageClass);
   }
   auto varPtrType = varType.cast<spirv::PointerType>();
   auto varPointeeType = varPtrType.getPointeeType().cast<spirv::StructType>();
 
for (auto argType : llvm::enumerate(funcOp.getType().getInputs())) {
     Value replacement =
         rewriter.create<spirv::AddressOfOp>(funcOp.getLoc(), var);
     // Check if the arg is a scalar or vector type. In that case, the value
     // needs to be loaded into registers.
     // TODO(ravishankarm) : This is loading value of the scalar into registers
     // at the start of the function. It is probably better to do the load just
     // before the use. There might be multiple loads and currently there is no
     // easy way to replace all uses with a sequence of operations.
-    if (isScalarOrVectorType(argType.value())) {
+    if (argType.value().cast<spirv::SPIRVType>().isScalarOrVector()) {
       auto indexType = SPIRVTypeConverter::getIndexType(funcOp.getContext());
       auto zero =
           spirv::ConstantOp::getZero(indexType, funcOp.getLoc(), &rewriter);
       auto loadPtr = rewriter.create<spirv::AccessChainOp>(
           funcOp.getLoc(), replacement, zero.constant());
       replacement = rewriter.create<spirv::LoadOp>(funcOp.getLoc(), loadPtr);
     }
     signatureConverter.remapInput(argType.index(), replacement);