[Mlir-commits] [mlir] [mlir][SPIR-V] Add ComplexToSPIRV lowering for complex.sign (PR #216722)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Wed Oct 14 02:59:21 PDT 2026
https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/216722
>From 856b2034dc703ee225c6f78569ad3efc3e924572 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 5 Aug 2026 11:46:58 +0200
Subject: [PATCH 1/2] [mlir][SPIR-V] Add ComplexToSPIRV lowering for
complex.sign
Computes complex.sign as z / abs(z), using a component-wise
spirv.Select to pick a zero result when z is zero without requiring
SPIR-V 1.4 (a vector-result Select would need a scalar condition
paired with a composite result, which is only legal from 1.4 on).
---
.../ComplexToSPIRV/ComplexToSPIRV.cpp | 58 ++++++++++++++++++-
.../ComplexToSPIRV/complex-to-spirv.mlir | 26 +++++++++
2 files changed, 83 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
index 884df06a4e705..e7bfb5548b24f 100644
--- a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
+++ b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
@@ -26,6 +26,13 @@ using namespace mlir;
namespace {
+/// Creates a scalar floating-point constant of the given value.
+Value createFPConstant(OpBuilder &builder, Location loc, Type type,
+ double value) {
+ return spirv::ConstantOp::create(
+ builder, loc, type, builder.getFloatAttr(cast<FloatType>(type), value));
+}
+
struct ConstantOpPattern final : OpConversionPattern<complex::ConstantOp> {
using Base::Base;
@@ -309,6 +316,54 @@ struct AngleOpPattern final : OpConversionPattern<complex::AngleOp> {
}
};
+/// Computes complex.sign as z / abs(z), selecting a zero result when z is
+/// zero to avoid a division by zero.
+template <typename SqrtOp>
+struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
+ using OpConversionPattern<complex::SignOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(complex::SignOp 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});
+
+ Value reSq = spirv::FMulOp::create(rewriter, loc, re, re);
+ Value imSq = spirv::FMulOp::create(rewriter, loc, im, im);
+ Value sum = spirv::FAddOp::create(rewriter, loc, reSq, imSq);
+ Value abs = SqrtOp::create(rewriter, loc, sum);
+
+ Value signRe = spirv::FDivOp::create(rewriter, loc, re, abs);
+ Value signIm = spirv::FDivOp::create(rewriter, loc, im, abs);
+
+ Value zero = createFPConstant(rewriter, loc, re.getType(), 0.0);
+ Value reIsZero = spirv::FOrdEqualOp::create(rewriter, loc, re, zero);
+ Value imIsZero = spirv::FOrdEqualOp::create(rewriter, loc, im, zero);
+ Value isZero =
+ spirv::LogicalAndOp::create(rewriter, loc, reIsZero, imIsZero);
+
+ // Select per component: spirv.Select with a scalar condition and a
+ // composite result requires SPIR-V 1.4, so operate on scalars here and
+ // assemble the result afterwards.
+ Value resultRe = spirv::SelectOp::create(rewriter, loc, isZero, re, signRe);
+ Value resultIm = spirv::SelectOp::create(rewriter, loc, isZero, im, signIm);
+
+ rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
+ op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
+ return success();
+ }
+};
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -331,6 +386,7 @@ void mlir::populateComplexToSPIRVPatterns(
NegationOpPattern<complex::NegOp, /*NegateReal=*/true>,
NegationOpPattern<complex::ConjOp, /*NegateReal=*/false>,
AbsOpPattern<spirv::GLSqrtOp>, AbsOpPattern<spirv::CLSqrtOp>,
- AngleOpPattern<spirv::GLAtan2Op>, AngleOpPattern<spirv::CLAtan2Op>>(
+ AngleOpPattern<spirv::GLAtan2Op>, AngleOpPattern<spirv::CLAtan2Op>,
+ SignOpPattern<spirv::GLSqrtOp>, SignOpPattern<spirv::CLSqrtOp>>(
typeConverter, context);
}
diff --git a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
index e060c1c0ad44a..10ee62d043e83 100644
--- a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
+++ b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
@@ -268,3 +268,29 @@ func.func @complex_angle_opencl(%arg: complex<f32>) -> f32 {
// CHECK: spirv.CL.atan2
}
+
+// -----
+
+func.func @complex_sign(%arg: complex<f32>) -> complex<f32> {
+ %sign = complex.sign %arg : complex<f32>
+ return %sign : complex<f32>
+}
+
+// CHECK-LABEL: func.func @complex_sign
+// 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: %[[RESQ:.+]] = spirv.FMul %[[RE]], %[[RE]] : f32
+// CHECK: %[[IMSQ:.+]] = spirv.FMul %[[IM]], %[[IM]] : f32
+// CHECK: %[[SUM:.+]] = spirv.FAdd %[[RESQ]], %[[IMSQ]] : f32
+// CHECK: %[[ABS:.+]] = spirv.GL.Sqrt %[[SUM]] : f32
+// CHECK: %[[SIGNRE:.+]] = spirv.FDiv %[[RE]], %[[ABS]] : f32
+// CHECK: %[[SIGNIM:.+]] = spirv.FDiv %[[IM]], %[[ABS]] : f32
+// CHECK: %[[ZERO:.+]] = spirv.Constant 0.000000e+00 : f32
+// CHECK: %[[REZ:.+]] = spirv.FOrdEqual %[[RE]], %[[ZERO]] : f32
+// CHECK: %[[IMZ:.+]] = spirv.FOrdEqual %[[IM]], %[[ZERO]] : f32
+// CHECK: %[[ISZERO:.+]] = spirv.LogicalAnd %[[REZ]], %[[IMZ]] : i1
+// CHECK: %[[SELRE:.+]] = spirv.Select %[[ISZERO]], %[[RE]], %[[SIGNRE]] : i1, f32
+// CHECK: %[[SELIM:.+]] = spirv.Select %[[ISZERO]], %[[IM]], %[[SIGNIM]] : i1, f32
+// CHECK: %[[RESULT:.+]] = spirv.CompositeConstruct %[[SELRE]], %[[SELIM]] : (f32, f32) -> vector<2xf32>
>From a6830c2870521525d4f27c84b7cb7cb00d9b222b Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 14 Oct 2026 11:58:58 +0200
Subject: [PATCH 2/2] Avoid overflow
---
.../ComplexToSPIRV/ComplexToSPIRV.cpp | 26 ++++++++++++-------
.../ComplexToSPIRV/complex-to-spirv.mlir | 14 +++++++---
2 files changed, 27 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
index e7bfb5548b24f..32a39186ab0f6 100644
--- a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
+++ b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
@@ -317,8 +317,9 @@ struct AngleOpPattern final : OpConversionPattern<complex::AngleOp> {
};
/// Computes complex.sign as z / abs(z), selecting a zero result when z is
-/// zero to avoid a division by zero.
-template <typename SqrtOp>
+/// zero to avoid a division by zero. abs(z) is computed as
+/// max * sqrt(1 + (min / max)^2) to avoid overflow in squaring.
+template <typename SqrtOp, typename FAbsOp, typename FMaxOp, typename FMinOp>
struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
using OpConversionPattern<complex::SignOp>::OpConversionPattern;
@@ -338,10 +339,16 @@ struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
Value im =
spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
- Value reSq = spirv::FMulOp::create(rewriter, loc, re, re);
- Value imSq = spirv::FMulOp::create(rewriter, loc, im, im);
- Value sum = spirv::FAddOp::create(rewriter, loc, reSq, imSq);
- Value abs = SqrtOp::create(rewriter, loc, sum);
+ Value absRe = FAbsOp::create(rewriter, loc, re);
+ Value absIm = FAbsOp::create(rewriter, loc, im);
+ Value max = FMaxOp::create(rewriter, loc, absRe, absIm);
+ Value min = FMinOp::create(rewriter, loc, absRe, absIm);
+ Value ratio = spirv::FDivOp::create(rewriter, loc, min, max);
+ Value ratioSq = spirv::FMulOp::create(rewriter, loc, ratio, ratio);
+ Value one = createFPConstant(rewriter, loc, re.getType(), 1.0);
+ Value ratioSqPlusOne = spirv::FAddOp::create(rewriter, loc, ratioSq, one);
+ Value sqrt = SqrtOp::create(rewriter, loc, ratioSqPlusOne);
+ Value abs = spirv::FMulOp::create(rewriter, loc, max, sqrt);
Value signRe = spirv::FDivOp::create(rewriter, loc, re, abs);
Value signIm = spirv::FDivOp::create(rewriter, loc, im, abs);
@@ -363,7 +370,6 @@ struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
return success();
}
};
-
} // namespace
//===----------------------------------------------------------------------===//
@@ -387,6 +393,8 @@ void mlir::populateComplexToSPIRVPatterns(
NegationOpPattern<complex::ConjOp, /*NegateReal=*/false>,
AbsOpPattern<spirv::GLSqrtOp>, AbsOpPattern<spirv::CLSqrtOp>,
AngleOpPattern<spirv::GLAtan2Op>, AngleOpPattern<spirv::CLAtan2Op>,
- SignOpPattern<spirv::GLSqrtOp>, SignOpPattern<spirv::CLSqrtOp>>(
- typeConverter, context);
+ SignOpPattern<spirv::GLSqrtOp, spirv::GLFAbsOp, spirv::GLFMaxOp,
+ spirv::GLFMinOp>,
+ SignOpPattern<spirv::CLSqrtOp, spirv::CLFAbsOp, spirv::CLFMaxOp,
+ spirv::CLFMinOp>>(typeConverter, context);
}
diff --git a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
index 10ee62d043e83..1b39062ab985a 100644
--- a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
+++ b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
@@ -281,10 +281,16 @@ func.func @complex_sign(%arg: complex<f32>) -> 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: %[[RESQ:.+]] = spirv.FMul %[[RE]], %[[RE]] : f32
-// CHECK: %[[IMSQ:.+]] = spirv.FMul %[[IM]], %[[IM]] : f32
-// CHECK: %[[SUM:.+]] = spirv.FAdd %[[RESQ]], %[[IMSQ]] : f32
-// CHECK: %[[ABS:.+]] = spirv.GL.Sqrt %[[SUM]] : f32
+// CHECK: %[[ABSRE:.+]] = spirv.GL.FAbs %[[RE]] : f32
+// CHECK: %[[ABSIM:.+]] = spirv.GL.FAbs %[[IM]] : f32
+// CHECK: %[[MAX:.+]] = spirv.GL.FMax %[[ABSRE]], %[[ABSIM]] : f32
+// CHECK: %[[MIN:.+]] = spirv.GL.FMin %[[ABSRE]], %[[ABSIM]] : f32
+// CHECK: %[[RATIO:.+]] = spirv.FDiv %[[MIN]], %[[MAX]] : f32
+// CHECK: %[[RATIOSQ:.+]] = spirv.FMul %[[RATIO]], %[[RATIO]] : f32
+// CHECK: %[[ONE:.+]] = spirv.Constant 1.000000e+00 : f32
+// CHECK: %[[SUM:.+]] = spirv.FAdd %[[RATIOSQ]], %[[ONE]] : f32
+// CHECK: %[[SQRT:.+]] = spirv.GL.Sqrt %[[SUM]] : f32
+// CHECK: %[[ABS:.+]] = spirv.FMul %[[MAX]], %[[SQRT]] : f32
// CHECK: %[[SIGNRE:.+]] = spirv.FDiv %[[RE]], %[[ABS]] : f32
// CHECK: %[[SIGNIM:.+]] = spirv.FDiv %[[IM]], %[[ABS]] : f32
// CHECK: %[[ZERO:.+]] = spirv.Constant 0.000000e+00 : f32
More information about the Mlir-commits
mailing list