[InstCombine] Determine demanded and known bits for funnel shifts

[nikic]

Summary

Support funnel shifts in InstCombine demanded bits simplification. If the shift amount is constant, we can determine both the demanded bits of the operands, as well as the known bits of the result.

If one of the operands has no demanded bits, it will be replaced by undef and the funnel shift will be simplified into a simple shift due to the simplifications added in D54778.

---

This change has been split off from D54666.

Another place where support for computing fsh demanded bits can be added is the DemandedBits analysis driving BDCE. This will allow handling funnel shifts with multiple uses. Unfortunately DemandedBits only tracks bits on a per-instruction level, so that the different demanded bits of the two operands cannot be leveraged.

[nikic]

test/Transforms/InstCombine/fsh.ll

292 ↗

(On Diff #175159)

This first simplifies to (x << 26) >> 30, which is then reduced to (x >> 4) & 7.

[spatel]

LGTM. I don't know if vector types will show the same simplifications, but it would be nice to have at least one of the tests repeated with a vector type, so we know if that is missing or works as expected.

Also, I know we're not getting the motivating case due to multi-use, and I agree with @nagisa that it should be filed as a bug report because I'm also not sure how to solve it.

[nikic]

I've added two tests for vectors. It works for the constant splat case, and fails if it's only constant modulo the bitwidth. That may be something we want to handle in the future (also for the zero shift case).

I've opened a bug report for the multi-use case here: https://bugs.llvm.org/show_bug.cgi?id=39771

I think the next step here would be to add support for funnel shifts to DemandedBits analysis and ValueTracking. It won't solve the motivating case either, but may help for other things.

llvm/trunk/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;
+
+        // Normalize to funnel shift left. APInt shifts of BitWidth are well-
+        // defined, so no need to special-case zero shifts here.
+        uint64_t ShiftAmt = SA->urem(BitWidth);
+        if (II->getIntrinsicID() == Intrinsic::fshr)
+          ShiftAmt = BitWidth - ShiftAmt;
+
+        APInt DemandedMaskLHS(DemandedMask.lshr(ShiftAmt));
+        APInt DemandedMaskRHS(DemandedMask.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

llvm/trunk/test/Transforms/InstCombine/fsh.ll

;
   %r = call <2 x i32> @llvm.fshr.v2i32(<2 x i32> %x, <2 x i32> zeroinitializer, <2 x i32> <i32 7, i32 7>)
   ret <2 x i32> %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:    [[Z:%.*]] = call i32 @llvm.fshl.i32(i32 [[X:%.*]], i32 [[Y:%.*]], i32 7)
+; CHECK-NEXT:    [[Z:%.*]] = shl i32 [[X:%.*]], 7
 ; CHECK-NEXT:    [[R:%.*]] = and i32 [[Z]], 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:    [[Z:%.*]] = call i32 @llvm.fshl.i32(i32 [[X:%.*]], i32 [[Y:%.*]], i32 7)
+; CHECK-NEXT:    [[Z:%.*]] = lshr i32 [[Y:%.*]], 25
 ; CHECK-NEXT:    [[R:%.*]] = and i32 [[Z]], 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:    [[Z:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 7)
+define i33 @fshr_only_op1_demanded(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_only_op1_demanded(
+; CHECK-NEXT:    [[Z:%.*]] = lshr i33 [[Y:%.*]], 7
 ; CHECK-NEXT:    [[R:%.*]] = and i33 [[Z]], 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:    [[Z:%.*]] = call i33 @llvm.fshr.i33(i33 [[X:%.*]], i33 [[Y:%.*]], i33 7)
-; CHECK-NEXT:    [[R:%.*]] = lshr i33 [[Z]], 30
+define i33 @fshr_only_op0_demanded(i33 %x, i33 %y) {
+; CHECK-LABEL: @fshr_only_op0_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
 }
 
+define <2 x i31> @fshl_only_op1_demanded_vec_splat(<2 x i31> %x, <2 x i31> %y) {
+; CHECK-LABEL: @fshl_only_op1_demanded_vec_splat(
+; CHECK-NEXT:    [[Z:%.*]] = lshr <2 x i31> [[Y:%.*]], <i31 24, i31 24>
+; CHECK-NEXT:    [[R:%.*]] = and <2 x i31> [[Z]], <i31 63, i31 31>
+; CHECK-NEXT:    ret <2 x i31> [[R]]
+;
+  %z = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 7, i31 7>)
+  %r = and <2 x i31> %z, <i31 63, i31 31>
+  ret <2 x i31> %r
+}
+
+; The shift modulo bitwidth is the same for all vector elements, but this is not simplified yet.
+define <2 x i31> @fshl_only_op1_demanded_vec_nonsplat(<2 x i31> %x, <2 x i31> %y) {
+; CHECK-LABEL: @fshl_only_op1_demanded_vec_nonsplat(
+; CHECK-NEXT:    [[Z:%.*]] = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> [[X:%.*]], <2 x i31> [[Y:%.*]], <2 x i31> <i31 7, i31 38>)
+; CHECK-NEXT:    [[R:%.*]] = and <2 x i31> [[Z]], <i31 63, i31 31>
+; CHECK-NEXT:    ret <2 x i31> [[R]]
+;
+  %z = call <2 x i31> @llvm.fshl.v2i31(<2 x i31> %x, <2 x i31> %y, <2 x i31> <i31 7, i31 38>)
+  %r = and <2 x i31> %z, <i31 63, i31 31>
+  ret <2 x i31> %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]]
 ;
;
   %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:    [[X2:%.*]] = or i32 [[X:%.*]], 1
-; CHECK-NEXT:    [[Y2:%.*]] = lshr i32 [[Y:%.*]], 1
-; CHECK-NEXT:    [[Z:%.*]] = call i32 @llvm.fshl.i32(i32 [[X2]], i32 [[Y2]], i32 7)
-; CHECK-NEXT:    [[R:%.*]] = and i32 [[Z]], 192
-; CHECK-NEXT:    ret i32 [[R]]
+; 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:    [[X2:%.*]] = or i33 [[X:%.*]], 1
-; CHECK-NEXT:    [[Y2:%.*]] = lshr i33 [[Y:%.*]], 1
-; CHECK-NEXT:    [[Z:%.*]] = call i33 @llvm.fshr.i33(i33 [[X2]], i33 [[Y2]], i33 26)
-; CHECK-NEXT:    [[R:%.*]] = and i33 [[Z]], 192
-; CHECK-NEXT:    ret i33 [[R]]
+; 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
 }