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 
9 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
10 
11 #include <memory>
12 #include <type_traits>
13 
14 #include "../PassDetail.h"
15 #include "mlir/Dialect/Complex/IR/Complex.h"
16 #include "mlir/Dialect/Math/IR/Math.h"
17 #include "mlir/Dialect/StandardOps/IR/Ops.h"
18 #include "mlir/IR/ImplicitLocOpBuilder.h"
19 #include "mlir/IR/PatternMatch.h"
20 #include "mlir/Transforms/DialectConversion.h"
21 
22 using namespace mlir;
23 
24 namespace {
25 struct AbsOpConversion : public OpConversionPattern<complex::AbsOp> {
26   using OpConversionPattern<complex::AbsOp>::OpConversionPattern;
27 
28   LogicalResult
29   matchAndRewrite(complex::AbsOp op, ArrayRef<Value> operands,
30                   ConversionPatternRewriter &rewriter) const override {
31     complex::AbsOp::Adaptor transformed(operands);
32     auto loc = op.getLoc();
33     auto type = op.getType();
34 
35     Value real =
36         rewriter.create<complex::ReOp>(loc, type, transformed.complex());
37     Value imag =
38         rewriter.create<complex::ImOp>(loc, type, transformed.complex());
39     Value realSqr = rewriter.create<MulFOp>(loc, real, real);
40     Value imagSqr = rewriter.create<MulFOp>(loc, imag, imag);
41     Value sqNorm = rewriter.create<AddFOp>(loc, realSqr, imagSqr);
42 
43     rewriter.replaceOpWithNewOp<math::SqrtOp>(op, sqNorm);
44     return success();
45   }
46 };
47 
48 template <typename ComparisonOp, CmpFPredicate p>
49 struct ComparisonOpConversion : public OpConversionPattern<ComparisonOp> {
50   using OpConversionPattern<ComparisonOp>::OpConversionPattern;
51   using ResultCombiner =
52       std::conditional_t<std::is_same<ComparisonOp, complex::EqualOp>::value,
53                          AndOp, OrOp>;
54 
55   LogicalResult
56   matchAndRewrite(ComparisonOp op, ArrayRef<Value> operands,
57                   ConversionPatternRewriter &rewriter) const override {
58     typename ComparisonOp::Adaptor transformed(operands);
59     auto loc = op.getLoc();
60     auto type = transformed.lhs()
61                     .getType()
62                     .template cast<ComplexType>()
63                     .getElementType();
64 
65     Value realLhs =
66         rewriter.create<complex::ReOp>(loc, type, transformed.lhs());
67     Value imagLhs =
68         rewriter.create<complex::ImOp>(loc, type, transformed.lhs());
69     Value realRhs =
70         rewriter.create<complex::ReOp>(loc, type, transformed.rhs());
71     Value imagRhs =
72         rewriter.create<complex::ImOp>(loc, type, transformed.rhs());
73     Value realComparison = rewriter.create<CmpFOp>(loc, p, realLhs, realRhs);
74     Value imagComparison = rewriter.create<CmpFOp>(loc, p, imagLhs, imagRhs);
75 
76     rewriter.replaceOpWithNewOp<ResultCombiner>(op, realComparison,
77                                                 imagComparison);
78     return success();
79   }
80 };
81 
82 struct DivOpConversion : public OpConversionPattern<complex::DivOp> {
83   using OpConversionPattern<complex::DivOp>::OpConversionPattern;
84 
85   LogicalResult
86   matchAndRewrite(complex::DivOp op, ArrayRef<Value> operands,
87                   ConversionPatternRewriter &rewriter) const override {
88     complex::DivOp::Adaptor transformed(operands);
89     auto loc = op.getLoc();
90     auto type = transformed.lhs().getType().cast<ComplexType>();
91     auto elementType = type.getElementType().cast<FloatType>();
92 
93     Value lhsReal =
94         rewriter.create<complex::ReOp>(loc, elementType, transformed.lhs());
95     Value lhsImag =
96         rewriter.create<complex::ImOp>(loc, elementType, transformed.lhs());
97     Value rhsReal =
98         rewriter.create<complex::ReOp>(loc, elementType, transformed.rhs());
99     Value rhsImag =
100         rewriter.create<complex::ImOp>(loc, elementType, transformed.rhs());
101 
102     // Smith's algorithm to divide complex numbers. It is just a bit smarter
103     // way to compute the following formula:
104     //  (lhsReal + lhsImag * i) / (rhsReal + rhsImag * i)
105     //    = (lhsReal + lhsImag * i) (rhsReal - rhsImag * i) /
106     //          ((rhsReal + rhsImag * i)(rhsReal - rhsImag * i))
107     //    = ((lhsReal * rhsReal + lhsImag * rhsImag) +
108     //          (lhsImag * rhsReal - lhsReal * rhsImag) * i) / ||rhs||^2
109     //
110     // Depending on whether |rhsReal| < |rhsImag| we compute either
111     //   rhsRealImagRatio = rhsReal / rhsImag
112     //   rhsRealImagDenom = rhsImag + rhsReal * rhsRealImagRatio
113     //   resultReal = (lhsReal * rhsRealImagRatio + lhsImag) / rhsRealImagDenom
114     //   resultImag = (lhsImag * rhsRealImagRatio - lhsReal) / rhsRealImagDenom
115     //
116     // or
117     //
118     //   rhsImagRealRatio = rhsImag / rhsReal
119     //   rhsImagRealDenom = rhsReal + rhsImag * rhsImagRealRatio
120     //   resultReal = (lhsReal + lhsImag * rhsImagRealRatio) / rhsImagRealDenom
121     //   resultImag = (lhsImag - lhsReal * rhsImagRealRatio) / rhsImagRealDenom
122     //
123     // See https://dl.acm.org/citation.cfm?id=368661 for more details.
124     Value rhsRealImagRatio = rewriter.create<DivFOp>(loc, rhsReal, rhsImag);
125     Value rhsRealImagDenom = rewriter.create<AddFOp>(
126         loc, rhsImag, rewriter.create<MulFOp>(loc, rhsRealImagRatio, rhsReal));
127     Value realNumerator1 = rewriter.create<AddFOp>(
128         loc, rewriter.create<MulFOp>(loc, lhsReal, rhsRealImagRatio), lhsImag);
129     Value resultReal1 =
130         rewriter.create<DivFOp>(loc, realNumerator1, rhsRealImagDenom);
131     Value imagNumerator1 = rewriter.create<SubFOp>(
132         loc, rewriter.create<MulFOp>(loc, lhsImag, rhsRealImagRatio), lhsReal);
133     Value resultImag1 =
134         rewriter.create<DivFOp>(loc, imagNumerator1, rhsRealImagDenom);
135 
136     Value rhsImagRealRatio = rewriter.create<DivFOp>(loc, rhsImag, rhsReal);
137     Value rhsImagRealDenom = rewriter.create<AddFOp>(
138         loc, rhsReal, rewriter.create<MulFOp>(loc, rhsImagRealRatio, rhsImag));
139     Value realNumerator2 = rewriter.create<AddFOp>(
140         loc, lhsReal, rewriter.create<MulFOp>(loc, lhsImag, rhsImagRealRatio));
141     Value resultReal2 =
142         rewriter.create<DivFOp>(loc, realNumerator2, rhsImagRealDenom);
143     Value imagNumerator2 = rewriter.create<SubFOp>(
144         loc, lhsImag, rewriter.create<MulFOp>(loc, lhsReal, rhsImagRealRatio));
145     Value resultImag2 =
146         rewriter.create<DivFOp>(loc, imagNumerator2, rhsImagRealDenom);
147 
148     // Consider corner cases.
149     // Case 1. Zero denominator, numerator contains at most one NaN value.
150     Value zero = rewriter.create<ConstantOp>(loc, elementType,
151                                              rewriter.getZeroAttr(elementType));
152     Value rhsRealAbs = rewriter.create<AbsFOp>(loc, rhsReal);
153     Value rhsRealIsZero =
154         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, rhsRealAbs, zero);
155     Value rhsImagAbs = rewriter.create<AbsFOp>(loc, rhsImag);
156     Value rhsImagIsZero =
157         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, rhsImagAbs, zero);
158     Value lhsRealIsNotNaN =
159         rewriter.create<CmpFOp>(loc, CmpFPredicate::ORD, lhsReal, zero);
160     Value lhsImagIsNotNaN =
161         rewriter.create<CmpFOp>(loc, CmpFPredicate::ORD, lhsImag, zero);
162     Value lhsContainsNotNaNValue =
163         rewriter.create<OrOp>(loc, lhsRealIsNotNaN, lhsImagIsNotNaN);
164     Value resultIsInfinity = rewriter.create<AndOp>(
165         loc, lhsContainsNotNaNValue,
166         rewriter.create<AndOp>(loc, rhsRealIsZero, rhsImagIsZero));
167     Value inf = rewriter.create<ConstantOp>(
168         loc, elementType,
169         rewriter.getFloatAttr(
170             elementType, APFloat::getInf(elementType.getFloatSemantics())));
171     Value infWithSignOfRhsReal = rewriter.create<CopySignOp>(loc, inf, rhsReal);
172     Value infinityResultReal =
173         rewriter.create<MulFOp>(loc, infWithSignOfRhsReal, lhsReal);
174     Value infinityResultImag =
175         rewriter.create<MulFOp>(loc, infWithSignOfRhsReal, lhsImag);
176 
177     // Case 2. Infinite numerator, finite denominator.
178     Value rhsRealFinite =
179         rewriter.create<CmpFOp>(loc, CmpFPredicate::ONE, rhsRealAbs, inf);
180     Value rhsImagFinite =
181         rewriter.create<CmpFOp>(loc, CmpFPredicate::ONE, rhsImagAbs, inf);
182     Value rhsFinite = rewriter.create<AndOp>(loc, rhsRealFinite, rhsImagFinite);
183     Value lhsRealAbs = rewriter.create<AbsFOp>(loc, lhsReal);
184     Value lhsRealInfinite =
185         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, lhsRealAbs, inf);
186     Value lhsImagAbs = rewriter.create<AbsFOp>(loc, lhsImag);
187     Value lhsImagInfinite =
188         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, lhsImagAbs, inf);
189     Value lhsInfinite =
190         rewriter.create<OrOp>(loc, lhsRealInfinite, lhsImagInfinite);
191     Value infNumFiniteDenom =
192         rewriter.create<AndOp>(loc, lhsInfinite, rhsFinite);
193     Value one = rewriter.create<ConstantOp>(
194         loc, elementType, rewriter.getFloatAttr(elementType, 1));
195     Value lhsRealIsInfWithSign = rewriter.create<CopySignOp>(
196         loc, rewriter.create<SelectOp>(loc, lhsRealInfinite, one, zero),
197         lhsReal);
198     Value lhsImagIsInfWithSign = rewriter.create<CopySignOp>(
199         loc, rewriter.create<SelectOp>(loc, lhsImagInfinite, one, zero),
200         lhsImag);
201     Value lhsRealIsInfWithSignTimesRhsReal =
202         rewriter.create<MulFOp>(loc, lhsRealIsInfWithSign, rhsReal);
203     Value lhsImagIsInfWithSignTimesRhsImag =
204         rewriter.create<MulFOp>(loc, lhsImagIsInfWithSign, rhsImag);
205     Value resultReal3 = rewriter.create<MulFOp>(
206         loc, inf,
207         rewriter.create<AddFOp>(loc, lhsRealIsInfWithSignTimesRhsReal,
208                                 lhsImagIsInfWithSignTimesRhsImag));
209     Value lhsRealIsInfWithSignTimesRhsImag =
210         rewriter.create<MulFOp>(loc, lhsRealIsInfWithSign, rhsImag);
211     Value lhsImagIsInfWithSignTimesRhsReal =
212         rewriter.create<MulFOp>(loc, lhsImagIsInfWithSign, rhsReal);
213     Value resultImag3 = rewriter.create<MulFOp>(
214         loc, inf,
215         rewriter.create<SubFOp>(loc, lhsImagIsInfWithSignTimesRhsReal,
216                                 lhsRealIsInfWithSignTimesRhsImag));
217 
218     // Case 3: Finite numerator, infinite denominator.
219     Value lhsRealFinite =
220         rewriter.create<CmpFOp>(loc, CmpFPredicate::ONE, lhsRealAbs, inf);
221     Value lhsImagFinite =
222         rewriter.create<CmpFOp>(loc, CmpFPredicate::ONE, lhsImagAbs, inf);
223     Value lhsFinite = rewriter.create<AndOp>(loc, lhsRealFinite, lhsImagFinite);
224     Value rhsRealInfinite =
225         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, rhsRealAbs, inf);
226     Value rhsImagInfinite =
227         rewriter.create<CmpFOp>(loc, CmpFPredicate::OEQ, rhsImagAbs, inf);
228     Value rhsInfinite =
229         rewriter.create<OrOp>(loc, rhsRealInfinite, rhsImagInfinite);
230     Value finiteNumInfiniteDenom =
231         rewriter.create<AndOp>(loc, lhsFinite, rhsInfinite);
232     Value rhsRealIsInfWithSign = rewriter.create<CopySignOp>(
233         loc, rewriter.create<SelectOp>(loc, rhsRealInfinite, one, zero),
234         rhsReal);
235     Value rhsImagIsInfWithSign = rewriter.create<CopySignOp>(
236         loc, rewriter.create<SelectOp>(loc, rhsImagInfinite, one, zero),
237         rhsImag);
238     Value rhsRealIsInfWithSignTimesLhsReal =
239         rewriter.create<MulFOp>(loc, lhsReal, rhsRealIsInfWithSign);
240     Value rhsImagIsInfWithSignTimesLhsImag =
241         rewriter.create<MulFOp>(loc, lhsImag, rhsImagIsInfWithSign);
242     Value resultReal4 = rewriter.create<MulFOp>(
243         loc, zero,
244         rewriter.create<AddFOp>(loc, rhsRealIsInfWithSignTimesLhsReal,
245                                 rhsImagIsInfWithSignTimesLhsImag));
246     Value rhsRealIsInfWithSignTimesLhsImag =
247         rewriter.create<MulFOp>(loc, lhsImag, rhsRealIsInfWithSign);
248     Value rhsImagIsInfWithSignTimesLhsReal =
249         rewriter.create<MulFOp>(loc, lhsReal, rhsImagIsInfWithSign);
250     Value resultImag4 = rewriter.create<MulFOp>(
251         loc, zero,
252         rewriter.create<SubFOp>(loc, rhsRealIsInfWithSignTimesLhsImag,
253                                 rhsImagIsInfWithSignTimesLhsReal));
254 
255     Value realAbsSmallerThanImagAbs = rewriter.create<CmpFOp>(
256         loc, CmpFPredicate::OLT, rhsRealAbs, rhsImagAbs);
257     Value resultReal = rewriter.create<SelectOp>(loc, realAbsSmallerThanImagAbs,
258                                                  resultReal1, resultReal2);
259     Value resultImag = rewriter.create<SelectOp>(loc, realAbsSmallerThanImagAbs,
260                                                  resultImag1, resultImag2);
261     Value resultRealSpecialCase3 = rewriter.create<SelectOp>(
262         loc, finiteNumInfiniteDenom, resultReal4, resultReal);
263     Value resultImagSpecialCase3 = rewriter.create<SelectOp>(
264         loc, finiteNumInfiniteDenom, resultImag4, resultImag);
265     Value resultRealSpecialCase2 = rewriter.create<SelectOp>(
266         loc, infNumFiniteDenom, resultReal3, resultRealSpecialCase3);
267     Value resultImagSpecialCase2 = rewriter.create<SelectOp>(
268         loc, infNumFiniteDenom, resultImag3, resultImagSpecialCase3);
269     Value resultRealSpecialCase1 = rewriter.create<SelectOp>(
270         loc, resultIsInfinity, infinityResultReal, resultRealSpecialCase2);
271     Value resultImagSpecialCase1 = rewriter.create<SelectOp>(
272         loc, resultIsInfinity, infinityResultImag, resultImagSpecialCase2);
273 
274     Value resultRealIsNaN =
275         rewriter.create<CmpFOp>(loc, CmpFPredicate::UNO, resultReal, zero);
276     Value resultImagIsNaN =
277         rewriter.create<CmpFOp>(loc, CmpFPredicate::UNO, resultImag, zero);
278     Value resultIsNaN =
279         rewriter.create<AndOp>(loc, resultRealIsNaN, resultImagIsNaN);
280     Value resultRealWithSpecialCases = rewriter.create<SelectOp>(
281         loc, resultIsNaN, resultRealSpecialCase1, resultReal);
282     Value resultImagWithSpecialCases = rewriter.create<SelectOp>(
283         loc, resultIsNaN, resultImagSpecialCase1, resultImag);
284 
285     rewriter.replaceOpWithNewOp<complex::CreateOp>(
286         op, type, resultRealWithSpecialCases, resultImagWithSpecialCases);
287     return success();
288   }
289 };
290 
291 struct ExpOpConversion : public OpConversionPattern<complex::ExpOp> {
292   using OpConversionPattern<complex::ExpOp>::OpConversionPattern;
293 
294   LogicalResult
295   matchAndRewrite(complex::ExpOp op, ArrayRef<Value> operands,
296                   ConversionPatternRewriter &rewriter) const override {
297     complex::ExpOp::Adaptor transformed(operands);
298     auto loc = op.getLoc();
299     auto type = transformed.complex().getType().cast<ComplexType>();
300     auto elementType = type.getElementType().cast<FloatType>();
301 
302     Value real =
303         rewriter.create<complex::ReOp>(loc, elementType, transformed.complex());
304     Value imag =
305         rewriter.create<complex::ImOp>(loc, elementType, transformed.complex());
306     Value expReal = rewriter.create<math::ExpOp>(loc, real);
307     Value cosImag = rewriter.create<math::CosOp>(loc, imag);
308     Value resultReal = rewriter.create<MulFOp>(loc, expReal, cosImag);
309     Value sinImag = rewriter.create<math::SinOp>(loc, imag);
310     Value resultImag = rewriter.create<MulFOp>(loc, expReal, sinImag);
311 
312     rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
313                                                    resultImag);
314     return success();
315   }
316 };
317 
318 struct NegOpConversion : public OpConversionPattern<complex::NegOp> {
319   using OpConversionPattern<complex::NegOp>::OpConversionPattern;
320 
321   LogicalResult
322   matchAndRewrite(complex::NegOp op, ArrayRef<Value> operands,
323                   ConversionPatternRewriter &rewriter) const override {
324     complex::NegOp::Adaptor transformed(operands);
325     auto loc = op.getLoc();
326     auto type = transformed.complex().getType().cast<ComplexType>();
327     auto elementType = type.getElementType().cast<FloatType>();
328 
329     Value real =
330         rewriter.create<complex::ReOp>(loc, elementType, transformed.complex());
331     Value imag =
332         rewriter.create<complex::ImOp>(loc, elementType, transformed.complex());
333     Value negReal = rewriter.create<NegFOp>(loc, real);
334     Value negImag = rewriter.create<NegFOp>(loc, imag);
335     rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, negReal, negImag);
336     return success();
337   }
338 };
339 
340 struct SignOpConversion : public OpConversionPattern<complex::SignOp> {
341   using OpConversionPattern<complex::SignOp>::OpConversionPattern;
342 
343   LogicalResult
344   matchAndRewrite(complex::SignOp op, ArrayRef<Value> operands,
345                   ConversionPatternRewriter &rewriter) const override {
346     complex::SignOp::Adaptor transformed(operands);
347     auto type = transformed.complex().getType().cast<ComplexType>();
348     auto elementType = type.getElementType().cast<FloatType>();
349     mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
350 
351     Value real = b.create<complex::ReOp>(elementType, transformed.complex());
352     Value imag = b.create<complex::ImOp>(elementType, transformed.complex());
353     Value zero = b.create<ConstantOp>(elementType, b.getZeroAttr(elementType));
354     Value realIsZero = b.create<CmpFOp>(CmpFPredicate::OEQ, real, zero);
355     Value imagIsZero = b.create<CmpFOp>(CmpFPredicate::OEQ, imag, zero);
356     Value isZero = b.create<AndOp>(realIsZero, imagIsZero);
357     auto abs = b.create<complex::AbsOp>(elementType, transformed.complex());
358     Value realSign = b.create<DivFOp>(real, abs);
359     Value imagSign = b.create<DivFOp>(imag, abs);
360     Value sign = b.create<complex::CreateOp>(type, realSign, imagSign);
361     rewriter.replaceOpWithNewOp<SelectOp>(op, isZero, transformed.complex(),
362                                           sign);
363     return success();
364   }
365 };
366 } // namespace
367 
368 void mlir::populateComplexToStandardConversionPatterns(
369     RewritePatternSet &patterns) {
370   // clang-format off
371   patterns.add<
372       AbsOpConversion,
373       ComparisonOpConversion<complex::EqualOp, CmpFPredicate::OEQ>,
374       ComparisonOpConversion<complex::NotEqualOp, CmpFPredicate::UNE>,
375       DivOpConversion,
376       ExpOpConversion,
377       NegOpConversion,
378       SignOpConversion>(patterns.getContext());
379   // clang-format on
380 }
381 
382 namespace {
383 struct ConvertComplexToStandardPass
384     : public ConvertComplexToStandardBase<ConvertComplexToStandardPass> {
385   void runOnFunction() override;
386 };
387 
388 void ConvertComplexToStandardPass::runOnFunction() {
389   auto function = getFunction();
390 
391   // Convert to the Standard dialect using the converter defined above.
392   RewritePatternSet patterns(&getContext());
393   populateComplexToStandardConversionPatterns(patterns);
394 
395   ConversionTarget target(getContext());
396   target.addLegalDialect<StandardOpsDialect, math::MathDialect,
397                          complex::ComplexDialect>();
398   target.addIllegalOp<complex::AbsOp, complex::DivOp, complex::EqualOp,
399                       complex::ExpOp, complex::NotEqualOp, complex::NegOp,
400                       complex::SignOp>();
401   if (failed(applyPartialConversion(function, target, std::move(patterns))))
402     signalPassFailure();
403 }
404 } // namespace
405 
406 std::unique_ptr<OperationPass<FuncOp>>
407 mlir::createConvertComplexToStandardPass() {
408   return std::make_unique<ConvertComplexToStandardPass>();
409 }
410