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

[Graham Hunter via llvm-commits]

================
@@ -0,0 +1,344 @@
+//===------- 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 zero extended.
+//
+// Something like the following:
+//   %bgra = call ... @llvm.masked.load
+//   %deinterleave = call ... @llvm.vector.deinterleave4(%bgra)
+// If the load was of a <vscale x 8 x i16>, we now have 4 deinterleaved
+// <vscale x 2 x i16> values.
+//   %b.i16 = extractvalue %deinterleave, 0
+//   %b.i64 = zext <vscale x 2 x i16> %b.i16 to <vscale x 2 x i64>
+//   %acc.b.next = add <vscale x 2 x i64> %acc.b, %b.i64
+//   <repeat for the other 3 subvectors>
+//
+// 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.
+// 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.
+//
+// This basic transform could be performed in CodeGenPrepare (as the equivalent
+// for NEON is), or in a DAG Combine. However, we hope to extend it to detect
+// other shuffles that we can fold into the tbl. Extending the above example,
+// if instead of directly adding to the accumulator we multiplied it by a
+// common term for all 4 components that had been reversed:
+//   %common.load = call @llvm.masked.load
+//   %common.reverse = call @llvm.vector.reverse
+// These would be loaded at the extended size, <vscale x 2 x i64> in our
+// example.
+//   %b.mul = mul <vscale x 2 x i64> %b.i64, %common.reverse
+//   %acc.b.next = add <vscale x 2 x i64> %acc.b, %b.mul
+//   <repeat for the other 3 subvectors, using %common.reverse for each)
+//
+// In this case, the reverse isn't applied to the deinterleaved data in the
+// original IR, but to the common term multiplied by the individual bgra
+// elements. If the order of the elements in the accumulator is important, we
+// cannot change that. If, however, we know that the accumulator is reduced to
+// a single scalar after the loop and the data is either integers or floating
+// point with reassociation allowed, we could instead choose a different mask
+// for the tbls to reverse the individual bgra elements instead, removing an
+// additional instruction from the loop. This does require looking beyond the
+// blocks in the loop, so DAGCombine won't help.
+//
+// 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.
+//
+// This pass may end up being a temporary solution that is removed if we can
+// create a generic vector shuffle intrinsic and move this feature to
+// LoopVectorize itself, as that would allow for better cost modelling.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "Utils/AArch64BaseInfo.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/MemorySSA.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 <array>
+#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 *, std::array<CastInst *, 4>>;
+
+/// Evaluate a deinterleave and see what the uses are. If we find other
+/// operations that we can combine into a tbl shuffle, add the deinterleave and
+/// the operations (currently only zext or uitofp) to the candidates map.
+static void evaluateDeinterleave(IntrinsicInst *I, DeinterleaveMap &Candidates,
+                                 Loop &L) {
+  // This pass currently only handles legal SVE vector types.
+  if (I->getOperand(0)
+          ->getType()
+          ->getPrimitiveSizeInBits()
+          .getKnownMinValue() != AArch64::SVEBitsPerBlock)
+    return;
+
+  unsigned IntId = I->getIntrinsicID();
+  assert(IntId == Intrinsic::vector_deinterleave4 &&
+         "Only deinterleave4 supported currently");
+  std::array<CastInst *, 4> Extends = {};
+  unsigned Opcode = 0;
+  Type *DestTy = nullptr;
+  for (User *U : I->users()) {
+    auto *Extract = dyn_cast<ExtractValueInst>(U);
+    if (!Extract)
+      return;
+
+    // We expect only a single cast instruction as a user for the extract.
+    auto *Extend =
+        dyn_cast_if_present<CastInst>(Extract->getUniqueUndroppableUser());
+    if (!Extend || (!isa<ZExtInst>(Extend) && !isa<UIToFPInst>(Extend)))
+      return;
+
+    // We're only interested if the uses are in the loop. This is almost
+    // certainly the case.
+    if (!L.contains(Extract) || !L.contains(Extend))
+      return;
+
+    Opcode = Extend->getOpcode();
+    DestTy = Extend->getDestTy();
+    Type *SrcTy = Extend->getSrcTy();
+
+    // For now, we only want to handle scalable vectors here.
+    if (!DestTy->isScalableTy())
+      return;
+
+    unsigned SrcBits = SrcTy->getScalarSizeInBits();
+    unsigned DestBits = DestTy->getScalarSizeInBits();
+
+    // Looking to match the deinterleave factor.
+    if (DestBits / SrcBits != 4)
+      return;
+
+    // TBL zeroes elements with an out-of-bounds index, but for the largest
+    // possible SVE vector (2048b) a maximum value for i8 elements (256) is not
+    // large enough to encode an 'out of bounds' value.
+    // TODO: If we know vscale is 8 or less, then we could use tbls for bytes.
+    if (SrcBits <= 8)
+      return;
+
+    Extends[Extract->getIndices().front()] = Extend;
+  }
+
+  // Check that all extracted values are being extended the same way, and that
+  // we have the expected number of extensions.
+  if (!all_of(Extends, [DestTy, Opcode](CastInst *CI) {
+        return CI && CI->getDestTy() == DestTy && CI->getOpcode() == Opcode;
+      }))
+    return;
+
+  Candidates.try_emplace(I, Extends);
----------------
huntergr-arm wrote:

Checked that the deinterleave has the expected number of uses.

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