[SimplifyLibCalls] Simplify cabs libcall if real or imaginary part of input is zero (#97976)

cabs(a + i0) -> abs(a)
cabs(0 +ib) -> abs(b)

Closes #97336

Author: HendrikHuebner

llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp

@@ -27,6 +27,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
@@ -1958,25 +1959,53 @@ static Value *optimizeBinaryDoubleFP(CallInst *CI, IRBuilderBase &B,
 
 // cabs(z) -> sqrt((creal(z)*creal(z)) + (cimag(z)*cimag(z)))
 Value *LibCallSimplifier::optimizeCAbs(CallInst *CI, IRBuilderBase &B) {
-  if (!CI->isFast())
-    return nullptr;
-
-  // Propagate fast-math flags from the existing call to new instructions.
-  IRBuilderBase::FastMathFlagGuard Guard(B);
-  B.setFastMathFlags(CI->getFastMathFlags());
-
   Value *Real, *Imag;
+
   if (CI->arg_size() == 1) {
+
+    if (!CI->isFast())
+      return nullptr;
+
     Value *Op = CI->getArgOperand(0);
     assert(Op->getType()->isArrayTy() && "Unexpected signature for cabs!");
+
     Real = B.CreateExtractValue(Op, 0, "real");
     Imag = B.CreateExtractValue(Op, 1, "imag");
+
   } else {
     assert(CI->arg_size() == 2 && "Unexpected signature for cabs!");
+
     Real = CI->getArgOperand(0);
     Imag = CI->getArgOperand(1);
+
+    // if real or imaginary part is zero, simplify to abs(cimag(z))
+    // or abs(creal(z))
+    Value *AbsOp = nullptr;
+    if (ConstantFP *ConstReal = dyn_cast<ConstantFP>(Real)) {
+      if (ConstReal->isZero())
+        AbsOp = Imag;
+
+    } else if (ConstantFP *ConstImag = dyn_cast<ConstantFP>(Imag)) {
+      if (ConstImag->isZero())
+        AbsOp = Real;
+    }
+
+    if (AbsOp) {
+      IRBuilderBase::FastMathFlagGuard Guard(B);
+      B.setFastMathFlags(CI->getFastMathFlags());
+
+      return copyFlags(
+          *CI, B.CreateUnaryIntrinsic(Intrinsic::fabs, AbsOp, nullptr, "cabs"));
+    }
+
+    if (!CI->isFast())
+      return nullptr;
   }
 
+  // Propagate fast-math flags from the existing call to new instructions.
+  IRBuilderBase::FastMathFlagGuard Guard(B);
+  B.setFastMathFlags(CI->getFastMathFlags());
+
   Value *RealReal = B.CreateFMul(Real, Real);
   Value *ImagImag = B.CreateFMul(Imag, Imag);
 

llvm/test/Transforms/InstCombine/cabs-discrete.ll

@@ -40,6 +40,33 @@ define double @fast_cabs(double %real, double %imag) {
   ret double %call
 }
 
+define double @cabs_zero_real(double %imag) {
+; CHECK-LABEL: @cabs_zero_real(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call double @llvm.fabs.f64(double [[IMAG:%.*]])
+; CHECK-NEXT:    ret double [[CABS]]
+;
+  %call = tail call double @cabs(double 0.0, double %imag)
+  ret double %call
+}
+
+define double @fast_cabs_neg_zero_real(double %imag) {
+; CHECK-LABEL: @fast_cabs_neg_zero_real(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call fast double @llvm.fabs.f64(double [[IMAG:%.*]])
+; CHECK-NEXT:    ret double [[CABS]]
+;
+  %call = tail call fast double @cabs(double -0.0, double %imag)
+  ret double %call
+}
+
+define double @cabs_zero_imag(double %real) {
+; CHECK-LABEL: @cabs_zero_imag(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call double @llvm.fabs.f64(double [[REAL:%.*]])
+; CHECK-NEXT:    ret double [[CABS]]
+;
+  %call = tail call double @cabs(double %real, double 0.0)
+  ret double %call
+}
+
 define float @fast_cabsf(float %real, float %imag) {
 ; CHECK-LABEL: @fast_cabsf(
 ; CHECK-NEXT:    [[TMP1:%.*]] = fmul fast float [[REAL:%.*]], [[REAL]]
@@ -52,6 +79,33 @@ define float @fast_cabsf(float %real, float %imag) {
   ret float %call
 }
 
+define float @cabsf_zero_real(float %imag) {
+; CHECK-LABEL: @cabsf_zero_real(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call float @llvm.fabs.f32(float [[IMAG:%.*]])
+; CHECK-NEXT:    ret float [[CABS]]
+;
+  %call = tail call float @cabsf(float 0.0, float %imag)
+  ret float %call
+}
+
+define float @fast_cabsf_neg_zero_real(float %imag) {
+; CHECK-LABEL: @fast_cabsf_neg_zero_real(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call fast float @llvm.fabs.f32(float [[IMAG:%.*]])
+; CHECK-NEXT:    ret float [[CABS]]
+;
+  %call = tail call fast float @cabsf(float -0.0, float %imag)
+  ret float %call
+}
+
+define float @cabsf_zero_imag(float %real) {
+; CHECK-LABEL: @cabsf_zero_imag(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call float @llvm.fabs.f32(float [[REAL:%.*]])
+; CHECK-NEXT:    ret float [[CABS]]
+;
+  %call = tail call float @cabsf(float %real, float 0.0)
+  ret float %call
+}
+
 define fp128 @fast_cabsl(fp128 %real, fp128 %imag) {
 ; CHECK-LABEL: @fast_cabsl(
 ; CHECK-NEXT:    [[TMP1:%.*]] = fmul fast fp128 [[REAL:%.*]], [[REAL]]
@@ -64,6 +118,33 @@ define fp128 @fast_cabsl(fp128 %real, fp128 %imag) {
   ret fp128 %call
 }
 
+define fp128 @cabsl_zero_real(fp128 %imag) {
+; CHECK-LABEL: @cabsl_zero_real(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call fp128 @llvm.fabs.f128(fp128 [[IMAG:%.*]])
+; CHECK-NEXT:    ret fp128 [[CABS]]
+;
+  %call = tail call fp128 @cabsl(fp128 0xL00000000000000000000000000000000, fp128 %imag)
+  ret fp128 %call
+}
+
+define fp128 @cabsl_zero_imag(fp128 %real) {
+; CHECK-LABEL: @cabsl_zero_imag(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call fp128 @llvm.fabs.f128(fp128 [[REAL:%.*]])
+; CHECK-NEXT:    ret fp128 [[CABS]]
+;
+  %call = tail call fp128 @cabsl(fp128 %real, fp128 0xL00000000000000000000000000000000)
+  ret fp128 %call
+}
+
+define fp128 @fast_cabsl_neg_zero_imag(fp128 %real) {
+; CHECK-LABEL: @fast_cabsl_neg_zero_imag(
+; CHECK-NEXT:    [[CABS:%.*]] = tail call fast fp128 @llvm.fabs.f128(fp128 [[REAL:%.*]])
+; CHECK-NEXT:    ret fp128 [[CABS]]
+;
+  %call = tail call fast fp128 @cabsl(fp128 %real, fp128 0xL00000000000000008000000000000000)
+  ret fp128 %call
+}
+
 declare double @cabs(double %real, double %imag)
 declare float @cabsf(float %real, float %imag)
 declare fp128 @cabsl(fp128 %real, fp128 %imag)