[llvm] [LoopIdiomRecognize] Enable clmul optimization for CRC loops (PR #203405)

[Sean Clarke via llvm-commits]

================
@@ -1549,7 +1554,152 @@ bool LoopIdiomRecognize::avoidLIRForMultiBlockLoop(bool IsMemset,
   return false;
 }
 
-bool LoopIdiomRecognize::optimizeCRCLoop(const PolynomialInfo &Info) {
+bool LoopIdiomRecognize::optimizeCRCLoopUsingClmul(const PolynomialInfo &Info) {
+  Type *CRCTy = Info.LHS->getType();
+  LLVMContext &Ctx = CRCTy->getContext();
+  unsigned CRCBW = CRCTy->getIntegerBitWidth();
+  // The TripCount determines how many bits of data are processed, regardless of
+  // whether the actual data bit width matches (if auxiliary data is even used
+  // at all).
+  unsigned TC = Info.TripCount;
+  // The first clmul uses 2*TC bits, and the second clmul uses CRCBW+TC bits.
+  // For simplicity, have both operate on the same bit width.
+  unsigned ClmulBW = std::max(2 * TC, CRCBW + TC);
+  auto *ClmulTy = IntegerType::get(Ctx, ClmulBW);
+
+  // This optimization should only be applied if clmul for the required width is
+  // a fast operation on the target.
+  // TODO: If TC > CRCBW, then the data could probably be split into multiple
+  // chunks and processed in a loop.
+  if (!TTI->haveFastClmul(ClmulTy))
+    return false;
+
+  // First, generate the constants required for GF(2) Barrett reduction.
+  auto [Mu, FullGenPoly] =
+      HashRecognize::genBarrettConstants(Info.RHS, TC, Info.IsBigEndian);
+  Value *MuConst = ConstantInt::get(Ctx, Mu.zext(ClmulBW));
+  Value *GenPolyConst = ConstantInt::get(Ctx, FullGenPoly.zext(ClmulBW));
+
+  IRBuilder<> Builder(CurLoop->getLoopPreheader()->getTerminator());
+
+  auto ShiftNetAmt = [&](Value *Op, unsigned LShrAmt, unsigned ShlAmt,
+                         const Twine &Name) {
+    if (LShrAmt > ShlAmt)
+      return Builder.CreateLShr(Op, LShrAmt - ShlAmt, Name);
+    if (ShlAmt > LShrAmt)
+      return Builder.CreateShl(Op, ShlAmt - LShrAmt, Name);
+    return Op;
+  };
+
+  auto LoTCBits = [&](Value *Op, const Twine &Name) {
+    unsigned OpBW = Op->getType()->getIntegerBitWidth();
+    assert(OpBW >= TC && "Bit width should be at least TripCount");
+    auto *Mask = ConstantInt::get(Ctx, APInt::getLowBitsSet(OpBW, TC));
+    return Builder.CreateAnd(Op, Mask, Name);
+  };
+
+  auto MostSignificantTCBits = [&](Value *Op, unsigned BW, const Twine &Name) {
+    assert(BW >= TC && "Bit width should be at least TripCount");
+    return Info.IsBigEndian ? Builder.CreateLShr(Op, BW - TC, Name + ".be.lshr")
+                            : LoTCBits(Op, Name + ".le.mask");
+  };
+
+  Value *LHS = Builder.CreateZExt(Info.LHS, ClmulTy, "crc.cast");
+
+  // Based on the Intel white paper, in our case, we have
+  // R(x) = (LHS*x^TC) xor (LHSAux ? getTCBits(LHSAux)*x^CRCBW : 0)
+  // since the CRC loop multiplies LHS by x each iteration, and the x^CRCBW term
+  // of getTCBits(LHSAux) is XORed in for the significant bit check.
+  // Rather than compute the full R(x), we can split it in two where the most
+  // significant part is used in step 1 (floor(R(x)/x^CRCBW)) and the least
+  // significant part is used in step 3 (R(x) mod x^CRCBW).
+  // ClmulMuInput is an evolving variable that will eventually become the part
+  // used in step 1, which can be simplified to
+  // (LHS*x^(TC-CRCBW)) xor (LHSAux ? getTCBits(LHSAux) : 0). However, due to a
+  // quirk in HashRecognize, getTCBits(LHSAux) = LHSAux*x^(TC-CRCBW), so this
+  // can be further simplified to (LHS xor (LHSAux ? LHSAux : 0))*x^(TC-CRCBW).
+  Value *ClmulMuInput = LHS;
+
+  // If auxiliary data is present, XOR it in with the CRC.
+  if (Value *Data = Info.LHSAux) {
+    // The reason for the HashRecognize quirk mentioned above is that it detects
+    // (CastOrSelf LHS) xor (CastOrSelf LHSAux), which is incorrect for
+    // big-endian CRCs. This mostly allows us to handle LHS and LHSAux in the
+    // same way, regardless of bit widths, but there is an exception here.
+    // If DataBW < CRCBW, then LHSAux will always be zexted before being XORed,
+    // and the significant bit check extracts the (CRCBW-1) bit of LHSAux, which
+    // will always be zero. XORing in the data in this case gives an incorrect
+    // result, so just skip the step entirely since the XOR is with zero anyway.
+    if (!Info.IsBigEndian || Data->getType()->getIntegerBitWidth() >= CRCBW) {
+      // This is usually a zext, but DataBW may exceed ClmulBW if both CRCBW and
+      // TC are small enough.
+      Data = Builder.CreateZExtOrTrunc(Data, ClmulTy, "data.cast");
+
+      ClmulMuInput = Builder.CreateXor(ClmulMuInput, Data, "xor.crc.data");
+    }
----------------
xarkenz wrote:

Sure, I can make a separate PR for that.

I don't think it's necessary to forbid this case, unless that's what we want to do. It should be enough to just null LHSAux, no? If ignoring the dead auxiliary data allows it to be optimized without any special logic, I don't see why that would be undesirable.

https://github.com/llvm/llvm-project/pull/203405