[InstCombine] form copysign from select of FP constants (PR44153)

[spatel]

Summary

This should be the last step needed to solve the problem in the description of PR44153:
https://bugs.llvm.org/show_bug.cgi?id=44153

If we're casting an FP value to int, testing its signbit, and then choosing between a value and its negated value, that's a complicated way of saying "copysign":

(bitcast X) <  0 ? -TC :  TC --> copysign(TC,  X)

[nick]

I cannot be a reviewer, just added some dummies questions. And, @spatel, thanks for working on this!

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

I would never expect TC == FC to happen here. Are branches folded later than this transform?

2338

Is the one use check necessary? I would imagine that this transform will unblock other one use transforms.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

We have to look at this 'select' independently of any other instruction because of the way instcombine processes. But I think you're correct that we can assume TC != FC because that should be simplified before we reach here. I will change that to an assert.

2338

It's not strictly necessary, but the transform becomes less clearly canonical in the case of >1 use because we can end up with more instructions than we started with (and I'm not sure what that will do in the backend/codegen). For those reasons, I think we should defer that change to its own minimal follow-on patch if we want to try it.

[xbolva00]

Looks ok

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

+1 for assert

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

After I wrote that comment, I thought more about the vector case (undef lanes are always a pain!), and we can't assert yet because we're missing a simplify:
rG26d2ace9e230

And that's the motivation for:
D72784

So let me mark this patch as on hold for the moment and try to make progress on the underlying analysis.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

FWIW the m_APInt and m_APFloat matchers don't allow undefs. Only m_Zero and friends allow undefs.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2331

This is correct. We'll have an option after D72975, but there really wasn't a reason to hold this up currently. I can change to an assert and push this.

[spatel]

Patch updated:
Changed check for equal constants to assert (I changed one of the tests to vectors with undef, so we can enhance that as a follow-up).

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2358

@spatel https://github.com/llvm/llvm-project/issues/54077 suggests that we shouldn't pass forward the FMF

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2358

Thanks for tracking it down!
Should be fixed with:
84812b9b072a

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

static Instruction *foldSelectRotate(SelectInst &Sel) {
   // Convert to funnel shift intrinsic.
   bool IsFshl = (ShAmt == SA0 && ShiftOpcode0 == BinaryOperator::Shl) ||
                 (ShAmt == SA1 && ShiftOpcode1 == BinaryOperator::Shl);
   Intrinsic::ID IID = IsFshl ? Intrinsic::fshl : Intrinsic::fshr;
   Function *F = Intrinsic::getDeclaration(Sel.getModule(), IID, Sel.getType());
   return IntrinsicInst::Create(F, { TVal, TVal, ShAmt });
 }
 
+static Instruction *foldSelectToCopysign(SelectInst &Sel,
+                                         InstCombiner::BuilderTy &Builder) {
+  Value *Cond = Sel.getCondition();
+  Value *TVal = Sel.getTrueValue();
+  Value *FVal = Sel.getFalseValue();
+  Type *SelType = Sel.getType();
+
+  // Match select ?, TC, FC where the constants are equal but negated.
+  // TODO: Generalize to handle a negated variable operand?
+  const APFloat *TC, *FC;
+  if (!match(TVal, m_APFloat(TC)) || !match(FVal, m_APFloat(FC)) ||
+      !abs(*TC).bitwiseIsEqual(abs(*FC)))
+    return nullptr;
+
+  assert(TC != FC && "Expected equal select arms to simplify");
+
+  Value *X;
+  const APInt *C;
+  bool IsTrueIfSignSet;
+  ICmpInst::Predicate Pred;
+  if (!match(Cond, m_OneUse(m_ICmp(Pred, m_BitCast(m_Value(X)), m_APInt(C)))) ||
+      !isSignBitCheck(Pred, *C, IsTrueIfSignSet) || X->getType() != SelType)
+    return nullptr;
+
+  // If needed, negate the value that will be the sign argument of the copysign:
+  // (bitcast X) <  0 ? -TC :  TC --> copysign(TC,  X)
+  // (bitcast X) <  0 ?  TC : -TC --> copysign(TC, -X)
+  // (bitcast X) >= 0 ? -TC :  TC --> copysign(TC, -X)
+  // (bitcast X) >= 0 ?  TC : -TC --> copysign(TC,  X)
+  if (IsTrueIfSignSet ^ TC->isNegative())
+    X = Builder.CreateFNegFMF(X, &Sel);
+
+  // Canonicalize the magnitude argument as the positive constant since we do
+  // not care about its sign.
+  Value *MagArg = TC->isNegative() ? FVal : TVal;
+  Function *F = Intrinsic::getDeclaration(Sel.getModule(), Intrinsic::copysign,
+                                          Sel.getType());
+  Instruction *CopySign = IntrinsicInst::Create(F, { MagArg, X });
+  CopySign->setFastMathFlags(Sel.getFastMathFlags());
+  return CopySign;
+}
+
 Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   // FIXME: Remove this workaround when freeze related patches are done.
   // For select with undef operand which feeds into an equality comparison,
if (Instruction *Select = foldSelectCmpXchg(SI))
     return Select;
 
   if (Instruction *Select = foldSelectBinOpIdentity(SI, TLI))
     return Select;
 
   if (Instruction *Rot = foldSelectRotate(SI))
     return Rot;
 
+  if (Instruction *Copysign = foldSelectToCopysign(SI, Builder))
+    return Copysign;
+
   return nullptr;
 }

llvm/test/Transforms/InstCombine/select.ll

 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %r = select <2 x i1> %t, <2 x i32> <i32 4, i32 8>, <2 x i32> zeroinitializer
   ret <2 x i32> %r
 }
 
 define float @copysign1(float %x) {
 ; CHECK-LABEL: @copysign1(
-; CHECK-NEXT:    [[I:%.*]] = bitcast float [[X:%.*]] to i32
-; CHECK-NEXT:    [[ISPOS:%.*]] = icmp sgt i32 [[I]], -1
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISPOS]], float 1.000000e+00, float -1.000000e+00
+; CHECK-NEXT:    [[R:%.*]] = call float @llvm.copysign.f32(float 1.000000e+00, float [[X:%.*]])
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast float %x to i32
   %ispos = icmp sgt i32 %i, -1
   %r = select i1 %ispos, float 1.0, float -1.0
   ret float %r
 }
 
 define <2 x float> @copysign2(<2 x float> %x) {
 ; CHECK-LABEL: @copysign2(
-; CHECK-NEXT:    [[I:%.*]] = bitcast <2 x float> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[I]], zeroinitializer
-; CHECK-NEXT:    [[R:%.*]] = select nsz <2 x i1> [[ISNEG]], <2 x float> <float 4.200000e+01, float 4.200000e+01>, <2 x float> <float -4.200000e+01, float -4.200000e+01>
+; CHECK-NEXT:    [[TMP1:%.*]] = fneg nsz <2 x float> [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = call nsz <2 x float> @llvm.copysign.v2f32(<2 x float> <float 4.200000e+01, float 4.200000e+01>, <2 x float> [[TMP1]])
 ; CHECK-NEXT:    ret <2 x float> [[R]]
 ;
   %i = bitcast <2 x float> %x to <2 x i32>
   %isneg = icmp slt <2 x i32> %i, zeroinitializer
   %r = select nsz <2 x i1> %isneg, <2 x float> <float 42.0, float 42.0>, <2 x float> <float -42.0, float -42.0>
   ret <2 x float> %r
 }
 
 define float @copysign3(float %x) {
 ; CHECK-LABEL: @copysign3(
-; CHECK-NEXT:    [[I:%.*]] = bitcast float [[X:%.*]] to i32
-; CHECK-NEXT:    [[ISPOS:%.*]] = icmp sgt i32 [[I]], -1
-; CHECK-NEXT:    [[R:%.*]] = select fast i1 [[ISPOS]], float -4.300000e+01, float 4.300000e+01
+; CHECK-NEXT:    [[TMP1:%.*]] = fneg fast float [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = call fast float @llvm.copysign.f32(float 4.300000e+01, float [[TMP1]])
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast float %x to i32
   %ispos = icmp ult i32 %i, 2147483648
   %r = select fast i1 %ispos, float -43.0, float 43.0
   ret float %r
 }
 
+; TODO: Allow undefs when matching vectors.
+
 define <2 x float> @copysign4(<2 x float> %x) {
 ; CHECK-LABEL: @copysign4(
 ; CHECK-NEXT:    [[I:%.*]] = bitcast <2 x float> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt <2 x i32> [[I]], zeroinitializer
 ; CHECK-NEXT:    [[R:%.*]] = select nnan arcp <2 x i1> [[ISNEG]], <2 x float> <float 4.200000e+01, float undef>, <2 x float> <float -4.200000e+01, float -4.200000e+01>
 ; CHECK-NEXT:    ret <2 x float> [[R]]
 ;
   %i = bitcast <2 x float> %x to <2 x i32>
   %isneg = icmp ugt <2 x i32> %i, <i32 2147483647, i32 2147483647>
   %r = select arcp nnan <2 x i1> %isneg, <2 x float> <float 42.0, float undef>, <2 x float> <float -42.0, float -42.0>
   ret <2 x float> %r
 }
 
 declare void @use1(i1)
 
+; Negative test
+
 define float @copysign_extra_use(float %x) {
 ; CHECK-LABEL: @copysign_extra_use(
 ; CHECK-NEXT:    [[I:%.*]] = bitcast float [[X:%.*]] to i32
 ; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[I]], 0
 ; CHECK-NEXT:    call void @use1(i1 [[ISNEG]])
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISNEG]], float -4.400000e+01, float 4.400000e+01
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast float %x to i32
   %isneg = icmp ugt i32 %i, 2147483647
   call void @use1(i1 %isneg)
   %r = select i1 %isneg, float -44.0, float 44.0
   ret float %r
 }
 
+; Negative test
+
 define float @copysign_type_mismatch(double %x) {
 ; CHECK-LABEL: @copysign_type_mismatch(
 ; CHECK-NEXT:    [[I:%.*]] = bitcast double [[X:%.*]] to i64
 ; CHECK-NEXT:    [[ISPOS:%.*]] = icmp sgt i64 [[I]], -1
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISPOS]], float 1.000000e+00, float -1.000000e+00
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast double %x to i64
   %ispos = icmp sgt i64 %i, -1
   %r = select i1 %ispos, float 1.0, float -1.0
   ret float %r
 }
 
+; Negative test
+
 define float @copysign_wrong_cmp(float %x) {
 ; CHECK-LABEL: @copysign_wrong_cmp(
 ; CHECK-NEXT:    [[I:%.*]] = bitcast float [[X:%.*]] to i32
 ; CHECK-NEXT:    [[ISPOS:%.*]] = icmp sgt i32 [[I]], 0
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISPOS]], float 1.000000e+00, float -1.000000e+00
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast float %x to i32
   %ispos = icmp sgt i32 %i, 0
   %r = select i1 %ispos, float 1.0, float -1.0
   ret float %r
 }
 
+; Negative test
+
 define float @copysign_wrong_const(float %x) {
 ; CHECK-LABEL: @copysign_wrong_const(
 ; CHECK-NEXT:    [[I:%.*]] = bitcast float [[X:%.*]] to i32
 ; CHECK-NEXT:    [[ISPOS:%.*]] = icmp sgt i32 [[I]], -1
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[ISPOS]], float 2.000000e+00, float -1.000000e+00
 ; CHECK-NEXT:    ret float [[R]]
 ;
   %i = bitcast float %x to i32
   %ispos = icmp sgt i32 %i, -1
   %r = select i1 %ispos, float 2.0, float -1.0
   ret float %r
 }