From 7d2d8e2a7245e4e64da22cb3c422ea3be5a0bf0a Mon Sep 17 00:00:00 2001 From: Kai Sasaki Date: Mon, 25 Mar 2024 10:59:42 +0900 Subject: [PATCH] [mlir][complex] Fastmath flag for the trigonometric ops in complex (#85563) Support Fastmath flag to convert trigonometric ops in the complex dialect. See: https://discourse.llvm.org/t/rfc-fastmath-flags-support-in-complex-dialect/71981 --- .../ComplexToStandard/ComplexToStandard.cpp | 50 +++++++++++-------- .../convert-to-standard.mlir | 46 +++++++++++++++++ 2 files changed, 75 insertions(+), 21 deletions(-) diff --git a/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp b/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp index 76729278ec1b..17f64f1b65b7 100644 --- a/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp +++ b/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp @@ -196,6 +196,7 @@ struct TrigonometricOpConversion : public OpConversionPattern { auto loc = op.getLoc(); auto type = cast(adaptor.getComplex().getType()); auto elementType = cast(type.getElementType()); + arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr(); Value real = rewriter.create(loc, elementType, adaptor.getComplex()); @@ -207,14 +208,14 @@ struct TrigonometricOpConversion : public OpConversionPattern { // implementation in the subclass to combine them. Value half = rewriter.create( loc, elementType, rewriter.getFloatAttr(elementType, 0.5)); - Value exp = rewriter.create(loc, imag); - Value scaledExp = rewriter.create(loc, half, exp); - Value reciprocalExp = rewriter.create(loc, half, exp); - Value sin = rewriter.create(loc, real); - Value cos = rewriter.create(loc, real); + Value exp = rewriter.create(loc, imag, fmf); + Value scaledExp = rewriter.create(loc, half, exp, fmf); + Value reciprocalExp = rewriter.create(loc, half, exp, fmf); + Value sin = rewriter.create(loc, real, fmf); + Value cos = rewriter.create(loc, real, fmf); auto resultPair = - combine(loc, scaledExp, reciprocalExp, sin, cos, rewriter); + combine(loc, scaledExp, reciprocalExp, sin, cos, rewriter, fmf); rewriter.replaceOpWithNewOp(op, type, resultPair.first, resultPair.second); @@ -223,15 +224,17 @@ struct TrigonometricOpConversion : public OpConversionPattern { virtual std::pair combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin, - Value cos, ConversionPatternRewriter &rewriter) const = 0; + Value cos, ConversionPatternRewriter &rewriter, + arith::FastMathFlagsAttr fmf) const = 0; }; struct CosOpConversion : public TrigonometricOpConversion { using TrigonometricOpConversion::TrigonometricOpConversion; - std::pair - combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin, - Value cos, ConversionPatternRewriter &rewriter) const override { + std::pair combine(Location loc, Value scaledExp, + Value reciprocalExp, Value sin, Value cos, + ConversionPatternRewriter &rewriter, + arith::FastMathFlagsAttr fmf) const override { // Complex cosine is defined as; // cos(x + iy) = 0.5 * (exp(i(x + iy)) + exp(-i(x + iy))) // Plugging in: @@ -241,10 +244,12 @@ struct CosOpConversion : public TrigonometricOpConversion { // We get: // Re(cos(x + iy)) = (0.5/t + 0.5*t) * cos x // Im(cos(x + iy)) = (0.5/t - 0.5*t) * sin x - Value sum = rewriter.create(loc, reciprocalExp, scaledExp); - Value resultReal = rewriter.create(loc, sum, cos); - Value diff = rewriter.create(loc, reciprocalExp, scaledExp); - Value resultImag = rewriter.create(loc, diff, sin); + Value sum = + rewriter.create(loc, reciprocalExp, scaledExp, fmf); + Value resultReal = rewriter.create(loc, sum, cos, fmf); + Value diff = + rewriter.create(loc, reciprocalExp, scaledExp, fmf); + Value resultImag = rewriter.create(loc, diff, sin, fmf); return {resultReal, resultImag}; } }; @@ -813,9 +818,10 @@ struct NegOpConversion : public OpConversionPattern { struct SinOpConversion : public TrigonometricOpConversion { using TrigonometricOpConversion::TrigonometricOpConversion; - std::pair - combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin, - Value cos, ConversionPatternRewriter &rewriter) const override { + std::pair combine(Location loc, Value scaledExp, + Value reciprocalExp, Value sin, Value cos, + ConversionPatternRewriter &rewriter, + arith::FastMathFlagsAttr fmf) const override { // Complex sine is defined as; // sin(x + iy) = -0.5i * (exp(i(x + iy)) - exp(-i(x + iy))) // Plugging in: @@ -825,10 +831,12 @@ struct SinOpConversion : public TrigonometricOpConversion { // We get: // Re(sin(x + iy)) = (0.5*t + 0.5/t) * sin x // Im(cos(x + iy)) = (0.5*t - 0.5/t) * cos x - Value sum = rewriter.create(loc, scaledExp, reciprocalExp); - Value resultReal = rewriter.create(loc, sum, sin); - Value diff = rewriter.create(loc, scaledExp, reciprocalExp); - Value resultImag = rewriter.create(loc, diff, cos); + Value sum = + rewriter.create(loc, scaledExp, reciprocalExp, fmf); + Value resultReal = rewriter.create(loc, sum, sin, fmf); + Value diff = + rewriter.create(loc, scaledExp, reciprocalExp, fmf); + Value resultImag = rewriter.create(loc, diff, cos, fmf); return {resultReal, resultImag}; } }; diff --git a/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir b/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir index 5918ff2e0f36..bac94aae6b74 100644 --- a/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir +++ b/mlir/test/Conversion/ComplexToStandard/convert-to-standard.mlir @@ -1834,3 +1834,49 @@ func.func @complex_sqrt_with_fmf(%arg: complex) -> complex { // CHECK: %[[VAR40:.*]] = arith.select %[[VAR38]], %cst, %[[VAR32]] : f32 // CHECK: %[[VAR41:.*]] = complex.create %[[VAR39]], %[[VAR40]] : complex // CHECK: return %[[VAR41]] : complex + +// ----- + +// CHECK-LABEL: func @complex_cos_with_fmf +// CHECK-SAME: %[[ARG:.*]]: complex +func.func @complex_cos_with_fmf(%arg: complex) -> complex { + %cos = complex.cos %arg fastmath : complex + return %cos : complex +} +// CHECK-DAG: %[[REAL:.*]] = complex.re %[[ARG]] +// CHECK-DAG: %[[IMAG:.*]] = complex.im %[[ARG]] +// CHECK-DAG: %[[HALF:.*]] = arith.constant 5.000000e-01 : f32 +// CHECK-DAG: %[[EXP:.*]] = math.exp %[[IMAG]] fastmath : f32 +// CHECK-DAG: %[[HALF_EXP:.*]] = arith.mulf %[[HALF]], %[[EXP]] fastmath +// CHECK-DAG: %[[HALF_REXP:.*]] = arith.divf %[[HALF]], %[[EXP]] fastmath +// CHECK-DAG: %[[SIN:.*]] = math.sin %[[REAL]] fastmath : f32 +// CHECK-DAG: %[[COS:.*]] = math.cos %[[REAL]] fastmath : f32 +// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_REXP]], %[[HALF_EXP]] fastmath +// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[COS]] fastmath +// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_REXP]], %[[HALF_EXP]] fastmath +// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[SIN]] fastmath +// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex +// CHECK: return %[[RESULT]] + +// ----- + +// CHECK-LABEL: func @complex_sin_with_fmf +// CHECK-SAME: %[[ARG:.*]]: complex +func.func @complex_sin_with_fmf(%arg: complex) -> complex { + %cos = complex.sin %arg fastmath : complex + return %cos : complex +} +// CHECK-DAG: %[[REAL:.*]] = complex.re %[[ARG]] +// CHECK-DAG: %[[IMAG:.*]] = complex.im %[[ARG]] +// CHECK-DAG: %[[HALF:.*]] = arith.constant 5.000000e-01 : f32 +// CHECK-DAG: %[[EXP:.*]] = math.exp %[[IMAG]] fastmath : f32 +// CHECK-DAG: %[[HALF_EXP:.*]] = arith.mulf %[[HALF]], %[[EXP]] fastmath +// CHECK-DAG: %[[HALF_REXP:.*]] = arith.divf %[[HALF]], %[[EXP]] fastmath +// CHECK-DAG: %[[SIN:.*]] = math.sin %[[REAL]] fastmath : f32 +// CHECK-DAG: %[[COS:.*]] = math.cos %[[REAL]] fastmath : f32 +// CHECK-DAG: %[[EXP_SUM:.*]] = arith.addf %[[HALF_EXP]], %[[HALF_REXP]] fastmath +// CHECK-DAG: %[[RESULT_REAL:.*]] = arith.mulf %[[EXP_SUM]], %[[SIN]] fastmath +// CHECK-DAG: %[[EXP_DIFF:.*]] = arith.subf %[[HALF_EXP]], %[[HALF_REXP]] fastmath +// CHECK-DAG: %[[RESULT_IMAG:.*]] = arith.mulf %[[EXP_DIFF]], %[[COS]] fastmath +// CHECK-DAG: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex +// CHECK: return %[[RESULT]]