[mlir][complex] complex.arg op to calculate the angle of complex number

[Lewuathe]

Summary

Add complex.arg op which calculates the angle of complex number. The op name is inspired by the function carg in libm.

See: https://sourceware.org/newlib/libm.html#carg

[Lewuathe]

Add new line.

[pifon2a]

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

1061

nit: sort the patterns

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

703

bikeshedding: let's call it complex.angle. I find complex.arg a bit confusing, even though carg exists in libm.

https://www.geeksforgeeks.org/numpy-angle-in-python/#:~:text=angle()%20function%20is%20used,1(x%2Fy).

[Lewuathe]

Rename complex.arg to complex.angle

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

let description = [{
     ```mlir
     %a = complex.conj %b: complex<f32>
     ```
   }];
 
   let results = (outs Complex<AnyFloat>:$result);
 }
 
+//===----------------------------------------------------------------------===//
+// AngleOp
+//===----------------------------------------------------------------------===//
+
+def AngleOp : ComplexUnaryOp<"angle",
+                           [TypesMatchWith<"complex element type matches result type",
+                                           "complex", "result",
+                                           "$_self.cast<ComplexType>().getElementType()">]> {
+  let summary = "computes argument value of a complex number";
+  let description = [{
+    The `angle` op takes a single complex number and computes its argument value with a branch cut along the negative real axis.
+
+    Example:
+
+    ```mlir
+         %a = complex.angle %b : complex<f32>
+    ```
+  }];
+  let results = (outs AnyFloat:$result);
+}
+
 #endif // COMPLEX_OPS

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

matchAndRewrite(complex::RsqrtOp op, OpAdaptor adaptor,
     Value d = builder.create<arith::ConstantOp>(
         elementType, builder.getFloatAttr(elementType, 0));
 
     rewriter.replaceOp(op, {powOpConversionImpl(builder, type, a, b, c, d)});
     return success();
   }
 };
 
+struct AngleOpConversion : public OpConversionPattern<complex::AngleOp> {
+  using OpConversionPattern<complex::AngleOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::AngleOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    auto type = op.getType();
+
+    Value real =
+        rewriter.create<complex::ReOp>(loc, type, adaptor.getComplex());
+    Value imag =
+        rewriter.create<complex::ImOp>(loc, type, adaptor.getComplex());
+
+    rewriter.replaceOpWithNewOp<math::Atan2Op>(op, imag, real);
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::populateComplexToStandardConversionPatterns(
     RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
       AbsOpConversion,
+      AngleOpConversion,
       Atan2OpConversion,
       BinaryComplexOpConversion<complex::AddOp, arith::AddFOp>,
       BinaryComplexOpConversion<complex::SubOp, arith::SubFOp>,
       ComparisonOpConversion<complex::EqualOp, arith::CmpFPredicate::OEQ>,
       ComparisonOpConversion<complex::NotEqualOp, arith::CmpFPredicate::UNE>,
       ConjOpConversion,
       CosOpConversion,
       DivOpConversion,
       ExpOpConversion,
       Expm1OpConversion,
       Log1pOpConversion,
       LogOpConversion,
       MulOpConversion,
       NegOpConversion,
       SignOpConversion,
       SinOpConversion,
       SqrtOpConversion,
       TanOpConversion,
       TanhOpConversion,
       PowOpConversion,
       RsqrtOpConversion
   >(patterns.getContext());
   // clang-format on
 }
 
 namespace {
 struct ConvertComplexToStandardPass
     : public ConvertComplexToStandardBase<ConvertComplexToStandardPass> {
   void runOnOperation() override;
 };

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

 
 // -----
 
 // CHECK-LABEL:   func.func @complex_rsqrt
 func.func @complex_rsqrt(%arg: complex<f32>) -> complex<f32> {
   %rsqrt = complex.rsqrt %arg : complex<f32>
   return %rsqrt : complex<f32>
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @complex_angle
+// CHECK-SAME: %[[ARG:.*]]: complex<f32>
+func.func @complex_angle(%arg: complex<f32>) -> f32 {
+  %angle = complex.angle %arg : complex<f32>
+  return %angle : f32
+}
+// CHECK: %[[REAL:.*]] = complex.re %[[ARG]] : complex<f32>
+// CHECK: %[[IMAG:.*]] = complex.im %[[ARG]] : complex<f32>
+// CHECK: %[[RESULT:.*]] = math.atan2 %[[IMAG]], %[[REAL]] : f32
+// CHECK: return %[[RESULT]] : f32

mlir/test/Integration/Dialect/Complex/CPU/correctness.mlir

func.func @atan2(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
   func.return %atan2 : complex<f32>
 }
 
 func.func @pow(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
   %pow = complex.pow %lhs, %rhs : complex<f32>
   func.return %pow : complex<f32>
 }
 
+func.func @test_element(%input: tensor<?xcomplex<f32>>,
+                      %func: (complex<f32>) -> f32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>
+
+  scf.for %i = %c0 to %size step %c1 {
+    %elem = tensor.extract %input[%i]: tensor<?xcomplex<f32>>
+
+    %val = func.call_indirect %func(%elem) : (complex<f32>) -> f32
+    vector.print %val : f32
+    scf.yield
+  }
+  func.return
+}
+
+func.func @angle(%arg: complex<f32>) -> f32 {
+  %angle = complex.angle %arg : complex<f32>
+  func.return %angle : f32
+}
+
 func.func @entry() {
   // complex.sqrt test
   %sqrt_test = arith.constant dense<[
     (-1.0, -1.0),
     // CHECK:       0.455
     // CHECK-NEXT: -1.098
     (-1.0, 1.0),
     // CHECK-NEXT:  0.455
%conj_test = arith.constant dense<[
     // CHECK-NEXT:  -1
   ]> : tensor<7xcomplex<f32>>
   %conj_test_cast = tensor.cast %conj_test
     :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
 
   %conj_func = func.constant @conj : (complex<f32>) -> complex<f32>
   call @test_unary(%conj_test_cast, %conj_func)
     : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
-    
+  // complex.angle test
+  %angle_test = arith.constant dense<[
+    (-1.0, -1.0),
+    // CHECK:      -2.356
+    (-1.0, 1.0),
+    // CHECK-NEXT:  2.356
+    (0.0, 0.0),
+    // CHECK-NEXT:  0
+    (0.0, 1.0),
+    // CHECK-NEXT:  1.570
+    (1.0, -1.0),
+    // CHECK-NEXT:  -0.785
+    (1.0, 0.0),
+    // CHECK-NEXT:  0
+    (1.0, 1.0)
+    // CHECK-NEXT:  0.785
+  ]> : tensor<7xcomplex<f32>>
+  %angle_test_cast = tensor.cast %angle_test
+    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
+
+  %angle_func = func.constant @angle : (complex<f32>) -> f32
+  call @test_element(%angle_test_cast, %angle_func)
+    : (tensor<?xcomplex<f32>>, (complex<f32>) -> f32) -> ()
+
   func.return
 }