[Mlir-commits] [mlir] [mlir][spirv] Update math.powf lowering (PR #111388)

[Dmitriy Smirnov]

https://github.com/d-smirnov created https://github.com/llvm/llvm-project/pull/111388

The PR updates math.powf lowering to produce NaN result for a negative base with a fractional exponent which matches the actual behaviour of the C/C++ implementation.

>From 54191573a1c056856f8969375b301b7fdc22079b Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Mon, 7 Oct 2024 14:29:11 +0100
Subject: [PATCH] [mlir][spirv] Update math.powf lowering

math.powf lowering now produces NaN result for negative base with fractonal exponent
which matches the actual behavior of the C/C++ implementation
---
 .../Conversion/MathToSPIRV/MathToSPIRV.cpp    | 26 ++++++++++++++++++-
 .../MathToSPIRV/math-to-gl-spirv.mlir         | 12 ++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index 6f948e80d5af8f..fb7816b2fc2473 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -389,7 +389,31 @@ struct PowFOpPattern final : public OpConversionPattern<math::PowFOp> {
         spirv::ConstantOp::getZero(adaptor.getLhs().getType(), loc, rewriter);
     Value lessThan =
         rewriter.create<spirv::FOrdLessThanOp>(loc, adaptor.getLhs(), zero);
-    Value abs = rewriter.create<spirv::GLFAbsOp>(loc, adaptor.getLhs());
+
+    // Per C/CPP spec:
+    // "pow(base, exponent) returns NaN (and raises FE_INVALID) if base is "
+    // " finite and negative and exponent is finite and non-integer. "
+    // Calculae calc reminder from exponent and check whether it is zero
+    Value floatOne =
+        spirv::ConstantOp::getOne(adaptor.getRhs().getType(), loc, rewriter);
+    Value expRem =
+        rewriter.create<spirv::FRemOp>(loc, adaptor.getRhs(), floatOne);
+    Value expRemNonZero =
+        rewriter.create<spirv::FOrdNotEqualOp>(loc, expRem, zero);
+    Value cmpNegativeWithFractionalExp =
+        rewriter.create<spirv::LogicalAndOp>(loc, expRemNonZero, lessThan);
+    // Create NaN result and replace base value if conditions meet
+    const auto &floatSemantics = scalarFloatType.getFloatSemantics();
+    const auto nan = APFloat::getNaN(floatSemantics);
+    Attribute nanAttr = rewriter.getFloatAttr(scalarFloatType, nan);
+    if (auto vectorType = dyn_cast<VectorType>(adaptor.getRhs().getType()))
+      nanAttr = DenseElementsAttr::get(vectorType, nan);
+
+    Value NanValue = rewriter.create<spirv::ConstantOp>(
+        loc, adaptor.getRhs().getType(), nanAttr);
+    Value lhs = rewriter.create<spirv::SelectOp>(
+        loc, cmpNegativeWithFractionalExp, NanValue, adaptor.getLhs());
+    Value abs = rewriter.create<spirv::GLFAbsOp>(loc, lhs);
 
     // TODO: The following just forcefully casts y into an integer value in
     // order to properly propagate the sign, assuming integer y cases. It
diff --git a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
index ecbd59e54971ef..5c6561c1043892 100644
--- a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
+++ b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
@@ -156,7 +156,13 @@ func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
 func.func @powf_scalar(%lhs: f32, %rhs: f32) -> f32 {
   // CHECK: %[[F0:.+]] = spirv.Constant 0.000000e+00 : f32
   // CHECK: %[[LT:.+]] = spirv.FOrdLessThan %[[LHS]], %[[F0]] : f32
-  // CHECK: %[[ABS:.+]] = spirv.GL.FAbs %[[LHS]] : f32
+  // CHECK: %[[F1:.+]] = spirv.Constant 1.000000e+00 : f32
+  // CHECK: %[[REM:.+]] = spirv.FRem %[[RHS]], %[[F1]] : f32
+  // CHECK: %[[IS_FRACTION:.+]] = spirv.FOrdNotEqual %[[REM]], %[[F0]] : f32
+  // CHECK: %[[AND:.+]] = spirv.LogicalAnd %[[IS_FRACTION]], %[[LT]] : i1
+  // CHECK: %[[NAN:.+]] = spirv.Constant 0x7FC00000 : f32
+  // CHECK: %[[NEW_LHS:.+]] = spirv.Select %[[AND]], %[[NAN]], %[[LHS]] : i1, f32
+  // CHECK: %[[ABS:.+]] = spirv.GL.FAbs %[[NEW_LHS]] : f32
   // CHECK: %[[IRHS:.+]] = spirv.ConvertFToS
   // CHECK: %[[CST1:.+]] = spirv.Constant 1 : i32
   // CHECK: %[[REM:.+]] = spirv.BitwiseAnd %[[IRHS]]
@@ -173,6 +179,10 @@ func.func @powf_scalar(%lhs: f32, %rhs: f32) -> f32 {
 // CHECK-LABEL: @powf_vector
 func.func @powf_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> vector<4xf32> {
   // CHECK: spirv.FOrdLessThan
+  // CHECK: spirv.FRem
+  // CHECK: spirv.FOrdNotEqual
+  // CHECK: spirv.LogicalAnd
+  // CHECK: spirv.Select
   // CHECK: spirv.GL.FAbs
   // CHECK: spirv.BitwiseAnd %{{.*}} : vector<4xi32>
   // CHECK: spirv.IEqual %{{.*}} : vector<4xi32>