[InstCombine] foldSelectRotate - generalize to foldSelectFunnelShift

[RKSimon]

Summary

This is the last of the rotate->funnel shift InstCombine generalizations for PR46896

We still have foldGuardedRotateToFunnelShift to deal with in AggressiveInstCombine

[nikic]

LGTM

[nlopes]

Alive2 says this test is incorrect (because select blocks poison and funnel shift doesn't):

define i8 @fshr_select(i8 %x, i8 %y, i8 %shamt) {
%0:
  %cmp = icmp eq i8 %shamt, 0
  %sub = sub i8 8, %shamt
  %shr = lshr i8 %y, %shamt
  %shl = shl i8 %x, %sub
  %or = or i8 %shl, %shr
  %r = select i1 %cmp, i8 %y, i8 %or
  ret i8 %r
}
=>
define i8 @fshr_select(i8 %x, i8 %y, i8 %shamt) {
%0:
  %r = fshr i8 %x, i8 %y, i8 %shamt
  ret i8 %r
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
i8 %x = poison
i8 %y = any
i8 %shamt = #x00 (0)

Source:
i1 %cmp = #x1 (1)
i8 %sub = #x08 (8)
i8 %shr = any
i8 %shl = poison
i8 %or = poison
i8 %r = any

Target:
i8 %r = poison
Source value: any
Target value: poison

https://web.ist.utl.pt/nuno.lopes/alive2/index.php?hash=bdd9860ccefdbd55&test=Transforms%2FInstCombine%2Ffunnel.ll

[nikic]

@nlopes I think we should adjust the funnel shift definition to say that it blocks poison on one operand if the shift amount is zero. Basically the poison semantics should be "as if" the funnel shift were expanded, which does include an explicit select for the zero shift amount case.

[RKSimon]

Thanks @nlopes - I assume a suitably positioned isGuaranteedNotToBeUndefOrPoison guard should be enough to help with this?

[nlopes]

In D90382#2375905, @nikic wrote:

@nlopes I think we should adjust the funnel shift definition to say that it blocks poison on one operand if the shift amount is zero. Basically the poison semantics should be "as if" the funnel shift were expanded, which does include an explicit select for the zero shift amount case.

The issue with having instructions block poison is that then you can't remove those instruction. Right now only select and freeze block poison. So I would try to stay away of making funnel shift block poison.

In D90382#2375906, @RKSimon wrote:

Thanks @nlopes - I assume a suitably positioned isGuaranteedNotToBeUndefOrPoison guard should be enough to help with this?

Right. Or use freeze if the transformation is quite important for your benchmarks.

[nikic]

In D90382#2375931, @nlopes wrote:

In D90382#2375905, @nikic wrote:

@nlopes I think we should adjust the funnel shift definition to say that it blocks poison on one operand if the shift amount is zero. Basically the poison semantics should be "as if" the funnel shift were expanded, which does include an explicit select for the zero shift amount case.

The issue with having instructions block poison is that then you can't remove those instruction. Right now only select and freeze block poison. So I would try to stay away of making funnel shift block poison.

That only applies to basic instructions. Here we're talking about a function call. The poison semantics of the call depend on the implementation of the called function. If we implemented llvm.fshr.i8 in LLVM IR, it would look something like...

define i8 @llvm.fshr.i8(i8 %x, i8 %y, i8 %shamt) {
  %xyz = ...
  %ret = select i1 %cmp, i8 %y, i8 %xyz
  ret i8 %ret
}

... and would have corresponding poison semantics. Now, because it is an intrinsic, we have the technical capability of giving it different semantics, but I don't see why we should go out of the way to specify something non-standard here.

Note that we currently don't do any special modelling of poison wrt calls (apart from noundef attributes), so we are currently not making use of the fact that for most intrinsics "poison in poison out" holds.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

static Instruction *factorizeMinMaxTree(SelectPatternFlavor SPF, Value *LHS,
   if (!MinMaxOp || !ThirdOp)
     return nullptr;
 
   CmpInst::Predicate P = getMinMaxPred(SPF);
   Value *CmpABC = Builder.CreateICmp(P, MinMaxOp, ThirdOp);
   return SelectInst::Create(CmpABC, MinMaxOp, ThirdOp);
 }
 
-/// Try to reduce a rotate pattern that includes a compare and select into a
-/// funnel shift intrinsic. Example:
+/// Try to reduce a funnel/rotate pattern that includes a compare and select
+/// into a funnel shift intrinsic. Example:
 /// rotl32(a, b) --> (b == 0 ? a : ((a >> (32 - b)) | (a << b)))
 ///              --> call llvm.fshl.i32(a, a, b)
-static Instruction *foldSelectRotate(SelectInst &Sel,
+/// fshl32(a, b, c) --> (c == 0 ? a : ((b >> (32 - c)) | (a << c)))
+///                 --> call llvm.fshl.i32(a, b, c)
+/// fshr32(a, b, c) --> (c == 0 ? b : ((a >> (32 - c)) | (b << c)))
+///                 --> call llvm.fshr.i32(a, b, c)
+static Instruction *foldSelectFunnelShift(SelectInst &Sel,
                                           InstCombiner::BuilderTy &Builder) {
-  // The false value of the select must be a rotate of the true value.
+  // This must be a power-of-2 type for a bitmasking transform to be valid.
+  unsigned Width = Sel.getType()->getScalarSizeInBits();
+  if (!isPowerOf2_32(Width))
+    return nullptr;
+
   BinaryOperator *Or0, *Or1;
   if (!match(Sel.getFalseValue(), m_OneUse(m_Or(m_BinOp(Or0), m_BinOp(Or1)))))
     return nullptr;
 
-  Value *TVal = Sel.getTrueValue();
-  Value *SA0, *SA1;
-  if (!match(Or0, m_OneUse(m_LogicalShift(m_Specific(TVal),
+  Value *SV0, *SV1, *SA0, *SA1;
+  if (!match(Or0, m_OneUse(m_LogicalShift(m_Value(SV0),
                                           m_ZExtOrSelf(m_Value(SA0))))) ||
-      !match(Or1, m_OneUse(m_LogicalShift(m_Specific(TVal),
+      !match(Or1, m_OneUse(m_LogicalShift(m_Value(SV1),
                                           m_ZExtOrSelf(m_Value(SA1))))) ||
       Or0->getOpcode() == Or1->getOpcode())
     return nullptr;
 
-  // Canonicalize to or(shl(TVal, SA0), lshr(TVal, SA1)).
+  // Canonicalize to or(shl(SV0, SA0), lshr(SV1, SA1)).
   if (Or0->getOpcode() == BinaryOperator::LShr) {
     std::swap(Or0, Or1);
+    std::swap(SV0, SV1);
     std::swap(SA0, SA1);
   }
   assert(Or0->getOpcode() == BinaryOperator::Shl &&
          Or1->getOpcode() == BinaryOperator::LShr &&
          "Illegal or(shift,shift) pair");
 
-  // We should now have this pattern:
-  // select ?, TVal, (or (shl TVal, SA0), (lshr TVal, SA1))
-  // This must be a power-of-2 rotate for a bitmasking transform to be valid.
-  unsigned Width = Sel.getType()->getScalarSizeInBits();
-  if (!isPowerOf2_32(Width))
-    return nullptr;
-
   // Check the shift amounts to see if they are an opposite pair.
   Value *ShAmt;
   if (match(SA1, m_OneUse(m_Sub(m_SpecificInt(Width), m_Specific(SA0)))))
     ShAmt = SA0;
   else if (match(SA0, m_OneUse(m_Sub(m_SpecificInt(Width), m_Specific(SA1)))))
     ShAmt = SA1;
   else
     return nullptr;
 
+  // We should now have this pattern:
+  // select ?, TVal, (or (shl SV0, SA0), (lshr SV1, SA1))
+  // The false value of the select must be a funnel-shift of the true value:
+  // IsFShl -> TVal must be SV0 else TVal must be SV1.
+  bool IsFshl = (ShAmt == SA0);
+  Value *TVal = Sel.getTrueValue();
+  if ((IsFshl && TVal != SV0) || (!IsFshl && TVal != SV1))
+    return nullptr;
+
   // Finally, see if the select is filtering out a shift-by-zero.
   Value *Cond = Sel.getCondition();
   ICmpInst::Predicate Pred;
   if (!match(Cond, m_OneUse(m_ICmp(Pred, m_Specific(ShAmt), m_ZeroInt()))) ||
       Pred != ICmpInst::ICMP_EQ)
     return nullptr;
 
-  // This is a rotate that avoids shift-by-bitwidth UB in a suboptimal way.
+  // This is a funnel/rotate that avoids shift-by-bitwidth UB in a suboptimal way.
   // Convert to funnel shift intrinsic.
-  bool IsFshl = (ShAmt == SA0);
   Intrinsic::ID IID = IsFshl ? Intrinsic::fshl : Intrinsic::fshr;
   Function *F = Intrinsic::getDeclaration(Sel.getModule(), IID, Sel.getType());
   ShAmt = Builder.CreateZExt(ShAmt, Sel.getType());
-  return IntrinsicInst::Create(F, { TVal, TVal, ShAmt });
+  return IntrinsicInst::Create(F, { SV0, SV1, 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();
Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
 
   // Simplify selects that test the returned flag of cmpxchg instructions.
   if (Value *V = foldSelectCmpXchg(SI))
     return replaceInstUsesWith(SI, V);
 
   if (Instruction *Select = foldSelectBinOpIdentity(SI, TLI, *this))
     return Select;
 
-  if (Instruction *Rot = foldSelectRotate(SI, Builder))
-    return Rot;
+  if (Instruction *Funnel = foldSelectFunnelShift(SI, Builder))
+    return Funnel;
 
   if (Instruction *Copysign = foldSelectToCopysign(SI, Builder))
     return Copysign;
 
   if (Instruction *PN = foldSelectToPhi(SI, DT, Builder))
     return replaceInstUsesWith(SI, PN);
 
   if (Value *Fr = foldSelectWithFrozenICmp(SI, Builder))
     return replaceInstUsesWith(SI, Fr);
 
   return nullptr;
 }

llvm/test/Transforms/InstCombine/funnel.ll

;
   %conv2 = trunc i32 %or to i8
   ret i8 %conv2
 }
 
 ; Convert select pattern to funnel shift that ends in 'or'.
 
 define i8 @fshr_select(i8 %x, i8 %y, i8 %shamt) {
 ; CHECK-LABEL: @fshr_select(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[SHAMT:%.*]], 0
-; CHECK-NEXT:    [[SUB:%.*]] = sub i8 8, [[SHAMT]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i8 [[Y:%.*]], [[SHAMT]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i8 [[X:%.*]], [[SUB]]
-; CHECK-NEXT:    [[OR:%.*]] = or i8 [[SHL]], [[SHR]]
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i8 [[Y]], i8 [[OR]]
+; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.fshr.i8(i8 [[X:%.*]], i8 [[Y:%.*]], i8 [[SHAMT:%.*]])
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %cmp = icmp eq i8 %shamt, 0
   %sub = sub i8 8, %shamt
   %shr = lshr i8 %y, %shamt
   %shl = shl i8 %x, %sub
   %or = or i8 %shl, %shr
   %r = select i1 %cmp, i8 %y, i8 %or
   ret i8 %r
 }
 
 ; Convert select pattern to funnel shift that ends in 'or'.
 
 define i16 @fshl_select(i16 %x, i16 %y, i16 %shamt) {
 ; CHECK-LABEL: @fshl_select(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i16 [[SHAMT:%.*]], 0
-; CHECK-NEXT:    [[SUB:%.*]] = sub i16 16, [[SHAMT]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i16 [[Y:%.*]], [[SUB]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i16 [[X:%.*]], [[SHAMT]]
-; CHECK-NEXT:    [[OR:%.*]] = or i16 [[SHR]], [[SHL]]
-; CHECK-NEXT:    [[R:%.*]] = select i1 [[CMP]], i16 [[X]], i16 [[OR]]
+; CHECK-NEXT:    [[R:%.*]] = call i16 @llvm.fshl.i16(i16 [[X:%.*]], i16 [[Y:%.*]], i16 [[SHAMT:%.*]])
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %cmp = icmp eq i16 %shamt, 0
   %sub = sub i16 16, %shamt
   %shr = lshr i16 %y, %sub
   %shl = shl i16 %x, %shamt
   %or = or i16 %shr, %shl
   %r = select i1 %cmp, i16 %x, i16 %or
   ret i16 %r
 }
 
 ; Convert select pattern to funnel shift that ends in 'or'.
 
 define <2 x i64> @fshl_select_vector(<2 x i64> %x, <2 x i64> %y, <2 x i64> %shamt) {
 ; CHECK-LABEL: @fshl_select_vector(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq <2 x i64> [[SHAMT:%.*]], zeroinitializer
-; CHECK-NEXT:    [[SUB:%.*]] = sub <2 x i64> <i64 64, i64 64>, [[SHAMT]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <2 x i64> [[X:%.*]], [[SUB]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i64> [[Y:%.*]], [[SHAMT]]
-; CHECK-NEXT:    [[OR:%.*]] = or <2 x i64> [[SHL]], [[SHR]]
-; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP]], <2 x i64> [[Y]], <2 x i64> [[OR]]
+; CHECK-NEXT:    [[R:%.*]] = call <2 x i64> @llvm.fshl.v2i64(<2 x i64> [[Y:%.*]], <2 x i64> [[X:%.*]], <2 x i64> [[SHAMT:%.*]])
 ; CHECK-NEXT:    ret <2 x i64> [[R]]
 ;
   %cmp = icmp eq <2 x i64> %shamt, zeroinitializer
   %sub = sub <2 x i64> <i64 64, i64 64>, %shamt
   %shr = lshr <2 x i64> %x, %sub
   %shl = shl <2 x i64> %y, %shamt
   %or = or <2 x i64> %shl, %shr
   %r = select <2 x i1> %cmp, <2 x i64> %y, <2 x i64> %or
   ret <2 x i64> %r
 }