Skip to content

Commit 5c9315f

Browse files
jreiffersjreiffers
authored
Fix complex log1p accuracy with large abs values. (#88364)
This ports openxla/xla#10503 by @pearu. In addition to the filecheck test here, the accuracy was tested with XLA's complex_unary_op_test and its MLIR emitters. This is a fixed version of #88260. The previous version relied on implementation-specific behavior in the order of evaluation of maxAbsOfRealPlusOneAndImagMinusOne's operands.
1 parent fe3b20d commit 5c9315f

File tree

2 files changed

+58
-41
lines changed

2 files changed

+58
-41
lines changed

‎mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp‎

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -570,37 +570,40 @@ struct Log1pOpConversion : public OpConversionPattern<complex::Log1pOp> {
570570
ConversionPatternRewriter &rewriter) const override {
571571
auto type = cast<ComplexType>(adaptor.getComplex().getType());
572572
auto elementType = cast<FloatType>(type.getElementType());
573-
arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
573+
arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
574574
mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
575575

576-
Value real = b.create<complex::ReOp>(elementType, adaptor.getComplex());
577-
Value imag = b.create<complex::ImOp>(elementType, adaptor.getComplex());
576+
Value real = b.create<complex::ReOp>(adaptor.getComplex());
577+
Value imag = b.create<complex::ImOp>(adaptor.getComplex());
578578

579579
Value half = b.create<arith::ConstantOp>(elementType,
580580
b.getFloatAttr(elementType, 0.5));
581581
Value one = b.create<arith::ConstantOp>(elementType,
582582
b.getFloatAttr(elementType, 1));
583-
Value two = b.create<arith::ConstantOp>(elementType,
584-
b.getFloatAttr(elementType, 2));
585-
586-
// log1p(a+bi) = .5*log((a+1)^2+b^2) + i*atan2(b, a + 1)
587-
// log((a+1)+bi) = .5*log(a*a + 2*a + 1 + b*b) + i*atan2(b, a+1)
588-
// log((a+1)+bi) = .5*log1p(a*a + 2*a + b*b) + i*atan2(b, a+1)
589-
Value sumSq = b.create<arith::MulFOp>(real, real, fmf.getValue());
590-
sumSq = b.create<arith::AddFOp>(
591-
sumSq, b.create<arith::MulFOp>(real, two, fmf.getValue()),
592-
fmf.getValue());
593-
sumSq = b.create<arith::AddFOp>(
594-
sumSq, b.create<arith::MulFOp>(imag, imag, fmf.getValue()),
595-
fmf.getValue());
596-
Value logSumSq =
597-
b.create<math::Log1pOp>(elementType, sumSq, fmf.getValue());
598-
Value resultReal = b.create<arith::MulFOp>(logSumSq, half, fmf.getValue());
599-
600-
Value realPlusOne = b.create<arith::AddFOp>(real, one, fmf.getValue());
601-
602-
Value resultImag =
603-
b.create<math::Atan2Op>(elementType, imag, realPlusOne, fmf.getValue());
583+
Value realPlusOne = b.create<arith::AddFOp>(real, one, fmf);
584+
Value absRealPlusOne = b.create<math::AbsFOp>(realPlusOne, fmf);
585+
Value absImag = b.create<math::AbsFOp>(imag, fmf);
586+
587+
Value maxAbs = b.create<arith::MaximumFOp>(absRealPlusOne, absImag, fmf);
588+
Value minAbs = b.create<arith::MinimumFOp>(absRealPlusOne, absImag, fmf);
589+
590+
Value useReal = b.create<arith::CmpFOp>(arith::CmpFPredicate::OGT,
591+
realPlusOne, absImag, fmf);
592+
Value maxMinusOne = b.create<arith::SubFOp>(maxAbs, one, fmf);
593+
Value maxAbsOfRealPlusOneAndImagMinusOne =
594+
b.create<arith::SelectOp>(useReal, real, maxMinusOne);
595+
Value minMaxRatio = b.create<arith::DivFOp>(minAbs, maxAbs, fmf);
596+
Value logOfMaxAbsOfRealPlusOneAndImag =
597+
b.create<math::Log1pOp>(maxAbsOfRealPlusOneAndImagMinusOne, fmf);
598+
Value logOfSqrtPart = b.create<math::Log1pOp>(
599+
b.create<arith::MulFOp>(minMaxRatio, minMaxRatio, fmf), fmf);
600+
Value r = b.create<arith::AddFOp>(
601+
b.create<arith::MulFOp>(half, logOfSqrtPart, fmf),
602+
logOfMaxAbsOfRealPlusOneAndImag, fmf);
603+
Value resultReal = b.create<arith::SelectOp>(
604+
b.create<arith::CmpFOp>(arith::CmpFPredicate::UNO, r, r, fmf), minAbs,
605+
r);
606+
Value resultImag = b.create<math::Atan2Op>(imag, realPlusOne, fmf);
604607
rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
605608
resultImag);
606609
return success();

‎mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir‎

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -300,15 +300,22 @@ func.func @complex_log1p(%arg: complex<f32>) -> complex<f32> {
300300
// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
301301
// CHECK: %[[ONE_HALF:.*]] = arith.constant 5.000000e-01 : f32
302302
// CHECK: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
303-
// CHECK: %[[TWO:.*]] = arith.constant 2.000000e+00 : f32
304-
// CHECK: %[[SQ_SUM_0:.*]] = arith.mulf %[[REAL]], %[[REAL]] : f32
305-
// CHECK: %[[TWO_REAL:.*]] = arith.mulf %[[REAL]], %[[TWO]] : f32
306-
// CHECK: %[[SQ_SUM_1:.*]] = arith.addf %[[SQ_SUM_0]], %[[TWO_REAL]] : f32
307-
// CHECK: %[[SQ_IMAG:.*]] = arith.mulf %[[IMAG]], %[[IMAG]] : f32
308-
// CHECK: %[[SQ_SUM_2:.*]] = arith.addf %[[SQ_SUM_1]], %[[SQ_IMAG]] : f32
309-
// CHECK: %[[LOG_SQ_SUM:.*]] = math.log1p %[[SQ_SUM_2]] : f32
310-
// CHECK: %[[RESULT_REAL:.*]] = arith.mulf %[[LOG_SQ_SUM]], %[[ONE_HALF]] : f32
311303
// CHECK: %[[REAL_PLUS_ONE:.*]] = arith.addf %[[REAL]], %[[ONE]] : f32
304+
// CHECK: %[[ABS_REAL_PLUS_ONE:.*]] = math.absf %[[REAL_PLUS_ONE]] : f32
305+
// CHECK: %[[ABS_IMAG:.*]] = math.absf %[[IMAG]] : f32
306+
// CHECK: %[[MAX:.*]] = arith.maximumf %[[ABS_REAL_PLUS_ONE]], %[[ABS_IMAG]] : f32
307+
// CHECK: %[[MIN:.*]] = arith.minimumf %[[ABS_REAL_PLUS_ONE]], %[[ABS_IMAG]] : f32
308+
// CHECK: %[[CMPF:.*]] = arith.cmpf ogt, %[[REAL_PLUS_ONE]], %[[ABS_IMAG]] : f32
309+
// CHECK: %[[MAX_MINUS_ONE:.*]] = arith.subf %[[MAX]], %[[ONE]] : f32
310+
// CHECK: %[[SELECT:.*]] = arith.select %[[CMPF]], %0, %[[MAX_MINUS_ONE]] : f32
311+
// CHECK: %[[MIN_MAX_RATIO:.*]] = arith.divf %[[MIN]], %[[MAX]] : f32
312+
// CHECK: %[[LOG_1:.*]] = math.log1p %[[SELECT]] : f32
313+
// CHECK: %[[RATIO_SQ:.*]] = arith.mulf %[[MIN_MAX_RATIO]], %[[MIN_MAX_RATIO]] : f32
314+
// CHECK: %[[LOG_SQ:.*]] = math.log1p %[[RATIO_SQ]] : f32
315+
// CHECK: %[[HALF_LOG_SQ:.*]] = arith.mulf %cst, %[[LOG_SQ]] : f32
316+
// CHECK: %[[R:.*]] = arith.addf %[[HALF_LOG_SQ]], %[[LOG_1]] : f32
317+
// CHECK: %[[ISNAN:.*]] = arith.cmpf uno, %[[R]], %[[R]] : f32
318+
// CHECK: %[[RESULT_REAL:.*]] = arith.select %[[ISNAN]], %[[MIN]], %[[R]] : f32
312319
// CHECK: %[[RESULT_IMAG:.*]] = math.atan2 %[[IMAG]], %[[REAL_PLUS_ONE]] : f32
313320
// CHECK: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
314321
// CHECK: return %[[RESULT]] : complex<f32>
@@ -963,15 +970,22 @@ func.func @complex_log1p_with_fmf(%arg: complex<f32>) -> complex<f32> {
963970
// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
964971
// CHECK: %[[ONE_HALF:.*]] = arith.constant 5.000000e-01 : f32
965972
// CHECK: %[[ONE:.*]] = arith.constant 1.000000e+00 : f32
966-
// CHECK: %[[TWO:.*]] = arith.constant 2.000000e+00 : f32
967-
// CHECK: %[[SQ_SUM_0:.*]] = arith.mulf %[[REAL]], %[[REAL]] fastmath<nnan,contract> : f32
968-
// CHECK: %[[TWO_REAL:.*]] = arith.mulf %[[REAL]], %[[TWO]] fastmath<nnan,contract> : f32
969-
// CHECK: %[[SQ_SUM_1:.*]] = arith.addf %[[SQ_SUM_0]], %[[TWO_REAL]] fastmath<nnan,contract> : f32
970-
// CHECK: %[[SQ_IMAG:.*]] = arith.mulf %[[IMAG]], %[[IMAG]] fastmath<nnan,contract> : f32
971-
// CHECK: %[[SQ_SUM_2:.*]] = arith.addf %[[SQ_SUM_1]], %[[SQ_IMAG]] fastmath<nnan,contract> : f32
972-
// CHECK: %[[LOG_SQ_SUM:.*]] = math.log1p %[[SQ_SUM_2]] fastmath<nnan,contract> : f32
973-
// CHECK: %[[RESULT_REAL:.*]] = arith.mulf %[[LOG_SQ_SUM]], %[[ONE_HALF]] fastmath<nnan,contract> : f32
974-
// CHECK: %[[REAL_PLUS_ONE:.*]] = arith.addf %[[REAL]], %[[ONE]] fastmath<nnan,contract> : f32
973+
// CHECK: %[[REAL_PLUS_ONE:.*]] = arith.addf %[[REAL]], %[[ONE]] fastmath<nnan,contract> : f32
974+
// CHECK: %[[ABS_REAL_PLUS_ONE:.*]] = math.absf %[[REAL_PLUS_ONE]] fastmath<nnan,contract> : f32
975+
// CHECK: %[[ABS_IMAG:.*]] = math.absf %[[IMAG]] fastmath<nnan,contract> : f32
976+
// CHECK: %[[MAX:.*]] = arith.maximumf %[[ABS_REAL_PLUS_ONE]], %[[ABS_IMAG]] fastmath<nnan,contract> : f32
977+
// CHECK: %[[MIN:.*]] = arith.minimumf %[[ABS_REAL_PLUS_ONE]], %[[ABS_IMAG]] fastmath<nnan,contract> : f32
978+
// CHECK: %[[CMPF:.*]] = arith.cmpf ogt, %[[REAL_PLUS_ONE]], %[[ABS_IMAG]] fastmath<nnan,contract> : f32
979+
// CHECK: %[[MAX_MINUS_ONE:.*]] = arith.subf %[[MAX]], %[[ONE]] fastmath<nnan,contract> : f32
980+
// CHECK: %[[SELECT:.*]] = arith.select %[[CMPF]], %0, %[[MAX_MINUS_ONE]] : f32
981+
// CHECK: %[[MIN_MAX_RATIO:.*]] = arith.divf %[[MIN]], %[[MAX]] fastmath<nnan,contract> : f32
982+
// CHECK: %[[LOG_1:.*]] = math.log1p %[[SELECT]] fastmath<nnan,contract> : f32
983+
// CHECK: %[[RATIO_SQ:.*]] = arith.mulf %[[MIN_MAX_RATIO]], %[[MIN_MAX_RATIO]] fastmath<nnan,contract> : f32
984+
// CHECK: %[[LOG_SQ:.*]] = math.log1p %[[RATIO_SQ]] fastmath<nnan,contract> : f32
985+
// CHECK: %[[HALF_LOG_SQ:.*]] = arith.mulf %cst, %[[LOG_SQ]] fastmath<nnan,contract> : f32
986+
// CHECK: %[[R:.*]] = arith.addf %[[HALF_LOG_SQ]], %[[LOG_1]] fastmath<nnan,contract> : f32
987+
// CHECK: %[[ISNAN:.*]] = arith.cmpf uno, %[[R]], %[[R]] fastmath<nnan,contract> : f32
988+
// CHECK: %[[RESULT_REAL:.*]] = arith.select %[[ISNAN]], %[[MIN]], %[[R]] : f32
975989
// CHECK: %[[RESULT_IMAG:.*]] = math.atan2 %[[IMAG]], %[[REAL_PLUS_ONE]] fastmath<nnan,contract> : f32
976990
// CHECK: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
977991
// CHECK: return %[[RESULT]] : complex<f32>

0 commit comments

Comments
 (0)