[mlir][math] Expand coverage of atan2 expansion

[jpienaar]

Summary

Reuse the higher precision F32 approximation for the F16 one (by expanding and
truncating).

[ezhulenev]

I think it's ok as a fallback for f16/f32, but maybe @mehdi_amini has a different opinion.

[rsuderman]

Mind changing the atan2 test to run on f16? It should handle all possible cases that way.

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

360

Shouldn't this be > 32?

[jpienaar]

Update test instance and use helper pattern

[mehdi_amini]

That's fine with me, I'd be wary that just using f16 blindly in the f32 approximation wouldn't be good enough in terms of precision and that a different expansion would be needed there. This seems safer and we can optimize this for f16 if needed in the future.

[jpienaar]

Generalize helper a bit so that it could be reused for other fallback approximations.

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

360

Indeed :) (I moved these into a helper pattern as that could be made into something a bit more reusable)

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

 };
 } // namespace
 
 LogicalResult
 Atan2Approximation::matchAndRewrite(math::Atan2Op op,
                                     PatternRewriter &rewriter) const {
   auto y = op.getOperand(0);
   auto x = op.getOperand(1);
-  if (!getElementTypeOrSelf(x).isF32())
+  // The expansion below is for F32. For lower precision floats, we can use
+  // the F32 approximation and truncate.
+  if (!getElementTypeOrSelf(x).isF32()) {
+    if (x.getType().getIntOrFloatBitWidth() < 32)
       return rewriter.notifyMatchFailure(op, "unsupported operand type");
+    auto f32Type = x.getType().cast<ShapedType>().clone(rewriter.getF32Type());
+    x = rewriter.create<arith::ExtFOp>(op.getLoc(), f32Type, x);
+    y = rewriter.create<arith::ExtFOp>(op.getLoc(), f32Type, y);
+  }
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   ArrayRef<int64_t> shape = vectorShape(op.getResult());
 
   // Compute atan in the valid range.
   auto div = builder.create<arith::DivFOp>(y, x);
   auto atan = builder.create<math::AtanOp>(div);
 
Atan2Approximation::matchAndRewrite(math::Atan2Op op,
 
   // Handle x = 0, y = 0;
   Value yZero =
       builder.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, y, zero);
   Value isNan = builder.create<arith::AndIOp>(xZero, yZero);
   Value cstNan = broadcast(builder, f32FromBits(builder, 0x7fc00000), shape);
   result = builder.create<arith::SelectOp>(isNan, cstNan, result);
 
+  // Truncate if needed.
+  if (!getElementTypeOrSelf(op.getResult()).isF32()) {
+    auto f16Type = op.getResult().getType().cast<ShapedType>().clone(
+        rewriter.getF16Type());
+    result = rewriter.create<arith::TruncFOp>(op.getLoc(), f16Type, result);
+  }
+
   rewriter.replaceOp(op, result);
   return success();
 }
 
 //----------------------------------------------------------------------------//
 // TanhOp approximation.
 //----------------------------------------------------------------------------//