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

[Sean Clarke via llvm-commits]

================
@@ -1549,7 +1552,151 @@ bool LoopIdiomRecognize::avoidLIRForMultiBlockLoop(bool IsMemset,
   return false;
 }
 
-bool LoopIdiomRecognize::optimizeCRCLoop(const PolynomialInfo &Info) {
+bool LoopIdiomRecognize::optimizeCRCLoopToClmul(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);
+  Type *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.
+  CRCBarrettConstants Constants =
+      HashRecognize::genBarrettConstants(Info.RHS, TC, Info.ByteOrderSwapped);
+  Value *Mu = ConstantInt::get(Ctx, Constants.Mu.zext(ClmulBW));
+  Value *FullGenPoly =
+      ConstantInt::get(Ctx, Constants.FullGenPoly.zext(ClmulBW));
+
+  // Mark all PHIs for removal since we're getting rid of the loop.
+  SmallVector<PHINode *, 2> Cleanup;
+  for (PHINode &PN : CurLoop->getHeader()->phis()) {
+    PN.replaceAllUsesWith(PoisonValue::get(PN.getType()));
+    Cleanup.push_back(&PN);
+  }
+
+  IRBuilder<> Builder(CurLoop->getLoopPreheader()->getTerminator());
+
+  Value *CRCExt = Builder.CreateZExt(Info.LHS, ClmulTy, "crc.ext");
+
+  // For the big-endian case, align the leftmost bit of the CRC with the
+  // leftmost bit of the data which is used. For the little-endian case, align
+  // the rightmost bits (nothing to do).
+  Value *CRCAlignTC;
+  if (!Info.ByteOrderSwapped)
+    CRCAlignTC = CRCExt;
+  else if (CRCBW > TC)
+    CRCAlignTC = Builder.CreateLShr(CRCExt, CRCBW - TC, "crc.be.lshr");
+  else if (TC > CRCBW)
+    CRCAlignTC = Builder.CreateShl(CRCExt, TC - CRCBW, "crc.be.shl");
+  else
+    CRCAlignTC = CRCExt;
+
+  // If auxiliary data is present, XOR it in with the CRC.
+  Value *ClmulMuInput = CRCAlignTC;
+  if (Value *Data = Info.LHSAux) {
+    unsigned DataBW = Data->getType()->getIntegerBitWidth();
+    // For big-endian CRC loops where auxiliary data is XORed with the CRC
+    // inside the loop, the bits won't be aligned properly if the bit widths
+    // don't match, and thus the CRC computation is incorrect, but HashRecognize
+    // will still detect the loop. To handle the case where the data is zexted
+    // before XORing with the CRC, just ignore the auxiliary data entirely,
+    // because the extracted bit will always be zero.
+    if (!(Info.ByteOrderSwapped && DataBW < CRCBW)) {
+      // For the aforementioned HashRecognize quirk, to handle the case where
+      // the data is truncated before XORing with the CRC, shift the data so the
+      // CRCBW-1 bit becomes the leftmost bit, and then the remaining logic
+      // treats the DataBW-1 bit as the first bit to be processed.
+      if (Info.ByteOrderSwapped && DataBW > CRCBW)
+        Data = Builder.CreateShl(Data, DataBW - CRCBW, "data.be.shl");
+      if (DataBW > TC) {
+        // Extract the useful bits of the data and discard the rest.
+        if (Info.ByteOrderSwapped) {
+          Data = Builder.CreateLShr(Data, DataBW - TC, "data.be.lshr");
+        } else {
+          ConstantInt *Mask =
+              ConstantInt::get(Ctx, APInt::getLowBitsSet(DataBW, TC));
+          Data = Builder.CreateAnd(Data, Mask, "data.le.mask");
+        }
+      }
+      // This is always a zext since TripCount <= DataBW < ClmulBW.
+      Value *DataExt = Builder.CreateZExt(Data, ClmulTy, "data.ext");
+
+      ClmulMuInput = Builder.CreateXor(CRCAlignTC, DataExt, "xor.crc.data");
+    }
+  }
+
+  // Perform the first clmul operation with the mu/mu' constant. Input is TC
+  // bits and Mu is TC+1 bits, so the result will be 2*TC bits.
+  Value *ClmulMu = Builder.CreateBinaryIntrinsic(Intrinsic::clmul, ClmulMuInput,
+                                                 Mu, {}, "clmul.mu");
+
+  // Extract the relevant bits from the result.
+  Value *ClmulGPInput;
+  if (Info.ByteOrderSwapped) {
+    ClmulGPInput = Builder.CreateLShr(ClmulMu, TC, "quot.be.lshr");
+  } else {
+    ConstantInt *Mask =
+        ConstantInt::get(Ctx, APInt::getLowBitsSet(ClmulBW, TC));
+    ClmulGPInput = Builder.CreateAnd(ClmulMu, Mask, "quot.le.mask");
+  }
+
+  // Perform the second clmul operation with the P(x)/P(x)' constant. Input is
+  // TC bits and GP is CRCBW+1 bits, so the result will be CRCBW+TC bits.
+  Value *ClmulGP = Builder.CreateBinaryIntrinsic(Intrinsic::clmul, ClmulGPInput,
+                                                 FullGenPoly, {}, "clmul.gp");
+
+  // For the big-endian case, align the leftmost bit of the CRC with the
+  // leftmost bit of the clmul result. For the little-endian case, align the
+  // rightmost bits (nothing to do).
+  Value *CRCAlignClmul = CRCExt;
+  if (Info.ByteOrderSwapped)
+    CRCAlignClmul = Builder.CreateShl(CRCExt, TC, "crc.be.shl");
+
+  // Get the remainder by subtracting (XORing) the calculated multiple of
+  // GenPoly from the CRC.
+  Value *CRCNext = Builder.CreateXor(CRCAlignClmul, ClmulGP, "xor.crc.mult");
+
+  // For the little-endian case, the leftmost bits of the XOR are relevant.
+  if (!Info.ByteOrderSwapped)
+    CRCNext = Builder.CreateLShr(CRCNext, TC, "crc.le.lshr");
+  CRCNext = Builder.CreateTrunc(CRCNext, CRCTy, "crc.next");
+
+  // Replace the result of the loop with the new computed CRC value.
+  Info.ComputedValue->replaceUsesOutsideBlock(CRCNext, CurLoop->getLoopLatch());
+
+  // Clean up the loop as much as possible so it can be trivially deleted.
+  for (PHINode *PN : Cleanup)
+    RecursivelyDeleteDeadPHINode(PN);
+  deleteDeadInstruction(CurLoop->getLatchCmpInst());
+  // Make the conditional branch always go to the exit block without changing
+  // the loop successors.
+  CondBrInst *BrInst =
+      cast<CondBrInst>(CurLoop->getLoopLatch()->getTerminator());
+  BasicBlock *ExitBlk = CurLoop->getExitBlock();
+  BasicBlock *OtherSucc = BrInst->getSuccessor(0);
+  if (OtherSucc == ExitBlk)
+    OtherSucc = BrInst->getSuccessor(1);
+  Builder.SetInsertPoint(BrInst);
+  Builder.CreateCondBr(ConstantInt::getBool(Ctx, true), ExitBlk, OtherSucc);
+  BrInst->eraseFromParent();
----------------
xarkenz wrote:

I'm going to modify this to use either true or false depending on which successor is the exit block, but this is definitely a lot simpler, thanks.

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