[mlir][arith][spirv] Account for possible type conversion failures

[kuhar]

Summary

Check results of all type conversions in --convert-arith-to-spirv.

Fixes: https://github.com/llvm/llvm-project/issues/59496

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

if (status != APFloat::opOK || losesInfo) {
     return FloatAttr();
   }
 
   return builder.getF32FloatAttr(dstVal.convertToFloat());
 }
 
 /// Returns true if the given `type` is a boolean scalar or vector type.
 static bool isBoolScalarOrVector(Type type) {
+  assert(type && "Not a valid type");
   if (type.isInteger(1))
     return true;
+
   if (auto vecType = type.dyn_cast<VectorType>())
     return vecType.getElementType().isInteger(1);
+
   return false;
 }
 
 /// Returns true if scalar/vector type `a` and `b` have the same number of
 /// bitwidth.
 static bool hasSameBitwidth(Type a, Type b) {
   auto getNumBitwidth = [](Type type) {
     unsigned bw = 0;
     if (type.isIntOrFloat())
       bw = type.getIntOrFloatBitWidth();
     else if (auto vecType = type.dyn_cast<VectorType>())
       bw = vecType.getElementTypeBitWidth() * vecType.getNumElements();
     return bw;
   };
   unsigned aBW = getNumBitwidth(a);
   unsigned bBW = getNumBitwidth(b);
   return aBW != 0 && bBW != 0 && aBW == bBW;
 }
 
+/// Returns a source type conversion failure for `srcType` and operation `op`.
+static LogicalResult
+getTypeConversionFailure(ConversionPatternRewriter &rewriter, Operation *op,
+                         Type srcType) {
+  return rewriter.notifyMatchFailure(
+      op->getLoc(),
+      llvm::formatv("failed to convert source type '{0}'", srcType));
+}
+
+/// Returns a source type conversion failure for the result type of `op`.
+static LogicalResult
+getTypeConversionFailure(ConversionPatternRewriter &rewriter, Operation *op) {
+  assert(op->getNumResults() == 1);
+  return getTypeConversionFailure(rewriter, op, op->getResultTypes().front());
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp with composite type
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantCompositeOpPattern::matchAndRewrite(
     arith::ConstantOp constOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   auto srcType = constOp.getType().dyn_cast<ShapedType>();

 //===----------------------------------------------------------------------===//
 
 template <typename Op, typename SPIRVLogicalOp, typename SPIRVBitwiseOp>
 LogicalResult
 BitwiseOpPattern<Op, SPIRVLogicalOp, SPIRVBitwiseOp>::matchAndRewrite(
     Op op, typename Op::Adaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   assert(adaptor.getOperands().size() == 2);
-  auto dstType =
-      this->getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = this->getTypeConverter()->convertType(op.getType());
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op);
+
   if (isBoolScalarOrVector(adaptor.getOperands().front().getType())) {
     rewriter.template replaceOpWithNewOp<SPIRVLogicalOp>(op, dstType,
                                                          adaptor.getOperands());
   } else {
     rewriter.template replaceOpWithNewOp<SPIRVBitwiseOp>(op, dstType,
                                                          adaptor.getOperands());
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // XOrIOpLogicalPattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult XOrIOpLogicalPattern::matchAndRewrite(
     arith::XOrIOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   assert(adaptor.getOperands().size() == 2);
 
   if (isBoolScalarOrVector(adaptor.getOperands().front().getType()))
     return failure();
 
   Type dstType = getTypeConverter()->convertType(op.getType());
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op);
+
   rewriter.replaceOpWithNewOp<spirv::BitwiseXorOp>(op, dstType,
                                                    adaptor.getOperands());
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // XOrIOpBooleanPattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult XOrIOpBooleanPattern::matchAndRewrite(
     arith::XOrIOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   assert(adaptor.getOperands().size() == 2);
 
   if (!isBoolScalarOrVector(adaptor.getOperands().front().getType()))
     return failure();
 
   Type dstType = getTypeConverter()->convertType(op.getType());
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op);
+
   rewriter.replaceOpWithNewOp<spirv::LogicalNotEqualOp>(op, dstType,
                                                         adaptor.getOperands());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // UIToFPI1Pattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 UIToFPI1Pattern::matchAndRewrite(arith::UIToFPOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   Type srcType = adaptor.getOperands().front().getType();
   if (!isBoolScalarOrVector(srcType))
     return failure();
 
-  Type dstType = getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = getTypeConverter()->convertType(op.getType());
+  if (!dstType)
+    return getTypeConversionFailure(rewriter, op);
+
   Location loc = op.getLoc();
   Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
   Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
   rewriter.replaceOpWithNewOp<spirv::SelectOp>(
       op, dstType, adaptor.getOperands().front(), one, zero);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ExtSII1Pattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ExtSII1Pattern::matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   Value operand = adaptor.getIn();
   if (!isBoolScalarOrVector(operand.getType()))
     return failure();
 
   Location loc = op.getLoc();
-  Type dstType = getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = getTypeConverter()->convertType(op.getType());
+  if (!dstType)
+    return getTypeConversionFailure(rewriter, op);
 
   Value allOnes;
   if (auto intTy = dstType.dyn_cast<IntegerType>()) {
     unsigned componentBitwidth = intTy.getWidth();
     allOnes = rewriter.create<spirv::ConstantOp>(
         loc, intTy,
         rewriter.getIntegerAttr(intTy, APInt::getAllOnes(componentBitwidth)));
   } else if (auto vectorTy = dstType.dyn_cast<VectorType>()) {

 
 LogicalResult
 ExtUII1Pattern::matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   Type srcType = adaptor.getOperands().front().getType();
   if (!isBoolScalarOrVector(srcType))
     return failure();
 
-  Type dstType = getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = getTypeConverter()->convertType(op.getType());
+  if (!dstType)
+    return getTypeConversionFailure(rewriter, op);
+
   Location loc = op.getLoc();
   Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
   Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
   rewriter.replaceOpWithNewOp<spirv::SelectOp>(
       op, dstType, adaptor.getOperands().front(), one, zero);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TruncII1Pattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 TruncII1Pattern::matchAndRewrite(arith::TruncIOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
-  Type dstType = getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = getTypeConverter()->convertType(op.getType());
+  if (!dstType)
+    return getTypeConversionFailure(rewriter, op);
+
   if (!isBoolScalarOrVector(dstType))
     return failure();
 
   Location loc = op.getLoc();
   auto srcType = adaptor.getOperands().front().getType();
   // Check if (x & 1) == 1.
   Value mask = spirv::ConstantOp::getOne(srcType, loc, rewriter);
   Value maskedSrc = rewriter.create<spirv::BitwiseAndOp>(

 //===----------------------------------------------------------------------===//
 
 template <typename Op, typename SPIRVOp>
 LogicalResult TypeCastingOpPattern<Op, SPIRVOp>::matchAndRewrite(
     Op op, typename Op::Adaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   assert(adaptor.getOperands().size() == 1);
   Type srcType = adaptor.getOperands().front().getType();
-  Type dstType =
-      this->getTypeConverter()->convertType(op.getResult().getType());
+  Type dstType = this->getTypeConverter()->convertType(op.getType());
+  if (!dstType)
+    return getTypeConversionFailure(rewriter, op);
+
   if (isBoolScalarOrVector(srcType) || isBoolScalarOrVector(dstType))
     return failure();
+
   if (dstType == srcType) {
     // Due to type conversion, we are seeing the same source and target type.
     // Then we can just erase this operation by forwarding its operand.
     rewriter.replaceOp(op, adaptor.getOperands().front());
   } else {
     rewriter.template replaceOpWithNewOp<SPIRVOp>(op, dstType,
                                                   adaptor.getOperands());
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOpBooleanPattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult CmpIOpBooleanPattern::matchAndRewrite(
     arith::CmpIOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Type srcType = op.getLhs().getType();
   if (!isBoolScalarOrVector(srcType))
     return failure();
   Type dstType = getTypeConverter()->convertType(srcType);
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op, srcType);
 
   switch (op.getPredicate()) {
   case arith::CmpIPredicate::eq: {
     rewriter.replaceOpWithNewOp<spirv::LogicalEqualOp>(op, adaptor.getLhs(),
                                                        adaptor.getRhs());
     return success();
   }
   case arith::CmpIPredicate::ne: {

 LogicalResult
 CmpIOpPattern::matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
   Type srcType = op.getLhs().getType();
   if (isBoolScalarOrVector(srcType))
     return failure();
   Type dstType = getTypeConverter()->convertType(srcType);
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op, srcType);
 
   switch (op.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
     if (spirvOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&            \
         srcType != dstType && !hasSameBitwidth(srcType, dstType)) {            \
       return op.emitError(                                                     \
           "bitwidth emulation is not implemented yet on unsigned op");         \

 
 template <typename Op, typename SPIRVOp>
 LogicalResult MinMaxFOpPattern<Op, SPIRVOp>::matchAndRewrite(
     Op op, typename Op::Adaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   auto *converter = this->template getTypeConverter<SPIRVTypeConverter>();
   Type dstType = converter->convertType(op.getType());
   if (!dstType)
-    return failure();
+    return getTypeConversionFailure(rewriter, op);
 
   // arith.maxf/minf:
   //   "if one of the arguments is NaN, then the result is also NaN."
   // spirv.GL.FMax/FMin
   //   "which operand is the result is undefined if one of the operands
   //   is a NaN."
   // spirv.CL.fmax/fmin:
   //   "If one argument is a NaN, Fmin returns the other argument."

mlir/test/Conversion/ArithToSPIRV/arith-to-spirv-unsupported.mlir

 // Check that we do not emualte f64 by truncating to i32.
 func.func @unsupported_f64(%arg0: f64) {
   // expected-error@+1 {{failed to legalize operation 'arith.addf'}}
   %2 = arith.addf %arg0, %arg0: f64
   return
 }
 
 } // end module
+
+// -----
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
+} {
+
+// i64 is not a valid result type in this target env.
+func.func @type_conversion_failure(%arg0: i32) {
+  // expected-error@+1 {{failed to legalize operation 'arith.extsi'}}
+  %2 = arith.extsi %arg0 : i32 to i64
+  return
+}
+
+} // end module