[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:

I tried to address this in the comment above the condition, but here is an example of what I mean based on the `crc16.be.data8` test in llvm-test-suite:
```c
uint16_t crc_loop(uint16_t crc_initval, uint8_t data) {
  uint16_t crc = crc_initval;
  for (size_t i = 0; i < BW(data); ++i) {
    uint16_t xor_crc_data = crc ^ data;
    uint16_t crc_shl = crc << 1;
    crc = (xor_crc_data & SMIN(uint16_t)) ? (crc_shl ^ GENPOLY) : crc_shl;
    data <<= 1;
  }
  return crc;
}
```
In computing `xor_crc_data`, `data` gets zexted to 16 bits, then the conditional below is dependent on bit 15 of the result. No matter how many times `data` is shifted, it never gets used here. In fact, during testing, some cases such as `crc24.be.data16` compute an incorrect result if this guard is removed. (I can add a case like this to the regression tests if needed.)

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