[mlir][spirv] Add path for math.round to spirv for OCL and GLSL

[rsuderman]

Summary

OpenCL's round function matches math.round so we can directly lower to
the op, this includes adding the op definition to the SPIRV OCL ops.
GLSL does not guarantee rounding direction so we include custom rounding
code to guarantee correct rounding direction.

[rsuderman]

Removed unintentional inclusion.

[antiagainst]

Nice, thanks for supporting this, @rsuderman!

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

237

Nit: ... to GLSL SPIR-V extended ops.

260

Hmm, this is using expensive FMul and FDiv ops. What about using FAdd like the expansion pattern in the below?

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

57

Missing a test for this?

68

This should be math::FloorOp?

133

This should be exposed in the .h file?

[rsuderman]

Moving extraneous code.

[rsuderman]

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

57

This was some extraneous changes that leaked in. Just removed it entirely.

68

Ditto on removed entirely.

133

Ditto on the removal.

[rsuderman]

Updated glsl comment

[rsuderman]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

237

Fixed.

260

I discovered an error with the add that causes unintentional double rounding. Specifically the case of 0.49999997f + 0,5f -> 1.0.

I can see if there is an alternative solution? One option is to do more bit magic instead but my f32 mantissa manipulations are not fantastic.

[antiagainst]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

260

Ha, interesting. Okay that makes sense then. Could you add comments in the code explaining this so later when we come back to this we can still understand why doing mul & div?

[rsuderman]

Updated to use select op instead

[rsuderman]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

260

Changed to use a select operation. It is slightly more ops but should be computationally faster.

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

 class SPV_OCLTernaryArithmeticOp<string mnemonic, int opcode, Type type,
                                 list<Trait> traits = []> :
   SPV_OCLTernaryOp<mnemonic, opcode, type, type,
                   traits # [SameOperandsAndResultType]> {
   let assemblyFormat = "operands attr-dict `:` type($result)";
 }
 
 
-// -----
-
-def SPV_OCLFmaOp : SPV_OCLTernaryArithmeticOp<"fma", 26, SPV_Float> {
-  let summary = [{
-    Compute the correctly rounded floating-point representation of the sum
-    of c with the infinitely precise product of a and b. Rounding of
-    intermediate products shall not occur. Edge case results are per the
-    IEEE 754-2008 standard.
-  }];
-
-  let description = [{
-    Result Type, a, b and c 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 -->
-
-    ```
-    fma-op ::= ssa-id `=` `spv.OCL.fma` ssa-use, ssa-use, ssa-use `:`
-               float-scalar-vector-type
-    ```mlir
-
-    ```
-    %0 = spv.OCL.fma %a, %b, %c : f32
-    %1 = spv.OCL.fma %a, %b, %c : vector<3xf16>
-    ```
-  }];
-}
 
 // -----
 
 def SPV_OCLCeilOp : SPV_OCLUnaryArithmeticOp<"ceil", 12, SPV_Float> {
   let summary = [{
     Round x to integral value using the round to positive infinity rounding
     mode.
   }];
let description = [{
     %2 = spv.OCL.floor %0 : f32
     %3 = spv.OCL.ceifloorl %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPV_OCLFmaOp : SPV_OCLTernaryArithmeticOp<"fma", 26, SPV_Float> {
+  let summary = [{
+    Compute the correctly rounded floating-point representation of the sum
+    of c with the infinitely precise product of a and b. Rounding of
+    intermediate products shall not occur. Edge case results are per the
+    IEEE 754-2008 standard.
+  }];
+
+  let description = [{
+    Result Type, a, b and c 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 -->
+
+    ```
+    fma-op ::= ssa-id `=` `spv.OCL.fma` ssa-use, ssa-use, ssa-use `:`
+               float-scalar-vector-type
+    ```mlir
+
+    ```
+    %0 = spv.OCL.fma %a, %b, %c : f32
+    %1 = spv.OCL.fma %a, %b, %c : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_OCLLogOp : SPV_OCLUnaryArithmeticOp<"log", 37, SPV_Float> {
   let summary = "Compute the natural logarithm 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 = [{
     %2 = spv.OCL.pow %0, %1 : f32
     %3 = spv.OCL.pow %0, %1 : vector<3xf16>
     ```
   }];
 }
 
 // -----
 
+def SPV_OCLRoundOp : SPV_OCLUnaryArithmeticOp<"round", 55, SPV_Float> {
+  let summary = [{
+    Return the integral value nearest to x rounding halfway cases away from
+    zero, regardless of the current rounding direction.
+  }];
+
+  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 `>`
+    round-op ::= ssa-id `=` `spv.OCL.round` ssa-use `:`
+               float-scalar-vector-type
+    ```
+    #### Example:
+
+    ```mlir
+    %2 = spv.OCL.round %0 : f32
+    %3 = spv.OCL.round %0 : vector<3xf16>
+    ```
+  }];
+}
+
+// -----
+
 def SPV_OCLRsqrtOp : SPV_OCLUnaryArithmeticOp<"rsqrt", 56, SPV_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

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

 //===----------------------------------------------------------------------===//
 
 #include "../SPIRVCommon/Pattern.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "math-to-spirv-pattern"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
matchAndRewrite(math::PowFOp powfOp, OpAdaptor adaptor,
     Value abs = rewriter.create<spirv::GLSLFAbsOp>(loc, adaptor.getLhs());
     Value pow = rewriter.create<spirv::GLSLPowOp>(loc, abs, adaptor.getRhs());
     Value negate = rewriter.create<spirv::FNegateOp>(loc, pow);
     rewriter.replaceOpWithNewOp<spirv::SelectOp>(powfOp, lessThan, negate, pow);
     return success();
   }
 };
 
+/// Converts math.round to SPIRV-Ops for GLSL.
+struct RoundOpPattern final : public OpConversionPattern<math::RoundOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(math::RoundOp roundOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = roundOp.getLoc();
+    auto operand = roundOp.getOperand();
+    auto ty = operand.getType();
+    auto ety = getElementTypeOrSelf(ty);
+
+    Value two;
+    if (VectorType vty = ty.dyn_cast<VectorType>()) {
+      two = rewriter.create<spirv::ConstantOp>(
+          loc, vty,
+          DenseElementsAttr::get(vty,
+                                 rewriter.getFloatAttr(ety, 2.0).getValue()));
+    } else {
+      two = rewriter.create<spirv::ConstantOp>(loc, ty,
+                                               rewriter.getFloatAttr(ety, 2.0));
+    }
+
+    auto abs = rewriter.create<spirv::GLSLFAbsOp>(loc, operand);
+    auto mul = rewriter.create<spirv::FMulOp>(loc, abs, two);
+    auto floor = rewriter.create<spirv::GLSLFloorOp>(loc, mul);
+    auto div = rewriter.create<spirv::FDivOp>(loc, floor, two);
+    auto ceil = rewriter.create<spirv::GLSLCeilOp>(loc, div);
+    rewriter.replaceOpWithNewOp<math::CopySignOp>(roundOp, ceil, operand);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                  RewritePatternSet &patterns) {
   // Core patterns
   patterns.add<CopySignPattern>(typeConverter, patterns.getContext());
 
   // GLSL patterns
   patterns
       .add<CountLeadingZerosPattern, Log1pOpPattern<spirv::GLSLLogOp>,
-           ExpM1OpPattern<spirv::GLSLExpOp>, PowFOpPattern,
+           ExpM1OpPattern<spirv::GLSLExpOp>, PowFOpPattern, RoundOpPattern,
            spirv::ElementwiseOpPattern<math::AbsOp, spirv::GLSLFAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLSLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLSLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLSLFloorOp>,
            spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLSLFmaOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLSLLogOp>,
            spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::GLSLInverseSqrtOp>,
            spirv::ElementwiseOpPattern<math::SinOp, spirv::GLSLSinOp>,
            spirv::ElementwiseOpPattern<math::SqrtOp, spirv::GLSLSqrtOp>,
            spirv::ElementwiseOpPattern<math::TanhOp, spirv::GLSLTanhOp>>(
           typeConverter, patterns.getContext());
 
   // OpenCL patterns
   patterns.add<Log1pOpPattern<spirv::OCLLogOp>, ExpM1OpPattern<spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::AbsOp, spirv::OCLFAbsOp>,
                spirv::ElementwiseOpPattern<math::CeilOp, spirv::OCLCeilOp>,
                spirv::ElementwiseOpPattern<math::CosOp, spirv::OCLCosOp>,
                spirv::ElementwiseOpPattern<math::ErfOp, spirv::OCLErfOp>,
                spirv::ElementwiseOpPattern<math::ExpOp, spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::FloorOp, spirv::OCLFloorOp>,
                spirv::ElementwiseOpPattern<math::FmaOp, spirv::OCLFmaOp>,
                spirv::ElementwiseOpPattern<math::LogOp, spirv::OCLLogOp>,
                spirv::ElementwiseOpPattern<math::PowFOp, spirv::OCLPowOp>,
+               spirv::ElementwiseOpPattern<math::RoundOp, spirv::OCLRoundOp>,
                spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::OCLRsqrtOp>,
                spirv::ElementwiseOpPattern<math::SinOp, spirv::OCLSinOp>,
                spirv::ElementwiseOpPattern<math::SqrtOp, spirv::OCLSqrtOp>,
                spirv::ElementwiseOpPattern<math::TanhOp, spirv::OCLTanhOp>>(
       typeConverter, patterns.getContext());
 }
 
 } // namespace mlir

mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp

static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
   Value zero = rewriter.create<arith::ConstantOp>(loc, floatZero);
   Value cmpRes = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OGE,
                                                 op.getOperand(), zero);
   rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cmpRes, positiveRes,
                                                negativeRes);
   return success();
 }
 
+static LogicalResult convertRoundOp(math::RoundOp op,
+                                    PatternRewriter &rewriter) {
+  auto operand = op.getOperand();
+  auto ty = operand.getType();
+  auto loc = operand.getLoc();
+
+  auto floatHalf = rewriter.getFloatAttr(ty, 0.5);
+  Value half = rewriter.create<arith::ConstantOp>(loc, floatHalf);
+
+  Value copysign = rewriter.create<math::CopySignOp>(loc, half, operand);
+  Value add = rewriter.create<arith::AddFOp>(loc, operand, copysign);
+  rewriter.replaceOpWithNewOp<math::CeilOp>(op, add);
+  return success();
+}
+
 static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
                                    PatternRewriter &rewriter) {
   auto operand = op.getOperand();
   auto elementTy = operand.getType();
   auto resultTy = op.getType();
   Location loc = op.getLoc();
 
   int bitWidth = elementTy.getIntOrFloatBitWidth();
static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
   rewriter.replaceOp(op, whileOp->getResult(1));
   return success();
 }
 
 void mlir::populateExpandCtlzPattern(RewritePatternSet &patterns) {
   patterns.add(convertCtlzOp);
 }
 
+void mlir::populateExpandRoundPattern(RewritePatternSet &patterns) {
+  patterns.add(convertRoundOp);
+}
+
 void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }

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

func.func @powf_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> vector<4xf32> {
   // CHEKC: spv.GLSL.FAbs
   // CHECK: spv.GLSL.Pow %{{.*}}: vector<4xf32>
   // CHECK: spv.FNegate
   // CHECK: spv.Select
   %0 = math.powf %lhs, %rhs : vector<4xf32>
   return %0: vector<4xf32>
 }
 
+// CHECK-LABEL: @round_scalar
+func.func @round_scalar(%x: f32) -> f32 {
+  // CHECK: %[[TWO:.+]] = spv.Constant 2.000000e+00
+  // CHECK: %[[ABS:.+]] = spv.GLSL.FAbs %arg0
+  // CHECK: %[[MUL:.+]] = spv.FMul %[[ABS]], %[[TWO]]
+  // CHECK: %[[FLOOR:.+]] = spv.GLSL.Floor %[[MUL]]
+  // CHECK: %[[DIV:.+]] = spv.FDiv %[[FLOOR]], %[[TWO]]
+  // CHECK: %[[CEIL:.+]] = spv.GLSL.Ceil %[[DIV]]
+  %0 = math.round %x : f32
+  return %0: f32
+}
+
+// CHECK-LABEL: @round_vector
+func.func @round_vector(%x: vector<4xf32>) -> vector<4xf32> {
+  // CHECK: %[[TWO:.+]] = spv.Constant dense<2.000000e+00>
+  // CHECK: %[[ABS:.+]] = spv.GLSL.FAbs %arg0
+  // CHECK: %[[MUL:.+]] = spv.FMul %[[ABS]], %[[TWO]]
+  // CHECK: %[[FLOOR:.+]] = spv.GLSL.Floor %[[MUL]]
+  // CHECK: %[[DIV:.+]] = spv.FDiv %[[FLOOR]], %[[TWO]]
+  // CHECK: %[[CEIL:.+]] = spv.GLSL.Ceil %[[DIV]]
+  %0 = math.round %x : vector<4xf32>
+  return %0: vector<4xf32>
+}
+
 } // end module
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader, Int64, Int16], []>, #spv.resource_limits<>>
 } {
 

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

func.func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spv.OCL.fabs %{{.*}}: f32
   %9 = math.abs %arg0 : f32
   // CHECK: spv.OCL.ceil %{{.*}}: f32
   %10 = math.ceil %arg0 : f32
   // CHECK: spv.OCL.floor %{{.*}}: f32
   %11 = math.floor %arg0 : f32
   // CHECK: spv.OCL.erf %{{.*}}: f32
   %12 = math.erf %arg0 : f32
+  // CHECK: spv.OCL.round %{{.*}}: f32
+  %13 = math.round %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func.func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spv.OCL.cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spv.OCL.exp %{{.*}}: vector<3xf32>