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

[Matt Arsenault via llvm-commits]

================
@@ -718,6 +720,288 @@ GCNTTIImpl::hoistLaneIntrinsicThroughOperand(InstCombiner &IC,
   return nullptr;
 }
 
+static std::optional<unsigned> evalLaneExpr(Value *V, unsigned Lane,
+                                            const GCNSubtarget &ST,
+                                            const DataLayout &DL,
+                                            unsigned Depth = 0) {
+  if (Depth > 16)
+    return std::nullopt;
+
+  if (isThreadID(ST, V))
+    return Lane;
+
+  if (const auto *CI = dyn_cast<ConstantInt>(V))
+    return CI->getZExtValue();
+
+  auto *BO = dyn_cast<BinaryOperator>(V);
+  if (!BO)
+    return std::nullopt;
+
+  auto LHS = evalLaneExpr(BO->getOperand(0), Lane, ST, DL, Depth + 1);
+  auto RHS = evalLaneExpr(BO->getOperand(1), Lane, ST, DL, Depth + 1);
+  if (!LHS || !RHS)
+    return std::nullopt;
+
+  Type *Ty = BO->getType();
+  Constant *LC = ConstantInt::get(Ty, *LHS);
+  Constant *RC = ConstantInt::get(Ty, *RHS);
+  Constant *Result = ConstantFoldBinaryOpOperands(BO->getOpcode(), LC, RC, DL);
+  auto *CI = dyn_cast_or_null<ConstantInt>(Result);
+  if (!CI)
+    return std::nullopt;
+  return CI->getZExtValue();
+}
+
+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 = 0; Lane < WaveSize; ++Lane) {
+    auto Val = evalLaneExpr(Index, Lane, ST, DL);
+    if (!Val || *Val >= WaveSize)
+      return false;
+    Ids[Lane] = static_cast<uint8_t>(*Val);
+  }
+  return true;
+}
+
+static bool isHalfRowPattern(ArrayRef<uint8_t> Ids) {
+  for (unsigned i = 0; i < 8; ++i)
+    if (Ids[i] >= 8)
+      return false;
+  for (unsigned i = 8, E = Ids.size(); i < E; ++i)
+    if (Ids[i] != Ids[i % 8] + (i & ~7u))
+      return false;
+  return true;
+}
+
+static bool isFullRowPattern(ArrayRef<uint8_t> Ids) {
+  for (unsigned i = 0; i < 16; ++i)
+    if (Ids[i] >= 16)
+      return false;
+  for (unsigned i = 16, E = Ids.size(); i < E; ++i)
+    if (Ids[i] != Ids[i % 16] + (i & ~15u))
+      return false;
+  return true;
+}
+
+static std::optional<unsigned> matchQuadPermPattern(ArrayRef<uint8_t> Ids) {
+  for (unsigned i = 0; i < 4; ++i)
+    if (Ids[i] >= 4)
+      return std::nullopt;
+  for (unsigned i = 4, E = Ids.size(); i < E; ++i)
+    if (Ids[i] != Ids[i % 4] + (i & ~3u))
+      return std::nullopt;
+  return (Ids[3] << 6) | (Ids[2] << 4) | (Ids[1] << 2) | Ids[0];
+}
+
+static bool matchHalfRowMirrorPattern(ArrayRef<uint8_t> Ids) {
+  if (!isHalfRowPattern(Ids))
+    return false;
+  for (unsigned j = 0; j < 8; ++j)
+    if (Ids[j] != 7 - j)
+      return false;
+  return true;
+}
+
+static bool matchFullRowMirrorPattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return false;
+  for (unsigned j = 0; j < 16; ++j)
+    if (Ids[j] != 15 - j)
+      return false;
+  return true;
+}
+
+static std::optional<unsigned> matchRowRotatePattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return std::nullopt;
+  if (Ids[0] == 0 || Ids[15] == 15)
+    return std::nullopt;
+  for (unsigned j = 1; j < 16; ++j) {
+    int Diff = static_cast<int>(Ids[j]) - static_cast<int>(Ids[j - 1]);
+    if (Diff != 1 && Diff != -15)
+      return std::nullopt;
+  }
+  return 16u - Ids[0];
+}
+
+static std::optional<unsigned> matchRowSharePattern(ArrayRef<uint8_t> Ids) {
+  if (!isFullRowPattern(Ids))
+    return std::nullopt;
+  for (unsigned j = 1; j < 16; ++j)
+    if (Ids[j] != Ids[0])
+      return std::nullopt;
+  return static_cast<unsigned>(Ids[0]);
+}
+
+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;
+}
+
+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 |= static_cast<unsigned>(Ids[j] & 0x7) << (j * 3);
+  return Selector;
+}
+
+static std::pair<uint32_t, uint32_t>
+computePermlane16Masks(ArrayRef<uint8_t> Ids) {
+  uint32_t Lo = 0, Hi = 0;
+  for (unsigned j = 0; j < 8; ++j)
+    Lo |= static_cast<uint32_t>(Ids[j] & 0xF) << (j * 4);
+  for (unsigned j = 8; j < 16; ++j)
+    Hi |= static_cast<uint32_t>(Ids[j] & 0xF) << ((j - 8) * 4);
+  return {Lo, Hi};
+}
+
+static std::optional<unsigned> matchHalfRowSharePattern(ArrayRef<uint8_t> Ids) {
+  if (!isHalfRowPattern(Ids))
+    return std::nullopt;
+  for (unsigned j = 1; j < 8; ++j)
+    if (Ids[j] != Ids[0])
+      return std::nullopt;
+  // Bitmask-mode: AND=0x18 (keep group-select bits), OR=ids[0], XOR=0
+  return (static_cast<unsigned>(Ids[0]) << AMDGPU::Swizzle::BITMASK_OR_SHIFT) |
+         0x18u;
+}
+
+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;
+}
+
+static Value *createUpdateDpp(IRBuilderBase &B, Value *Val, unsigned Ctrl,
+                              Module *M) {
+  Type *Ty = Val->getType();
+  Function *F =
+      Intrinsic::getOrInsertDeclaration(M, Intrinsic::amdgcn_update_dpp, {Ty});
+  return B.CreateCall(F, {PoisonValue::get(Ty), Val, B.getInt32(Ctrl),
+                          B.getInt32(0xF), B.getInt32(0xF), B.getTrue()});
+}
+
+static Value *createMovDpp8(IRBuilderBase &B, Value *Val, unsigned Selector,
+                            Module *M) {
+  Type *Ty = Val->getType();
+  Function *F =
+      Intrinsic::getOrInsertDeclaration(M, Intrinsic::amdgcn_mov_dpp8, {Ty});
+  return B.CreateCall(F, {Val, B.getInt32(Selector)});
+}
+
+static Value *createPermlane16(IRBuilderBase &B, Value *Val, uint32_t Lo,
+                               uint32_t Hi, Module *M) {
+  Type *Ty = Val->getType();
+  Function *F =
+      Intrinsic::getOrInsertDeclaration(M, Intrinsic::amdgcn_permlane16, {Ty});
+  return B.CreateCall(F, {PoisonValue::get(Ty), Val, B.getInt32(Lo),
+                          B.getInt32(Hi), B.getFalse(), B.getFalse()});
+}
+
+static Value *createDsSwizzle(IRBuilderBase &B, Value *Val, unsigned Offset,
+                              Module *M) {
+  Type *OrigTy = Val->getType();
+  Value *Src = Val;
+  if (!OrigTy->isIntegerTy(32))
+    Src = B.CreateBitCast(Src, B.getInt32Ty());
+  Function *F =
+      Intrinsic::getOrInsertDeclaration(M, Intrinsic::amdgcn_ds_swizzle);
+  Value *Result = B.CreateCall(F, {Src, B.getInt32(Offset)});
----------------
arsenm wrote:

CreateIntrinsic 

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