[llvm] [AArch64] Add SVE shuffle optimization pass (PR #193951)

[Graham Hunter via llvm-commits]

================
@@ -0,0 +1,300 @@
+//===------- SVEShuffleOpts - SVE Shuffle Optimization --------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Tries to pattern match and combine scalable vector shuffles that could
+// be more efficiently performed by tbl instructions.
+//
+// An example would be a loop with 4 multiply-accumulate reductions, where the
+// new data in each vector iterations comes from a 4-way deinterleaving of
+// smaller datatypes loaded from memory,  which are then extended and multiplied
+// by a common term loaded in reverse order from memory before being added to
+// the accumulator.
+//
+// If the initial load is a legal vector rather than 4x the size (generating a
+// structured ld4 instead), we would see multiple uunpkhi/lo instructions for
+// the extensions, followed by uzp1/2 instructions for the deinterleave, and rev
+// instructions for the common terms. Instead, we can replace all of those with
+// 4 tbl instructions. The tradeoff, of course, is that we now have 4 mask
+// values to maintain which increases register pressure.
+//
+// We should also be able to introduce new shuffles in order to balance out
+// SVE's bottom/top instruction pairs, which act on even/odd lanes instead of
+// the high or low half of a register.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "Utils/AArch64BaseInfo.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsAArch64.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/InitializePasses.h"
+#include <optional>
+
+using namespace llvm;
+using namespace llvm::PatternMatch;
+
+#define DEBUG_TYPE "aarch64-sve-shuffle-opts"
+
+namespace {
+
+class SVEShuffleImpl {
+  const AArch64TargetMachine *TM = nullptr;
+  const LoopInfo *LI = nullptr;
+
+public:
+  SVEShuffleImpl() {};
+  SVEShuffleImpl(const AArch64TargetMachine *TM) : TM(TM) {};
+
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+  bool runOnFunction(Function &F, Pass &P);
+
+private:
+  bool processLoop(Loop &L);
+};
+
+struct SVEShuffleOpts : public FunctionPass {
+  SVEShuffleImpl Impl;
+  static char ID; // Pass identification, replacement for typeid
+  SVEShuffleOpts() : FunctionPass(ID) {}
+
+  bool runOnFunction(Function &F) override {
+    if (skipFunction(F))
+      return false;
+
+    return Impl.runOnFunction(F, *this);
+  }
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+  StringRef getPassName() const override { return "SVE Tbl Folding Opts"; }
+
+private:
+};
+} // end anonymous namespace
+
+/// A mapping between a vector_deinterleaveN intrinsic and extending cast
+/// instructions used on the resulting subvectors.
+using DeinterleaveMap =
+    SmallDenseMap<CallInst *, SmallVector<std::pair<CastInst *, unsigned>, 4>>;
+
+static void evaluateDeinterleave(IntrinsicInst *I, DeinterleaveMap &Candidates,
+                                 Loop &L) {
+  // TODO: 'Legalize' if the input is wider than nx128b but not wide enough
+  //       to match a structured load?
+  if (I->getOperand(0)
+          ->getType()
+          ->getPrimitiveSizeInBits()
+          .getKnownMinValue() != AArch64::SVEBitsPerBlock)
+    return;
+
+  unsigned IntId = I->getIntrinsicID();
+  assert(IntId == Intrinsic::vector_deinterleave4 &&
+         "Only deinterleave4 supported currently");
+  SmallVector<std::pair<CastInst *, unsigned>, 4> Extends;
+  unsigned Opcode = 0;
+  Type *DestTy = nullptr;
+  for (User *U : I->users()) {
+    auto *Extract = dyn_cast<ExtractValueInst>(U);
+    // We expect only a single cast instruction as a user.
+    if (!Extract || Extract->getNumIndices() != 1)
+      return;
+
+    auto *Extend = dyn_cast<CastInst>(Extract->getUniqueUndroppableUser());
----------------
huntergr-arm wrote:

done

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