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

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.td

28

Could you make this fit within 80 chars.

[mravishankar]

Could you add some serialization tests as well. Thanks

[denis13]

@mravishankar thanks for review, added serialization tests.

[mravishankar]

Thanks Denis for taking these up and driving them

[antiagainst]

Thanks!

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.td

class BinaryOpPattern<Type type, Op src, Op tgt> :
       Pat<(src SPV_ScalarOrVectorOf<type>:$l, SPV_ScalarOrVectorOf<type>:$r),
           (tgt $l, $r)>;
 
 class UnaryOpPattern<Type type, Op src, Op tgt> :
       Pat<(src type:$input),
           (tgt $input)>;
 
 def : BinaryOpPattern<SPV_Bool, AndOp, SPV_LogicalAndOp>;
 def : BinaryOpPattern<SPV_Bool, OrOp, SPV_LogicalOrOp>;
+def : BinaryOpPattern<SPV_Integer, AndOp, SPV_BitwiseAndOp>;
+def : BinaryOpPattern<SPV_Integer, OrOp, SPV_BitwiseOrOp>;
+def : BinaryOpPattern<SPV_Integer, ShiftLeftOp, SPV_ShiftLeftLogicalOp>;
+def : BinaryOpPattern<SPV_Integer, SignedShiftRightOp,
+                      SPV_ShiftRightArithmeticOp>;
+def : BinaryOpPattern<SPV_Integer, UnsignedShiftRightOp,
+                      SPV_ShiftRightLogicalOp>;
+def : BinaryOpPattern<SPV_Integer, XOrOp, SPV_BitwiseXorOp>;
 def : BinaryOpPattern<SPV_Float, AddFOp, SPV_FAddOp>;
 def : BinaryOpPattern<SPV_Float, DivFOp, SPV_FDivOp>;
 def : BinaryOpPattern<SPV_Float, MulFOp, SPV_FMulOp>;
 def : BinaryOpPattern<SPV_Float, RemFOp, SPV_FRemOp>;
 def : BinaryOpPattern<SPV_Float, SubFOp, SPV_FSubOp>;
 
 def : UnaryOpPattern<SPV_Integer, SIToFPOp, SPV_ConvertSToFOp>;
 def : UnaryOpPattern<SPV_Float, FPExtOp, SPV_FConvertOp>;

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

func @div_rem(%arg0 : i32, %arg1 : i32) {
   // CHECK: spv.SDiv
   %0 = divi_signed %arg0, %arg1 : i32
   // CHECK: spv.SMod
   %1 = remi_signed %arg0, %arg1 : i32
   return
 }
 
 //===----------------------------------------------------------------------===//
+// std bit ops
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @bitwise_scalar
+func @bitwise_scalar(%arg0 : i32, %arg1 : i32) {
+  // CHECK: spv.BitwiseAnd
+  %0 = and %arg0, %arg1 : i32
+  // CHECK: spv.BitwiseOr
+  %1 = or %arg0, %arg1 : i32
+  // CHECK: spv.BitwiseXor
+  %2 = xor %arg0, %arg1 : i32
+  return
+}
+
+// CHECK-LABEL: @bitwise_vector
+func @bitwise_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>) {
+  // CHECK: spv.BitwiseAnd
+  %0 = and %arg0, %arg1 : vector<4xi32>
+  // CHECK: spv.BitwiseOr
+  %1 = or %arg0, %arg1 : vector<4xi32>
+  // CHECK: spv.BitwiseXor
+  %2 = xor %arg0, %arg1 : vector<4xi32>
+  return
+}
+
+// CHECK-LABEL: @shift_scalar
+func @shift_scalar(%arg0 : i32, %arg1 : i32) {
+  // CHECK: spv.ShiftLeftLogical
+  %0 = shift_left %arg0, %arg1 : i32
+  // CHECK: spv.ShiftRightArithmetic
+  %1 = shift_right_signed %arg0, %arg1 : i32
+  // CHECK: spv.ShiftRightLogical
+  %2 = shift_right_unsigned %arg0, %arg1 : i32
+  return
+}
+
+// CHECK-LABEL: @shift_vector
+func @shift_vector(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>) {
+  // CHECK: spv.ShiftLeftLogical
+  %0 = shift_left %arg0, %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.cmpi
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @cmpi
 func @cmpi(%arg0 : i32, %arg1 : i32) {
   // CHECK: spv.IEqual
   %0 = cmpi "eq", %arg0, %arg1 : i32
   // CHECK: spv.INotEqual

 func @logical_vector(%arg0 : vector<4xi1>, %arg1 : vector<4xi1>) {
   // CHECK: spv.LogicalAnd
   %0 = and %arg0, %arg1 : vector<4xi1>
   // CHECK: spv.LogicalOr
   %1 = or %arg0, %arg1 : vector<4xi1>
   return
 }
 
-// CHECK-LABEL: @logical_scalar_fail
-func @logical_scalar_fail(%arg0 : i32, %arg1 : i32) {
-  // CHECK-NOT: spv.LogicalAnd
-  %0 = and %arg0, %arg1 : i32
-  // CHECK-NOT: spv.LogicalOr
-  %1 = or %arg0, %arg1 : i32
-  return
-}
-
-// CHECK-LABEL: @logical_vector_fail
-func @logical_vector_fail(%arg0 : vector<4xi32>, %arg1 : vector<4xi32>) {
-  // CHECK-NOT: spv.LogicalAnd
-  %0 = and %arg0, %arg1 : vector<4xi32>
-  // CHECK-NOT: spv.LogicalOr
-  %1 = or %arg0, %arg1 : vector<4xi32>
-  return
-}
-
 //===----------------------------------------------------------------------===//
 // std.fpext
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @fpext
 func @fpext(%arg0 : f32) {
   // CHECK: spv.FConvert
   %0 = std.fpext %arg0 : f32 to f64

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

func @bitreverse(%arg: i32) -> i32 {
     %0 = spv.BitReverse %arg : i32
     spv.ReturnValue %0 : i32
   }
   func @not(%arg: i32) -> i32 {
     // CHECK: spv.Not {{%.*}} : i32
     %0 = spv.Not %arg : i32
     spv.ReturnValue %0 : i32
   }
+  func @bitwise_scalar(%arg0 : i32, %arg1 : i32) {
+    // CHECK: spv.BitwiseAnd
+    %0 = spv.BitwiseAnd %arg0, %arg1 : i32
+    // CHECK: spv.BitwiseOr
+    %1 = spv.BitwiseOr %arg0, %arg1 : i32
+    // CHECK: spv.BitwiseXor
+    %2 = spv.BitwiseXor %arg0, %arg1 : i32
+    spv.Return
+  }
   func @shift_left_logical(%arg0: i32, %arg1 : i16) -> i32 {
     // CHECK: {{%.*}} = spv.ShiftLeftLogical {{%.*}}, {{%.*}} : i32, i16
     %0 = spv.ShiftLeftLogical %arg0, %arg1: i32, i16
     spv.ReturnValue %0 : i32
   }
   func @shift_right_arithmetic(%arg0: vector<4xi32>, %arg1 : vector<4xi8>) -> vector<4xi32> {
     // CHECK: {{%.*}} = spv.ShiftRightArithmetic {{%.*}}, {{%.*}} : vector<4xi32>, vector<4xi8>
     %0 = spv.ShiftRightArithmetic %arg0, %arg1: vector<4xi32>, vector<4xi8>
     spv.ReturnValue %0 : vector<4xi32>
   }
   func @shift_right_logical(%arg0: vector<2xi32>, %arg1 : vector<2xi8>) -> vector<2xi32> {
     // CHECK: {{%.*}} = spv.ShiftRightLogical {{%.*}}, {{%.*}} : vector<2xi32>, vector<2xi8>
     %0 = spv.ShiftRightLogical %arg0, %arg1: vector<2xi32>, vector<2xi8>
     spv.ReturnValue %0 : vector<2xi32>
   }
 }