[mlir][spirv] Add lowering for std cmp ops.

[mravishankar]

Could you add some serialization tests as well?

[denis13]

Added serialization tests.

[mravishankar]

THanks Denis. Looks great!

[antiagainst]

Thanks!

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

 public:
   using SPIRVOpLowering<ConstantOp>::SPIRVOpLowering;
 
   PatternMatchResult
   matchAndRewrite(ConstantOp constIndexOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Convert floating-point comparison operations to SPIR-V dialect.
+class CmpFOpConversion final : public SPIRVOpLowering<CmpFOp> {
+public:
+  using SPIRVOpLowering<CmpFOp>::SPIRVOpLowering;
+
+  PatternMatchResult
+  matchAndRewrite(CmpFOp cmpFOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Convert compare operation to SPIR-V dialect.
 class CmpIOpConversion final : public SPIRVOpLowering<CmpIOp> {
 public:
   using SPIRVOpLowering<CmpIOp>::SPIRVOpLowering;
 
   PatternMatchResult
   matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
PatternMatchResult ConstantIndexOpConversion::matchAndRewrite(
   auto spirvConstVal =
       rewriter.getIntegerAttr(spirvConstType, constAttr.getInt());
   rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constIndexOp, spirvConstType,
                                                  spirvConstVal);
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
+// CmpFOp
+//===----------------------------------------------------------------------===//
+
+PatternMatchResult
+CmpFOpConversion::matchAndRewrite(CmpFOp cmpFOp, ArrayRef<Value> operands,
+                                  ConversionPatternRewriter &rewriter) const {
+  CmpFOpOperandAdaptor cmpFOpOperands(operands);
+
+  switch (cmpFOp.getPredicate()) {
+#define DISPATCH(cmpPredicate, spirvOp)                                        \
+  case cmpPredicate:                                                           \
+    rewriter.replaceOpWithNewOp<spirvOp>(                                      \
+        cmpFOp, cmpFOp.getResult()->getType(), cmpFOpOperands.lhs(),           \
+        cmpFOpOperands.rhs());                                                 \
+    return matchSuccess();
+
+    // Ordered.
+    DISPATCH(CmpFPredicate::OEQ, spirv::FOrdEqualOp);
+    DISPATCH(CmpFPredicate::OGT, spirv::FOrdGreaterThanOp);
+    DISPATCH(CmpFPredicate::OGE, spirv::FOrdGreaterThanEqualOp);
+    DISPATCH(CmpFPredicate::OLT, spirv::FOrdLessThanOp);
+    DISPATCH(CmpFPredicate::OLE, spirv::FOrdLessThanEqualOp);
+    DISPATCH(CmpFPredicate::ONE, spirv::FOrdNotEqualOp);
+    // Unordered.
+    DISPATCH(CmpFPredicate::UEQ, spirv::FUnordEqualOp);
+    DISPATCH(CmpFPredicate::UGT, spirv::FUnordGreaterThanOp);
+    DISPATCH(CmpFPredicate::UGE, spirv::FUnordGreaterThanEqualOp);
+    DISPATCH(CmpFPredicate::ULT, spirv::FUnordLessThanOp);
+    DISPATCH(CmpFPredicate::ULE, spirv::FUnordLessThanEqualOp);
+    DISPATCH(CmpFPredicate::UNE, spirv::FUnordNotEqualOp);
+
+#undef DISPATCH
+
+  default:
+    break;
+  }
+  return matchFailure();
+}
+
+//===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult
 CmpIOpConversion::matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
                                   ConversionPatternRewriter &rewriter) const {
   CmpIOpOperandAdaptor cmpIOpOperands(operands);
 
   switch (cmpIOp.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
     rewriter.replaceOpWithNewOp<spirvOp>(                                      \
         cmpIOp, cmpIOp.getResult()->getType(), cmpIOpOperands.lhs(),           \
         cmpIOpOperands.rhs());                                                 \
     return matchSuccess();
 
     DISPATCH(CmpIPredicate::eq, spirv::IEqualOp);
     DISPATCH(CmpIPredicate::ne, spirv::INotEqualOp);
     DISPATCH(CmpIPredicate::slt, spirv::SLessThanOp);
     DISPATCH(CmpIPredicate::sle, spirv::SLessThanEqualOp);
     DISPATCH(CmpIPredicate::sgt, spirv::SGreaterThanOp);
     DISPATCH(CmpIPredicate::sge, spirv::SGreaterThanEqualOp);
+    DISPATCH(CmpIPredicate::ult, spirv::ULessThanOp);
+    DISPATCH(CmpIPredicate::ule, spirv::ULessThanEqualOp);
+    DISPATCH(CmpIPredicate::ugt, spirv::UGreaterThanOp);
+    DISPATCH(CmpIPredicate::uge, spirv::UGreaterThanEqualOp);
 
 #undef DISPATCH
-
-  default:
-    break;
   }
   return matchFailure();
 }
 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 

 } // namespace
 
 namespace mlir {
 void populateStandardToSPIRVPatterns(MLIRContext *context,
                                      SPIRVTypeConverter &typeConverter,
                                      OwningRewritePatternList &patterns) {
   // Add patterns that lower operations into SPIR-V dialect.
   populateWithGenerated(context, &patterns);
-  patterns.insert<ConstantIndexOpConversion, CmpIOpConversion,
+  patterns.insert<ConstantIndexOpConversion, CmpFOpConversion, CmpIOpConversion,
                   IntegerOpConversion<AddIOp, spirv::IAddOp>,
                   IntegerOpConversion<MulIOp, spirv::IMulOp>,
                   IntegerOpConversion<SignedDivIOp, spirv::SDivOp>,
                   IntegerOpConversion<SignedRemIOp, spirv::SModOp>,
                   IntegerOpConversion<SubIOp, spirv::ISubOp>, LoadOpConversion,
                   ReturnOpConversion, SelectOpConversion, StoreOpConversion>(
       context, typeConverter);
 }
 } // namespace mlir

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

func @shift_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>) {
   // CHECK: spv.ShiftRightArithmetic
   %1 = shift_right_signed %arg0, %arg1 : vector<4xi32>
   // CHECK: spv.ShiftRightLogical
   %2 = shift_right_unsigned %arg0, %arg1 : vector<4xi32>
   return
 }
 
 //===----------------------------------------------------------------------===//
+// std.cmpf
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @cmpf
+func @cmpf(%arg0 : f32, %arg1 : f32) {
+  // CHECK: spv.FOrdEqual
+  %1 = cmpf "oeq", %arg0, %arg1 : f32
+  // CHECK: spv.FOrdGreaterThan
+  %2 = cmpf "ogt", %arg0, %arg1 : f32
+  // CHECK: spv.FOrdGreaterThanEqual
+  %3 = cmpf "oge", %arg0, %arg1 : f32
+  // CHECK: spv.FOrdLessThan
+  %4 = cmpf "olt", %arg0, %arg1 : f32
+  // CHECK: spv.FOrdLessThanEqual
+  %5 = cmpf "ole", %arg0, %arg1 : f32
+  // CHECK: spv.FOrdNotEqual
+  %6 = cmpf "one", %arg0, %arg1 : f32
+  // CHECK: spv.FUnordEqual
+  %7 = cmpf "ueq", %arg0, %arg1 : f32
+  // CHECK: spv.FUnordGreaterThan
+  %8 = cmpf "ugt", %arg0, %arg1 : f32
+  // CHECK: spv.FUnordGreaterThanEqual
+  %9 = cmpf "uge", %arg0, %arg1 : f32
+  // CHECK: spv.FUnordLessThan
+  %10 = cmpf "ult", %arg0, %arg1 : f32
+  // CHECK: FUnordLessThanEqual
+  %11 = cmpf "ule", %arg0, %arg1 : f32
+  // CHECK: spv.FUnordNotEqual
+  %12 = cmpf "une", %arg0, %arg1 : f32
+  return
+}
+
+//===----------------------------------------------------------------------===//
 // std.cmpi
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @cmpi
 func @cmpi(%arg0 : i32, %arg1 : i32) {
   // CHECK: spv.IEqual
   %0 = cmpi "eq", %arg0, %arg1 : i32
   // CHECK: spv.INotEqual
   %1 = cmpi "ne", %arg0, %arg1 : i32
   // CHECK: spv.SLessThan
   %2 = cmpi "slt", %arg0, %arg1 : i32
   // CHECK: spv.SLessThanEqual
   %3 = cmpi "sle", %arg0, %arg1 : i32
   // CHECK: spv.SGreaterThan
   %4 = cmpi "sgt", %arg0, %arg1 : i32
   // CHECK: spv.SGreaterThanEqual
   %5 = cmpi "sge", %arg0, %arg1 : i32
+  // CHECK: spv.ULessThan
+  %6 = cmpi "ult", %arg0, %arg1 : i32
+  // CHECK: spv.ULessThanEqual
+  %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
 }
 
 //===----------------------------------------------------------------------===//
 // std.constant
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @constant

mlir/test/Dialect/SPIRV/Serialization/logical-ops.mlir

func @ult_vector(%arg0: vector<4xi32>, %arg1: vector<4xi32>) {
     %0 = spv.ULessThan %arg0, %arg1 : vector<4xi32>
     spv.Return
   }
   func @ulte_vector(%arg0: vector<4xi32>, %arg1: vector<4xi32>)  {
     // CHECK: {{.*}} = spv.ULessThanEqual {{.*}}, {{.*}} : vector<4xi32>
     %0 = spv.ULessThanEqual %arg0, %arg1 : vector<4xi32>
     spv.Return
   }
+  func @cmpf(%arg0 : f32, %arg1 : f32) {
+    // CHECK: spv.FOrdEqual
+    %1 = spv.FOrdEqual %arg0, %arg1 : f32
+    // CHECK: spv.FOrdGreaterThan
+    %2 = spv.FOrdGreaterThan %arg0, %arg1 : f32
+    // CHECK: spv.FOrdGreaterThanEqual
+    %3 = spv.FOrdGreaterThanEqual %arg0, %arg1 : f32
+    // CHECK: spv.FOrdLessThan
+    %4 = spv.FOrdLessThan %arg0, %arg1 : f32
+    // CHECK: spv.FOrdLessThanEqual
+    %5 = spv.FOrdLessThanEqual %arg0, %arg1 : f32
+    // CHECK: spv.FOrdNotEqual
+    %6 = spv.FOrdNotEqual %arg0, %arg1 : f32
+    // CHECK: spv.FUnordEqual
+    %7 = spv.FUnordEqual %arg0, %arg1 : f32
+    // CHECK: spv.FUnordGreaterThan
+    %8 = spv.FUnordGreaterThan %arg0, %arg1 : f32
+    // CHECK: spv.FUnordGreaterThanEqual
+    %9 = spv.FUnordGreaterThanEqual %arg0, %arg1 : f32
+    // CHECK: spv.FUnordLessThan
+    %10 = spv.FUnordLessThan %arg0, %arg1 : f32
+    // CHECK: spv.FUnordLessThanEqual
+    %11 = spv.FUnordLessThanEqual %arg0, %arg1 : f32
+    // CHECK: spv.FUnordNotEqual
+    %12 = spv.FUnordNotEqual %arg0, %arg1 : f32
+    spv.Return
+  }
 }
 
 // -----
 
 spv.module "Logical" "GLSL450" {
   spv.specConstant @condition_scalar = true
   func @select() -> () {
     %0 = spv.constant 4.0 : f32