[llvm] [X86] Fold nested VGF2P8AFFINEQB instructions (PR #195210)

[Simon Pilgrim via llvm-commits]

================
@@ -62136,6 +62155,83 @@ static SDValue combineKSHIFT(SDNode *N, SelectionDAG &DAG,
   return SDValue();
 }
 
+// Fold: GF2P8AFFINEQB(GF2P8AFFINEQB(X, YSub), YSup)
+//    => GF2P8AFFINEQB(X, YFolded)
+// Permuting the sub-matrix by the super-matrix at a byte, rather than bit,
+// granularity produces a matrix that performs both permutations at once.
+static SDValue combineNestedGF2P8AFFINEQB(SDNode *N, const SDLoc &DL,
+                                          SelectionDAG &DAG, EVT VT) {
+  using namespace SDPatternMatch;
+
+  unsigned VecWidth = VT.getSizeInBits();
+  unsigned NumElts = VT.getVectorNumElements();
+  unsigned EltWidth = VT.getScalarSizeInBits();
+
+  SDValue X, YSub, YSup;
+  APInt ImmSub, ImmSup, ConstUndef;
+  SmallVector<APInt> YSubEltBits, YSupEltBits;
+
+  if (!(sd_match(N, m_TernaryOp(X86ISD::GF2P8AFFINEQB,
+                                m_TernaryOp(X86ISD::GF2P8AFFINEQB, m_Value(X),
+                                            m_Value(YSub), m_ConstInt(ImmSub)),
+                                m_Value(YSup), m_ConstInt(ImmSup))) &&
+        getTargetConstantBitsFromNode(YSub, EltWidth, ConstUndef, YSubEltBits,
+                                      /*AllowWholeUndefs=*/false) &&
+        getTargetConstantBitsFromNode(YSup, EltWidth, ConstUndef, YSupEltBits,
+                                      /*AllowWholeUndefs=*/false)))
+    return SDValue();
+
+  APInt SubM(VecWidth, 0);
+  APInt SupM(VecWidth, 0);
+  for (unsigned i = 0; i < NumElts; ++i) {
+    SubM.insertBits(YSubEltBits[i], i * EltWidth);
+    SupM.insertBits(YSupEltBits[i], i * EltWidth);
+  }
+
+  // Immediate is shared and needs to be permuted in the same manner
+  if (!SupM.isSplat(64) && ImmSub != 0)
+    return SDValue();
+
+  // Immediate permute
+  APInt FoldedImm = getGFNIByteAffine(ImmSub, SupM.trunc(64), ImmSup);
+
+  // Matrix permute
+  APInt FoldedMatrix = APInt(VecWidth, 0);
+  APInt LeastRowMask = APInt::getSplat(VecWidth, APInt(64, 0xFF));
+  APInt LeastBitInByte = APInt::getSplat(VecWidth, APInt(8, 0x01));
+  APInt RowSplatter = APInt(VecWidth, 0x0101010101010101ull);
+
+  for (unsigned Row = 0; Row < 8; ++Row) {
+    APInt RowSplat = (SubM & LeastRowMask) * RowSplatter;
+    SubM = SubM.lshr(EltWidth);
+
+    APInt ByteMaskIfSet = (SupM.lshr(7 - Row)) & LeastBitInByte;
+    ByteMaskIfSet *= 0xFF;
+
+    FoldedMatrix ^= RowSplat & ByteMaskIfSet;
+  }
+
+  SmallVector<SDValue> FoldedVector;
+  for (unsigned i = 0; i < NumElts; ++i) {
----------------
RKSimon wrote:

```suggestion
  for (unsigned I = 0; I != NumElts; ++I) {
```

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