MLIR  22.0.0git
ComplexToStandard.cpp
Go to the documentation of this file.
1 //===- ComplexToStandard.cpp - conversion from Complex to Standard dialect ===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 
15 #include "mlir/IR/PatternMatch.h"
17 #include <type_traits>
18 
19 namespace mlir {
20 #define GEN_PASS_DEF_CONVERTCOMPLEXTOSTANDARDPASS
21 #include "mlir/Conversion/Passes.h.inc"
22 } // namespace mlir
23 
24 using namespace mlir;
25 
26 namespace {
27 
28 enum class AbsFn { abs, sqrt, rsqrt };
29 
30 // Returns the absolute value, its square root or its reciprocal square root.
31 Value computeAbs(Value real, Value imag, arith::FastMathFlags fmf,
32  ImplicitLocOpBuilder &b, AbsFn fn = AbsFn::abs) {
33  Value one = arith::ConstantOp::create(b, real.getType(),
34  b.getFloatAttr(real.getType(), 1.0));
35 
36  Value absReal = math::AbsFOp::create(b, real, fmf);
37  Value absImag = math::AbsFOp::create(b, imag, fmf);
38 
39  Value max = arith::MaximumFOp::create(b, absReal, absImag, fmf);
40  Value min = arith::MinimumFOp::create(b, absReal, absImag, fmf);
41 
42  // The lowering below requires NaNs and infinities to work correctly.
43  arith::FastMathFlags fmfWithNaNInf = arith::bitEnumClear(
44  fmf, arith::FastMathFlags::nnan | arith::FastMathFlags::ninf);
45  Value ratio = arith::DivFOp::create(b, min, max, fmfWithNaNInf);
46  Value ratioSq = arith::MulFOp::create(b, ratio, ratio, fmfWithNaNInf);
47  Value ratioSqPlusOne = arith::AddFOp::create(b, ratioSq, one, fmfWithNaNInf);
48  Value result;
49 
50  if (fn == AbsFn::rsqrt) {
51  ratioSqPlusOne = math::RsqrtOp::create(b, ratioSqPlusOne, fmfWithNaNInf);
52  min = math::RsqrtOp::create(b, min, fmfWithNaNInf);
53  max = math::RsqrtOp::create(b, max, fmfWithNaNInf);
54  }
55 
56  if (fn == AbsFn::sqrt) {
57  Value quarter = arith::ConstantOp::create(
58  b, real.getType(), b.getFloatAttr(real.getType(), 0.25));
59  // sqrt(sqrt(a*b)) would avoid the pow, but will overflow more easily.
60  Value sqrt = math::SqrtOp::create(b, max, fmfWithNaNInf);
61  Value p025 =
62  math::PowFOp::create(b, ratioSqPlusOne, quarter, fmfWithNaNInf);
63  result = arith::MulFOp::create(b, sqrt, p025, fmfWithNaNInf);
64  } else {
65  Value sqrt = math::SqrtOp::create(b, ratioSqPlusOne, fmfWithNaNInf);
66  result = arith::MulFOp::create(b, max, sqrt, fmfWithNaNInf);
67  }
68 
69  Value isNaN = arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, result,
70  result, fmfWithNaNInf);
71  return arith::SelectOp::create(b, isNaN, min, result);
72 }
73 
74 struct AbsOpConversion : public OpConversionPattern<complex::AbsOp> {
76 
77  LogicalResult
78  matchAndRewrite(complex::AbsOp op, OpAdaptor adaptor,
79  ConversionPatternRewriter &rewriter) const override {
80  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
81 
82  arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
83 
84  Value real = complex::ReOp::create(b, adaptor.getComplex());
85  Value imag = complex::ImOp::create(b, adaptor.getComplex());
86  rewriter.replaceOp(op, computeAbs(real, imag, fmf, b));
87 
88  return success();
89  }
90 };
91 
92 // atan2(y,x) = -i * log((x + i * y)/sqrt(x**2+y**2))
93 struct Atan2OpConversion : public OpConversionPattern<complex::Atan2Op> {
95 
96  LogicalResult
97  matchAndRewrite(complex::Atan2Op op, OpAdaptor adaptor,
98  ConversionPatternRewriter &rewriter) const override {
99  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
100 
101  auto type = cast<ComplexType>(op.getType());
102  Type elementType = type.getElementType();
103  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
104 
105  Value lhs = adaptor.getLhs();
106  Value rhs = adaptor.getRhs();
107 
108  Value rhsSquared = complex::MulOp::create(b, type, rhs, rhs, fmf);
109  Value lhsSquared = complex::MulOp::create(b, type, lhs, lhs, fmf);
110  Value rhsSquaredPlusLhsSquared =
111  complex::AddOp::create(b, type, rhsSquared, lhsSquared, fmf);
112  Value sqrtOfRhsSquaredPlusLhsSquared =
113  complex::SqrtOp::create(b, type, rhsSquaredPlusLhsSquared, fmf);
114 
115  Value zero =
116  arith::ConstantOp::create(b, elementType, b.getZeroAttr(elementType));
117  Value one = arith::ConstantOp::create(b, elementType,
118  b.getFloatAttr(elementType, 1));
119  Value i = complex::CreateOp::create(b, type, zero, one);
120  Value iTimesLhs = complex::MulOp::create(b, i, lhs, fmf);
121  Value rhsPlusILhs = complex::AddOp::create(b, rhs, iTimesLhs, fmf);
122 
123  Value divResult = complex::DivOp::create(
124  b, rhsPlusILhs, sqrtOfRhsSquaredPlusLhsSquared, fmf);
125  Value logResult = complex::LogOp::create(b, divResult, fmf);
126 
127  Value negativeOne = arith::ConstantOp::create(
128  b, elementType, b.getFloatAttr(elementType, -1));
129  Value negativeI = complex::CreateOp::create(b, type, zero, negativeOne);
130 
131  rewriter.replaceOpWithNewOp<complex::MulOp>(op, negativeI, logResult, fmf);
132  return success();
133  }
134 };
135 
136 template <typename ComparisonOp, arith::CmpFPredicate p>
137 struct ComparisonOpConversion : public OpConversionPattern<ComparisonOp> {
139  using ResultCombiner =
140  std::conditional_t<std::is_same<ComparisonOp, complex::EqualOp>::value,
141  arith::AndIOp, arith::OrIOp>;
142 
143  LogicalResult
144  matchAndRewrite(ComparisonOp op, typename ComparisonOp::Adaptor adaptor,
145  ConversionPatternRewriter &rewriter) const override {
146  auto loc = op.getLoc();
147  auto type = cast<ComplexType>(adaptor.getLhs().getType()).getElementType();
148 
149  Value realLhs =
150  complex::ReOp::create(rewriter, loc, type, adaptor.getLhs());
151  Value imagLhs =
152  complex::ImOp::create(rewriter, loc, type, adaptor.getLhs());
153  Value realRhs =
154  complex::ReOp::create(rewriter, loc, type, adaptor.getRhs());
155  Value imagRhs =
156  complex::ImOp::create(rewriter, loc, type, adaptor.getRhs());
157  Value realComparison =
158  arith::CmpFOp::create(rewriter, loc, p, realLhs, realRhs);
159  Value imagComparison =
160  arith::CmpFOp::create(rewriter, loc, p, imagLhs, imagRhs);
161 
162  rewriter.replaceOpWithNewOp<ResultCombiner>(op, realComparison,
163  imagComparison);
164  return success();
165  }
166 };
167 
168 // Default conversion which applies the BinaryStandardOp separately on the real
169 // and imaginary parts. Can for example be used for complex::AddOp and
170 // complex::SubOp.
171 template <typename BinaryComplexOp, typename BinaryStandardOp>
172 struct BinaryComplexOpConversion : public OpConversionPattern<BinaryComplexOp> {
174 
175  LogicalResult
176  matchAndRewrite(BinaryComplexOp op, typename BinaryComplexOp::Adaptor adaptor,
177  ConversionPatternRewriter &rewriter) const override {
178  auto type = cast<ComplexType>(adaptor.getLhs().getType());
179  auto elementType = cast<FloatType>(type.getElementType());
180  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
181  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
182 
183  Value realLhs = complex::ReOp::create(b, elementType, adaptor.getLhs());
184  Value realRhs = complex::ReOp::create(b, elementType, adaptor.getRhs());
185  Value resultReal = BinaryStandardOp::create(b, elementType, realLhs,
186  realRhs, fmf.getValue());
187  Value imagLhs = complex::ImOp::create(b, elementType, adaptor.getLhs());
188  Value imagRhs = complex::ImOp::create(b, elementType, adaptor.getRhs());
189  Value resultImag = BinaryStandardOp::create(b, elementType, imagLhs,
190  imagRhs, fmf.getValue());
191  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
192  resultImag);
193  return success();
194  }
195 };
196 
197 template <typename TrigonometricOp>
198 struct TrigonometricOpConversion : public OpConversionPattern<TrigonometricOp> {
199  using OpAdaptor = typename OpConversionPattern<TrigonometricOp>::OpAdaptor;
200 
202 
203  LogicalResult
204  matchAndRewrite(TrigonometricOp op, OpAdaptor adaptor,
205  ConversionPatternRewriter &rewriter) const override {
206  auto loc = op.getLoc();
207  auto type = cast<ComplexType>(adaptor.getComplex().getType());
208  auto elementType = cast<FloatType>(type.getElementType());
209  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
210 
211  Value real =
212  complex::ReOp::create(rewriter, loc, elementType, adaptor.getComplex());
213  Value imag =
214  complex::ImOp::create(rewriter, loc, elementType, adaptor.getComplex());
215 
216  // Trigonometric ops use a set of common building blocks to convert to real
217  // ops. Here we create these building blocks and call into an op-specific
218  // implementation in the subclass to combine them.
219  Value half = arith::ConstantOp::create(
220  rewriter, loc, elementType, rewriter.getFloatAttr(elementType, 0.5));
221  Value exp = math::ExpOp::create(rewriter, loc, imag, fmf);
222  Value scaledExp = arith::MulFOp::create(rewriter, loc, half, exp, fmf);
223  Value reciprocalExp = arith::DivFOp::create(rewriter, loc, half, exp, fmf);
224  Value sin = math::SinOp::create(rewriter, loc, real, fmf);
225  Value cos = math::CosOp::create(rewriter, loc, real, fmf);
226 
227  auto resultPair =
228  combine(loc, scaledExp, reciprocalExp, sin, cos, rewriter, fmf);
229 
230  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultPair.first,
231  resultPair.second);
232  return success();
233  }
234 
235  virtual std::pair<Value, Value>
236  combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin,
237  Value cos, ConversionPatternRewriter &rewriter,
238  arith::FastMathFlagsAttr fmf) const = 0;
239 };
240 
241 struct CosOpConversion : public TrigonometricOpConversion<complex::CosOp> {
242  using TrigonometricOpConversion<complex::CosOp>::TrigonometricOpConversion;
243 
244  std::pair<Value, Value> combine(Location loc, Value scaledExp,
245  Value reciprocalExp, Value sin, Value cos,
246  ConversionPatternRewriter &rewriter,
247  arith::FastMathFlagsAttr fmf) const override {
248  // Complex cosine is defined as;
249  // cos(x + iy) = 0.5 * (exp(i(x + iy)) + exp(-i(x + iy)))
250  // Plugging in:
251  // exp(i(x+iy)) = exp(-y + ix) = exp(-y)(cos(x) + i sin(x))
252  // exp(-i(x+iy)) = exp(y + i(-x)) = exp(y)(cos(x) + i (-sin(x)))
253  // and defining t := exp(y)
254  // We get:
255  // Re(cos(x + iy)) = (0.5/t + 0.5*t) * cos x
256  // Im(cos(x + iy)) = (0.5/t - 0.5*t) * sin x
257  Value sum =
258  arith::AddFOp::create(rewriter, loc, reciprocalExp, scaledExp, fmf);
259  Value resultReal = arith::MulFOp::create(rewriter, loc, sum, cos, fmf);
260  Value diff =
261  arith::SubFOp::create(rewriter, loc, reciprocalExp, scaledExp, fmf);
262  Value resultImag = arith::MulFOp::create(rewriter, loc, diff, sin, fmf);
263  return {resultReal, resultImag};
264  }
265 };
266 
267 struct DivOpConversion : public OpConversionPattern<complex::DivOp> {
268  DivOpConversion(MLIRContext *context, complex::ComplexRangeFlags target)
269  : OpConversionPattern<complex::DivOp>(context), complexRange(target) {}
270 
272 
273  LogicalResult
274  matchAndRewrite(complex::DivOp op, OpAdaptor adaptor,
275  ConversionPatternRewriter &rewriter) const override {
276  auto loc = op.getLoc();
277  auto type = cast<ComplexType>(adaptor.getLhs().getType());
278  auto elementType = cast<FloatType>(type.getElementType());
279  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
280 
281  Value lhsReal =
282  complex::ReOp::create(rewriter, loc, elementType, adaptor.getLhs());
283  Value lhsImag =
284  complex::ImOp::create(rewriter, loc, elementType, adaptor.getLhs());
285  Value rhsReal =
286  complex::ReOp::create(rewriter, loc, elementType, adaptor.getRhs());
287  Value rhsImag =
288  complex::ImOp::create(rewriter, loc, elementType, adaptor.getRhs());
289 
290  Value resultReal, resultImag;
291 
292  if (complexRange == complex::ComplexRangeFlags::basic ||
293  complexRange == complex::ComplexRangeFlags::none) {
295  rewriter, loc, lhsReal, lhsImag, rhsReal, rhsImag, fmf, &resultReal,
296  &resultImag);
297  } else if (complexRange == complex::ComplexRangeFlags::improved) {
299  rewriter, loc, lhsReal, lhsImag, rhsReal, rhsImag, fmf, &resultReal,
300  &resultImag);
301  }
302 
303  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
304  resultImag);
305 
306  return success();
307  }
308 
309 private:
310  complex::ComplexRangeFlags complexRange;
311 };
312 
313 struct ExpOpConversion : public OpConversionPattern<complex::ExpOp> {
315 
316  LogicalResult
317  matchAndRewrite(complex::ExpOp op, OpAdaptor adaptor,
318  ConversionPatternRewriter &rewriter) const override {
319  auto loc = op.getLoc();
320  auto type = cast<ComplexType>(adaptor.getComplex().getType());
321  auto elementType = cast<FloatType>(type.getElementType());
322  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
323 
324  Value real =
325  complex::ReOp::create(rewriter, loc, elementType, adaptor.getComplex());
326  Value imag =
327  complex::ImOp::create(rewriter, loc, elementType, adaptor.getComplex());
328  Value expReal = math::ExpOp::create(rewriter, loc, real, fmf.getValue());
329  Value cosImag = math::CosOp::create(rewriter, loc, imag, fmf.getValue());
330  Value resultReal =
331  arith::MulFOp::create(rewriter, loc, expReal, cosImag, fmf.getValue());
332  Value sinImag = math::SinOp::create(rewriter, loc, imag, fmf.getValue());
333  Value resultImag =
334  arith::MulFOp::create(rewriter, loc, expReal, sinImag, fmf.getValue());
335 
336  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
337  resultImag);
338  return success();
339  }
340 };
341 
342 Value evaluatePolynomial(ImplicitLocOpBuilder &b, Value arg,
343  ArrayRef<double> coefficients,
344  arith::FastMathFlagsAttr fmf) {
345  auto argType = mlir::cast<FloatType>(arg.getType());
346  Value poly =
347  arith::ConstantOp::create(b, b.getFloatAttr(argType, coefficients[0]));
348  for (unsigned i = 1; i < coefficients.size(); ++i) {
349  poly = math::FmaOp::create(
350  b, poly, arg,
351  arith::ConstantOp::create(b, b.getFloatAttr(argType, coefficients[i])),
352  fmf);
353  }
354  return poly;
355 }
356 
357 struct Expm1OpConversion : public OpConversionPattern<complex::Expm1Op> {
359 
360  // e^(a+bi)-1 = (e^a*cos(b)-1)+e^a*sin(b)i
361  // [handle inaccuracies when a and/or b are small]
362  // = ((e^a - 1) * cos(b) + cos(b) - 1) + e^a*sin(b)i
363  // = (expm1(a) * cos(b) + cosm1(b)) + e^a*sin(b)i
364  LogicalResult
365  matchAndRewrite(complex::Expm1Op op, OpAdaptor adaptor,
366  ConversionPatternRewriter &rewriter) const override {
367  auto type = op.getType();
368  auto elemType = mlir::cast<FloatType>(type.getElementType());
369 
370  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
371  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
372  Value real = complex::ReOp::create(b, adaptor.getComplex());
373  Value imag = complex::ImOp::create(b, adaptor.getComplex());
374 
375  Value zero = arith::ConstantOp::create(b, b.getFloatAttr(elemType, 0.0));
376  Value one = arith::ConstantOp::create(b, b.getFloatAttr(elemType, 1.0));
377 
378  Value expm1Real = math::ExpM1Op::create(b, real, fmf);
379  Value expReal = arith::AddFOp::create(b, expm1Real, one, fmf);
380 
381  Value sinImag = math::SinOp::create(b, imag, fmf);
382  Value cosm1Imag = emitCosm1(imag, fmf, b);
383  Value cosImag = arith::AddFOp::create(b, cosm1Imag, one, fmf);
384 
385  Value realResult = arith::AddFOp::create(
386  b, arith::MulFOp::create(b, expm1Real, cosImag, fmf), cosm1Imag, fmf);
387 
388  Value imagIsZero = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, imag,
389  zero, fmf.getValue());
390  Value imagResult = arith::SelectOp::create(
391  b, imagIsZero, zero, arith::MulFOp::create(b, expReal, sinImag, fmf));
392 
393  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, realResult,
394  imagResult);
395  return success();
396  }
397 
398 private:
399  Value emitCosm1(Value arg, arith::FastMathFlagsAttr fmf,
400  ImplicitLocOpBuilder &b) const {
401  auto argType = mlir::cast<FloatType>(arg.getType());
402  auto negHalf = arith::ConstantOp::create(b, b.getFloatAttr(argType, -0.5));
403  auto negOne = arith::ConstantOp::create(b, b.getFloatAttr(argType, -1.0));
404 
405  // Algorithm copied from cephes cosm1.
406  SmallVector<double, 7> kCoeffs{
407  4.7377507964246204691685E-14, -1.1470284843425359765671E-11,
408  2.0876754287081521758361E-9, -2.7557319214999787979814E-7,
409  2.4801587301570552304991E-5, -1.3888888888888872993737E-3,
410  4.1666666666666666609054E-2,
411  };
412  Value cos = math::CosOp::create(b, arg, fmf);
413  Value forLargeArg = arith::AddFOp::create(b, cos, negOne, fmf);
414 
415  Value argPow2 = arith::MulFOp::create(b, arg, arg, fmf);
416  Value argPow4 = arith::MulFOp::create(b, argPow2, argPow2, fmf);
417  Value poly = evaluatePolynomial(b, argPow2, kCoeffs, fmf);
418 
419  auto forSmallArg =
420  arith::AddFOp::create(b, arith::MulFOp::create(b, argPow4, poly, fmf),
421  arith::MulFOp::create(b, negHalf, argPow2, fmf));
422 
423  // (pi/4)^2 is approximately 0.61685
424  Value piOver4Pow2 =
425  arith::ConstantOp::create(b, b.getFloatAttr(argType, 0.61685));
426  Value cond = arith::CmpFOp::create(b, arith::CmpFPredicate::OGE, argPow2,
427  piOver4Pow2, fmf.getValue());
428  return arith::SelectOp::create(b, cond, forLargeArg, forSmallArg);
429  }
430 };
431 
432 struct LogOpConversion : public OpConversionPattern<complex::LogOp> {
434 
435  LogicalResult
436  matchAndRewrite(complex::LogOp op, OpAdaptor adaptor,
437  ConversionPatternRewriter &rewriter) const override {
438  auto type = cast<ComplexType>(adaptor.getComplex().getType());
439  auto elementType = cast<FloatType>(type.getElementType());
440  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
441  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
442 
443  Value abs = complex::AbsOp::create(b, elementType, adaptor.getComplex(),
444  fmf.getValue());
445  Value resultReal = math::LogOp::create(b, elementType, abs, fmf.getValue());
446  Value real = complex::ReOp::create(b, elementType, adaptor.getComplex());
447  Value imag = complex::ImOp::create(b, elementType, adaptor.getComplex());
448  Value resultImag =
449  math::Atan2Op::create(b, elementType, imag, real, fmf.getValue());
450  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
451  resultImag);
452  return success();
453  }
454 };
455 
456 struct Log1pOpConversion : public OpConversionPattern<complex::Log1pOp> {
458 
459  LogicalResult
460  matchAndRewrite(complex::Log1pOp op, OpAdaptor adaptor,
461  ConversionPatternRewriter &rewriter) const override {
462  auto type = cast<ComplexType>(adaptor.getComplex().getType());
463  auto elementType = cast<FloatType>(type.getElementType());
464  arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
465  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
466 
467  Value real = complex::ReOp::create(b, adaptor.getComplex());
468  Value imag = complex::ImOp::create(b, adaptor.getComplex());
469 
470  Value half = arith::ConstantOp::create(b, elementType,
471  b.getFloatAttr(elementType, 0.5));
472  Value one = arith::ConstantOp::create(b, elementType,
473  b.getFloatAttr(elementType, 1));
474  Value realPlusOne = arith::AddFOp::create(b, real, one, fmf);
475  Value absRealPlusOne = math::AbsFOp::create(b, realPlusOne, fmf);
476  Value absImag = math::AbsFOp::create(b, imag, fmf);
477 
478  Value maxAbs = arith::MaximumFOp::create(b, absRealPlusOne, absImag, fmf);
479  Value minAbs = arith::MinimumFOp::create(b, absRealPlusOne, absImag, fmf);
480 
481  Value useReal = arith::CmpFOp::create(b, arith::CmpFPredicate::OGT,
482  realPlusOne, absImag, fmf);
483  Value maxMinusOne = arith::SubFOp::create(b, maxAbs, one, fmf);
484  Value maxAbsOfRealPlusOneAndImagMinusOne =
485  arith::SelectOp::create(b, useReal, real, maxMinusOne);
486  arith::FastMathFlags fmfWithNaNInf = arith::bitEnumClear(
487  fmf, arith::FastMathFlags::nnan | arith::FastMathFlags::ninf);
488  Value minMaxRatio = arith::DivFOp::create(b, minAbs, maxAbs, fmfWithNaNInf);
489  Value logOfMaxAbsOfRealPlusOneAndImag =
490  math::Log1pOp::create(b, maxAbsOfRealPlusOneAndImagMinusOne, fmf);
491  Value logOfSqrtPart = math::Log1pOp::create(
492  b, arith::MulFOp::create(b, minMaxRatio, minMaxRatio, fmfWithNaNInf),
493  fmfWithNaNInf);
494  Value r = arith::AddFOp::create(
495  b, arith::MulFOp::create(b, half, logOfSqrtPart, fmfWithNaNInf),
496  logOfMaxAbsOfRealPlusOneAndImag, fmfWithNaNInf);
497  Value resultReal = arith::SelectOp::create(
498  b,
499  arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, r, r,
500  fmfWithNaNInf),
501  minAbs, r);
502  Value resultImag = math::Atan2Op::create(b, imag, realPlusOne, fmf);
503  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
504  resultImag);
505  return success();
506  }
507 };
508 
509 struct MulOpConversion : public OpConversionPattern<complex::MulOp> {
511 
512  LogicalResult
513  matchAndRewrite(complex::MulOp op, OpAdaptor adaptor,
514  ConversionPatternRewriter &rewriter) const override {
515  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
516  auto type = cast<ComplexType>(adaptor.getLhs().getType());
517  auto elementType = cast<FloatType>(type.getElementType());
518  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
519  auto fmfValue = fmf.getValue();
520  Value lhsReal = complex::ReOp::create(b, elementType, adaptor.getLhs());
521  Value lhsImag = complex::ImOp::create(b, elementType, adaptor.getLhs());
522  Value rhsReal = complex::ReOp::create(b, elementType, adaptor.getRhs());
523  Value rhsImag = complex::ImOp::create(b, elementType, adaptor.getRhs());
524  Value lhsRealTimesRhsReal =
525  arith::MulFOp::create(b, lhsReal, rhsReal, fmfValue);
526  Value lhsImagTimesRhsImag =
527  arith::MulFOp::create(b, lhsImag, rhsImag, fmfValue);
528  Value real = arith::SubFOp::create(b, lhsRealTimesRhsReal,
529  lhsImagTimesRhsImag, fmfValue);
530  Value lhsImagTimesRhsReal =
531  arith::MulFOp::create(b, lhsImag, rhsReal, fmfValue);
532  Value lhsRealTimesRhsImag =
533  arith::MulFOp::create(b, lhsReal, rhsImag, fmfValue);
534  Value imag = arith::AddFOp::create(b, lhsImagTimesRhsReal,
535  lhsRealTimesRhsImag, fmfValue);
536  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, real, imag);
537  return success();
538  }
539 };
540 
541 struct NegOpConversion : public OpConversionPattern<complex::NegOp> {
543 
544  LogicalResult
545  matchAndRewrite(complex::NegOp op, OpAdaptor adaptor,
546  ConversionPatternRewriter &rewriter) const override {
547  auto loc = op.getLoc();
548  auto type = cast<ComplexType>(adaptor.getComplex().getType());
549  auto elementType = cast<FloatType>(type.getElementType());
550 
551  Value real =
552  complex::ReOp::create(rewriter, loc, elementType, adaptor.getComplex());
553  Value imag =
554  complex::ImOp::create(rewriter, loc, elementType, adaptor.getComplex());
555  Value negReal = arith::NegFOp::create(rewriter, loc, real);
556  Value negImag = arith::NegFOp::create(rewriter, loc, imag);
557  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, negReal, negImag);
558  return success();
559  }
560 };
561 
562 struct SinOpConversion : public TrigonometricOpConversion<complex::SinOp> {
563  using TrigonometricOpConversion<complex::SinOp>::TrigonometricOpConversion;
564 
565  std::pair<Value, Value> combine(Location loc, Value scaledExp,
566  Value reciprocalExp, Value sin, Value cos,
567  ConversionPatternRewriter &rewriter,
568  arith::FastMathFlagsAttr fmf) const override {
569  // Complex sine is defined as;
570  // sin(x + iy) = -0.5i * (exp(i(x + iy)) - exp(-i(x + iy)))
571  // Plugging in:
572  // exp(i(x+iy)) = exp(-y + ix) = exp(-y)(cos(x) + i sin(x))
573  // exp(-i(x+iy)) = exp(y + i(-x)) = exp(y)(cos(x) + i (-sin(x)))
574  // and defining t := exp(y)
575  // We get:
576  // Re(sin(x + iy)) = (0.5*t + 0.5/t) * sin x
577  // Im(cos(x + iy)) = (0.5*t - 0.5/t) * cos x
578  Value sum =
579  arith::AddFOp::create(rewriter, loc, scaledExp, reciprocalExp, fmf);
580  Value resultReal = arith::MulFOp::create(rewriter, loc, sum, sin, fmf);
581  Value diff =
582  arith::SubFOp::create(rewriter, loc, scaledExp, reciprocalExp, fmf);
583  Value resultImag = arith::MulFOp::create(rewriter, loc, diff, cos, fmf);
584  return {resultReal, resultImag};
585  }
586 };
587 
588 // The algorithm is listed in https://dl.acm.org/doi/pdf/10.1145/363717.363780.
589 struct SqrtOpConversion : public OpConversionPattern<complex::SqrtOp> {
591 
592  LogicalResult
593  matchAndRewrite(complex::SqrtOp op, OpAdaptor adaptor,
594  ConversionPatternRewriter &rewriter) const override {
595  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
596 
597  auto type = cast<ComplexType>(op.getType());
598  auto elementType = cast<FloatType>(type.getElementType());
599  arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
600 
601  auto cst = [&](APFloat v) {
602  return arith::ConstantOp::create(b, elementType,
603  b.getFloatAttr(elementType, v));
604  };
605  const auto &floatSemantics = elementType.getFloatSemantics();
606  Value zero = cst(APFloat::getZero(floatSemantics));
607  Value half = arith::ConstantOp::create(b, elementType,
608  b.getFloatAttr(elementType, 0.5));
609 
610  Value real = complex::ReOp::create(b, elementType, adaptor.getComplex());
611  Value imag = complex::ImOp::create(b, elementType, adaptor.getComplex());
612  Value absSqrt = computeAbs(real, imag, fmf, b, AbsFn::sqrt);
613  Value argArg = math::Atan2Op::create(b, imag, real, fmf);
614  Value sqrtArg = arith::MulFOp::create(b, argArg, half, fmf);
615  Value cos = math::CosOp::create(b, sqrtArg, fmf);
616  Value sin = math::SinOp::create(b, sqrtArg, fmf);
617  // sin(atan2(0, inf)) = 0, sqrt(abs(inf)) = inf, but we can't multiply
618  // 0 * inf.
619  Value sinIsZero =
620  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, sin, zero, fmf);
621 
622  Value resultReal = arith::MulFOp::create(b, absSqrt, cos, fmf);
623  Value resultImag = arith::SelectOp::create(
624  b, sinIsZero, zero, arith::MulFOp::create(b, absSqrt, sin, fmf));
625  if (!arith::bitEnumContainsAll(fmf, arith::FastMathFlags::nnan |
626  arith::FastMathFlags::ninf)) {
627  Value inf = cst(APFloat::getInf(floatSemantics));
628  Value negInf = cst(APFloat::getInf(floatSemantics, true));
629  Value nan = cst(APFloat::getNaN(floatSemantics));
630  Value absImag = math::AbsFOp::create(b, elementType, imag, fmf);
631 
632  Value absImagIsInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
633  absImag, inf, fmf);
634  Value absImagIsNotInf = arith::CmpFOp::create(
635  b, arith::CmpFPredicate::ONE, absImag, inf, fmf);
636  Value realIsInf =
637  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, real, inf, fmf);
638  Value realIsNegInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
639  real, negInf, fmf);
640 
641  resultReal = arith::SelectOp::create(
642  b, arith::AndIOp::create(b, realIsNegInf, absImagIsNotInf), zero,
643  resultReal);
644  resultReal = arith::SelectOp::create(
645  b, arith::OrIOp::create(b, absImagIsInf, realIsInf), inf, resultReal);
646 
647  Value imagSignInf = math::CopySignOp::create(b, inf, imag, fmf);
648  resultImag = arith::SelectOp::create(
649  b,
650  arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, absSqrt, absSqrt),
651  nan, resultImag);
652  resultImag = arith::SelectOp::create(
653  b, arith::OrIOp::create(b, absImagIsInf, realIsNegInf), imagSignInf,
654  resultImag);
655  }
656 
657  Value resultIsZero =
658  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, absSqrt, zero, fmf);
659  resultReal = arith::SelectOp::create(b, resultIsZero, zero, resultReal);
660  resultImag = arith::SelectOp::create(b, resultIsZero, zero, resultImag);
661 
662  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
663  resultImag);
664  return success();
665  }
666 };
667 
668 struct SignOpConversion : public OpConversionPattern<complex::SignOp> {
670 
671  LogicalResult
672  matchAndRewrite(complex::SignOp op, OpAdaptor adaptor,
673  ConversionPatternRewriter &rewriter) const override {
674  auto type = cast<ComplexType>(adaptor.getComplex().getType());
675  auto elementType = cast<FloatType>(type.getElementType());
676  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
677  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
678 
679  Value real = complex::ReOp::create(b, elementType, adaptor.getComplex());
680  Value imag = complex::ImOp::create(b, elementType, adaptor.getComplex());
681  Value zero =
682  arith::ConstantOp::create(b, elementType, b.getZeroAttr(elementType));
683  Value realIsZero =
684  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, real, zero);
685  Value imagIsZero =
686  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, imag, zero);
687  Value isZero = arith::AndIOp::create(b, realIsZero, imagIsZero);
688  auto abs =
689  complex::AbsOp::create(b, elementType, adaptor.getComplex(), fmf);
690  Value realSign = arith::DivFOp::create(b, real, abs, fmf);
691  Value imagSign = arith::DivFOp::create(b, imag, abs, fmf);
692  Value sign = complex::CreateOp::create(b, type, realSign, imagSign);
693  rewriter.replaceOpWithNewOp<arith::SelectOp>(op, isZero,
694  adaptor.getComplex(), sign);
695  return success();
696  }
697 };
698 
699 template <typename Op>
700 struct TanTanhOpConversion : public OpConversionPattern<Op> {
702 
703  LogicalResult
704  matchAndRewrite(Op op, typename Op::Adaptor adaptor,
705  ConversionPatternRewriter &rewriter) const override {
706  ImplicitLocOpBuilder b(op.getLoc(), rewriter);
707  auto loc = op.getLoc();
708  auto type = cast<ComplexType>(adaptor.getComplex().getType());
709  auto elementType = cast<FloatType>(type.getElementType());
710  arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
711  const auto &floatSemantics = elementType.getFloatSemantics();
712 
713  Value real =
714  complex::ReOp::create(b, loc, elementType, adaptor.getComplex());
715  Value imag =
716  complex::ImOp::create(b, loc, elementType, adaptor.getComplex());
717  Value negOne = arith::ConstantOp::create(b, elementType,
718  b.getFloatAttr(elementType, -1.0));
719 
720  if constexpr (std::is_same_v<Op, complex::TanOp>) {
721  // tan(x+yi) = -i*tanh(-y + xi)
722  std::swap(real, imag);
723  real = arith::MulFOp::create(b, real, negOne, fmf);
724  }
725 
726  auto cst = [&](APFloat v) {
727  return arith::ConstantOp::create(b, elementType,
728  b.getFloatAttr(elementType, v));
729  };
730  Value inf = cst(APFloat::getInf(floatSemantics));
731  Value four = arith::ConstantOp::create(b, elementType,
732  b.getFloatAttr(elementType, 4.0));
733  Value twoReal = arith::AddFOp::create(b, real, real, fmf);
734  Value negTwoReal = arith::MulFOp::create(b, negOne, twoReal, fmf);
735 
736  Value expTwoRealMinusOne = math::ExpM1Op::create(b, twoReal, fmf);
737  Value expNegTwoRealMinusOne = math::ExpM1Op::create(b, negTwoReal, fmf);
738  Value realNum = arith::SubFOp::create(b, expTwoRealMinusOne,
739  expNegTwoRealMinusOne, fmf);
740 
741  Value cosImag = math::CosOp::create(b, imag, fmf);
742  Value cosImagSq = arith::MulFOp::create(b, cosImag, cosImag, fmf);
743  Value twoCosTwoImagPlusOne = arith::MulFOp::create(b, cosImagSq, four, fmf);
744  Value sinImag = math::SinOp::create(b, imag, fmf);
745 
746  Value imagNum = arith::MulFOp::create(
747  b, four, arith::MulFOp::create(b, cosImag, sinImag, fmf), fmf);
748 
749  Value expSumMinusTwo = arith::AddFOp::create(b, expTwoRealMinusOne,
750  expNegTwoRealMinusOne, fmf);
751  Value denom =
752  arith::AddFOp::create(b, expSumMinusTwo, twoCosTwoImagPlusOne, fmf);
753 
754  Value isInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
755  expSumMinusTwo, inf, fmf);
756  Value realLimit = math::CopySignOp::create(b, negOne, real, fmf);
757 
758  Value resultReal = arith::SelectOp::create(
759  b, isInf, realLimit, arith::DivFOp::create(b, realNum, denom, fmf));
760  Value resultImag = arith::DivFOp::create(b, imagNum, denom, fmf);
761 
762  if (!arith::bitEnumContainsAll(fmf, arith::FastMathFlags::nnan |
763  arith::FastMathFlags::ninf)) {
764  Value absReal = math::AbsFOp::create(b, real, fmf);
765  Value zero = arith::ConstantOp::create(b, elementType,
766  b.getFloatAttr(elementType, 0.0));
767  Value nan = cst(APFloat::getNaN(floatSemantics));
768 
769  Value absRealIsInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
770  absReal, inf, fmf);
771  Value imagIsZero =
772  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, imag, zero, fmf);
773  Value absRealIsNotInf = arith::XOrIOp::create(
774  b, absRealIsInf, arith::ConstantIntOp::create(b, true, /*width=*/1));
775 
776  Value imagNumIsNaN = arith::CmpFOp::create(b, arith::CmpFPredicate::UNO,
777  imagNum, imagNum, fmf);
778  Value resultRealIsNaN =
779  arith::AndIOp::create(b, imagNumIsNaN, absRealIsNotInf);
780  Value resultImagIsZero = arith::OrIOp::create(
781  b, imagIsZero, arith::AndIOp::create(b, absRealIsInf, imagNumIsNaN));
782 
783  resultReal = arith::SelectOp::create(b, resultRealIsNaN, nan, resultReal);
784  resultImag =
785  arith::SelectOp::create(b, resultImagIsZero, zero, resultImag);
786  }
787 
788  if constexpr (std::is_same_v<Op, complex::TanOp>) {
789  // tan(x+yi) = -i*tanh(-y + xi)
790  std::swap(resultReal, resultImag);
791  resultImag = arith::MulFOp::create(b, resultImag, negOne, fmf);
792  }
793 
794  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
795  resultImag);
796  return success();
797  }
798 };
799 
800 struct ConjOpConversion : public OpConversionPattern<complex::ConjOp> {
802 
803  LogicalResult
804  matchAndRewrite(complex::ConjOp op, OpAdaptor adaptor,
805  ConversionPatternRewriter &rewriter) const override {
806  auto loc = op.getLoc();
807  auto type = cast<ComplexType>(adaptor.getComplex().getType());
808  auto elementType = cast<FloatType>(type.getElementType());
809  Value real =
810  complex::ReOp::create(rewriter, loc, elementType, adaptor.getComplex());
811  Value imag =
812  complex::ImOp::create(rewriter, loc, elementType, adaptor.getComplex());
813  Value negImag = arith::NegFOp::create(rewriter, loc, elementType, imag);
814 
815  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, real, negImag);
816 
817  return success();
818  }
819 };
820 
821 /// Converts lhs^y = (a+bi)^(c+di) to
822 /// (a*a+b*b)^(0.5c) * exp(-d*atan2(b,a)) * (cos(q) + i*sin(q)),
823 /// where q = c*atan2(b,a)+0.5d*ln(a*a+b*b)
824 static Value powOpConversionImpl(mlir::ImplicitLocOpBuilder &builder,
825  ComplexType type, Value lhs, Value c, Value d,
826  arith::FastMathFlags fmf) {
827  auto elementType = cast<FloatType>(type.getElementType());
828 
829  Value a = complex::ReOp::create(builder, lhs);
830  Value b = complex::ImOp::create(builder, lhs);
831 
832  Value abs = complex::AbsOp::create(builder, lhs, fmf);
833  Value absToC = math::PowFOp::create(builder, abs, c, fmf);
834 
835  Value negD = arith::NegFOp::create(builder, d, fmf);
836  Value argLhs = math::Atan2Op::create(builder, b, a, fmf);
837  Value negDArgLhs = arith::MulFOp::create(builder, negD, argLhs, fmf);
838  Value expNegDArgLhs = math::ExpOp::create(builder, negDArgLhs, fmf);
839 
840  Value coeff = arith::MulFOp::create(builder, absToC, expNegDArgLhs, fmf);
841  Value lnAbs = math::LogOp::create(builder, abs, fmf);
842  Value cArgLhs = arith::MulFOp::create(builder, c, argLhs, fmf);
843  Value dLnAbs = arith::MulFOp::create(builder, d, lnAbs, fmf);
844  Value q = arith::AddFOp::create(builder, cArgLhs, dLnAbs, fmf);
845  Value cosQ = math::CosOp::create(builder, q, fmf);
846  Value sinQ = math::SinOp::create(builder, q, fmf);
847 
848  Value inf = arith::ConstantOp::create(
849  builder, elementType,
850  builder.getFloatAttr(elementType,
851  APFloat::getInf(elementType.getFloatSemantics())));
852  Value zero = arith::ConstantOp::create(
853  builder, elementType, builder.getFloatAttr(elementType, 0.0));
854  Value one = arith::ConstantOp::create(builder, elementType,
855  builder.getFloatAttr(elementType, 1.0));
856  Value complexOne = complex::CreateOp::create(builder, type, one, zero);
857  Value complexZero = complex::CreateOp::create(builder, type, zero, zero);
858  Value complexInf = complex::CreateOp::create(builder, type, inf, zero);
859 
860  // Case 0:
861  // d^c is 0 if d is 0 and c > 0. 0^0 is defined to be 1.0, see
862  // Branch Cuts for Complex Elementary Functions or Much Ado About
863  // Nothing's Sign Bit, W. Kahan, Section 10.
864  Value absEqZero =
865  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, abs, zero, fmf);
866  Value dEqZero =
867  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, d, zero, fmf);
868  Value cEqZero =
869  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, c, zero, fmf);
870  Value bEqZero =
871  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, b, zero, fmf);
872 
873  Value zeroLeC =
874  arith::CmpFOp::create(builder, arith::CmpFPredicate::OLE, zero, c, fmf);
875  Value coeffCosQ = arith::MulFOp::create(builder, coeff, cosQ, fmf);
876  Value coeffSinQ = arith::MulFOp::create(builder, coeff, sinQ, fmf);
877  Value complexOneOrZero =
878  arith::SelectOp::create(builder, cEqZero, complexOne, complexZero);
879  Value coeffCosSin =
880  complex::CreateOp::create(builder, type, coeffCosQ, coeffSinQ);
881  Value cutoff0 = arith::SelectOp::create(
882  builder,
883  arith::AndIOp::create(
884  builder, arith::AndIOp::create(builder, absEqZero, dEqZero), zeroLeC),
885  complexOneOrZero, coeffCosSin);
886 
887  // Case 1:
888  // x^0 is defined to be 1 for any x, see
889  // Branch Cuts for Complex Elementary Functions or Much Ado About
890  // Nothing's Sign Bit, W. Kahan, Section 10.
891  Value rhsEqZero = arith::AndIOp::create(builder, cEqZero, dEqZero);
892  Value cutoff1 =
893  arith::SelectOp::create(builder, rhsEqZero, complexOne, cutoff0);
894 
895  // Case 2:
896  // 1^(c + d*i) = 1 + 0*i
897  Value lhsEqOne = arith::AndIOp::create(
898  builder,
899  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, a, one, fmf),
900  bEqZero);
901  Value cutoff2 =
902  arith::SelectOp::create(builder, lhsEqOne, complexOne, cutoff1);
903 
904  // Case 3:
905  // inf^(c + 0*i) = inf + 0*i, c > 0
906  Value lhsEqInf = arith::AndIOp::create(
907  builder,
908  arith::CmpFOp::create(builder, arith::CmpFPredicate::OEQ, a, inf, fmf),
909  bEqZero);
910  Value rhsGt0 = arith::AndIOp::create(
911  builder, dEqZero,
912  arith::CmpFOp::create(builder, arith::CmpFPredicate::OGT, c, zero, fmf));
913  Value cutoff3 = arith::SelectOp::create(
914  builder, arith::AndIOp::create(builder, lhsEqInf, rhsGt0), complexInf,
915  cutoff2);
916 
917  // Case 4:
918  // inf^(c + 0*i) = 0 + 0*i, c < 0
919  Value rhsLt0 = arith::AndIOp::create(
920  builder, dEqZero,
921  arith::CmpFOp::create(builder, arith::CmpFPredicate::OLT, c, zero, fmf));
922  Value cutoff4 = arith::SelectOp::create(
923  builder, arith::AndIOp::create(builder, lhsEqInf, rhsLt0), complexZero,
924  cutoff3);
925 
926  return cutoff4;
927 }
928 
929 struct PowOpConversion : public OpConversionPattern<complex::PowOp> {
931 
932  LogicalResult
933  matchAndRewrite(complex::PowOp op, OpAdaptor adaptor,
934  ConversionPatternRewriter &rewriter) const override {
935  mlir::ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
936  auto type = cast<ComplexType>(adaptor.getLhs().getType());
937  auto elementType = cast<FloatType>(type.getElementType());
938 
939  Value c = complex::ReOp::create(builder, elementType, adaptor.getRhs());
940  Value d = complex::ImOp::create(builder, elementType, adaptor.getRhs());
941 
942  rewriter.replaceOp(op, {powOpConversionImpl(builder, type, adaptor.getLhs(),
943  c, d, op.getFastmath())});
944  return success();
945  }
946 };
947 
948 struct RsqrtOpConversion : public OpConversionPattern<complex::RsqrtOp> {
950 
951  LogicalResult
952  matchAndRewrite(complex::RsqrtOp op, OpAdaptor adaptor,
953  ConversionPatternRewriter &rewriter) const override {
954  mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
955  auto type = cast<ComplexType>(adaptor.getComplex().getType());
956  auto elementType = cast<FloatType>(type.getElementType());
957 
958  arith::FastMathFlags fmf = op.getFastMathFlagsAttr().getValue();
959 
960  auto cst = [&](APFloat v) {
961  return arith::ConstantOp::create(b, elementType,
962  b.getFloatAttr(elementType, v));
963  };
964  const auto &floatSemantics = elementType.getFloatSemantics();
965  Value zero = cst(APFloat::getZero(floatSemantics));
966  Value inf = cst(APFloat::getInf(floatSemantics));
967  Value negHalf = arith::ConstantOp::create(
968  b, elementType, b.getFloatAttr(elementType, -0.5));
969  Value nan = cst(APFloat::getNaN(floatSemantics));
970 
971  Value real = complex::ReOp::create(b, elementType, adaptor.getComplex());
972  Value imag = complex::ImOp::create(b, elementType, adaptor.getComplex());
973  Value absRsqrt = computeAbs(real, imag, fmf, b, AbsFn::rsqrt);
974  Value argArg = math::Atan2Op::create(b, imag, real, fmf);
975  Value rsqrtArg = arith::MulFOp::create(b, argArg, negHalf, fmf);
976  Value cos = math::CosOp::create(b, rsqrtArg, fmf);
977  Value sin = math::SinOp::create(b, rsqrtArg, fmf);
978 
979  Value resultReal = arith::MulFOp::create(b, absRsqrt, cos, fmf);
980  Value resultImag = arith::MulFOp::create(b, absRsqrt, sin, fmf);
981 
982  if (!arith::bitEnumContainsAll(fmf, arith::FastMathFlags::nnan |
983  arith::FastMathFlags::ninf)) {
984  Value negOne = arith::ConstantOp::create(b, elementType,
985  b.getFloatAttr(elementType, -1));
986 
987  Value realSignedZero = math::CopySignOp::create(b, zero, real, fmf);
988  Value imagSignedZero = math::CopySignOp::create(b, zero, imag, fmf);
989  Value negImagSignedZero =
990  arith::MulFOp::create(b, negOne, imagSignedZero, fmf);
991 
992  Value absReal = math::AbsFOp::create(b, real, fmf);
993  Value absImag = math::AbsFOp::create(b, imag, fmf);
994 
995  Value absImagIsInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
996  absImag, inf, fmf);
997  Value realIsNan =
998  arith::CmpFOp::create(b, arith::CmpFPredicate::UNO, real, real, fmf);
999  Value realIsInf = arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ,
1000  absReal, inf, fmf);
1001  Value inIsNanInf = arith::AndIOp::create(b, absImagIsInf, realIsNan);
1002 
1003  Value resultIsZero = arith::OrIOp::create(b, inIsNanInf, realIsInf);
1004 
1005  resultReal =
1006  arith::SelectOp::create(b, resultIsZero, realSignedZero, resultReal);
1007  resultImag = arith::SelectOp::create(b, resultIsZero, negImagSignedZero,
1008  resultImag);
1009  }
1010 
1011  Value isRealZero =
1012  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, real, zero, fmf);
1013  Value isImagZero =
1014  arith::CmpFOp::create(b, arith::CmpFPredicate::OEQ, imag, zero, fmf);
1015  Value isZero = arith::AndIOp::create(b, isRealZero, isImagZero);
1016 
1017  resultReal = arith::SelectOp::create(b, isZero, inf, resultReal);
1018  resultImag = arith::SelectOp::create(b, isZero, nan, resultImag);
1019 
1020  rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
1021  resultImag);
1022  return success();
1023  }
1024 };
1025 
1026 struct AngleOpConversion : public OpConversionPattern<complex::AngleOp> {
1028 
1029  LogicalResult
1030  matchAndRewrite(complex::AngleOp op, OpAdaptor adaptor,
1031  ConversionPatternRewriter &rewriter) const override {
1032  auto loc = op.getLoc();
1033  auto type = op.getType();
1034  arith::FastMathFlagsAttr fmf = op.getFastMathFlagsAttr();
1035 
1036  Value real =
1037  complex::ReOp::create(rewriter, loc, type, adaptor.getComplex());
1038  Value imag =
1039  complex::ImOp::create(rewriter, loc, type, adaptor.getComplex());
1040 
1041  rewriter.replaceOpWithNewOp<math::Atan2Op>(op, imag, real, fmf);
1042 
1043  return success();
1044  }
1045 };
1046 
1047 } // namespace
1048 
1050  RewritePatternSet &patterns, complex::ComplexRangeFlags complexRange) {
1051  // clang-format off
1052  patterns.add<
1053  AbsOpConversion,
1054  AngleOpConversion,
1055  Atan2OpConversion,
1056  BinaryComplexOpConversion<complex::AddOp, arith::AddFOp>,
1057  BinaryComplexOpConversion<complex::SubOp, arith::SubFOp>,
1058  ComparisonOpConversion<complex::EqualOp, arith::CmpFPredicate::OEQ>,
1059  ComparisonOpConversion<complex::NotEqualOp, arith::CmpFPredicate::UNE>,
1060  ConjOpConversion,
1061  CosOpConversion,
1062  ExpOpConversion,
1063  Expm1OpConversion,
1064  Log1pOpConversion,
1065  LogOpConversion,
1066  MulOpConversion,
1067  NegOpConversion,
1068  SignOpConversion,
1069  SinOpConversion,
1070  SqrtOpConversion,
1071  TanTanhOpConversion<complex::TanOp>,
1072  TanTanhOpConversion<complex::TanhOp>,
1073  PowOpConversion,
1074  RsqrtOpConversion
1075  >(patterns.getContext());
1076 
1077  patterns.add<DivOpConversion>(patterns.getContext(), complexRange);
1078 
1079  // clang-format on
1080 }
1081 
1082 namespace {
1083 struct ConvertComplexToStandardPass
1084  : public impl::ConvertComplexToStandardPassBase<
1085  ConvertComplexToStandardPass> {
1086  using Base::Base;
1087 
1088  void runOnOperation() override;
1089 };
1090 
1091 void ConvertComplexToStandardPass::runOnOperation() {
1092  // Convert to the Standard dialect using the converter defined above.
1095 
1096  ConversionTarget target(getContext());
1097  target.addLegalDialect<arith::ArithDialect, math::MathDialect>();
1098  target.addLegalOp<complex::CreateOp, complex::ImOp, complex::ReOp>();
1099  if (failed(
1100  applyPartialConversion(getOperation(), target, std::move(patterns))))
1101  signalPassFailure();
1102 }
1103 } // namespace
static Value getZero(OpBuilder &b, Location loc, Type elementType)
Get zero value for an element type.
static MLIRContext * getContext(OpFoldResult val)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
FloatAttr getFloatAttr(Type type, double value)
Definition: Builders.cpp:249
TypedAttr getZeroAttr(Type type)
Definition: Builders.cpp:319
This class implements a pattern rewriter for use with ConversionPatterns.
void replaceOp(Operation *op, ValueRange newValues) override
Replace the given operation with the new values.
This class describes a specific conversion target.
ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...
Definition: Builders.h:621
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
typename SourceOp::Adaptor OpAdaptor
Location getLoc()
The source location the operation was defined or derived from.
Definition: OpDefinition.h:129
This provides public APIs that all operations should have.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:519
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
static ConstantIntOp create(OpBuilder &builder, Location location, int64_t value, unsigned width)
Definition: ArithOps.cpp:258
void convertDivToStandardUsingAlgebraic(ConversionPatternRewriter &rewriter, Location loc, Value lhsRe, Value lhsIm, Value rhsRe, Value rhsIm, arith::FastMathFlagsAttr fmf, Value *resultRe, Value *resultIm)
convert a complex division to the arith/math dialects using algebraic method
void convertDivToStandardUsingRangeReduction(ConversionPatternRewriter &rewriter, Location loc, Value lhsRe, Value lhsIm, Value rhsRe, Value rhsIm, arith::FastMathFlagsAttr fmf, Value *resultRe, Value *resultIm)
convert a complex division to the arith/math dialects using Smith's method
Fraction abs(const Fraction &f)
Definition: Fraction.h:107
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
OwningOpRef< spirv::ModuleOp > combine(ArrayRef< spirv::ModuleOp > inputModules, OpBuilder &combinedModuleBuilder, SymbolRenameListener symRenameListener)
Combines a list of SPIR-V inputModules into one.
Include the generated interface declarations.
void populateComplexToStandardConversionPatterns(RewritePatternSet &patterns, mlir::complex::ComplexRangeFlags complexRange=mlir::complex::ComplexRangeFlags::improved)
Populate the given list with patterns that convert from Complex to Standard.
const FrozenRewritePatternSet & patterns
LogicalResult applyPartialConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Below we define several entry points for operation conversion.