[mlir][spirv] Fix result type for arith.cmpi/cmpf conversion

[antiagainst]

Summary

We cannot directly use the original result type; instead we need
to deduce it from the converted operand type. This addresses
invalid ops generated from converting single element vectors.

Depends on D127573

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

 static bool isBoolScalarOrVector(Type 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;
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp with composite type
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantCompositeOpPattern::matchAndRewrite(
     arith::ConstantOp constOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   auto srcType = constOp.getType().dyn_cast<ShapedType>();
LogicalResult CmpIOpBooleanPattern::matchAndRewrite(
     arith::CmpIOp op, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Type operandType = op.getLhs().getType();
   if (!isBoolScalarOrVector(operandType))
     return failure();
 
   switch (op.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
-  case cmpPredicate:                                                           \
-    rewriter.replaceOpWithNewOp<spirvOp>(op, op.getResult().getType(),         \
-                                         adaptor.getLhs(), adaptor.getRhs());  \
-    return success();
+  case cmpPredicate: {                                                         \
+    rewriter.replaceOpWithNewOp<spirvOp>(op, adaptor.getLhs(),                 \
+                                         adaptor.getRhs());                    \
+    return success();                                                          \
+  }
 
     DISPATCH(arith::CmpIPredicate::eq, spirv::LogicalEqualOp);
     DISPATCH(arith::CmpIPredicate::ne, spirv::LogicalNotEqualOp);
 
 #undef DISPATCH
   default:;
   }
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOpPattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 CmpIOpPattern::matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
-  Type operandType = op.getLhs().getType();
-  if (isBoolScalarOrVector(operandType))
+  Type srcType = op.getLhs().getType();
+  if (isBoolScalarOrVector(srcType))
+    return failure();
+  Type dstType = getTypeConverter()->convertType(srcType);
+  if (!dstType)
     return failure();
 
   switch (op.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
     if (spirvOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&            \
-        operandType != this->getTypeConverter()->convertType(operandType)) {   \
+        srcType != dstType && !hasSameBitwidth(srcType, dstType)) {            \
       return op.emitError(                                                     \
           "bitwidth emulation is not implemented yet on unsigned op");         \
     }                                                                          \
-    rewriter.replaceOpWithNewOp<spirvOp>(op, op.getResult().getType(),         \
-                                         adaptor.getLhs(), adaptor.getRhs());  \
+    rewriter.replaceOpWithNewOp<spirvOp>(op, adaptor.getLhs(),                 \
+                                         adaptor.getRhs());                    \
     return success();
 
     DISPATCH(arith::CmpIPredicate::eq, spirv::IEqualOp);
     DISPATCH(arith::CmpIPredicate::ne, spirv::INotEqualOp);
     DISPATCH(arith::CmpIPredicate::slt, spirv::SLessThanOp);
     DISPATCH(arith::CmpIPredicate::sle, spirv::SLessThanEqualOp);
     DISPATCH(arith::CmpIPredicate::sgt, spirv::SGreaterThanOp);
     DISPATCH(arith::CmpIPredicate::sge, spirv::SGreaterThanEqualOp);

 //===----------------------------------------------------------------------===//
 
 LogicalResult
 CmpFOpPattern::matchAndRewrite(arith::CmpFOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
   switch (op.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
-    rewriter.replaceOpWithNewOp<spirvOp>(op, op.getResult().getType(),         \
-                                         adaptor.getLhs(), adaptor.getRhs());  \
+    rewriter.replaceOpWithNewOp<spirvOp>(op, adaptor.getLhs(),                 \
+                                         adaptor.getRhs());                    \
     return success();
 
     // Ordered.
     DISPATCH(arith::CmpFPredicate::OEQ, spirv::FOrdEqualOp);
     DISPATCH(arith::CmpFPredicate::OGT, spirv::FOrdGreaterThanOp);
     DISPATCH(arith::CmpFPredicate::OGE, spirv::FOrdGreaterThanEqualOp);
     DISPATCH(arith::CmpFPredicate::OLT, spirv::FOrdLessThanOp);
     DISPATCH(arith::CmpFPredicate::OLE, spirv::FOrdLessThanEqualOp);

mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir

func.func @cmpi(%arg0 : i32, %arg1 : i32) {
   %7 = arith.cmpi ule, %arg0, %arg1 : i32
   // CHECK: spv.UGreaterThan
   %8 = arith.cmpi ugt, %arg0, %arg1 : i32
   // CHECK: spv.UGreaterThanEqual
   %9 = arith.cmpi uge, %arg0, %arg1 : i32
   return
 }
 
+// CHECK-LABEL: @vec1cmpi
+func.func @vec1cmpi(%arg0 : vector<1xi32>, %arg1 : vector<1xi32>) {
+  // CHECK: spv.ULessThan
+  %0 = arith.cmpi ult, %arg0, %arg1 : vector<1xi32>
+  // CHECK: spv.SGreaterThan
+  %1 = arith.cmpi sgt, %arg0, %arg1 : vector<1xi32>
+  return
+}
+
 // CHECK-LABEL: @boolcmpi
 func.func @boolcmpi(%arg0 : i1, %arg1 : i1) {
   // CHECK: spv.LogicalEqual
   %0 = arith.cmpi eq, %arg0, %arg1 : i1
   // CHECK: spv.LogicalNotEqual
   %1 = arith.cmpi ne, %arg0, %arg1 : i1
   return
 }
 
+// CHECK-LABEL: @vec1boolcmpi
+func.func @vec1boolcmpi(%arg0 : vector<1xi1>, %arg1 : vector<1xi1>) {
+  // CHECK: spv.LogicalEqual
+  %0 = arith.cmpi eq, %arg0, %arg1 : vector<1xi1>
+  // CHECK: spv.LogicalNotEqual
+  %1 = arith.cmpi ne, %arg0, %arg1 : vector<1xi1>
+  return
+}
+
 // CHECK-LABEL: @vecboolcmpi
 func.func @vecboolcmpi(%arg0 : vector<4xi1>, %arg1 : vector<4xi1>) {
   // CHECK: spv.LogicalEqual
   %0 = arith.cmpi eq, %arg0, %arg1 : vector<4xi1>
   // CHECK: spv.LogicalNotEqual
   %1 = arith.cmpi ne, %arg0, %arg1 : vector<4xi1>
   return
 }
func.func @cmpf(%arg0 : f32, %arg1 : f32) {
   %10 = arith.cmpf ult, %arg0, %arg1 : f32
   // CHECK: FUnordLessThanEqual
   %11 = arith.cmpf ule, %arg0, %arg1 : f32
   // CHECK: spv.FUnordNotEqual
   %12 = arith.cmpf une, %arg0, %arg1 : f32
   return
 }
 
+// CHECK-LABEL: @vec1cmpf
+func.func @vec1cmpf(%arg0 : vector<1xf32>, %arg1 : vector<1xf32>) {
+  // CHECK: spv.FOrdGreaterThan
+  %0 = arith.cmpf ogt, %arg0, %arg1 : vector<1xf32>
+  // CHECK: spv.FUnordLessThan
+  %1 = arith.cmpf ult, %arg0, %arg1 : vector<1xf32>
+  return
+}
+
 } // end module
 
 // -----
 
 // With Kernel capability, we can convert NaN check to spv.Ordered/spv.Unordered.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Kernel], []>, #spv.resource_limits<>>
 } {