[llvm] [AMDGPU] Add shuffle optimizer pass (PR #155824)

[Matt Arsenault via llvm-commits]

================
@@ -0,0 +1,475 @@
+//===-- AMDGPUShuffleOptimizer.cpp - Optimize shuffle patterns -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass optimizes generic shuffle intrinsics by detecting constant
+// patterns and replacing them with efficient hardware-specific instructions
+// (DPP, PERMLANE*, etc.) when beneficial, falling back to DS_BPERMUTE_B32
+// for unmatched patterns.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "GCNSubtarget.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Target/TargetMachine.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "amdgpu-shuffle-optimizer"
+
+static cl::opt<bool>
+    EnableShuffleOptimization("amdgpu-enable-shuffle-optimization",
+                              cl::desc("Enable AMDGPU shuffle optimization"),
+                              cl::init(true), cl::Hidden);
+
+namespace {
+
+// Represents a detected shuffle pattern that can be optimized
+struct ShufflePattern {
+  enum PatternKind {
+    DPP_QUAD_PERM, // DPP quad permutation
+    DPP_ROW_SHL,   // DPP row shift left
+    DPP_ROW_SHR,   // DPP row shift right
+    DPP_WAVE_SHL,  // DPP wave shift left
+    DPP_WAVE_SHR,  // DPP wave shift right
+    PERMLANE16,    // V_PERMLANE16_B32
+    PERMLANEX16,   // V_PERMLANEX16_B32
+    PERMLANE64,    // V_PERMLANE64_B32
+    DS_BPERMUTE,   // Fallback to DS_BPERMUTE_B32
+    UNSUPPORTED    // Cannot be optimized
+  };
+
+  PatternKind Kind;
+  uint32_t DPPCtrl = 0;    // DPP control value
+  uint32_t RowMask = 0xf;  // DPP row mask
+  uint32_t BankMask = 0xf; // DPP bank mask
+  bool BoundCtrl = false;  // DPP bound control
+
+  ShufflePattern() : Kind(UNSUPPORTED) {}
+};
+
+class AMDGPUShuffleOptimizer : public FunctionPass {
+public:
+  static char ID;
+
+  AMDGPUShuffleOptimizer() : FunctionPass(ID) {}
+
+  bool runOnFunction(Function &F) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetPassConfig>();
+    AU.setPreservesCFG();
+  }
+
+  static bool runShuffleOptimizer(Function &F, const GCNSubtarget &ST);
+
+private:
+  const GCNSubtarget *ST = nullptr;
+
+  bool optimizeShuffleIntrinsic(CallInst *CI);
+  ShufflePattern analyzeShufflePattern(CallInst *CI);
+  ShufflePattern analyzeShuffleIdx(int Width, int Offset);
+  ShufflePattern analyzeShuffleUp(int Width, int Delta);
+  ShufflePattern analyzeShuffleDown(int Width, int Delta);
+  ShufflePattern analyzeShuffleXor(int Width, int Mask);
+
+  bool tryOptimizeToDPP(CallInst *CI, const ShufflePattern &Pattern);
+  bool tryOptimizeToPermlane(CallInst *CI, const ShufflePattern &Pattern);
+  bool fallbackToBpermute(CallInst *CI);
+
+  Value *createDPPIntrinsic(IRBuilder<> &Builder, Value *OldVal, Value *SrcVal,
+                            const ShufflePattern &Pattern);
+  Value *createPermlaneIntrinsic(IRBuilder<> &Builder, Value *Val,
+                                 const ShufflePattern &Pattern);
+  Value *createBpermuteIntrinsic(IRBuilder<> &Builder, Value *Val,
+                                 Value *Index);
+
+  bool processShuffleIntrinsics(Function &F);
+};
+
+char AMDGPUShuffleOptimizer::ID = 0;
+
+bool AMDGPUShuffleOptimizer::runOnFunction(Function &F) {
+  if (!EnableShuffleOptimization)
+    return false;
+
+  auto &TPC = getAnalysis<TargetPassConfig>();
+  const TargetMachine &TM = TPC.getTM<TargetMachine>();
+  ST = &TM.getSubtarget<GCNSubtarget>(F);
+
+  return processShuffleIntrinsics(F);
+}
+
+bool AMDGPUShuffleOptimizer::optimizeShuffleIntrinsic(CallInst *CI) {
+  ShufflePattern Pattern = analyzeShufflePattern(CI);
+
+  if (Pattern.Kind == ShufflePattern::UNSUPPORTED)
+    return fallbackToBpermute(CI);
+
+  switch (Pattern.Kind) {
+  case ShufflePattern::DPP_QUAD_PERM:
+  case ShufflePattern::DPP_ROW_SHL:
+  case ShufflePattern::DPP_ROW_SHR:
+  case ShufflePattern::DPP_WAVE_SHL:
+  case ShufflePattern::DPP_WAVE_SHR:
+    return tryOptimizeToDPP(CI, Pattern);
+
+  case ShufflePattern::PERMLANE16:
+  case ShufflePattern::PERMLANEX16:
+  case ShufflePattern::PERMLANE64:
+    return tryOptimizeToPermlane(CI, Pattern);
+
+  case ShufflePattern::DS_BPERMUTE:
+    return fallbackToBpermute(CI);
+
+  default:
+    return false;
+  }
+}
+
+ShufflePattern AMDGPUShuffleOptimizer::analyzeShufflePattern(CallInst *CI) {
+  auto *II = cast<IntrinsicInst>(CI);
+
+  // Get width parameter (must be constant)
+  auto *WidthConst = dyn_cast<ConstantInt>(CI->getArgOperand(2));
+  if (!WidthConst)
+    return ShufflePattern();
+
+  int Width = WidthConst->getSExtValue();
+
+  // Get offset/delta/mask parameter (must be constant for pattern optimization)
+  auto *ParamConst = dyn_cast<ConstantInt>(CI->getArgOperand(1));
+  if (!ParamConst)
+    return ShufflePattern();
+
+  int Param = ParamConst->getSExtValue();
+
+  switch (II->getIntrinsicID()) {
+  case Intrinsic::amdgcn_generic_shuffle:
+    return analyzeShuffleIdx(Width, Param);
+  case Intrinsic::amdgcn_generic_shuffle_up:
+    return analyzeShuffleUp(Width, Param);
+  case Intrinsic::amdgcn_generic_shuffle_down:
+    return analyzeShuffleDown(Width, Param);
+  case Intrinsic::amdgcn_generic_shuffle_xor:
+    return analyzeShuffleXor(Width, Param);
+  default:
+    return ShufflePattern();
+  }
+}
+
+ShufflePattern AMDGPUShuffleOptimizer::analyzeShuffleIdx(int Width,
+                                                         int Offset) {
+  // For idx shuffle, all lanes read from the same offset
+  if (Offset == 0) {
+    // Broadcast from lane 0 - use DPP if supported, otherwise bpermute
+    if (ST->hasDPP() && Width == 16) {
+      // Use DPP quad permutation to broadcast lane 0 within each group of 4
+      ShufflePattern Pattern;
+      Pattern.Kind = ShufflePattern::DPP_QUAD_PERM;
+      Pattern.DPPCtrl = 0x00; // [0,0,0,0] quad perm - broadcast lane 0
+      return Pattern;
+    } else {
----------------
arsenm wrote:

No else after return 

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