[mlir] Add sin & cos ops to complex dialect

[gflegar]

Summary

Also add conversions for those ops to math + arith.

[frgossen]

Awesome! Thank you!

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

116

It's correct but it got me wondering if the base classes could use some cleaning up, i.e. if we want to have something like ComplexArithmeticOp for unary and binary. Could be a separate CL

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

124

Very nice to extract these building blocks.
This is a bit risky though in that you create these ops no matter if the values are used or not.
You could extract these as functions to be used on demand.

153

nit: maybe indent the equations

664

nit: maybe indentation

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

38

Where the order of the ops does not matter, you can use CHECK-DAG. Here, you can use it everywhee except for the return op.
see https://llvm.org/docs/CommandGuide/FileCheck.html#the-check-dag-directive

51

nit: in the interest of readability, you don't need to check the types everywhere but only where it is non-obvious.

384

Same as above

[frgossen]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

135

Can we keep things simple and use a "normal" virtual function here to avoid the templating?

[gflegar]

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

116

I guess we could remove the repetitions of SameOperandsAndResultType, and let results = .... Could be a minor cleanup, but would definitely better fit into a separate CL.

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

124

I did have that consideration as well, but the thing is that so far we are using all of these components in both cases. We can always refactor later if we need something different, especially since this is not a part of a public interface.

With splitting this up into separate functions we are risking one of the following:

1. ending up just creating simple wrappers around rewriter.create<Op>, which defeats the whole purpose, and ends up having the same amount of code duplication, or
2. duplicating the same instructions multiple times, since the end products we care about (scaledExp, reciprocalExp, sin, cos) share the same intermediate values, or
3. ending up with more complex code that tries to avoid generating intermediate values multiple times, which I again think is not worth the effort.

I guess you can also make the case that if we went either for the current approach, or (2), the inefficiencies generated by both (though, I want to note again, that for our current use-case, there is no inefficiency with the current approach) could easily be eliminated by subsequent passes (DCE, and CSE respectively), and (3) would just implement a limited subset of such a pass.

[gflegar]

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

38

Synced offline with frgossen about this, so documenting the summary here in case anyone else has the same question as I did:

Why use CHECK-DAG for everything else, but not for return.

The rule of thumb we came up with would be:

Prefer CHECK, unless the order of operations could reasonably change without impacting the correctness of the test. In that case, use CHECK-DAG instead.

Note that applying this rule would also support using CHECK for the final complex.create op, but that starts requiring thinking about the dependencies, so a good heuristic without going that deep is:

Use CHECK-DAG for "expressions", and CHECK for "statements".

[frgossen]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

124

Didn't see the reuse of half. Let's keep it simple like this not to rely on CSE.

[gflegar]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

135

My original thinking was to allow the compiler to inline this call and optimize it better. Though that probably does not make a big difference here.

[gflegar]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

124

exp as well

[gflegar]

Make combine a regular function, use CHECK-DAG, formatting

[frgossen]

Awesome, thanks!!

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

155

super-nit: redundant brackets?

663

super-nit: brackets?

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

51

super-nit: adding some spaces between CHECK and return can make this easier to read

[gflegar]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

155

I don't see any redundant ones, which ones do you think are redundant?

663

same as above

[gflegar]

Minor formatting

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

def ConstantOp : Complex_Op<"constant", [
   let extraClassDeclaration = [{
     /// Returns true if a constant operation can be built with the given value
     /// and result type.
     static bool isBuildableWith(Attribute value, Type type);
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// CosOp
+//===----------------------------------------------------------------------===//
+
+def CosOp : ComplexUnaryOp<"cos", [SameOperandsAndResultType]> {
+  let summary = "computes cosine of a complex number";
+  let description = [{
+    The `cos` op takes a single complex number and computes the cosine of
+    it, i.e. `cos(x)`, where `x` is the input value.
+
+    Example:
+
+    ```mlir
+    %a = complex.cos %b : complex<f32>
+    ```
+  }];
+
+  let results = (outs Complex<AnyFloat>:$result);
+}
+
+//===----------------------------------------------------------------------===//
 // CreateOp
 //===----------------------------------------------------------------------===//
 
 def CreateOp : Complex_Op<"create",
     [NoSideEffect,
      AllTypesMatch<["real", "imaginary"]>,
      TypesMatchWith<"complex element type matches real operand type",
                     "complex", "real",
let description = [{
     %a = complex.sign %b : complex<f32>
     ```
   }];
 
   let results = (outs Complex<AnyFloat>:$result);
 }
 
 //===----------------------------------------------------------------------===//
+// SinOp
+//===----------------------------------------------------------------------===//
+
+def SinOp : ComplexUnaryOp<"sin", [SameOperandsAndResultType]> {
+  let summary = "computes sine of a complex number";
+  let description = [{
+    The `sin` op takes a single complex number and computes the sine of
+    it, i.e. `sin(x)`, where `x` is the input value.
+
+    Example:
+
+    ```mlir
+    %a = complex.sin %b : complex<f32>
+    ```
+  }];
+
+  let results = (outs Complex<AnyFloat>:$result);
+}
+
+//===----------------------------------------------------------------------===//
 // SubOp
 //===----------------------------------------------------------------------===//
 
 def SubOp : ComplexArithmeticOp<"sub"> {
   let summary = "complex subtraction";
   let description = [{
     The `sub` operation takes two complex numbers and returns their difference.
 

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

matchAndRewrite(BinaryComplexOp op, typename BinaryComplexOp::Adaptor adaptor,
     Value resultImag =
         b.create<BinaryStandardOp>(elementType, imagLhs, imagRhs);
     rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultReal,
                                                    resultImag);
     return success();
   }
 };
 
+template <typename TrigonometricOp>
+struct TrigonometricOpConversion : public OpConversionPattern<TrigonometricOp> {
+  using OpAdaptor = typename OpConversionPattern<TrigonometricOp>::OpAdaptor;
+
+  using OpConversionPattern<TrigonometricOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(TrigonometricOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto type = adaptor.getComplex().getType().template cast<ComplexType>();
+    auto elementType = type.getElementType().template cast<FloatType>();
+
+    Value real =
+        rewriter.create<complex::ReOp>(loc, elementType, adaptor.getComplex());
+    Value imag =
+        rewriter.create<complex::ImOp>(loc, elementType, adaptor.getComplex());
+
+    // Trigonometric ops use a set of common building blocks to convert to real
+    // ops. Here we create these building blocks and call into an op-specific
+    // implementation in the subclass to combine them.
+    Value half = rewriter.create<arith::ConstantOp>(
+        loc, elementType, rewriter.getFloatAttr(elementType, 0.5));
+    Value exp = rewriter.create<math::ExpOp>(loc, imag);
+    Value scaledExp = rewriter.create<arith::MulFOp>(loc, half, exp);
+    Value reciprocalExp = rewriter.create<arith::DivFOp>(loc, half, exp);
+    Value sin = rewriter.create<math::SinOp>(loc, real);
+    Value cos = rewriter.create<math::CosOp>(loc, real);
+
+    auto resultPair =
+        combine(loc, scaledExp, reciprocalExp, sin, cos, rewriter);
+
+    rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, resultPair.first,
+                                                   resultPair.second);
+    return success();
+  }
+
+  virtual std::pair<Value, Value>
+  combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin,
+          Value cos, ConversionPatternRewriter &rewriter) const = 0;
+};
+
+struct CosOpConversion : public TrigonometricOpConversion<complex::CosOp> {
+  using TrigonometricOpConversion<complex::CosOp>::TrigonometricOpConversion;
+
+  std::pair<Value, Value>
+  combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin,
+          Value cos, ConversionPatternRewriter &rewriter) const override {
+    // Complex cosine is defined as;
+    //   cos(x + iy) = 0.5 * (exp(i(x + iy)) + exp(-i(x + iy)))
+    // Plugging in:
+    //   exp(i(x+iy)) = exp(-y + ix) = exp(-y)(cos(x) + i sin(x))
+    //   exp(-i(x+iy)) = exp(y + i(-x)) = exp(y)(cos(x) + i (-sin(x)))
+    // and defining t := exp(y)
+    // 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<arith::AddFOp>(loc, reciprocalExp, scaledExp);
+    Value resultReal = rewriter.create<arith::MulFOp>(loc, sum, cos);
+    Value diff = rewriter.create<arith::SubFOp>(loc, reciprocalExp, scaledExp);
+    Value resultImag = rewriter.create<arith::MulFOp>(loc, diff, sin);
+    return {resultReal, resultImag};
+  }
+};
+
 struct DivOpConversion : public OpConversionPattern<complex::DivOp> {
   using OpConversionPattern<complex::DivOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(complex::DivOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     auto type = adaptor.getLhs().getType().cast<ComplexType>();
Value imag =
         rewriter.create<complex::ImOp>(loc, elementType, adaptor.getComplex());
     Value negReal = rewriter.create<arith::NegFOp>(loc, real);
     Value negImag = rewriter.create<arith::NegFOp>(loc, imag);
     rewriter.replaceOpWithNewOp<complex::CreateOp>(op, type, negReal, negImag);
     return success();
   }
 };
 
+struct SinOpConversion : public TrigonometricOpConversion<complex::SinOp> {
+  using TrigonometricOpConversion<complex::SinOp>::TrigonometricOpConversion;
+
+  std::pair<Value, Value>
+  combine(Location loc, Value scaledExp, Value reciprocalExp, Value sin,
+          Value cos, ConversionPatternRewriter &rewriter) const override {
+    // Complex sine is defined as;
+    //   sin(x + iy) = -0.5i * (exp(i(x + iy)) - exp(-i(x + iy)))
+    // Plugging in:
+    //   exp(i(x+iy)) = exp(-y + ix) = exp(-y)(cos(x) + i sin(x))
+    //   exp(-i(x+iy)) = exp(y + i(-x)) = exp(y)(cos(x) + i (-sin(x)))
+    // and defining t := exp(y)
+    // 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<arith::AddFOp>(loc, scaledExp, reciprocalExp);
+    Value resultReal = rewriter.create<arith::MulFOp>(loc, sum, sin);
+    Value diff = rewriter.create<arith::SubFOp>(loc, scaledExp, reciprocalExp);
+    Value resultImag = rewriter.create<arith::MulFOp>(loc, diff, cos);
+    return {resultReal, resultImag};
+  }
+};
+
 struct SignOpConversion : public OpConversionPattern<complex::SignOp> {
   using OpConversionPattern<complex::SignOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(complex::SignOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto type = adaptor.getComplex().getType().cast<ComplexType>();
     auto elementType = type.getElementType().cast<FloatType>();
void mlir::populateComplexToStandardConversionPatterns(
     RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
       AbsOpConversion,
       ComparisonOpConversion<complex::EqualOp, arith::CmpFPredicate::OEQ>,
       ComparisonOpConversion<complex::NotEqualOp, arith::CmpFPredicate::UNE>,
       BinaryComplexOpConversion<complex::AddOp, arith::AddFOp>,
       BinaryComplexOpConversion<complex::SubOp, arith::SubFOp>,
+      CosOpConversion,
       DivOpConversion,
       ExpOpConversion,
       LogOpConversion,
       Log1pOpConversion,
       MulOpConversion,
       NegOpConversion,
-      SignOpConversion>(patterns.getContext());
+      SignOpConversion,
+      SinOpConversion>(patterns.getContext());
   // clang-format on
 }
 
 namespace {
 struct ConvertComplexToStandardPass
     : public ConvertComplexToStandardBase<ConvertComplexToStandardPass> {
   void runOnOperation() override;
 };

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

 // CHECK: %[[REAL_RHS:.*]] = complex.re %[[RHS]] : complex<f32>
 // CHECK: %[[RESULT_REAL:.*]] = arith.addf %[[REAL_LHS]], %[[REAL_RHS]] : f32
 // CHECK: %[[IMAG_LHS:.*]] = complex.im %[[LHS]] : complex<f32>
 // CHECK: %[[IMAG_RHS:.*]] = complex.im %[[RHS]] : complex<f32>
 // CHECK: %[[RESULT_IMAG:.*]] = arith.addf %[[IMAG_LHS]], %[[IMAG_RHS]] : f32
 // CHECK: %[[RESULT:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
 // CHECK: return %[[RESULT]] : complex<f32>
 
+// CHECK-LABEL: func @complex_cos
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_cos(%arg: complex<f32>) -> complex<f32> {
+  %cos = complex.cos %arg : complex<f32>
+  return %cos : complex<f32>
+}
+// 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]] : f32
+// CHECK-DAG: %[[HALF_EXP:.*]] = arith.mulf %[[HALF]], %[[EXP]]
+// CHECK-DAG: %[[HALF_REXP:.*]] = arith.divf %[[HALF]], %[[EXP]]
+// CHECK-DAG: %[[SIN:.*]] = math.sin %[[REAL]] : f32
+// CHECK-DAG: %[[COS:.*]] = math.cos %[[REAL]] : 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:.*]] = complex.create %[[RESULT_REAL]], %[[RESULT_IMAG]] : complex<f32>
+// CHECK:     return %[[RESULT]]
+
 // CHECK-LABEL: func @complex_div
 // CHECK-SAME: (%[[LHS:.*]]: complex<f32>, %[[RHS:.*]]: complex<f32>)
 func.func @complex_div(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
   %div = complex.div %lhs, %rhs : complex<f32>
   return %div : complex<f32>
 }
 // CHECK: %[[LHS_REAL:.*]] = complex.re %[[LHS]] : complex<f32>
 // CHECK: %[[LHS_IMAG:.*]] = complex.im %[[LHS]] : complex<f32>

 // CHECK: %[[IMAG_LHS:.*]] = complex.im %[[LHS]] : complex<f32>
 // CHECK: %[[REAL_RHS:.*]] = complex.re %[[RHS]] : complex<f32>
 // CHECK: %[[IMAG_RHS:.*]] = complex.im %[[RHS]] : complex<f32>
 // CHECK-DAG: %[[REAL_NOT_EQUAL:.*]] = arith.cmpf une, %[[REAL_LHS]], %[[REAL_RHS]] : f32
 // CHECK-DAG: %[[IMAG_NOT_EQUAL:.*]] = arith.cmpf une, %[[IMAG_LHS]], %[[IMAG_RHS]] : f32
 // CHECK: %[[NOT_EQUAL:.*]] = arith.ori %[[REAL_NOT_EQUAL]], %[[IMAG_NOT_EQUAL]] : i1
 // CHECK: return %[[NOT_EQUAL]] : i1
 
+// CHECK-LABEL: func @complex_sin
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_sin(%arg: complex<f32>) -> complex<f32> {
+  %sin = complex.sin %arg : complex<f32>
+  return %sin : complex<f32>
+}
+// 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]] : f32
+// CHECK-DAG: %[[HALF_EXP:.*]] = arith.mulf %[[HALF]], %[[EXP]]
+// CHECK-DAG: %[[HALF_REXP:.*]] = arith.divf %[[HALF]], %[[EXP]]
+// 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: %[[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>
+// CHECK:     return %[[RESULT]]
+
 // CHECK-LABEL: func @complex_sign
 // CHECK-SAME: %[[ARG:.*]]: complex<f32>
 func.func @complex_sign(%arg: complex<f32>) -> complex<f32> {
   %sign = complex.sign %arg: complex<f32>
   return %sign : complex<f32>
 }
 // CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
 // CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>

mlir/test/Dialect/Complex/ops.mlir

func.func @ops(%f: f32) {
   %imag = complex.im %complex : complex<f32>
 
   // CHECK: complex.abs %[[C]] : complex<f32>
   %abs = complex.abs %complex : complex<f32>
 
   // CHECK: complex.add %[[C]], %[[C]] : complex<f32>
   %sum = complex.add %complex, %complex : complex<f32>
 
+  // CHECK: complex.cos %[[C]] : complex<f32>
+  %cos = complex.cos %complex : complex<f32>
+
   // CHECK: complex.div %[[C]], %[[C]] : complex<f32>
   %div = complex.div %complex, %complex : complex<f32>
 
   // CHECK: complex.eq %[[C]], %[[C]] : complex<f32>
   %eq = complex.eq %complex, %complex : complex<f32>
 
   // CHECK: complex.exp %[[C]] : complex<f32>
   %exp = complex.exp %complex : complex<f32>
func.func @ops(%f: f32) {
   %neg = complex.neg %complex : complex<f32>
 
   // CHECK: complex.neq %[[C]], %[[C]] : complex<f32>
   %neq = complex.neq %complex, %complex : complex<f32>
 
   // CHECK: complex.sign %[[C]] : complex<f32>
   %sign = complex.sign %complex : complex<f32>
 
+  // CHECK: complex.sin %[[C]] : complex<f32>
+  %sin = complex.sin %complex : complex<f32>
+
   // CHECK: complex.sub %[[C]], %[[C]] : complex<f32>
   %diff = complex.sub %complex, %complex : complex<f32>
   return
 }