[Mlir-commits] [mlir] [mlir][complex] Fix sin and cos of +-0 + +-inf*i (PR #233040)
Chang Yong Lik
llvmlistbot at llvm.org
Sat Oct 17 00:27:51 PDT 2026
https://github.com/ahlag created https://github.com/llvm/llvm-project/pull/233040
The `complex.sin` and `complex.cos` lowerings in `ComplexToStandard` return NaN instead of a signed zero for the real part of `sin(±0 ± inf·i)` and the imaginary part of `cos(±0 ± inf·i)`. In both cases the coefficient `(exp(y) ± exp(-y)) / 2` is inf and is multiplied by `sin(x) = ±0`, giving `inf * 0 = NaN`. C99 Annex G gives `sin(+0 + inf·i) = +0 + inf·i` and `cos(+0 + inf·i) = inf - 0·i`.
The fix returns a zero with the sign of the product when `sin(x)` is zero. Like the exp fix in #232791, the extra ops are skipped under `ninf`.
This fixes [tensorflow#128985](https://github.com/tensorflow/tensorflow/issues/128985) for XLA CPU, which lowers HLO `sine`/`cosine` on complex numbers through this pass. With tf-nightly 2.23.0-dev20261014 on macOS arm64:
```
sin eager: [ 0.+infj -0.+infj 0.-infj 0.+nanj] XLA: [nan+infj nan+infj nan-infj nan+nanj]
cos eager: [inf+0.j inf-0.j inf+0.j nan+0.j] XLA: [inf+nanj inf+nanj inf+nanj nan+nanj]
```
With this patch, the lowered code gives `0+inf·i, -0+inf·i, 0-inf·i, 0+nan·i` for sin and `inf-0·i, inf+0·i, inf+0·i, nan+0·i` for cos on these inputs. These follow C99 and match NumPy. TF eager on macOS (libc++) differs only in the sign of the zero imaginary part for cos.
Tests:
- `convert-to-standard.mlir`: updated the sin/cos checks and added `ninf` cases.
- `Integration/Dialect/Complex/CPU/correctness.mlir`: added sin/cos cases. Integration tests are not enabled in my build, so I ran its RUN pipeline by hand: it passes with the fix and fails without it.
AI use: I used Claude to help write the patch and tests; I reviewed and tested the change.
>From da18068157fa5b59bdfc236e06e41275fc05aaad Mon Sep 17 00:00:00 2001
From: ahlag <theeahlag at gmail.com>
Date: Sat, 17 Oct 2026 16:25:13 +0900
Subject: [PATCH] [mlir][complex] Fix sin and cos of +-0 + +-inf*i
complex.sin and complex.cos lowerings returned NaN instead of a signed
zero for sin(+-0 + +-inf*i) (real part) and cos(+-0 + +-inf*i) (imaginary
part), because the coefficient (exp(y) +- exp(-y)) / 2 is inf there and
it is multiplied by sin(x) = +-0. C99 Annex G gives
sin(+0 + inf*i) = +0 + inf*i and cos(+0 + inf*i) = inf - 0*i, so use a
zero with the sign of the product when sin(x) is zero. The extra ops are
not emitted under ninf.
Fixes https://github.com/tensorflow/tensorflow/issues/128985 for XLA CPU,
which uses this lowering.
Assisted-by: Claude
---
.../ComplexToStandard/ComplexToStandard.cpp | 27 ++++++++-
.../convert-to-standard.mlir | 56 +++++++++++++++++--
.../Dialect/Complex/CPU/correctness.mlir | 54 ++++++++++++++++++
3 files changed, 131 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp b/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp
index d1cdd8ae17e3f4..7ddf545ccbb0f5 100644
--- a/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp
+++ b/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp
@@ -236,6 +236,29 @@ struct TrigonometricOpConversion : public OpConversionPattern<TrigonometricOp> {
combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin,
Value cos, ConversionPatternRewriter &rewriter,
arith::FastMathFlagsAttr fmf) const = 0;
+
+ // Returns coeff * sin. When x is +-0 and y is +-inf, this is inf * 0 = NaN,
+ // but C99 Annex G gives sin(+0 + I*inf) = +0 + inf*I and
+ // cos(+0 + I*inf) = inf - 0*I, so use a zero with the sign of the product.
+ static Value mulSin(Location loc, Value coeff, Value sin,
+ ConversionPatternRewriter &rewriter,
+ arith::FastMathFlagsAttr fmf) {
+ Value product = arith::MulFOp::create(rewriter, loc, coeff, sin, fmf);
+ if (arith::bitEnumContainsAll(fmf.getValue(), arith::FastMathFlags::ninf))
+ return product;
+ Type elementType = sin.getType();
+ Value zero = arith::ConstantOp::create(
+ rewriter, loc, elementType, rewriter.getFloatAttr(elementType, 0.0));
+ Value one = arith::ConstantOp::create(
+ rewriter, loc, elementType, rewriter.getFloatAttr(elementType, 1.0));
+ Value sinIsZero = arith::CmpFOp::create(
+ rewriter, loc, arith::CmpFPredicate::OEQ, sin, zero, fmf.getValue());
+ Value coeffSign = math::CopySignOp::create(rewriter, loc, one, coeff, fmf);
+ Value signedZero =
+ arith::MulFOp::create(rewriter, loc, coeffSign, sin, fmf);
+ return arith::SelectOp::create(rewriter, loc, sinIsZero, signedZero,
+ product);
+ }
};
struct CosOpConversion : public TrigonometricOpConversion<complex::CosOp> {
@@ -259,7 +282,7 @@ struct CosOpConversion : public TrigonometricOpConversion<complex::CosOp> {
Value resultReal = arith::MulFOp::create(rewriter, loc, sum, cos, fmf);
Value diff =
arith::SubFOp::create(rewriter, loc, reciprocalExp, scaledExp, fmf);
- Value resultImag = arith::MulFOp::create(rewriter, loc, diff, sin, fmf);
+ Value resultImag = mulSin(loc, diff, sin, rewriter, fmf);
return {resultReal, resultImag};
}
};
@@ -623,7 +646,7 @@ struct SinOpConversion : public TrigonometricOpConversion<complex::SinOp> {
// Im(sin(x + iy)) = (0.5*t - 0.5/t) * cos x
Value sum =
arith::AddFOp::create(rewriter, loc, scaledExp, reciprocalExp, fmf);
- Value resultReal = arith::MulFOp::create(rewriter, loc, sum, sin, fmf);
+ Value resultReal = mulSin(loc, sum, sin, rewriter, fmf);
Value diff =
arith::SubFOp::create(rewriter, loc, scaledExp, reciprocalExp, fmf);
Value resultImag = arith::MulFOp::create(rewriter, loc, diff, cos, fmf);
diff --git a/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir b/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir
index 448d82df503202..350f28792d6171 100644
--- a/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir
+++ b/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir
@@ -69,7 +69,13 @@ func.func @complex_cos(%arg: complex<f32>) -> complex<f32> {
// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_REXP]], %[[HALF_EXP]]
// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[COS]]
// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_REXP]], %[[HALF_EXP]]
-// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[SIN]]
+// CHECK-DAG: %[[RESULT_IMAG_PRODUCT:.*]] = arith.mulf %[[EXP_DIFF]], %[[SIN]]
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[SIN_IS_ZERO:.*]] = arith.cmpf oeq, %[[SIN]], %[[ZERO]]
+// CHECK-DAG: %[[COEFF_SIGN:.*]] = math.copysign %[[ONE]], %[[EXP_DIFF]]
+// CHECK-DAG: %[[SIGNED_ZERO:.*]] = arith.mulf %[[COEFF_SIGN]], %[[SIN]]
+// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.select %[[SIN_IS_ZERO]], %[[SIGNED_ZERO]], %[[RESULT_IMAG_PRODUCT]]
// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
// CHECK: return %[[RESULT]]
@@ -416,7 +422,13 @@ func.func @complex_sin(%arg: complex<f32>) -> complex<f32> {
// CHECK-DAG: %[[SIN:.*]] = math.sin %[[REAL]] : f32
// CHECK-DAG: %[[COS:.*]] = math.cos %[[REAL]] : f32
// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_EXP]], %[[HALF_REXP]]
-// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[SIN]]
+// CHECK-DAG: %[[RESULT_REAL_PRODUCT:.*]] = arith.mulf %[[EXP_SUM]], %[[SIN]]
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[SIN_IS_ZERO:.*]] = arith.cmpf oeq, %[[SIN]], %[[ZERO]]
+// CHECK-DAG: %[[COEFF_SIGN:.*]] = math.copysign %[[ONE]], %[[EXP_SUM]]
+// CHECK-DAG: %[[SIGNED_ZERO:.*]] = arith.mulf %[[COEFF_SIGN]], %[[SIN]]
+// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.select %[[SIN_IS_ZERO]], %[[SIGNED_ZERO]], %[[RESULT_REAL_PRODUCT]]
// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_EXP]], %[[HALF_REXP]]
// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[COS]]
// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
@@ -1367,7 +1379,13 @@ func.func @complex_cos_with_fmf(%arg: complex<f32>) -> complex<f32> {
// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_REXP]], %[[HALF_EXP]] fastmath<nnan,contract>
// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[COS]] fastmath<nnan,contract>
// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_REXP]], %[[HALF_EXP]] fastmath<nnan,contract>
-// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[RESULT_IMAG_PRODUCT:.*]] = arith.mulf %[[EXP_DIFF]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[SIN_IS_ZERO:.*]] = arith.cmpf oeq, %[[SIN]], %[[ZERO]] fastmath<nnan,contract>
+// CHECK-DAG: %[[COEFF_SIGN:.*]] = math.copysign %[[ONE]], %[[EXP_DIFF]] fastmath<nnan,contract>
+// CHECK-DAG: %[[SIGNED_ZERO:.*]] = arith.mulf %[[COEFF_SIGN]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.select %[[SIN_IS_ZERO]], %[[SIGNED_ZERO]], %[[RESULT_IMAG_PRODUCT]]
// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
// CHECK: return %[[RESULT]]
@@ -1388,7 +1406,13 @@ func.func @complex_sin_with_fmf(%arg: complex<f32>) -> complex<f32> {
// CHECK-DAG: %[[SIN:.*]] = math.sin %[[REAL]] fastmath<nnan,contract> : f32
// CHECK-DAG: %[[COS:.*]] = math.cos %[[REAL]] fastmath<nnan,contract> : f32
// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_EXP]], %[[HALF_REXP]] fastmath<nnan,contract>
-// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[RESULT_REAL_PRODUCT:.*]] = arith.mulf %[[EXP_SUM]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[SIN_IS_ZERO:.*]] = arith.cmpf oeq, %[[SIN]], %[[ZERO]] fastmath<nnan,contract>
+// CHECK-DAG: %[[COEFF_SIGN:.*]] = math.copysign %[[ONE]], %[[EXP_SUM]] fastmath<nnan,contract>
+// CHECK-DAG: %[[SIGNED_ZERO:.*]] = arith.mulf %[[COEFF_SIGN]], %[[SIN]] fastmath<nnan,contract>
+// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.select %[[SIN_IS_ZERO]], %[[SIGNED_ZERO]], %[[RESULT_REAL_PRODUCT]]
// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_EXP]], %[[HALF_REXP]] fastmath<nnan,contract>
// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[COS]] fastmath<nnan,contract>
// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
@@ -1396,6 +1420,30 @@ func.func @complex_sin_with_fmf(%arg: complex<f32>) -> complex<f32> {
// -----
+// CHECK-LABEL: func @complex_sin_ninf
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_sin_ninf(%arg: complex<f32>) -> complex<f32> {
+ %sin = complex.sin %arg fastmath<ninf> : complex<f32>
+ return %sin : complex<f32>
+}
+// CHECK-NOT: arith.select
+// CHECK: %[[RESULT:.*]] = complex.create
+// CHECK: return %[[RESULT]] : complex<f32>
+
+// -----
+
+// CHECK-LABEL: func @complex_cos_ninf
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_cos_ninf(%arg: complex<f32>) -> complex<f32> {
+ %cos = complex.cos %arg fastmath<ninf> : complex<f32>
+ return %cos : complex<f32>
+}
+// CHECK-NOT: arith.select
+// CHECK: %[[RESULT:.*]] = complex.create
+// CHECK: return %[[RESULT]] : complex<f32>
+
+// -----
+
// CHECK-LABEL: func @complex_sign_with_fmf
// CHECK-SAME: %[[ARG:.*]]: complex<f32>
func.func @complex_sign_with_fmf(%arg: complex<f32>) -> complex<f32> {
diff --git a/mlir/test/Integration/Dialect/Complex/CPU/correctness.mlir b/mlir/test/Integration/Dialect/Complex/CPU/correctness.mlir
index 29ad9603d1b9b2..feacf8bdf393a3 100644
--- a/mlir/test/Integration/Dialect/Complex/CPU/correctness.mlir
+++ b/mlir/test/Integration/Dialect/Complex/CPU/correctness.mlir
@@ -54,6 +54,16 @@ func.func @exp(%arg: complex<f32>) -> complex<f32> {
func.return %exp : complex<f32>
}
+func.func @sin(%arg: complex<f32>) -> complex<f32> {
+ %sin = complex.sin %arg : complex<f32>
+ func.return %sin : complex<f32>
+}
+
+func.func @cos(%arg: complex<f32>) -> complex<f32> {
+ %cos = complex.cos %arg : complex<f32>
+ func.return %cos : complex<f32>
+}
+
// %input contains pairs of lhs, rhs, i.e. [lhs_0, rhs_0, lhs_1, rhs_1,...]
func.func @test_binary(%input: tensor<?xcomplex<f32>>,
%func: (complex<f32>, complex<f32>) -> complex<f32>) {
@@ -389,5 +399,49 @@ func.func @entry() {
call @test_unary(%exp_test_cast, %exp_func)
: (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
+ // complex.sin and complex.cos tests. For x = +-0 and y = +-inf, computing
+ // exp(y) * sin(x) directly gives inf * 0 = NaN instead of a signed zero.
+ %trig_test = arith.constant dense<[
+ (1.0, 2.0),
+ // (0, inf)
+ (0.0, 0x7F800000),
+ // (-0, inf)
+ (0x80000000, 0x7F800000),
+ // (0, -inf)
+ (0.0, 0xFF800000),
+ // (0, nan)
+ (0.0, 0x7FC00000)
+ ]> : tensor<5xcomplex<f32>>
+ %trig_test_cast = tensor.cast %trig_test
+ : tensor<5xcomplex<f32>> to tensor<?xcomplex<f32>>
+
+ %sin_func = func.constant @sin : (complex<f32>) -> complex<f32>
+ call @test_unary(%trig_test_cast, %sin_func)
+ : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
+ // CHECK-NEXT: 3.1657
+ // CHECK-NEXT: 1.9596
+ // CHECK-NEXT: {{^}}0{{$}}
+ // CHECK-NEXT: inf
+ // CHECK-NEXT: {{^}}-0{{$}}
+ // CHECK-NEXT: inf
+ // CHECK-NEXT: {{^}}0{{$}}
+ // CHECK-NEXT: -inf
+ // CHECK-NEXT: {{^}}0{{$}}
+ // CHECK-NEXT: nan
+
+ %cos_func = func.constant @cos : (complex<f32>) -> complex<f32>
+ call @test_unary(%trig_test_cast, %cos_func)
+ : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
+ // CHECK-NEXT: 2.0327
+ // CHECK-NEXT: -3.0519
+ // CHECK-NEXT: inf
+ // CHECK-NEXT: {{^}}-0{{$}}
+ // CHECK-NEXT: inf
+ // CHECK-NEXT: {{^}}0{{$}}
+ // CHECK-NEXT: inf
+ // CHECK-NEXT: {{^}}0{{$}}
+ // CHECK-NEXT: nan
+ // CHECK-NEXT: {{^}}0{{$}}
+
func.return
}
More information about the Mlir-commits
mailing list