[Mlir-commits] [mlir] e3c417d - [mlir][SPIR-V] Extract complex abs and zero-check helpers in ComplexToSPIRV (NFC) (#236851)

[llvmlistbot at llvm.org]

Author: Arseniy Obolenskiy
Date: 2026-10-30T10:58:12+01:00
New Revision: e3c417dcd3cbdd57007cda1df9c9d2c9fb679d4b

URL: https://github.com/llvm/llvm-project/commit/e3c417dcd3cbdd57007cda1df9c9d2c9fb679d4b
DIFF: https://github.com/llvm/llvm-project/commit/e3c417dcd3cbdd57007cda1df9c9d2c9fb679d4b.diff

LOG: [mlir][SPIR-V] Extract complex abs and zero-check helpers in ComplexToSPIRV (NFC) (#236851)

Move the overflow safe magnitude and zero check out of SignOpPattern so
other lowerings can reuse them

Required for other ops lowering (e.g. sqrt/rsqrt)

Added: 
    

Modified: 
    mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
index 07c0378b86299..2ed827e4b5850 100644
--- a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
+++ b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
@@ -32,6 +32,30 @@ Value createFPConstant(OpBuilder &builder, Location loc, Type type,
       builder, loc, type, builder.getFloatAttr(cast<FloatType>(type), value));
 }
 
+/// Computes abs(re + i*im) as max * sqrt(1 + (min / max)^2) to avoid overflow
+/// in squaring. The result is NaN when both components are zero.
+template <typename SqrtOp, typename FAbsOp, typename FMaxOp, typename FMinOp>
+Value createComplexAbs(OpBuilder &builder, Location loc, Value re, Value im) {
+  Value absRe = FAbsOp::create(builder, loc, re);
+  Value absIm = FAbsOp::create(builder, loc, im);
+  Value max = FMaxOp::create(builder, loc, absRe, absIm);
+  Value min = FMinOp::create(builder, loc, absRe, absIm);
+  Value ratio = spirv::FDivOp::create(builder, loc, min, max);
+  Value ratioSq = spirv::FMulOp::create(builder, loc, ratio, ratio);
+  Value one = createFPConstant(builder, loc, re.getType(), 1.0);
+  Value ratioSqPlusOne = spirv::FAddOp::create(builder, loc, ratioSq, one);
+  Value sqrt = SqrtOp::create(builder, loc, ratioSqPlusOne);
+  return spirv::FMulOp::create(builder, loc, max, sqrt);
+}
+
+/// Returns true when both components are zero.
+Value createIsZero(OpBuilder &builder, Location loc, Value re, Value im) {
+  Value zero = createFPConstant(builder, loc, re.getType(), 0.0);
+  Value reIsZero = spirv::FOrdEqualOp::create(builder, loc, re, zero);
+  Value imIsZero = spirv::FOrdEqualOp::create(builder, loc, im, zero);
+  return spirv::LogicalAndOp::create(builder, loc, reIsZero, imIsZero);
+}
+
 struct ConstantOpPattern final : OpConversionPattern<complex::ConstantOp> {
   using Base::Base;
 
@@ -338,25 +362,13 @@ struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
     Value im =
         spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
 
-    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 abs =
+        createComplexAbs<SqrtOp, FAbsOp, FMaxOp, FMinOp>(rewriter, loc, re, im);
 
     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);
+    Value isZero = createIsZero(rewriter, loc, re, im);
 
     // Select per component: spirv.Select with a scalar condition and a
     // composite result requires SPIR-V 1.4, so operate on scalars here and