[llvm] [AMDGPU][InstCombine] Optimize constant shuffle patterns (PR #192246)

[Jay Foad via llvm-commits]

================
@@ -718,6 +661,395 @@ GCNTTIImpl::hoistLaneIntrinsicThroughOperand(InstCombiner &IC,
   return nullptr;
 }
 
+/// Evaluate V as a function of the lane ID and return its value on Lane, or
+/// std::nullopt if V is not a closed-form expression of the lane ID.
+static std::optional<unsigned> evalLaneExpr(Value *V, unsigned Lane,
+                                            const GCNSubtarget &ST,
+                                            const DataLayout &DL,
+                                            unsigned Depth = 0) {
+  if (Depth >= MaxAnalysisRecursionDepth)
+    return std::nullopt;
+
+  // Poison/undef in the index expression: bail and let InstCombine fold the
+  // intrinsic the usual way.
+  if (isa<UndefValue>(V))
+    return std::nullopt;
+
+  if (const ConstantInt *CI = dyn_cast<ConstantInt>(V))
+    return CI->getZExtValue();
+
+  if (isThreadID(ST, V))
+    return Lane;
+
+  const BinaryOperator *BO = dyn_cast<BinaryOperator>(V);
+  if (!BO)
+    return std::nullopt;
+
+  std::optional<unsigned> LHS =
+      evalLaneExpr(BO->getOperand(0), Lane, ST, DL, Depth + 1);
+  if (!LHS)
+    return std::nullopt;
+  std::optional<unsigned> RHS =
+      evalLaneExpr(BO->getOperand(1), Lane, ST, DL, Depth + 1);
+  if (!RHS)
+    return std::nullopt;
+
+  Type *Ty = BO->getType();
+  Constant *Ops[] = {ConstantInt::get(Ty, *LHS), ConstantInt::get(Ty, *RHS)};
+  auto *CI =
+      dyn_cast_or_null<ConstantInt>(ConstantFoldInstOperands(BO, Ops, DL));
+  if (!CI)
+    return std::nullopt;
+  return CI->getZExtValue();
+}
+
+/// Build the per-lane shuffle map by evaluating Index for every lane in the
+/// wave. Returns false if any lane index is non-constant or out of range.
+static bool tryBuildShuffleMap(Value *Index, const GCNSubtarget &ST,
+                               SmallVectorImpl<uint8_t> &Ids,
+                               const DataLayout &DL) {
+  unsigned WaveSize = ST.getWavefrontSize();
+  Ids.resize(WaveSize);
+  for (unsigned Lane : seq(WaveSize)) {
+    std::optional<unsigned> Val = evalLaneExpr(Index, Lane, ST, DL);
+    if (!Val || *Val >= WaveSize)
+      return false;
+    Ids[Lane] = *Val;
+  }
+  return true;
+}
+
+/// Lanes are partitioned into groups of Period; each group is a translated
+/// copy of the first: Ids[I] = Ids[I % Period] + (I & ~(Period - 1)).
+template <unsigned Period>
+static bool hasPeriodicLayout(ArrayRef<uint8_t> Ids) {
+  static_assert(isPowerOf2_32(Period), "Period must be a power of two");
+  for (unsigned I = Period, E = Ids.size(); I < E; ++I)
+    if (Ids[I] != Ids[I % Period] + (I & ~(Period - 1)))
+      return false;
+  return true;
+}
+
+/// Match an N-lane row pattern: each lane in [0, N) reads from a source lane
+/// in the same N-lane row, and the pattern repeats periodically across rows.
+template <unsigned N> static bool isRowPattern(ArrayRef<uint8_t> Ids) {
+  for (unsigned I = 0; I < N; ++I)
+    if (Ids[I] >= N)
+      return false;
+  return hasPeriodicLayout<N>(Ids);
+}
+
+static constexpr auto isQuadPattern = isRowPattern<4>;
+static constexpr auto isHalfRowPattern = isRowPattern<8>;
+static constexpr auto isFullRowPattern = isRowPattern<16>;
+
+/// Match a 4-lane (quad) permutation, encoded as the v_mov_b32_dpp
+/// QUAD_PERM control word: bits[1:0]=Ids[0], [3:2]=Ids[1], [5:4]=Ids[2],
+/// [7:6]=Ids[3].
+static std::optional<unsigned> matchQuadPermPattern(ArrayRef<uint8_t> Ids) {
+  if (!isQuadPattern(Ids))
+    return std::nullopt;
+  return Ids[3] << 6 | Ids[2] << 4 | Ids[1] << 2 | Ids[0];
+}
+
+/// Match an N-lane reversal (mirror) pattern.
+template <unsigned N> static bool matchMirrorPattern(ArrayRef<uint8_t> Ids) {
+  if (!isRowPattern<N>(Ids))
+    return false;
+  for (unsigned J = 0; J < N; ++J)
+    if (Ids[J] != (N - 1) - J)
+      return false;
+  return true;
+}
+
+static constexpr auto matchHalfRowMirrorPattern = matchMirrorPattern<8>;
+static constexpr auto matchFullRowMirrorPattern = matchMirrorPattern<16>;
+
+/// Match a 16-lane cyclic rotation; returns the rotation amount in [1, 15].
+static std::optional<unsigned> matchRowRotatePattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return std::nullopt;
+  if (Ids[0] == 0)
+    return std::nullopt;
+  for (unsigned J = 1; J < 16; ++J)
+    if (Ids[J] != (Ids[0] + J) % 16)
+      return std::nullopt;
+  return 16u - Ids[0];
+}
+
+/// Match a row-share pattern: all 16 lanes of each row read the same source
+/// lane. Returns the shared source lane index in [0, 16).
+static std::optional<unsigned> matchRowSharePattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return std::nullopt;
+  if (!all_equal(Ids.take_front(16)))
+    return std::nullopt;
+  return Ids[0];
+}
+
+/// Match an XOR mask pattern within each 16-lane row: Ids[J] == Mask ^ J,
+/// with Mask in [1, 15].
+static std::optional<unsigned> matchRowXMaskPattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return std::nullopt;
+  unsigned Mask = Ids[0];
+  if (Mask == 0)
+    return std::nullopt;
+  for (unsigned J = 0; J < 16; ++J)
+    if (Ids[J] != (Mask ^ J))
+      return std::nullopt;
+  return Mask;
+}
+
+/// Match an 8-lane arbitrary permutation, encoded as the v_mov_b32_dpp8
+/// 24-bit selector (three bits per output lane).
+static std::optional<unsigned> matchHalfRowPermPattern(ArrayRef<uint8_t> Ids) {
+  if (!isHalfRowPattern(Ids))
+    return std::nullopt;
+  unsigned Selector = 0;
+  for (unsigned J = 0; J < 8; ++J)
+    Selector |= Ids[J] << (J * 3);
+  return Selector;
+}
+
+/// Pack a 16-lane permutation into a single 64-bit value: four bits per output
+/// lane, lane J in bits [J*4 + 3 : J*4]. The caller splits it into the low and
+/// high 32-bit selector operands of v_permlane16 / v_permlanex16.
+static uint64_t computePermlane16Masks(ArrayRef<uint8_t> Ids) {
+  uint64_t Sel = 0;
+  for (unsigned J = 0; J < 16; ++J)
+    Sel |= static_cast<uint64_t>(Ids[J] & 0xF) << (J * 4);
+  return Sel;
+}
+
+/// Match a half-wave swap: lane J reads from lane J ^ 32. Only meaningful on
+/// wave64 targets.
+static bool matchHalfWaveSwapPattern(ArrayRef<uint8_t> Ids) {
+  if (Ids.size() != 64)
+    return false;
+  for (unsigned J = 0; J < 64; ++J)
+    if (Ids[J] != (J ^ 32))
+      return false;
+  return true;
+}
+
+/// Match a cross-row permutation suitable for v_permlanex16: every lane in
+/// the low 16-lane half reads from the high half of its own row, and vice
+/// versa.
+static bool isCrossRowPattern(ArrayRef<uint8_t> Ids) {
+  if (!hasPeriodicLayout<32>(Ids))
+    return false;
+  for (unsigned J = 0; J < 16; ++J) {
+    if (Ids[J] < 16 || Ids[J] >= 32)
+      return false;
+    if (Ids[J + 16] != Ids[J] - 16)
+      return false;
+  }
+  return true;
+}
+
+/// Match a DS_SWIZZLE bitmask-mode permutation:
+///   dst_lane = ((src_lane & AND) | OR) ^ XOR
+/// with each mask being five bits. Returns the encoded swizzle immediate.
+/// The hardware applies the formula independently within each 32-lane group,
+/// so on wave64 the high group must replicate the low one (translated by 32).
+static std::optional<unsigned>
+matchDsSwizzleBitmaskPattern(ArrayRef<uint8_t> Ids) {
+  if (!hasPeriodicLayout<32>(Ids))
+    return std::nullopt;
+
+  // The formula is per-bit: output bit B depends only on input bit B. Probe
+  // each bit with src=0 and src=(1<<B); if the output bit flipped, AND[B]=1
+  // and XOR[B] carries the constant offset; otherwise it is a constant bit
+  // encoded in OR (with AND[B]=0, XOR[B]=0).
+  unsigned AndMask = 0, OrMask = 0, XorMask = 0;
+  for (unsigned B = 0; B < 5; ++B) {
+    unsigned Bit0 = (Ids[0] >> B) & 1;
+    unsigned Bit1 = (Ids[1u << B] >> B) & 1;
+    if (Bit0 != Bit1) {
+      AndMask |= 1u << B;
+      XorMask |= Bit0 << B;
+    } else {
+      OrMask |= Bit0 << B;
+    }
+  }
+
+  // The per-bit derivation assumes bit independence; verify the masks
+  // actually reproduce every lane in the 32-lane group.
+  for (unsigned I : seq(32u)) {
+    unsigned Expected = ((I & AndMask) | OrMask) ^ XorMask;
+    if (Ids[I] != Expected)
+      return std::nullopt;
+  }
+
+  return AMDGPU::Swizzle::BITMASK_PERM_ENC |
+         (AndMask << AMDGPU::Swizzle::BITMASK_AND_SHIFT) |
+         (OrMask << AMDGPU::Swizzle::BITMASK_OR_SHIFT) |
+         (XorMask << AMDGPU::Swizzle::BITMASK_XOR_SHIFT);
----------------
jayfoad wrote:

Nit: don't need parens

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