diff --git a/mlir/include/mlir/Dialect/Math/IR/MathOps.td b/mlir/include/mlir/Dialect/Math/IR/MathOps.td --- a/mlir/include/mlir/Dialect/Math/IR/MathOps.td +++ b/mlir/include/mlir/Dialect/Math/IR/MathOps.td @@ -144,6 +144,7 @@ %a = math.atan2 %b, %c : f32 ``` }]; + let hasFolder = 1; } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Math/IR/MathOps.cpp b/mlir/lib/Dialect/Math/IR/MathOps.cpp --- a/mlir/lib/Dialect/Math/IR/MathOps.cpp +++ b/mlir/lib/Dialect/Math/IR/MathOps.cpp @@ -50,6 +50,28 @@ }); } +//===----------------------------------------------------------------------===// +// Atan2Op folder +//===----------------------------------------------------------------------===// + +OpFoldResult math::Atan2Op::fold(ArrayRef operands) { + return constFoldBinaryOpConditional( + operands, [](const APFloat &a, const APFloat &b) -> Optional { + if (a.isZero() && b.isZero()) + return llvm::APFloat::getNaN(a.getSemantics()); + + if (a.getSizeInBits(a.getSemantics()) == 64 && + b.getSizeInBits(b.getSemantics()) == 64) + return APFloat(atan2(a.convertToDouble(), b.convertToDouble())); + + if (a.getSizeInBits(a.getSemantics()) == 32 && + b.getSizeInBits(b.getSemantics()) == 32) + return APFloat(atan2f(a.convertToFloat(), b.convertToFloat())); + + return {}; + }); +} + //===----------------------------------------------------------------------===// // CeilOp folder //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/Math/canonicalize.mlir b/mlir/test/Dialect/Math/canonicalize.mlir --- a/mlir/test/Dialect/Math/canonicalize.mlir +++ b/mlir/test/Dialect/Math/canonicalize.mlir @@ -337,3 +337,24 @@ %0 = math.atan %v1 : vector<4xf32> return %0 : vector<4xf32> } + +// CHECK-LABEL: @atan2_fold +// CHECK-NEXT: %[[cst:.+]] = arith.constant 0.000000e+00 : f32 +// CHECK-NEXT: return %[[cst]] +func.func @atan2_fold() -> f32 { + %c1 = arith.constant 0.0 : f32 + %c2 = arith.constant 1.0 : f32 + %r = math.atan2 %c1, %c2 : f32 + return %r : f32 +} + +// CHECK-LABEL: @atan2_fold_vec +// CHECK-NEXT: %[[cst:.+]] = arith.constant dense<[0.000000e+00, 0.000000e+00, 0.463647604, 0.463647604]> : vector<4xf32> +// CHECK-NEXT: return %[[cst]] +func.func @atan2_fold_vec() -> (vector<4xf32>) { + %v1 = arith.constant dense<[0.0, 0.0, 1.0, 1.0]> : vector<4xf32> + %v2 = arith.constant dense<[1.0, 1.0, 2.0, 2.0]> : vector<4xf32> + %0 = math.atan2 %v1, %v2 : vector<4xf32> + return %0 : vector<4xf32> +} + diff --git a/mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir b/mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir --- a/mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir +++ b/mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir @@ -474,7 +474,7 @@ %atan2_8 = math.atan2 %neg_two, %one : f32 vector.print %atan2_8 : f32 - // CHECK: 0.463643 + // CHECK: 0.463648 %atan2_9 = math.atan2 %one, %two : f32 vector.print %atan2_9 : f32 @@ -490,7 +490,7 @@ %atan2_11 = math.atan2 %neg_one, %neg_two : f32 vector.print %atan2_11 : f32 - // CHECK: -0.463643 + // CHECK: -0.463648 %atan2_12 = math.atan2 %neg_one, %two : f32 vector.print %atan2_12 : f32