[llvm] [AArch64] Add SVE shuffle optimization pass (PR #193951)
Graham Hunter via llvm-commits
llvm-commits at lists.llvm.org
Thu Jun 11 03:34:53 PDT 2026
================
@@ -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);
+}
+
+/// Given a map of deinterleaves to zext or uitofp casts, remove the operations
+/// and replace them with tbl shuffles.
+static void optimizeSVEDeinterleavedExtends(DeinterleaveMap Deinterleaves) {
+ // TODO: Cache tbl patterns and reuse, and abandon transforms for a particular
+ // deinterleave if it would introduce too many. We probably want a
+ // hardcoded number of tbls to start with, but if we can estimate
+ // register pressure then we could make better decisions.
+ for (auto &[Deinterleave, Extends] : Deinterleaves) {
+ VectorType *DestTy = cast<VectorType>(Extends[0]->getDestTy());
+ VectorType *SrcTy = cast<VectorType>(Extends[0]->getSrcTy());
+ unsigned DstBits = DestTy->getScalarSizeInBits();
+ unsigned SrcBits = SrcTy->getScalarSizeInBits();
+ bool IsUIToFP = isa<UIToFPInst>(Extends[0]);
+ VectorType *StepVecTy = VectorType::getInteger(DestTy);
+ Type *StepTy = StepVecTy->getScalarType();
+ Value *Input = Deinterleave->getOperand(0);
+ Type *InputTy = Input->getType();
+
+ APInt Invalid = APInt::getAllOnes(DstBits);
+ for (auto [Idx, Extend] : enumerate(Extends)) {
+ // Build the mask using stepvectors and casting.
+ // We want to select the Idx'th element, and every 4 elements after that.
+ // Each element needs to be zero extended; we can do that by providing
+ // tbl index values that are out of range. We can't do that nicely with
+ // a stepvector of the same element type as the input type, but we can
+ // do it with elements the size of the output type.
+ // E.g. for element 0 of a 16b -> 64b zext, we would start with a mask of
+ // 0xFFFF_FFFF_FFFF_0000 + Idx for the start of the stepvector, and use a
+ // step of 4. We then cast that back to an element size of 16b, yielding
+ // <0x0000 + Idx, 0xFFFF, 0xFFFF, 0xFFFF, 0x0004 + Idx, 0xFFFF...>.
+ APInt StartIdx = Invalid << SrcBits;
+ StartIdx += Idx;
+ IRBuilder<> Builder(Extend);
+ Value *StepVector = Builder.CreateStepVector(StepVecTy);
+ Value *ScaledSteps = Builder.CreateMul(
+ StepVector, Builder.CreateVectorSplat(StepVecTy->getElementCount(),
+ ConstantInt::get(StepTy, 4)));
+ Value *Start = ConstantInt::get(StepTy, StartIdx);
+ Value *ZextTbl = Builder.CreateAdd(
+ ScaledSteps,
+ Builder.CreateVectorSplat(StepVecTy->getElementCount(), Start));
+ Value *FinalMask = Builder.CreateBitCast(ZextTbl, InputTy);
+
+ // Replace the deinterleave, extractvalue, and extension chain with
+ // a tbl directly on the input value.
+ Value *Tbl = Builder.CreateIntrinsic(Intrinsic::aarch64_sve_tbl,
+ {InputTy}, {Input, FinalMask});
+ Value *Widen = Builder.CreateBitCast(Tbl, StepVecTy);
+ if (IsUIToFP)
+ Widen = Builder.CreateUIToFP(Widen, DestTy);
+ LLVM_DEBUG(dbgs() << "SVETBLOPT: Replaced " << *Extend << " with "
+ << *Widen << "\n");
+ Extend->replaceAllUsesWith(Widen);
+ Extend->eraseFromParent();
----------------
huntergr-arm wrote:
done.
https://github.com/llvm/llvm-project/pull/193951
More information about the llvm-commits
mailing list