[InstCombine] Simplify funnel shifts based on demanded bits

[nikic]

Summary

The goal here is to simplify funnel shifts based on demanded or known bits. There are three parts to this change:

• In InstCombineSimplifyDemanded, determine which bits in the operands are demanded based on the demanded bits of the result. Also determine which bits in the result are known based on known bits in the operands. This only works for known shift amount. This is the primary change.
• SimplifyDemanded may replace operands of the funnel shift with undef. As such InstCombineCalls is taught to replace funnel shifts with one undef operand with either a shl or lshr. This is also limited to known shift amounts. In principle this can also be applied to variable shamt, but as this would require handling modular reduction, as well as zero shifts, it doesn't seem like a clear win.
• Finally, the changes in InstructionSimplify are added to consistently handle all undef operands, including on shamt and on both inputs.

The background for this patch is https://github.com/rust-lang/rust/issues/56009, where a performance regression due to the switch to funnel shift intrinsics in Rust was reported. Unfortunately, this issue is not resolved by this patch (see final test cases in fsh.ll), because this would require simplifying a fsh with multiple users. There's special code in SimplifyMultipleUseDemandedBits that handles this for and/or/xor, but adding fsh to that list would require creating new shl/lshr instructions, which is not necessarily a win. Any ideas on how that case could be handled?

[nagisa]

lib/Transforms/InstCombine/InstCombineCalls.cpp

2006

This probably could use a in-code comment about the reason why this replacing the funnel shifts with one undef operand to just a shift, so that somebody else wouldn’t replace it with a straight fshl(X, undef, C) -> undef (which would make a lot of sense in isolation) down the road.

Ditto for similar replacement just below.

test/Transforms/InstCombine/fsh.ll

276

This seems like it would be more appropriate for a bug report than a test?

[spatel]

Thanks for working on this! As with the saturated series, I'd prefer to split this into separate reviews/commits and have baseline tests committed ahead of any code changes. I think we can start with the instsimplify part, then instcombine, and finally demanded bits.

[nikic]

Abandoning in favor of finer-grained patches.

lib/Analysis/InstructionSimplify.cpp

case Intrinsic::masked_load: {
     Value *PassthruArg = ArgBegin[3];
     // If the mask is all zeros or undef, the "passthru" argument is the result.
     if (maskIsAllZeroOrUndef(MaskArg))
       return PassthruArg;
     return nullptr;
   }
   case Intrinsic::fshl:
   case Intrinsic::fshr: {
-    Value *ShAmtArg = ArgBegin[2];
+    Value *Op0 = ArgBegin[0], *Op1 = ArgBegin[1], *ShAmtArg = ArgBegin[2];
+
+    // If both operands are undef, the result is undef.
+    if (match(Op0, m_Undef()) && match(Op1, m_Undef()))
+      return UndefValue::get(F->getReturnType());
+
+    // If shift amount is undef, assume it is zero.
+    if (match(ShAmtArg, m_Undef()))
+      return ArgBegin[IID == Intrinsic::fshl ? 0 : 1];
+
     const APInt *ShAmtC;
     if (match(ShAmtArg, m_APInt(ShAmtC))) {
       // If there's effectively no shift, return the 1st arg or 2nd arg.
       // TODO: For vectors, we could check each element of a non-splat constant.
       APInt BitWidth = APInt(ShAmtC->getBitWidth(), ShAmtC->getBitWidth());
       if (ShAmtC->urem(BitWidth).isNullValue())
         return ArgBegin[IID == Intrinsic::fshl ? 0 : 1];
     }

lib/Transforms/InstCombine/InstCombineCalls.cpp

Instruction *InstCombiner::visitCallInst(CallInst &CI) {
 
   case Intrinsic::ctpop:
     if (auto *I = foldCtpop(*II, *this))
       return I;
     break;
 
   case Intrinsic::fshl:
   case Intrinsic::fshr: {
+    const APInt *SA;
+    if (match(II->getArgOperand(2), m_APInt(SA))) {
+      unsigned BitWidth = SA->getBitWidth();
+      uint64_t ShiftAmt = SA->urem(BitWidth);
+      // Zero shift is already handled in simplification.
+      if (ShiftAmt != 0) {
+        Value *Op0 = II->getArgOperand(0), *Op1 = II->getArgOperand(1);
+        // Normalize to funnel shift left.
+        if (II->getIntrinsicID() == Intrinsic::fshr)
+          ShiftAmt = BitWidth - ShiftAmt;
+
+        // fshl(X, undef, C) -> shl X, C
+        if (match(Op1, m_Undef()))
+          return replaceInstUsesWith(*II, Builder.CreateShl(Op0, ShiftAmt));
+
+        // fshl(undef, X, C) -> lshr X, (BW-C)
+        if (match(Op0, m_Undef())) {
+          return replaceInstUsesWith(*II,
+            Builder.CreateLShr(Op1, BitWidth - ShiftAmt));
+        }
+      }
+    }
+
     // The shift amount (operand 2) of a funnel shift is modulo the bitwidth,
     // so only the low bits of the shift amount are demanded if the bitwidth is
     // a power-of-2.
     unsigned BitWidth = II->getType()->getScalarSizeInBits();
     if (!isPowerOf2_32(BitWidth))
       break;
     APInt Op2Demanded = APInt::getLowBitsSet(BitWidth, Log2_32_Ceil(BitWidth));
     KnownBits Op2Known(BitWidth);

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
                     ConstantInt::get(I->getType(), ResultBit-InputBit));
           NewVal->takeName(I);
           return InsertNewInstWith(NewVal, *I);
         }
 
         // TODO: Could compute known zero/one bits based on the input.
         break;
       }
+      case Intrinsic::fshr:
+      case Intrinsic::fshl: {
+        const APInt *SA;
+        if (!match(I->getOperand(2), m_APInt(SA)))
+          break;
+
+        // Zero shifts are handled during simplification.
+        uint64_t ShiftAmt = SA->urem(BitWidth);
+        if (ShiftAmt == 0)
+          break;
+
+        // Normalize to funnel shift left.
+        if (II->getIntrinsicID() == Intrinsic::fshr)
+          ShiftAmt = BitWidth - ShiftAmt;
+
+        APInt DemandedMaskLHS(DemandedMask.lshr(ShiftAmt));
+        APInt DemandedMaskRHS(
+          DemandedMask.getLoBits(ShiftAmt).shl(BitWidth - ShiftAmt));
+
+        if (SimplifyDemandedBits(I, 0, DemandedMaskLHS, LHSKnown, Depth + 1) ||
+            SimplifyDemandedBits(I, 1, DemandedMaskRHS, RHSKnown, Depth + 1))
+          return I;
+
+        Known.Zero = LHSKnown.Zero.shl(ShiftAmt)
+                   | RHSKnown.Zero.lshr(BitWidth - ShiftAmt);
+        Known.One = LHSKnown.One.shl(ShiftAmt)
+                  | RHSKnown.One.lshr(BitWidth - ShiftAmt);
+        break;
+      }
       case Intrinsic::x86_mmx_pmovmskb:
       case Intrinsic::x86_sse_movmsk_ps:
       case Intrinsic::x86_sse2_movmsk_pd:
       case Intrinsic::x86_sse2_pmovmskb_128:
       case Intrinsic::x86_avx_movmsk_ps_256:
       case Intrinsic::x86_avx_movmsk_pd_256:
       case Intrinsic::x86_avx2_pmovmskb: {
         // MOVMSK copies the vector elements' sign bits to the low bits

test/Transforms/InstCombine/fsh.ll

 ; CHECK-NEXT:    [[R:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> [[BOGUSBITS]])
 ; CHECK-NEXT:    ret <2 x i31> [[R]]
 ;
   %bogusbits = or <2 x i31> %sh, <i31 32, i31 32>
   %r = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> %bogusbits)
   ret <2 x i31> %r
 }
 
+; Simplify one undef operand and constant shift amount.
+
+define i32 @fshl_op0_undef(i32 %x) {
+; CHECK-LABEL: @fshl_op0_undef(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 25
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+  %r = call i32 @llvm.fshl.i32(i32 undef, i32 %x, i32 7)
+  ret i32 %r
+}
+
+define i33 @fshr_op0_undef(i33 %x) {
+; CHECK-LABEL: @fshr_op0_undef(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i33 [[X:%.*]], 7
+; CHECK-NEXT:    ret i33 [[TMP1]]
+;
+  %r = call i33 @llvm.fshr.i33(i33 undef, i33 %x, i33 7)
+  ret i33 %r
+}
+
+define i32 @fshl_op1_undef(i32 %x) {
+; CHECK-LABEL: @fshl_op1_undef(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 7
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+  %r = call i32 @llvm.fshl.i32(i32 %x, i32 undef, i32 7)
+  ret i32 %r
+}
+
+define i33 @fshr_op1_undef(i33 %x) {
+; CHECK-LABEL: @fshr_op1_undef(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i33 [[X:%.*]], 26
+; CHECK-NEXT:    ret i33 [[TMP1]]
+;
+  %r = call i33 @llvm.fshr.i33(i33 %x, i33 undef, i33 7)
+  ret i33 %r
+}
+
+; Only demand bits from one of the operands.
+
+define i32 @fshl_only_op0_demanded(i32 %x, i32 %y) {
+; CHECK-LABEL: @fshl_only_op0_demanded(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 7
+; CHECK-NEXT:    [[R:%.*]] = and i32 [[TMP1]], 128
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %z = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 7)
+  %r = and i32 %z, 128
+  ret i32 %r
+}
+
+define i32 @fshl_only_op1_demanded(i32 %x, i32 %y) {
+; CHECK-LABEL: @fshl_only_op1_demanded(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[Y:%.*]], 25
+; CHECK-NEXT:    [[R:%.*]] = and i32 [[TMP1]], 63
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %z = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 7)
+  %r = and i32 %z, 63
+  ret i32 %r
+}
+
+define i33 @fshr_only_op0_demanded(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_only_op0_demanded(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i33 [[Y:%.*]], 7
+; CHECK-NEXT:    [[R:%.*]] = and i33 [[TMP1]], 12392
+; CHECK-NEXT:    ret i33 [[R]]
+;
+  %z = call i33 @llvm.fshr.i33(i33 %x, i33 %y, i33 7)
+  %r = and i33 %z, 12392
+  ret i33 %r
+}
+
+define i33 @fshr_only_op1_demanded(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_only_op1_demanded(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i33 [[X:%.*]], 4
+; CHECK-NEXT:    [[R:%.*]] = and i33 [[TMP1]], 7
+; CHECK-NEXT:    ret i33 [[R]]
+;
+  %z = call i33 @llvm.fshr.i33(i33 %x, i33 %y, i33 7)
+  %r = lshr i33 %z, 30
+  ret i33 %r
+}
+
+; Demand bits from both operands -- cannot simplify.
+
+define i32 @fshl_both_ops_demanded(i32 %x, i32 %y) {
+; CHECK-LABEL: @fshl_both_ops_demanded(
+; CHECK-NEXT:    [[Z:%.*]] = call i32 @llvm.fshl.i32(i32 [[X:%.*]], i32 [[Y:%.*]], i32 7)
+; CHECK-NEXT:    [[R:%.*]] = and i32 [[Z]], 192
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %z = call i32 @llvm.fshl.i32(i32 %x, i32 %y, i32 7)
+  %r = and i32 %z, 192
+  ret i32 %r
+}
+
+define i33 @fshr_both_ops_demanded(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_both_ops_demanded(
+; CHECK-NEXT:    [[Z:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 26)
+; CHECK-NEXT:    [[R:%.*]] = and i33 [[Z]], 192
+; CHECK-NEXT:    ret i33 [[R]]
+;
+  %z = call i33 @llvm.fshr.i33(i33 %x, i33 %y, i33 26)
+  %r = and i33 %z, 192
+  ret i33 %r
+}
+
+; Both operands are demanded, but there are known bits.
+
+define i32 @fshl_known_bits(i32 %x, i32 %y) {
+; CHECK-LABEL: @fshl_known_bits(
+; CHECK-NEXT:    ret i32 128
+;
+  %x2 = or i32 %x, 1   ; lo bit set
+  %y2 = lshr i32 %y, 1 ; hi bit clear
+  %z = call i32 @llvm.fshl.i32(i32 %x2, i32 %y2, i32 7)
+  %r = and i32 %z, 192
+  ret i32 %r
+}
+
+define i33 @fshr_known_bits(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_known_bits(
+; CHECK-NEXT:    ret i33 128
+;
+  %x2 = or i33 %x, 1 ; lo bit set
+  %y2 = lshr i33 %y, 1 ; hi bit set
+  %z = call i33 @llvm.fshr.i33(i33 %x2, i33 %y2, i33 26)
+  %r = and i33 %z, 192
+  ret i33 %r
+}
+
+; This case fails to simplify due to multiple uses.
+
+define i33 @fshr_multi_use(i33 %a) {
+; CHECK-LABEL: @fshr_multi_use(
+; CHECK-NEXT:    [[B:%.*]] = tail call i33 @llvm.fshr.i33(i33 [[A:%.*]], i33 [[A]], i33 1)
+; CHECK-NEXT:    [[C:%.*]] = lshr i33 [[B]], 23
+; CHECK-NEXT:    [[D:%.*]] = xor i33 [[C]], [[B]]
+; CHECK-NEXT:    [[E:%.*]] = and i33 [[D]], 31
+; CHECK-NEXT:    ret i33 [[E]]
+;
+  %b = tail call i33 @llvm.fshr.i33(i33 %a, i33 %a, i33 1)
+  %c = lshr i33 %b, 23
+  %d = xor i33 %c, %b
+  %e = and i33 %d, 31
+  ret i33 %e
+}
+
+; This demonstrates the same simplification working if the fshr intrinsic
+; is expanded into shifts and or.
+
+define i33 @expanded_fshr_multi_use(i33 %a) {
+; CHECK-LABEL: @expanded_fshr_multi_use(
+; CHECK-NEXT:    [[TMP:%.*]] = lshr i33 [[A:%.*]], 1
+; CHECK-NEXT:    [[C:%.*]] = lshr i33 [[A]], 24
+; CHECK-NEXT:    [[D:%.*]] = xor i33 [[C]], [[TMP]]
+; CHECK-NEXT:    [[E:%.*]] = and i33 [[D]], 31
+; CHECK-NEXT:    ret i33 [[E]]
+;
+  %tmp = lshr i33 %a, 1
+  %tmp2 = shl i33 %a, 32
+  %b = or i33 %tmp, %tmp2
+  %c = lshr i33 %b, 23
+  %d = xor i33 %c, %b
+  %e = and i33 %d, 31
+  ret i33 %e
+}

test/Transforms/InstSimplify/call.ll

 ; CHECK-NEXT:    ret <2 x i8> [[F]]
 ;
   %c = icmp eq <2 x i8> %sh, zeroinitializer
   %f = call <2 x i8> @llvm.fshr.v2i8(<2 x i8> %x, <2 x i8> %y, <2 x i8> %sh)
   %s = select <2 x i1> %c, <2 x i8> %y, <2 x i8> %f
   ret <2 x i8> %s
 }
 
+; If first two operands of funnel shift are undef, the result is undef
+
+define i8 @fshl_ops_undef(i8 %shamt) {
+; CHECK-LABEL: @fshl_ops_undef(
+; CHECK-NEXT:    ret i8 undef
+;
+  %r = call i8 @llvm.fshl.i8(i8 undef, i8 undef, i8 %shamt)
+  ret i8 %r
+}
+
+define i9 @fshr_ops_undef(i9 %shamt) {
+; CHECK-LABEL: @fshr_ops_undef(
+; CHECK-NEXT:    ret i9 undef
+;
+  %r = call i9 @llvm.fshr.i9(i9 undef, i9 undef, i9 %shamt)
+  ret i9 %r
+}
+
+; If shift amount is undef, treat it as zero, returning operand 0 or 1
+
+define i8 @fshl_shift_undef(i8 %x, i8 %y) {
+; CHECK-LABEL: @fshl_shift_undef(
+; CHECK-NEXT:    ret i8 [[X:%.*]]
+;
+  %r = call i8 @llvm.fshl.i8(i8 %x, i8 %y, i8 undef)
+  ret i8 %r
+}
+
+define i9 @fshr_shift_undef(i9 %x, i9 %y) {
+; CHECK-LABEL: @fshr_shift_undef(
+; CHECK-NEXT:    ret i9 [[Y:%.*]]
+;
+  %r = call i9 @llvm.fshr.i9(i9 %x, i9 %y, i9 undef)
+  ret i9 %r
+}