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

[Graham Hunter via llvm-commits]

================
@@ -0,0 +1,242 @@
+//===------- 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.
+// TODO: Also handle bottom/top (de)interleaving.
+//
+//===----------------------------------------------------------------------===//
+
+#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
+
+// Optimize zext and uitofp from a 4-way deinterleaved load.
+static bool optimizeSVEUnsignedExtends(Instruction *I, Loop *L) {
+  assert(L && "No loop");
+  assert((isa<ZExtInst>(I) || isa<UIToFPInst>(I)) && "No conversion.");
+  Value *Src = I->getOperand(0);
+  VectorType *DstTy = cast<VectorType>(I->getType());
+  VectorType *SrcTy = cast<VectorType>(Src->getType());
+  unsigned DstBits = DstTy->getScalarSizeInBits();
+  unsigned SrcBits = SrcTy->getScalarSizeInBits();
+
+  // TODO: Would we want tbl-based deinterleaving even if the cast is a uitofp
+  //       at the same bitwidth? Maybe if we could combine with another shuffle?
+  if (SrcBits >= DstBits)
+    return false;
+
+  // We can't abuse the invalid index trick for tbls of bytes, since the
+  // largest possible SVE vector (2048b) would have 256 bytes, leaving no
+  // way of zeroing.
+  // TODO: If we know vscale is 8 or less, then we could use tbls for bytes.
+  if (SrcBits <= 8)
+    return false;
+
+  using namespace llvm::PatternMatch;
+  Value *Load = nullptr;
+  Value *Extract = nullptr;
+
+  bool IsUIToFP = isa<UIToFPInst>(I);
+  if (match(I->getOperand(0),
+            m_Value(Extract,
+                    m_ExtractValue(m_Intrinsic<Intrinsic::vector_deinterleave4>(
+                        m_Value(Load,
+                                m_CombineOr(m_MaskedLoad(m_Value(), m_Value(),
+                                                         m_CombineOr(m_Undef(),
+                                                                     m_Zero())),
+                                            m_Load(m_Value())))))))) {
----------------
huntergr-arm wrote:

Removed the restriction.

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