[mlir][spirv] Handle corner cases for math.powf conversion

[antiagainst]

Summary

Per GLSL Pow extended instruction spec: "Result is undefined if
x < 0. Result is undefined if x = 0 and y <= 0." So we need to
handle negative x values specifically.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

matchAndRewrite(math::Log1pOp operation, OpAdaptor adaptor,
     auto type = this->getTypeConverter()->convertType(operation.getType());
     auto one = spirv::ConstantOp::getOne(type, operation.getLoc(), rewriter);
     auto onePlus =
         rewriter.create<spirv::FAddOp>(loc, one, adaptor.getOperand());
     rewriter.replaceOpWithNewOp<LogOp>(operation, type, onePlus);
     return success();
   }
 };
+
+/// Converts math.powf to SPIRV-Ops.
+struct PowFOpPattern final : public OpConversionPattern<math::PowFOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(math::PowFOp powfOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = getTypeConverter()->convertType(powfOp.getType());
+    if (!dstType)
+      return failure();
+
+    // Per GLSL Pow extended instruction spec:
+    // "Result is undefined if x < 0. Result is undefined if x = 0 and y <= 0."
+    Location loc = powfOp.getLoc();
+    Value zero =
+        spirv::ConstantOp::getZero(adaptor.getLhs().getType(), loc, rewriter);
+    Value lessThan =
+        rewriter.create<spirv::FOrdLessThanOp>(loc, adaptor.getLhs(), zero);
+    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();
+  }
+};
+
 } // 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>,
+           ExpM1OpPattern<spirv::GLSLExpOp>, PowFOpPattern,
            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::PowFOp, spirv::GLSLPowOp>,
            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>,

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

func.func @float32_unary_vector(%arg0: vector<3xf32>) {
   %6 = math.sqrt %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Tanh %{{.*}}: vector<3xf32>
   %7 = math.tanh %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Sin %{{.*}}: vector<3xf32>
   %8 = math.sin %arg0 : vector<3xf32>
   return
 }
 
-// CHECK-LABEL: @float32_binary_scalar
-func.func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
-  // CHECK: spv.GLSL.Pow %{{.*}}: f32
-  %0 = math.powf %lhs, %rhs : f32
-  return
-}
-
-// CHECK-LABEL: @float32_binary_vector
-func.func @float32_binary_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) {
-  // CHECK: spv.GLSL.Pow %{{.*}}: vector<4xf32>
-  %0 = math.powf %lhs, %rhs : vector<4xf32>
-  return
-}
-
 // CHECK-LABEL: @float32_ternary_scalar
 func.func @float32_ternary_scalar(%a: f32, %b: f32, %c: f32) {
   // CHECK: spv.GLSL.Fma %{{.*}}: f32
   %0 = math.fma %a, %b, %c : f32
   return
 }
 
 // CHECK-LABEL: @float32_ternary_vector
func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
   // CHECK: %[[SUB:.+]] = spv.ISub %[[V31]], %[[MSB]] : vector<2xi32>
   // CHECK: %[[CMP:.+]] = spv.IEqual %[[MSB]], %[[MAX]] : vector<2xi32>
   // CHECK: %[[R:.+]] = spv.Select %[[CMP]], %[[V32]], %[[SUB]] : vector<2xi1>, vector<2xi32>
   // CHECK: return %[[R]]
   %0 = math.ctlz %val : vector<2xi32>
   return %0 : vector<2xi32>
 }
 
+// CHECK-LABEL: @powf_scalar
+//  CHECK-SAME: (%[[LHS:.+]]: f32, %[[RHS:.+]]: f32)
+func.func @powf_scalar(%lhs: f32, %rhs: f32) -> f32 {
+  // CHECK: %[[F0:.+]] = spv.Constant 0.000000e+00 : f32
+  // CHECK: %[[LT:.+]] = spv.FOrdLessThan %[[LHS]], %[[F0]] : f32
+  // CHECK: %[[ABS:.+]] = spv.GLSL.FAbs %[[LHS]] : f32
+  // CHECK: %[[POW:.+]] = spv.GLSL.Pow %[[ABS]], %[[RHS]] : f32
+  // CHECK: %[[NEG:.+]] = spv.FNegate %[[POW]] : f32
+  // CHECK: %[[SEL:.+]] = spv.Select %[[LT]], %[[NEG]], %[[POW]] : i1, f32
+  %0 = math.powf %lhs, %rhs : f32
+  // CHECK: return %[[SEL]]
+  return %0: f32
+}
+
+// CHECK-LABEL: @powf_vector
+func.func @powf_vector(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> vector<4xf32> {
+  // CHECK: spv.FOrdLessThan
+  // 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>
+}
+
 } // end module
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader, Int64, Int16], []>, #spv.resource_limits<>>
 } {