[llvm] [AArch64] Add predicate-as-counter loop rewrite pass (PR #220960)

[Benjamin Maxwell via llvm-commits]

================
@@ -0,0 +1,459 @@
+//===- AArch64PredicateAsCounterLoopRewrites.cpp --------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Rewrites IR for loop-carried wide masks that can be represented as
+// predicate-as-counter values. This applies when the mask is used by load/store
+// operations that can be mapped to multi-vector instructions (with +sve2p1).
+//
+// For example, a loop like:
+//
+//   entry:
+//     %step = vscale x 64
+//     %mask.entry = @get.active.lane.mask(0, %n)
+//
+//   loop:
+//     %iv   = phi i64 [0, entry], [%iv.next, loop]
+//     %mask = phi <vscale x 64 x i1> [%mask.entry, entry],
+//                                    [%mask.next, loop]
+//
+//     %load = load <vscale x 64 x i8> %src[%iv], %mask
+//     store <vscale x 64 x i8> %load, %dst[%iv], %mask
+//
+//     %iv.next   = %iv + %step
+//     %mask.next = @get.active.lane.mask(%iv.next, %n)
+//     br first.active(%mask.next), loop, exit
+//
+// Could be rewritten to:
+//
+//   entry:
+//     %step = vscale x 64
+//     %mask.entry = @whilelo.c8(0, %n, VLx4)
+//
+//   loop:
+//     %iv   = phi i64 [0, entry], [%iv.next, loop]
+//     %mask = phi target("aarch64.svcount") [%mask.entry, entry],
+//                                           [%mask.next, loop]
+//
+//     %load = @ld1.pn.x4 <4 x <vscale x 16 x i8>> %src[%iv], %mask
+//     @st1.pn.x4 <4 x <vscale x 16 x i8>> %load, %dest[%iv], %mask
+//
+//     %iv.next   = %iv + %step
+//     %mask.next = @whilelo.c8(%iv.next, %n, VLx4)
+//     br first.active(@pext(%mask.next, 0)), loop, exit
+//
+// This replaces the `get.active.lane.mask` intrinsics with AArch64
+// predicate-as-counter `whilelo` intrinsics and updates the mask phi to use the
+// `aarch64.svcount` target type. Within the loop, load/store users are mapped
+// to multi-vector load/store intrinsics where possible. Users that cannot be
+// mapped to multi-vector instructions materialize vector masks using the `pext`
+// intrinsic (which extracts vector predicates from a predicate-as-counter).
+//
+// This pass may be a temporary solution that is removed if we gain support
+// for target-specific VPlan transforms in the loop vectorizer.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AArch64.h"
+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/LoopPass.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsAArch64.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils.h"
+#include "llvm/Transforms/Utils/Local.h"
+#include <optional>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "aarch64-predicate-as-counter-loop-rewrites"
+namespace {
+
+STATISTIC(LoopsRewritten, "Number of loops rewritten");
+
+struct MaskRewriteCandidate {
+  /// The preheader block for the loop.
+  BasicBlock *Preheader = nullptr;
+  /// The latch block for the loop.
+  BasicBlock *Latch = nullptr;
+  /// The mask phi node (used by masked operations within the loop).
+  PHINode *MaskPhi = nullptr;
+  /// The initial value for the mask (incoming value from the preheader).
+  IntrinsicInst *StartMask = nullptr;
+  /// The updated value for the mask (incoming value from the loop latch).
+  IntrinsicInst *NextMask = nullptr;
+  /// The multi-vector scale for the predicate-as-counter (2 or 4).
+  unsigned VectorScale = 0;
+  /// The element size (in bits) for the predicate-as-counter.
+  unsigned ElementSizeInBits = 0;
+};
+
+static void logLoopBailout(const Loop &L, const Twine &Reason) {
+  LLVM_DEBUG({
+    dbgs() << "PAC loop rewrite: skipping loop with header ";
----------------
MacDue wrote:

I've gone with "PN" (which matches other predicate-as-counter intrinsics/types).  

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