[mlir][math][spirv] Add `math.roundeven` lowering to SPIR-V

[kuhar]

Summary

This has two lowering path, one for each extended instructions set:

• to OpenGL's RoundEven,
• to OpenCL's rint.

Implement those two ops and adds minimal tests.

[kuhar]

Add missing test

[antiagainst]

Nice, thanks!

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td

let description = [{
     %2 = spirv.CL.round %0 : f32
     %3 = spirv.CL.round %0 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPIRV_CLRintOp : SPIRV_CLUnaryArithmeticOp<"rint", 53, SPIRV_Float> {
+  let summary = [{
+    Round x to integral value (using round to nearest even rounding mode) in
+    floating-point format.
+  }];
+
+  let description = [{
+    Result Type and x must be floating-point or vector(2,3,4,8,16) of
+    floating-point values.
+
+    All of the operands, including the Result Type operand, must be of the
+    same type.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    rint-op ::= ssa-id `=` `spirv.CL.rint` ssa-use `:`
+               float-scalar-vector-type
+    ```
+
+    #### Example:
+
+    ```mlir
+    %0 = spirv.CL.rint %0 : f32
+    %1 = spirv.CL.rint %1 : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPIRV_CLRsqrtOp : SPIRV_CLUnaryArithmeticOp<"rsqrt", 56, SPIRV_Float> {
   let summary = "Compute inverse square root of x.";
 
   let description = [{
     Result Type and x must be floating-point or vector(2,3,4,8,16) of
     floating-point values.
 
     All of the operands, including the Result Type operand, must be of the
let description = [{
     #### Example:
     ```mlir
     %2 = spirv.CL.u_max %0, %1 : i32
     %3 = spirv.CL.u_max %0, %1 : vector<3xi16>
     ```
   }];
 }
 
+// -----
+
 def SPIRV_CLSMinOp : SPIRV_CLBinaryArithmeticOp<"s_min", 158, SPIRV_Integer> {
   let summary = "Return minimum of two signed integer operands";
 
   let description = [{
     Returns y if x < y, otherwise it returns x, where x and y are treated as signed integers.
 
     Result Type,x and y must be integer or vector(2,3,4,8,16) of integer values.
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLOps.td

let description = [{
     %3 = spirv.GL.Floor %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
 def SPIRV_GLRoundOp: SPIRV_GLUnaryArithmeticOp<"Round", 1, SPIRV_Float> {
-  let summary = "Rounds to the whole number";
+  let summary = "Rounds to the nearest whole number";
 
   let description = [{
-    Result is the value equal to the nearest whole number.
+    Result is the value equal to the nearest whole number to x. The fraction
+    0.5 will round in a direction chosen by the implementation, presumably
+    the direction that is fastest. This includes the possibility that
+    Round x is the same value as RoundEven x for all values of x.
 
     The operand x must be a scalar or vector whose component type is
     floating-point.
 
     Result Type and the type of x must be the same type. Results are computed
     per component.
 
     <!-- End of AutoGen section -->
     ```
     float-scalar-vector-type ::= float-type |
                                  `vector<` integer-literal `x` float-type `>`
-    floor-op ::= ssa-id `=` `spirv.GL.Round` ssa-use `:`
+    round-op ::= ssa-id `=` `spirv.GL.Round` ssa-use `:`
                 float-scalar-vector-type
     ```
     #### Example:
 
     ```mlir
     %2 = spirv.GL.Round %0 : f32
     %3 = spirv.GL.Round %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPIRV_GLRoundEvenOp: SPIRV_GLUnaryArithmeticOp<"RoundEven", 2, SPIRV_Float> {
+  let summary = "Rounds to the nearest even whole number";
+
+  let description = [{
+    Result is the value equal to the nearest whole number to x. A fractional
+    part of 0.5 will round toward the nearest even whole number. (Both 3.5 and
+    4.5 for x will be 4.0.)
+
+    The operand x must be a scalar or vector whose component type is
+    floating-point.
+
+    Result Type and the type of x must be the same type. Results are computed
+    per component.
+
+    <!-- End of AutoGen section -->
+    ```
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    round-even-op ::= ssa-id `=` `spirv.GL.RoundEven` ssa-use `:`
+                float-scalar-vector-type
+    ```
+    #### Example:
+
+    ```mlir
+    %2 = spirv.GL.RoundEven %0 : f32
+    %3 = spirv.GL.RoundEven %1 : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPIRV_GLInverseSqrtOp : SPIRV_GLUnaryArithmeticOp<"InverseSqrt", 32, SPIRV_Float> {
   let summary = "Reciprocal of sqrt(operand)";
 
   let description = [{
     Result is the reciprocal of sqrt x. Result is undefined if x <= 0.
 
     The operand x must be a scalar or vector whose component type is
     floating-point.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

patterns
            spirv::ElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
            spirv::ElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLExpOp>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLFloorOp>,
            spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLFmaOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLLogOp>,
+           spirv::ElementwiseOpPattern<math::RoundEvenOp, spirv::GLRoundEvenOp>,
            spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::GLInverseSqrtOp>,
            spirv::ElementwiseOpPattern<math::SinOp, spirv::GLSinOp>,
            spirv::ElementwiseOpPattern<math::SqrtOp, spirv::GLSqrtOp>,
            spirv::ElementwiseOpPattern<math::TanhOp, spirv::GLTanhOp>>(
           typeConverter, patterns.getContext());
 
   // OpenCL patterns
   patterns.add<Log1pOpPattern<spirv::CLLogOp>, ExpM1OpPattern<spirv::CLExpOp>,
                spirv::ElementwiseOpPattern<math::AbsFOp, spirv::CLFAbsOp>,
                spirv::ElementwiseOpPattern<math::CeilOp, spirv::CLCeilOp>,
                spirv::ElementwiseOpPattern<math::CosOp, spirv::CLCosOp>,
                spirv::ElementwiseOpPattern<math::ErfOp, spirv::CLErfOp>,
                spirv::ElementwiseOpPattern<math::ExpOp, spirv::CLExpOp>,
                spirv::ElementwiseOpPattern<math::FloorOp, spirv::CLFloorOp>,
                spirv::ElementwiseOpPattern<math::FmaOp, spirv::CLFmaOp>,
                spirv::ElementwiseOpPattern<math::LogOp, spirv::CLLogOp>,
                spirv::ElementwiseOpPattern<math::PowFOp, spirv::CLPowOp>,
+               spirv::ElementwiseOpPattern<math::RoundEvenOp, spirv::CLRintOp>,
                spirv::ElementwiseOpPattern<math::RoundOp, spirv::CLRoundOp>,
                spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::CLRsqrtOp>,
                spirv::ElementwiseOpPattern<math::SinOp, spirv::CLSinOp>,
                spirv::ElementwiseOpPattern<math::SqrtOp, spirv::CLSqrtOp>,
                spirv::ElementwiseOpPattern<math::TanhOp, spirv::CLTanhOp>>(
       typeConverter, patterns.getContext());
 }
 
 } // namespace mlir

mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir

func.func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spirv.FSub %[[EXP]], %[[ONE]]
   %2 = math.expm1 %arg0 : f32
   // CHECK: spirv.GL.Log %{{.*}}: f32
   %3 = math.log %arg0 : f32
   // CHECK: %[[ONE:.+]] = spirv.Constant 1.000000e+00 : f32
   // CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spirv.GL.Log %[[ADDONE]]
   %4 = math.log1p %arg0 : f32
+  // CHECK: spirv.GL.RoundEven %{{.*}}: f32
+  %5 = math.roundeven %arg0 : f32
   // CHECK: spirv.GL.InverseSqrt %{{.*}}: f32
-  %5 = math.rsqrt %arg0 : f32
+  %6 = math.rsqrt %arg0 : f32
   // CHECK: spirv.GL.Sqrt %{{.*}}: f32
-  %6 = math.sqrt %arg0 : f32
+  %7 = math.sqrt %arg0 : f32
   // CHECK: spirv.GL.Tanh %{{.*}}: f32
-  %7 = math.tanh %arg0 : f32
+  %8 = math.tanh %arg0 : f32
   // CHECK: spirv.GL.Sin %{{.*}}: f32
-  %8 = math.sin %arg0 : f32
+  %9 = math.sin %arg0 : f32
   // CHECK: spirv.GL.FAbs %{{.*}}: f32
-  %9 = math.absf %arg0 : f32
+  %10 = math.absf %arg0 : f32
   // CHECK: spirv.GL.Ceil %{{.*}}: f32
-  %10 = math.ceil %arg0 : f32
+  %11 = math.ceil %arg0 : f32
   // CHECK: spirv.GL.Floor %{{.*}}: f32
-  %11 = math.floor %arg0 : f32
+  %12 = math.floor %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func.func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spirv.GL.Cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spirv.GL.Exp %{{.*}}: vector<3xf32>
   %1 = math.exp %arg0 : vector<3xf32>
   // CHECK: %[[EXP:.+]] = spirv.GL.Exp %arg0
   // CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: spirv.FSub %[[EXP]], %[[ONE]]
   %2 = math.expm1 %arg0 : vector<3xf32>
   // CHECK: spirv.GL.Log %{{.*}}: vector<3xf32>
   %3 = math.log %arg0 : vector<3xf32>
   // CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spirv.GL.Log %[[ADDONE]]
   %4 = math.log1p %arg0 : vector<3xf32>
+  // CHECK: spirv.GL.RoundEven %{{.*}}: vector<3xf32>
+  %5 = math.roundeven %arg0 : vector<3xf32>
   // CHECK: spirv.GL.InverseSqrt %{{.*}}: vector<3xf32>
-  %5 = math.rsqrt %arg0 : vector<3xf32>
+  %6 = math.rsqrt %arg0 : vector<3xf32>
   // CHECK: spirv.GL.Sqrt %{{.*}}: vector<3xf32>
-  %6 = math.sqrt %arg0 : vector<3xf32>
+  %7 = math.sqrt %arg0 : vector<3xf32>
   // CHECK: spirv.GL.Tanh %{{.*}}: vector<3xf32>
-  %7 = math.tanh %arg0 : vector<3xf32>
+  %8 = math.tanh %arg0 : vector<3xf32>
   // CHECK: spirv.GL.Sin %{{.*}}: vector<3xf32>
-  %8 = math.sin %arg0 : vector<3xf32>
+  %9 = math.sin %arg0 : vector<3xf32>
   return
 }
 
 // CHECK-LABEL: @float32_ternary_scalar
 func.func @float32_ternary_scalar(%a: f32, %b: f32, %c: f32) {
   // CHECK: spirv.GL.Fma %{{.*}}: f32
   %0 = math.fma %a, %b, %c : f32
   return

mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir

func.func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spirv.FSub %[[EXP]], %[[ONE]]
   %2 = math.expm1 %arg0 : f32
   // CHECK: spirv.CL.log %{{.*}}: f32
   %3 = math.log %arg0 : f32
   // CHECK: %[[ONE:.+]] = spirv.Constant 1.000000e+00 : f32
   // CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spirv.CL.log %[[ADDONE]]
   %4 = math.log1p %arg0 : f32
+  // CHECK: spirv.CL.rint %{{.*}}: f32
+  %5 = math.roundeven %arg0 : f32
   // CHECK: spirv.CL.rsqrt %{{.*}}: f32
-  %5 = math.rsqrt %arg0 : f32
+  %6 = math.rsqrt %arg0 : f32
   // CHECK: spirv.CL.sqrt %{{.*}}: f32
-  %6 = math.sqrt %arg0 : f32
+  %7 = math.sqrt %arg0 : f32
   // CHECK: spirv.CL.tanh %{{.*}}: f32
-  %7 = math.tanh %arg0 : f32
+  %8 = math.tanh %arg0 : f32
   // CHECK: spirv.CL.sin %{{.*}}: f32
-  %8 = math.sin %arg0 : f32
+  %9 = math.sin %arg0 : f32
   // CHECK: spirv.CL.fabs %{{.*}}: f32
-  %9 = math.absf %arg0 : f32
+  %10 = math.absf %arg0 : f32
   // CHECK: spirv.CL.ceil %{{.*}}: f32
-  %10 = math.ceil %arg0 : f32
+  %11 = math.ceil %arg0 : f32
   // CHECK: spirv.CL.floor %{{.*}}: f32
-  %11 = math.floor %arg0 : f32
+  %12 = math.floor %arg0 : f32
   // CHECK: spirv.CL.erf %{{.*}}: f32
-  %12 = math.erf %arg0 : f32
+  %13 = math.erf %arg0 : f32
   // CHECK: spirv.CL.round %{{.*}}: f32
-  %13 = math.round %arg0 : f32
+  %14 = math.round %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func.func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spirv.CL.cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spirv.CL.exp %{{.*}}: vector<3xf32>
   %1 = math.exp %arg0 : vector<3xf32>
   // CHECK: %[[EXP:.+]] = spirv.CL.exp %arg0
   // CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: spirv.FSub %[[EXP]], %[[ONE]]
   %2 = math.expm1 %arg0 : vector<3xf32>
   // CHECK: spirv.CL.log %{{.*}}: vector<3xf32>
   %3 = math.log %arg0 : vector<3xf32>
   // CHECK: %[[ONE:.+]] = spirv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: %[[ADDONE:.+]] = spirv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spirv.CL.log %[[ADDONE]]
   %4 = math.log1p %arg0 : vector<3xf32>
+  // CHECK: spirv.CL.rint %{{.*}}: vector<3xf32>
+  %5 = math.roundeven %arg0 : vector<3xf32>
   // CHECK: spirv.CL.rsqrt %{{.*}}: vector<3xf32>
-  %5 = math.rsqrt %arg0 : vector<3xf32>
+  %6 = math.rsqrt %arg0 : vector<3xf32>
   // CHECK: spirv.CL.sqrt %{{.*}}: vector<3xf32>
-  %6 = math.sqrt %arg0 : vector<3xf32>
+  %7 = math.sqrt %arg0 : vector<3xf32>
   // CHECK: spirv.CL.tanh %{{.*}}: vector<3xf32>
-  %7 = math.tanh %arg0 : vector<3xf32>
+  %8 = math.tanh %arg0 : vector<3xf32>
   // CHECK: spirv.CL.sin %{{.*}}: vector<3xf32>
-  %8 = math.sin %arg0 : vector<3xf32>
+  %9 = math.sin %arg0 : vector<3xf32>
   return
 }
 
 // CHECK-LABEL: @float32_binary_scalar
 func.func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   // CHECK: spirv.CL.pow %{{.*}}: f32
   %0 = math.powf %lhs, %rhs : f32
   return

mlir/test/Dialect/SPIRV/IR/gl-ops.mlir

 }
 
 func.func @roundvec(%arg0 : vector<3xf16>) -> () {
   // CHECK: spirv.GL.Round {{%.*}} : vector<3xf16>
   %2 = spirv.GL.Round %arg0 : vector<3xf16>
   return
 }
 
+//===----------------------------------------------------------------------===//
+// spirv.GL.RoundEven
+//===----------------------------------------------------------------------===//
+
+func.func @round_even(%arg0 : f32) -> () {
+  // CHECK: spirv.GL.RoundEven {{%.*}} : f32
+  %2 = spirv.GL.RoundEven %arg0 : f32
+  return
+}
+
+func.func @round_even_vec(%arg0 : vector<3xf16>) -> () {
+  // CHECK: spirv.GL.RoundEven {{%.*}} : vector<3xf16>
+  %2 = spirv.GL.RoundEven %arg0 : vector<3xf16>
+  return
+}
+
 // -----
 
 //===----------------------------------------------------------------------===//
 // spirv.GL.FClamp
 //===----------------------------------------------------------------------===//
 
 func.func @fclamp(%arg0 : f32, %min : f32, %max : f32) -> () {
   // CHECK: spirv.GL.FClamp {{%[^,]*}}, {{%[^,]*}}, {{%[^,]*}} : f32

mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir

func.func @iminmax(%arg0: i32, %arg1: i32) {
   // CHECK: spirv.CL.u_max {{%.*}}, {{%.*}} : i32
   %2 = spirv.CL.u_max %arg0, %arg1 : i32
   // CHECK: spirv.CL.s_min {{%.*}}, {{%.*}} : i32
   %3 = spirv.CL.s_min %arg0, %arg1 : i32
   // CHECK: spirv.CL.u_min {{%.*}}, {{%.*}} : i32
   %4 = spirv.CL.u_min %arg0, %arg1 : i32
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.CL.rint
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @rint(
+func.func @rint(%arg0 : f32) -> () {
+  // CHECK: spirv.CL.rint {{%.*}} : f32
+  %0 = spirv.CL.rint %arg0 : f32
+  return
+}
+
+// CHECK-LABEL: func.func @rintvec(
+func.func @rintvec(%arg0 : vector<3xf16>) -> () {
+  // CHECK: spirv.CL.rint {{%.*}} : vector<3xf16>
+  %0 = spirv.CL.rint %arg0 : vector<3xf16>
+  return
+}