Add logic for math.powf -> spirv for some previously non covered cases

[qedawkins]

This will need a test in mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir.

[qedawkins]

A few additional comments. Otherwise looks good to me after updating the test.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

318

nit: add period to the end of comment.

319

I think it would make more sense to just convert to i32 because we don't know whether something like f16 -> int16 will be supported (@antiagainst to confirm).

320

nit: unused variable

338–349

It is probably cleaner to get the Attribute inside the if and then create the constant based on that attribute.

[antiagainst]

Thanks for the fix! A few comments in addition to the missing test.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

308

Nit: // Get the scalar float type.

309

Nit: Name this as scalarFloatType to be consistent with scalarIntType.

318

Nit: // Get int type of the same bitwidth and shape as the float type.

319

Nit: Type intType = scalarIntType to avoid duplication .

334

Could you add a TODO here to mention that "The following just forcefully casts y into an integer value in order to properly propagate the sign, assuming integer y cases. It doesn't cover other cases and should be fixed."

339

Can we repurpose getScalarOrVectorI32Constant to handle more bitwidths and use it here?

351

In Vulkan, "For the OpSRem and OpSMod instructions, if either operand is negative the result is undefined." https://registry.khronos.org/vulkan/specs/1.3-extensions/html/chap50.html Also, mod is an expensive operation. Here we just want to check whether it's odd or even. You can do bit operation to bitwise and 1 == 1.

[dan-garvey]

fixed the lit test

[antiagainst]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

319

yup good point. we can always convert to i32 for capability reasons. (the exact value doesn't matter that much given we only care about the odd/even part.)

[dan-garvey]

addressed comments

[dan-garvey]

clang-format

[dan-garvey]

added periods to some comments that I missed. Sorry for spamming your inboxes = (

[dan-garvey]

lit test fixes

[dan-garvey]

added changed to vector test and fixed both

[antiagainst]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

339

Remaining debugging dump(). :) I'll remove these before landing.

[dan-garvey]

removes debug dumps

[antiagainst]

Should I use Daniel Garvey <[e-mail]> as the author for landing the patch?

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

matchAndRewrite(math::PowFOp powfOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (LogicalResult res = checkSourceOpTypes(rewriter, powfOp); failed(res))
       return res;
 
     Type dstType = getTypeConverter()->convertType(powfOp.getType());
     if (!dstType)
       return failure();
 
+    // Get the scalar float type.
+    FloatType scalarFloatType;
+    if (auto scalarType = powfOp.getType().dyn_cast<FloatType>()) {
+      scalarFloatType = scalarType;
+    } else if (auto vectorType = powfOp.getType().dyn_cast<VectorType>()) {
+      scalarFloatType = vectorType.getElementType().cast<FloatType>();
+    } else {
+      return failure();
+    }
+
+    // Get int type of the same shape as the float type.
+    Type scalarIntType = rewriter.getIntegerType(32);
+    Type intType = scalarIntType;
+    if (auto vectorType = adaptor.getRhs().getType().dyn_cast<VectorType>()) {
+      auto shape = vectorType.getShape();
+      intType = VectorType::get(shape, scalarIntType);
+    }
+
     // Per GL 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::GLFAbsOp>(loc, adaptor.getLhs());
+
+    // TODO: The following just forcefully casts y into an integer value in
+    // order to properly propagate the sign, assuming integer y cases. It
+    // doesn't cover other cases and should be fixed.
+
+    // Cast exponent to integer and calculate exponent % 2 != 0.
+    Value intRhs =
+        rewriter.create<spirv::ConvertFToSOp>(loc, intType, adaptor.getRhs());
+    Value intOne = spirv::ConstantOp::getOne(intType, loc, rewriter);
+    Value bitwiseAndOne =
+        rewriter.create<spirv::BitwiseAndOp>(loc, intRhs, intOne);
+    Value isOdd = rewriter.create<spirv::IEqualOp>(loc, bitwiseAndOne, intOne);
+
+    // calculate pow based on abs(lhs)^rhs.
     Value pow = rewriter.create<spirv::GLPowOp>(loc, abs, adaptor.getRhs());
     Value negate = rewriter.create<spirv::FNegateOp>(loc, pow);
-    rewriter.replaceOpWithNewOp<spirv::SelectOp>(powfOp, lessThan, negate, pow);
+    // if the exponent is odd and lhs < 0, negate the result.
+    Value shouldNegate =
+        rewriter.create<spirv::LogicalAndOp>(loc, lessThan, isOdd);
+    rewriter.replaceOpWithNewOp<spirv::SelectOp>(powfOp, shouldNegate, negate,
+                                                 pow);
     return success();
   }
 };
 
 /// Converts math.round to GLSL SPIRV extended ops.
 struct RoundOpPattern final : public OpConversionPattern<math::RoundOp> {
   using OpConversionPattern::OpConversionPattern;
 

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

 }
 
 // CHECK-LABEL: @powf_scalar
 //  CHECK-SAME: (%[[LHS:.+]]: f32, %[[RHS:.+]]: f32)
 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: %[[IRHS:.+]] = spirv.ConvertFToS
+  // CHECK: %[[CST1:.+]] = spirv.Constant 1 : i32
+  // CHECK: %[[REM:.+]] = spirv.BitwiseAnd %[[IRHS]]
+  // CHECK: %[[ODD:.+]] = spirv.IEqual %[[REM]], %[[CST1]] : i32
   // CHECK: %[[POW:.+]] = spirv.GL.Pow %[[ABS]], %[[RHS]] : f32
   // CHECK: %[[NEG:.+]] = spirv.FNegate %[[POW]] : f32
-  // CHECK: %[[SEL:.+]] = spirv.Select %[[LT]], %[[NEG]], %[[POW]] : i1, f32
+  // CHECK: %[[SNEG:.+]] = spirv.LogicalAnd %[[LT]], %[[ODD]] : i1
+  // CHECK: %[[SEL:.+]] = spirv.Select %[[SNEG]], %[[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: spirv.FOrdLessThan
   // CHECK: spirv.GL.FAbs
+  // CHECK: spirv.BitwiseAnd %{{.*}} : vector<4xi32>
+  // CHECK: spirv.IEqual %{{.*}} : vector<4xi32>
   // CHECK: spirv.GL.Pow %{{.*}}: vector<4xf32>
   // CHECK: spirv.FNegate
   // CHECK: spirv.Select
   %0 = math.powf %lhs, %rhs : vector<4xf32>
   return %0: vector<4xf32>
 }
 
 // CHECK-LABEL: @round_scalar