[Mlir-commits] [mlir] [mlir][SPIR-V] Add ComplexToSPIRV lowering for complex.angle (PR #214172)

[Arseniy Obolenskiy]

https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/214172

Lower complex.angle to spirv.GL.Atan2/spirv.CL.Atan2 on the real and imaginary components, following the existing complex.abs pattern

>From 340cfe2602a5db2a27a5a7d2e3855fbf1c0011d4 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 5 Aug 2026 11:45:00 +0200
Subject: [PATCH] [mlir][SPIR-V] Add ComplexToSPIRV lowering for complex.angle

Lowers complex.angle to spirv.GL.Atan2/spirv.CL.Atan2 on the real
and imaginary components, following the existing complex.abs pattern.
---
 .../ComplexToSPIRV/ComplexToSPIRV.cpp         | 51 ++++++++++++++-----
 .../ComplexToSPIRV/complex-to-spirv.mlir      | 31 +++++++++++
 2 files changed, 70 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
index eb0f81dfecf30..884df06a4e705 100644
--- a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
+++ b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
@@ -284,6 +284,31 @@ struct DivOpPattern final : OpConversionPattern<complex::DivOp> {
   }
 };
 
+template <typename Atan2Op>
+struct AngleOpPattern final : OpConversionPattern<complex::AngleOp> {
+  using OpConversionPattern<complex::AngleOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::AngleOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type spirvType =
+        this->getTypeConverter()->convertType(op.getResult().getType());
+    if (!spirvType)
+      return rewriter.notifyMatchFailure(op, "unable to convert result type");
+
+    Location loc = op.getLoc();
+    Value complexVal = adaptor.getComplex();
+
+    Value re =
+        spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
+    Value im =
+        spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
+
+    rewriter.replaceOpWithNewOp<Atan2Op>(op, im, re);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -294,16 +319,18 @@ void mlir::populateComplexToSPIRVPatterns(
     const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) {
   MLIRContext *context = patterns.getContext();
 
-  patterns.add<ConstantOpPattern, CreateOpPattern, ReOpPattern, ImOpPattern,
-               ElementwiseBinaryOpPattern<complex::AddOp, spirv::FAddOp>,
-               ElementwiseBinaryOpPattern<complex::SubOp, spirv::FSubOp>,
-               ComparisonOpPattern<complex::EqualOp, spirv::FOrdEqualOp,
-                                   spirv::LogicalAndOp>,
-               ComparisonOpPattern<complex::NotEqualOp, spirv::FUnordNotEqualOp,
-                                   spirv::LogicalOrOp>,
-               MulOpPattern, DivOpPattern,
-               NegationOpPattern<complex::NegOp, /*NegateReal=*/true>,
-               NegationOpPattern<complex::ConjOp, /*NegateReal=*/false>,
-               AbsOpPattern<spirv::GLSqrtOp>, AbsOpPattern<spirv::CLSqrtOp>>(
-      typeConverter, context);
+  patterns
+      .add<ConstantOpPattern, CreateOpPattern, ReOpPattern, ImOpPattern,
+           ElementwiseBinaryOpPattern<complex::AddOp, spirv::FAddOp>,
+           ElementwiseBinaryOpPattern<complex::SubOp, spirv::FSubOp>,
+           ComparisonOpPattern<complex::EqualOp, spirv::FOrdEqualOp,
+                               spirv::LogicalAndOp>,
+           ComparisonOpPattern<complex::NotEqualOp, spirv::FUnordNotEqualOp,
+                               spirv::LogicalOrOp>,
+           MulOpPattern, DivOpPattern,
+           NegationOpPattern<complex::NegOp, /*NegateReal=*/true>,
+           NegationOpPattern<complex::ConjOp, /*NegateReal=*/false>,
+           AbsOpPattern<spirv::GLSqrtOp>, AbsOpPattern<spirv::CLSqrtOp>,
+           AngleOpPattern<spirv::GLAtan2Op>, AngleOpPattern<spirv::CLAtan2Op>>(
+          typeConverter, context);
 }
diff --git a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
index deb4eb6d9d08c..e060c1c0ad44a 100644
--- a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
+++ b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
@@ -237,3 +237,34 @@ func.func @complex_abs_opencl(%arg: complex<f32>) -> f32 {
 //       CHECK:   spirv.CL.sqrt %[[SUM]] : f32
 
 }
+
+// -----
+
+func.func @complex_angle(%arg: complex<f32>) -> f32 {
+  %angle = complex.angle %arg : complex<f32>
+  return %angle : f32
+}
+
+// CHECK-LABEL: func.func @complex_angle
+//  CHECK-SAME: %[[ARG:.+]]: complex<f32>
+//       CHECK:   %[[V:.+]] = builtin.unrealized_conversion_cast %[[ARG]] : complex<f32> to vector<2xf32>
+//       CHECK:   %[[RE:.+]] = spirv.CompositeExtract %[[V]][0 : i32] : vector<2xf32>
+//       CHECK:   %[[IM:.+]] = spirv.CompositeExtract %[[V]][1 : i32] : vector<2xf32>
+//       CHECK:   %[[ANGLE:.+]] = spirv.GL.Atan2 %[[IM]], %[[RE]] : f32
+//       CHECK:   return %[[ANGLE]] : f32
+
+// -----
+
+module attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel], []>, #spirv.resource_limits<>>
+} {
+
+func.func @complex_angle_opencl(%arg: complex<f32>) -> f32 {
+  %angle = complex.angle %arg : complex<f32>
+  return %angle : f32
+}
+
+// CHECK-LABEL: func.func @complex_angle_opencl
+//       CHECK:   spirv.CL.atan2
+
+}