[ValueTracking] Support funnel shifts in computeKnownBits()

[nikic]

Summary

This is a continuation of D54869, this time adding support for funnel shifts in ValueTracking computeKnownBits(). Implementation is essentially the same as for InstCombineSimplifyDemandedBits.

[nikic]

Apart from this, support for ROTL and ROTR could be added to the SelectionDAG variant of computeKnownBits.

[spatel]

LGTM

lib/Analysis/ValueTracking.cpp

if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
         // can't contribute to the population.
         unsigned BitsPossiblySet = Known2.countMaxPopulation();
         unsigned LowBits = Log2_32(BitsPossiblySet)+1;
         Known.Zero.setBitsFrom(LowBits);
         // TODO: we could bound KnownOne using the lower bound on the number
         // of bits which might be set provided by popcnt KnownOne2.
         break;
       }
+      case Intrinsic::fshr:
+      case Intrinsic::fshl: {
+        const APInt *SA;
+        if (!match(I->getOperand(2), m_APInt(SA)))
+          break;
+
+        // Normalize to funnel shift left.
+        uint64_t ShiftAmt = SA->urem(BitWidth);
+        if (II->getIntrinsicID() == Intrinsic::fshr)
+          ShiftAmt = BitWidth - ShiftAmt;
+
+        KnownBits Known3(Known);
+        computeKnownBits(I->getOperand(0), Known2, Depth + 1, Q);
+        computeKnownBits(I->getOperand(1), Known3, Depth + 1, Q);
+
+        Known.Zero =
+            Known2.Zero.shl(ShiftAmt) | Known3.Zero.lshr(BitWidth - ShiftAmt);
+        Known.One =
+            Known2.One.shl(ShiftAmt) | Known3.One.lshr(BitWidth - ShiftAmt);
+        break;
+      }
       case Intrinsic::x86_sse42_crc32_64_64:
         Known.Zero.setBitsFrom(32);
         break;
       }
     }
     break;
   case Instruction::ExtractElement:
     // Look through extract element. At the moment we keep this simple and skip

unittests/Analysis/ValueTrackingTest.cpp

parseAssembly(
       "  %bb = shl i32 %b, 5\n"
       "  %aaa = or i32 %aa, 24\n"
       "  %bbb = or i32 %bb, 28\n"
       "  %A = mul i32 %aaa, %bbb\n"
       "  ret i32 %A\n"
       "}\n");
   expectKnownBits(/*zero*/ 95u, /*one*/ 32u);
 }
+
+TEST_F(ComputeKnownBitsTest, ComputeKnownFshl) {
+  // fshl(....1111....0000, 00..1111........, 6)
+  // = 11....000000..11
+  parseAssembly(
+      "define i16 @test(i16 %a, i16 %b) {\n"
+      "  %aa = shl i16 %a, 4\n"
+      "  %bb = lshr i16 %b, 2\n"
+      "  %aaa = or i16 %aa, 3840\n"
+      "  %bbb = or i16 %bb, 3840\n"
+      "  %A = call i16 @llvm.fshl.i16(i16 %aaa, i16 %bbb, i16 6)\n"
+      "  ret i16 %A\n"
+      "}\n"
+      "declare i16 @llvm.fshl.i16(i16, i16, i16)\n");
+  expectKnownBits(/*zero*/ 1008u, /*one*/ 49155u);
+}
+
+TEST_F(ComputeKnownBitsTest, ComputeKnownFshr) {
+  // fshr(....1111....0000, 00..1111........, 26)
+  // = 11....000000..11
+  parseAssembly(
+      "define i16 @test(i16 %a, i16 %b) {\n"
+      "  %aa = shl i16 %a, 4\n"
+      "  %bb = lshr i16 %b, 2\n"
+      "  %aaa = or i16 %aa, 3840\n"
+      "  %bbb = or i16 %bb, 3840\n"
+      "  %A = call i16 @llvm.fshr.i16(i16 %aaa, i16 %bbb, i16 26)\n"
+      "  ret i16 %A\n"
+      "}\n"
+      "declare i16 @llvm.fshr.i16(i16, i16, i16)\n");
+  expectKnownBits(/*zero*/ 1008u, /*one*/ 49155u);
+}
+
+TEST_F(ComputeKnownBitsTest, ComputeKnownFshlZero) {
+  // fshl(....1111....0000, 00..1111........, 0)
+  // = ....1111....0000
+  parseAssembly(
+      "define i16 @test(i16 %a, i16 %b) {\n"
+      "  %aa = shl i16 %a, 4\n"
+      "  %bb = lshr i16 %b, 2\n"
+      "  %aaa = or i16 %aa, 3840\n"
+      "  %bbb = or i16 %bb, 3840\n"
+      "  %A = call i16 @llvm.fshl.i16(i16 %aaa, i16 %bbb, i16 0)\n"
+      "  ret i16 %A\n"
+      "}\n"
+      "declare i16 @llvm.fshl.i16(i16, i16, i16)\n");
+  expectKnownBits(/*zero*/ 15u, /*one*/ 3840u);
+}