[mlir][spirv] Timely fail type conversion

[hanchung]

Summary

Mirror of https://reviews.llvm.org/D100058

Per the TypeConverter API contract, returning llvm:None means
other conversion rules should be tried. But we only have one
rule per input type. So there is no need to try others and we can
just directly fail, which should return nullptr. This avoids
unnecessary checks.

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

static Optional<int64_t> getTypeNumBytes(Type t) {
   return llvm::None;
 }
 
 Optional<int64_t> SPIRVTypeConverter::getConvertedTypeNumBytes(Type t) {
   return getTypeNumBytes(t);
 }
 
 /// Converts a scalar `type` to a suitable type under the given `targetEnv`.
-static Optional<Type>
-convertScalarType(const spirv::TargetEnv &targetEnv, spirv::ScalarType type,
+static Type convertScalarType(const spirv::TargetEnv &targetEnv,
+                              spirv::ScalarType type,
                               Optional<spirv::StorageClass> storageClass = {}) {
   // Get extension and capability requirements for the given type.
   SmallVector<ArrayRef<spirv::Extension>, 1> extensions;
   SmallVector<ArrayRef<spirv::Capability>, 2> capabilities;
   type.getExtensions(extensions, storageClass);
   type.getCapabilities(capabilities, storageClass);
 
   // If all requirements are met, then we can accept this type as-is.
   if (succeeded(checkCapabilityRequirements(type, targetEnv, capabilities)) &&
static Type convertScalarType(const spirv::TargetEnv &targetEnv,
 
   auto intType = type.cast<IntegerType>();
   LLVM_DEBUG(llvm::dbgs() << type << " converted to 32-bit for SPIR-V\n");
   return IntegerType::get(targetEnv.getContext(), /*width=*/32,
                           intType.getSignedness());
 }
 
 /// Converts a vector `type` to a suitable type under the given `targetEnv`.
-static Optional<Type>
-convertVectorType(const spirv::TargetEnv &targetEnv, VectorType type,
+static Type convertVectorType(const spirv::TargetEnv &targetEnv,
+                              VectorType type,
                               Optional<spirv::StorageClass> storageClass = {}) {
   if (type.getRank() == 1 && type.getNumElements() == 1)
     return type.getElementType();
 
   if (!spirv::CompositeType::isValid(type)) {
     // TODO: Vector types with more than four elements can be translated into
     // array types.
     LLVM_DEBUG(llvm::dbgs() << type << " illegal: > 4-element unimplemented\n");
-    return llvm::None;
+    return nullptr;
   }
 
   // Get extension and capability requirements for the given type.
   SmallVector<ArrayRef<spirv::Extension>, 1> extensions;
   SmallVector<ArrayRef<spirv::Capability>, 2> capabilities;
   type.cast<spirv::CompositeType>().getExtensions(extensions, storageClass);
   type.cast<spirv::CompositeType>().getCapabilities(capabilities, storageClass);
 
   // If all requirements are met, then we can accept this type as-is.
   if (succeeded(checkCapabilityRequirements(type, targetEnv, capabilities)) &&
       succeeded(checkExtensionRequirements(type, targetEnv, extensions)))
     return type;
 
   auto elementType = convertScalarType(
       targetEnv, type.getElementType().cast<spirv::ScalarType>(), storageClass);
   if (elementType)
-    return VectorType::get(type.getShape(), *elementType);
-  return llvm::None;
+    return VectorType::get(type.getShape(), elementType);
+  return nullptr;
 }
 
 /// Converts a tensor `type` to a suitable type under the given `targetEnv`.
 ///
 /// Note that this is mainly for lowering constant tensors.In SPIR-V one can
 /// create composite constants with OpConstantComposite to embed relative large
 /// constant values and use OpCompositeExtract and OpCompositeInsert to
 /// manipulate, like what we do for vectors.
-static Optional<Type> convertTensorType(const spirv::TargetEnv &targetEnv,
-                                        TensorType type) {
+static Type convertTensorType(const spirv::TargetEnv &targetEnv,
   // TODO: Handle dynamic shapes.
   if (!type.hasStaticShape()) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: dynamic shape unimplemented\n");
-    return llvm::None;
+    return nullptr;
   }
 
   auto scalarType = type.getElementType().dyn_cast<spirv::ScalarType>();
   if (!scalarType) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert non-scalar element type\n");
-    return llvm::None;
+    return nullptr;
   }
 
   Optional<int64_t> scalarSize = getTypeNumBytes(scalarType);
   Optional<int64_t> tensorSize = getTypeNumBytes(type);
   if (!scalarSize || !tensorSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element count\n");
-    return llvm::None;
+    return nullptr;
   }
 
   auto arrayElemCount = *tensorSize / *scalarSize;
   auto arrayElemType = convertScalarType(targetEnv, scalarType);
   if (!arrayElemType)
-    return llvm::None;
-  Optional<int64_t> arrayElemSize = getTypeNumBytes(*arrayElemType);
+    return nullptr;
+  Optional<int64_t> arrayElemSize = getTypeNumBytes(arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
-    return llvm::None;
+    return nullptr;
   }
 
-  return spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
+  return spirv::ArrayType::get(arrayElemType, arrayElemCount, *arrayElemSize);
 }
 
-static Optional<Type> convertMemrefType(const spirv::TargetEnv &targetEnv,
-                                        MemRefType type) {
+static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
   Optional<spirv::StorageClass> storageClass =
       SPIRVTypeConverter::getStorageClassForMemorySpace(
           type.getMemorySpaceAsInt());
   if (!storageClass) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert memory space\n");
-    return llvm::None;
+    return nullptr;
   }
 
-  Optional<Type> arrayElemType;
+  Type arrayElemType;
   Type elementType = type.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>()) {
     arrayElemType = convertVectorType(targetEnv, vecType, storageClass);
   } else if (auto scalarType = elementType.dyn_cast<spirv::ScalarType>()) {
     arrayElemType = convertScalarType(targetEnv, scalarType, storageClass);
   } else {
     LLVM_DEBUG(
         llvm::dbgs()
         << type
         << " unhandled: can only convert scalar or vector element type\n");
-    return llvm::None;
+    return nullptr;
   }
   if (!arrayElemType)
-    return llvm::None;
+    return nullptr;
 
   Optional<int64_t> elementSize = getTypeNumBytes(elementType);
   if (!elementSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element size\n");
-    return llvm::None;
+    return nullptr;
   }
 
   if (!type.hasStaticShape()) {
-    auto arrayType = spirv::RuntimeArrayType::get(*arrayElemType, *elementSize);
+    auto arrayType = spirv::RuntimeArrayType::get(arrayElemType, *elementSize);
     // Wrap in a struct to satisfy Vulkan interface requirements.
     auto structType = spirv::StructType::get(arrayType, 0);
     return spirv::PointerType::get(structType, *storageClass);
   }
 
   Optional<int64_t> memrefSize = getTypeNumBytes(type);
   if (!memrefSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element count\n");
-    return llvm::None;
+    return nullptr;
   }
 
   auto arrayElemCount = *memrefSize / *elementSize;
 
-  Optional<int64_t> arrayElemSize = getTypeNumBytes(*arrayElemType);
+  Optional<int64_t> arrayElemSize = getTypeNumBytes(arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
-    return llvm::None;
+    return nullptr;
   }
 
   auto arrayType =
-      spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
+      spirv::ArrayType::get(arrayElemType, arrayElemCount, *arrayElemSize);
 
   // Wrap in a struct to satisfy Vulkan interface requirements. Memrefs with
   // workgroup storage class do not need the struct to be laid out explicitly.
   auto structType = *storageClass == spirv::StorageClass::Workgroup
                         ? spirv::StructType::get(arrayType)
                         : spirv::StructType::get(arrayType, 0);
   return spirv::PointerType::get(structType, *storageClass);
 }
 
 SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr)
     : targetEnv(targetAttr) {
   // Add conversions. The order matters here: later ones will be tried earlier.
 
-  // All other cases failed. Then we cannot convert this type.
-  addConversion([](Type type) { return llvm::None; });
-
   // Allow all SPIR-V dialect specific types. This assumes all builtin types
   // adopted in the SPIR-V dialect (i.e., IntegerType, FloatType, VectorType)
   // were tried before.
   //
   // TODO: this assumes that the SPIR-V types are valid to use in
   // the given target environment, which should be the case if the whole
   // pipeline is driven by the same target environment. Still, we probably still
   // want to validate and convert to be safe.
   addConversion([](spirv::SPIRVType type) { return type; });
 
   addConversion([](IndexType indexType) {
     return SPIRVTypeConverter::getIndexType(indexType.getContext());
   });
 
   addConversion([this](IntegerType intType) -> Optional<Type> {
     if (auto scalarType = intType.dyn_cast<spirv::ScalarType>())
       return convertScalarType(targetEnv, scalarType);
-    return llvm::None;
+    return Type();
   });
 
   addConversion([this](FloatType floatType) -> Optional<Type> {
     if (auto scalarType = floatType.dyn_cast<spirv::ScalarType>())
       return convertScalarType(targetEnv, scalarType);
-    return llvm::None;
+    return Type();
   });
 
   addConversion([this](VectorType vectorType) {
     return convertVectorType(targetEnv, vectorType);
   });
 
   addConversion([this](TensorType tensorType) {
     return convertTensorType(targetEnv, tensorType);