[InstCombine] Expand select+masked_load combine to include FP splats of -0.0

[david-arm]

Summary

When tail-folding the vectoriser sometimes generates a select instruction
that selects between data and a splat of -0.0. We already have a transform
in InstCombine that folds a select of a masked load and zeroinitializer
into a single masked load with a zeroinitializer passthru value. We can
do something similar for splats of -0.0 when the function has the
attribute no-signed-zeros-fp-math set to true. In this case we can
just replace a splat of -0.0 with zeroinitializer, i.e. a splat of 0.0.

[nikic]

Can this be based on nsz FMF instead?

[david-arm]

In D126692#3547055, @nikic wrote:

Can this be based on nsz FMF instead?

I was hoping to do that originally, but I think that select instructions do not have any fast math flags attached?

[nikic]

In D126692#3547061, @david-arm wrote:

In D126692#3547055, @nikic wrote:

Can this be based on nsz FMF instead?

I was hoping to do that originally, but I think that select instructions do not have any fast math flags attached?

Selects (and phis) do support FMF nowadays.

[david-arm]

In D126692#3547065, @nikic wrote:

In D126692#3547061, @david-arm wrote:

In D126692#3547055, @nikic wrote:

Can this be based on nsz FMF instead?

I was hoping to do that originally, but I think that select instructions do not have any fast math flags attached?

Selects (and phis) do support FMF nowadays.

OK. The loop vectoriser is kicking out selects that don't have any FMF and there is no other pass afterwards that adds them. So I guess an alternative would be to automatically add FMF flags to all select instructions even though it's actually a bitwise operation? If that's acceptable to everyone of course.

[paulwalker-arm]

In general how safe it is to use the function's no-signed-zeros-fp-math attribute to infer it's safe to convert one bit pattern to another? Fair enough if we're dealing with floating point operations then we can ignore the signedness of zero but this is just doing a select of some loaded data, which could be immediately stored out. So it feels dangerous to just change the data in this way.

[david-arm]

In D126692#3547103, @paulwalker-arm wrote:

In general how safe it is to use the function's no-signed-zeros-fp-math attribute to infer it's safe to convert one bit pattern to another? Fair enough if we're dealing with floating point operations then we can ignore the signedness of zero but this is just doing a select of some loaded data, which could be immediately stored out. So it feels dangerous to just change the data in this way.

This is a good point, so I wonder if I can restrict this to single uses of the select where the user is a FP operator? We can even check the user's FMF to see how the data is going to be interpreted.

[david-arm]

In D126692#3547055, @nikic wrote:

Can this be based on nsz FMF instead?

I guess it depends what other reviewers think about this patch and whether it's useful or not. Another way of achieving the same result is to ensure the vectoriser uses a splat of +0.0 for the select when the nsz flag is set.

[david-arm]

Hi @nikic, looking again at this I realised that it's not the vectoriser generating the splat of -0.0. I think it's instcombine itself! A sample output from the vectoriser looks like this:

%wide.masked.load = call <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0(ptr %11, i32 4, <vscale x 4 x i1> %active.lane.mask, <vscale x 4 x float> poison), !tbaa !10
%12 = fadd fast <vscale x 4 x float> %wide.masked.load, %vec.phi
%13 = select <vscale x 4 x i1> %active.lane.mask, <vscale x 4 x float> %12, <vscale x 4 x float> %vec.phi

and I think that instcombine moves the select between the fadd and masked.load. In doing so it creates a new select that selects between %wide.masked.load and a splat of -0.0. So perhaps the real solution here is to teach instcombine to create a sensible splat value?

[spatel]

In D126692#3547119, @david-arm wrote:

In D126692#3547055, @nikic wrote:

Can this be based on nsz FMF instead?

I guess it depends what other reviewers think about this patch and whether it's useful or not. Another way of achieving the same result is to ensure the vectoriser uses a splat of +0.0 for the select when the nsz flag is set.

If you can get the vectorizer to produce the better IR in the first place, that would be ideal. Function attributes are the legacy annotations for fast-math; the goal is to rely on FMF whenever possible in the optimizer and codegen.

We should probably canonicalize -0.0 to +0.0 in "select nsz":
https://alive2.llvm.org/ce/z/Z_-yGP

We already canonicalize -0.0 to 0.0 in fcmp (no FMF needed in that case).

[david-arm]

Based on the comments on this patch I am abandoning it in favour of a different approach (D126774) where we use the nsz flag on a SelectInst to control the identity value used in certain folds.

[david-arm]

I also created D126778 for the loop vectoriser to add the fast math flags to select instructions involved in tail-folding + reductions.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/OverflowInstAnalysis.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
foldRoundUpIntegerWithPow2Alignment(SelectInst &SI,
   Type *Ty = X->getType();
   Value *XOffset = Builder.CreateAdd(X, ConstantInt::get(Ty, *LowBitMaskCst),
                                      X->getName() + ".biased");
   Value *R = Builder.CreateAnd(XOffset, ConstantInt::get(Ty, *HighBitMaskCst));
   R->takeName(&SI);
   return R;
 }
 
+bool isEffectivelyZero(Function *F, Value *Val) {
+  if (match(Val, m_Zero()))
+    return true;
+
+  if (F) {
+    Value *SplatVal = llvm::getSplatValue(Val);
+    ConstantFP *SplatFPVal = dyn_cast_or_null<ConstantFP>(SplatVal);
+    if (SplatFPVal && SplatFPVal->isExactlyValue(-0.0) &&
+        F->getFnAttribute("no-signed-zeros-fp-math").getValueAsBool())
+      return true;
+  }
+
+  return false;
+}
+
 Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
   Type *SelType = SI.getType();
 
   if (Value *V = SimplifySelectInst(CondVal, TrueVal, FalseVal,
                                     SQ.getWithInstruction(&SI)))
Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (Value *Fr = foldSelectWithFrozenICmp(SI, Builder))
     return replaceInstUsesWith(SI, Fr);
 
   if (Value *V = foldRoundUpIntegerWithPow2Alignment(SI, Builder))
     return replaceInstUsesWith(SI, V);
 
   // select(mask, mload(,,mask,0), 0) -> mload(,,mask,0)
   // Load inst is intentionally not checked for hasOneUse()
-  if (match(FalseVal, m_Zero()) &&
+  if (isEffectivelyZero(SI.getParent()->getParent(), FalseVal) &&
       (match(TrueVal, m_MaskedLoad(m_Value(), m_Value(), m_Specific(CondVal),
                                    m_CombineOr(m_Undef(), m_Zero()))) ||
        match(TrueVal, m_MaskedGather(m_Value(), m_Value(), m_Specific(CondVal),
                                      m_CombineOr(m_Undef(), m_Zero()))))) {
     auto *MaskedInst = cast<IntrinsicInst>(TrueVal);
     if (isa<UndefValue>(MaskedInst->getArgOperand(3)))
-      MaskedInst->setArgOperand(3, FalseVal /* Zero */);
+      MaskedInst->setArgOperand(
+          3, ConstantAggregateZero::get(FalseVal->getType()));
     return replaceInstUsesWith(SI, MaskedInst);
   }
 
   Value *Mask;
   if (match(TrueVal, m_Zero()) &&
       (match(FalseVal, m_MaskedLoad(m_Value(), m_Value(), m_Value(Mask),
                                     m_CombineOr(m_Undef(), m_Zero()))) ||
        match(FalseVal, m_MaskedGather(m_Value(), m_Value(), m_Value(Mask),

llvm/test/Transforms/InstCombine/select-masked_load.ll

 ; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[COND:%.*]], <8 x float> zeroinitializer, <8 x float> [[TMP0]]
 ; CHECK-NEXT:    ret <8 x float> [[TMP1]]
 entry:
   %0 = call <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>* %ptr, i32 32, <8 x i1> %mask, <8 x float> undef)
   %1 = select i1 %cond, <8 x float> zeroinitializer, <8 x float> %0
   ret <8 x float> %1
 }
 
+define <4 x float> @masked_load_and_minus_zero_inactive_1(<4 x float>* %ptr, <4 x i1> %mask) #0 {
+; CHECK-LABEL: @masked_load_and_minus_zero_inactive_1(
+; CHECK: %load = call <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>* %ptr, i32 4, <4 x i1> %mask, <4 x float> zeroinitializer)
+; CHECK-NEXT: ret <4 x float> %load
+  %load = call <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>* %ptr, i32 4, <4 x i1> %mask, <4 x float> poison)
+  %masked = select <4 x i1> %mask, <4 x float> %load, <4 x float> <float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>
+  ret <4 x float> %masked
+}
+
+define <vscale x 4 x float> @masked_load_and_minus_zero_inactive_2(<vscale x 4 x float>* %ptr, <vscale x 4 x i1> %mask) #0 {
+; CHECK-LABEL: @masked_load_and_minus_zero_inactive_2(
+; CHECK: %load = call <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0nxv4f32(<vscale x 4 x float>* %ptr, i32 4, <vscale x 4 x i1> %mask, <vscale x 4 x float> zeroinitializer)
+; CHECK-NEXT: ret <vscale x 4 x float> %load
+  %load = call <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0nxv4f32(<vscale x 4 x float>* %ptr, i32 4, <vscale x 4 x i1> %mask, <vscale x 4 x float> poison)
+  %masked = select <vscale x 4 x i1> %mask, <vscale x 4 x float> %load, <vscale x 4 x float> shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float> poison, float -0.000000e+00, i32 0), <vscale x 4 x float> poison, <vscale x 4 x i32> zeroinitializer)
+  ret <vscale x 4 x float> %masked
+}
+
 declare <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>*, i32 immarg, <8 x i1>, <8 x float>)
 declare <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>*, i32 immarg, <4 x i1>, <4 x i32>)
 declare <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>*, i32 immarg, <4 x i1>, <4 x float>)
+declare <vscale x 4 x float> @llvm.masked.load.nxv4f32.p0nxv4f32(<vscale x 4 x float>*, i32 immarg, <vscale x 4 x i1>, <vscale x 4 x float>)
+
+attributes #0 = { "no-signed-zeros-fp-math"="true" }