[mlir][StandardToSPIRV] Handle conversion of cmpi operation with i1 type operands.

[mravishankar]

Summary

The instructions used to convert std.cmpi cannot have i1 types
according to SPIR-V specification. A different set of operations are
specified in the SPIR-V spec for comparing boolean types. Enhance the
StandardToSPIRV lowering to target these instructions when operands to
std.cmpi operation are of i1 type.

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

197

Update the comment?

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

 public:
   using SPIRVOpLowering<CmpFOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(CmpFOp cmpFOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts integer compare operation on i1 type opearnds to SPIR-V ops.
+class BoolCmpIOpPattern final : public SPIRVOpLowering<CmpIOp> {
+public:
+  using SPIRVOpLowering<CmpIOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts integer compare operation to SPIR-V ops.
 class CmpIOpPattern final : public SPIRVOpLowering<CmpIOp> {
 public:
   using SPIRVOpLowering<CmpIOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
#undef DISPATCH
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
+BoolCmpIOpPattern::matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
+                                   ConversionPatternRewriter &rewriter) const {
+  CmpIOpOperandAdaptor cmpIOpOperands(operands);
+
+  Type operandType = cmpIOp.lhs().getType();
+  if (!operandType.isa<IntegerType>() ||
+      operandType.cast<IntegerType>().getWidth() != 1)
+    return failure();
+
+  switch (cmpIOp.getPredicate()) {
+#define DISPATCH(cmpPredicate, spirvOp)                                        \
+  case cmpPredicate:                                                           \
+    rewriter.replaceOpWithNewOp<spirvOp>(cmpIOp, cmpIOp.getResult().getType(), \
+                                         cmpIOpOperands.lhs(),                 \
+                                         cmpIOpOperands.rhs());                \
+    return success();
+
+    DISPATCH(CmpIPredicate::eq, spirv::LogicalEqualOp);
+    DISPATCH(CmpIPredicate::ne, spirv::LogicalNotEqualOp);
+
+#undef DISPATCH
+  default:;
+  }
+  return failure();
+}
+
+LogicalResult
 CmpIOpPattern::matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
                                ConversionPatternRewriter &rewriter) const {
   CmpIOpOperandAdaptor cmpIOpOperands(operands);
 
+  Type operandType = cmpIOp.lhs().getType();
+  if (operandType.isa<IntegerType>() &&
+      operandType.cast<IntegerType>().getWidth() == 1)
+    return failure();
+
   switch (cmpIOp.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
     rewriter.replaceOpWithNewOp<spirvOp>(cmpIOp, cmpIOp.getResult().getType(), \
                                          cmpIOpOperands.lhs(),                 \
                                          cmpIOpOperands.rhs());                \
     return success();
 
patterns.insert<
       UnaryAndBinaryOpPattern<SubFOp, spirv::FSubOp>,
       UnaryAndBinaryOpPattern<SubIOp, spirv::ISubOp>,
       UnaryAndBinaryOpPattern<TanhOp, spirv::GLSLTanhOp>,
       UnaryAndBinaryOpPattern<UnsignedDivIOp, spirv::UDivOp>,
       UnaryAndBinaryOpPattern<UnsignedRemIOp, spirv::UModOp>,
       UnaryAndBinaryOpPattern<UnsignedShiftRightOp, spirv::ShiftRightLogicalOp>,
       BitwiseOpPattern<AndOp, spirv::LogicalAndOp, spirv::BitwiseAndOp>,
       BitwiseOpPattern<OrOp, spirv::LogicalOrOp, spirv::BitwiseOrOp>,
-      ConstantCompositeOpPattern, ConstantScalarOpPattern, CmpFOpPattern,
-      CmpIOpPattern, LoadOpPattern, ReturnOpPattern, SelectOpPattern,
-      StoreOpPattern, TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
+      BoolCmpIOpPattern, ConstantCompositeOpPattern, ConstantScalarOpPattern,
+      CmpFOpPattern, CmpIOpPattern, LoadOpPattern, ReturnOpPattern,
+      SelectOpPattern, StoreOpPattern,
+      TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
       TypeCastingOpPattern<FPExtOp, spirv::FConvertOp>,
       TypeCastingOpPattern<FPTruncOp, spirv::FConvertOp>, XOrOpPattern>(
       context, typeConverter);
 }
 } // namespace mlir

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

func @cmpi(%arg0 : i32, %arg1 : i32) {
   %7 = cmpi "ule", %arg0, %arg1 : i32
   // CHECK: spv.UGreaterThan
   %8 = cmpi "ugt", %arg0, %arg1 : i32
   // CHECK: spv.UGreaterThanEqual
   %9 = cmpi "uge", %arg0, %arg1 : i32
   return
 }
 
+// CHECK-LABEL: @boolcmpi
+func @boolcmpi(%arg0 : i1, %arg1 : i1) {
+  // CHECK: spv.LogicalEqual
+  %0 = cmpi "eq", %arg0, %arg1 : i1
+  // CHECK: spv.LogicalNotEqual
+  %1 = cmpi "ne", %arg0, %arg1 : i1
+  return
+}
+
 } // end module
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // std.constant
 //===----------------------------------------------------------------------===//