[llvm] [AArch64] SVE Shuffleopt: merge reduction reverse into tbl (PR #206047)

[Gaëtan Bossu via llvm-commits]

================
@@ -155,10 +155,130 @@ static void evaluateDeinterleave(IntrinsicInst *I, DeinterleaveMap &Candidates,
   Candidates.try_emplace(I, Extends);
 }
 
+/// Evaluate a reverse intrinsic to see what uses it. We want to find a reverse
+/// that pairs with an extract from a deinterleave, so that we can move the
+/// reverse into the tbl as well as deinterleave and extend. We also need to
+/// confirm that it's only used by single-use instructions, or instructions
+/// used by a phi and a reduction intrinsic outside the loop, where the
+/// reduction permits reassociation. Something like the following:
+///
+///  %acc.b.f64 = phi <vscale x 2 x double> [ splat(double 0.000000e+00),
+///                                          %entry ], [ %fadd.b.f64, %loop ]
+///  ...
+///  %rev.load = load <vscale x 2 x double>, ptr %rev.ptr
+///  %reversed = call <vscale x 2 x double> @llvm.vector.reverse.nxv2f64(
+///                                             <vscale x 2 x double> %rev.load)
+///  %bgra = call <vscale x 8 x i16> @llvm.masked.load(ptr %src.gep,
+///                 <vscale x 8 x i1> %mask, <vscale x 8 x i16> zeroinitializer)
+///  %deinterleave = tail call { <vscale x 2 x i16>, <vscale x 2 x i16>,
+///                               <vscale x 2 x i16>, <vscale x 2 x i16> }
+///                          @llvm.vector.deinterleave4(<vscale x 8 x i16>
+///                          %bgra)
+///  %b.i16 = extractvalue { <vscale x 2 x i16>, <vscale x 2 x i16>,
+///                   <vscale x 2 x i16>, <vscale x 2 x i16> } %deinterleave, 0
+///  %b.f64 = uitofp <vscale x 2 x i16> %b.i16 to <vscale x 2 x double>
+///  %b.mul.f64 = fmul <vscale x 2 x double> %b.f64, %reversed
+///  %fadd.b.f64 = fadd <vscale x 2 x double> %acc.b.f64, %b.mul.f64
+///  %iv.next = add nuw i64 %iv, %stride
+///  %ec = icmp eq i64 %iv.next, 2048
+///  br i1 %ec, label %exit, label %loop
+///  ...
+///  %b.acc = call fast double @llvm.vector.reduce.fadd.nxv2f64(double
+///                             0.000000e+00, <vscale x 2 x double> %fadd.b.f64)
+static void evaluateReverse(IntrinsicInst *Reverse,
+                            SmallVectorImpl<IntrinsicInst *> &Reverses, Loop &L,
+                            const AArch64TargetLowering &TL,
+                            const DataLayout DL) {
+  assert(Reverse->getIntrinsicID() == Intrinsic::vector_reverse &&
+         "Invalid intrinsic used");
+
+  // We want to check that for each use of the reverse, we eventually end up
+  // in a reassociable reduction, so that the ordering of elements does not
+  // matter.
+  for (User *U : Reverse->users()) {
+    Instruction *I = cast<Instruction>(U);
+
+    // For now, only iterate through normal operations until we hit a
+    // CallInst. If that's not a reduction, abandon.
+    // TODO: Handle more generic input when we have a use for it.
+    while (!isa<CallInst>(I)) {
+      // If there's only one use, proceeed to the next in the chain.
+      if (I->hasOneUse()) {
+        I = I->user_back();
+        continue;
+      }
+
+      // Otherwise, look for 2 users -- a phi in the loop, and an outside
+      // reduction intrinsic.
+      SmallVector<User *, 2> Users(I->users());
+      if (Users.size() != 2)
+        return;
+
+      Instruction *Phi = cast<Instruction>(Users[0]);
+      Instruction *Reduce = cast<Instruction>(Users[1]);
+      if (!isa<PHINode>(Phi))
+        std::swap(Phi, Reduce);
+
+      if (!isa<PHINode>(Phi) || !isa<IntrinsicInst>(Reduce))
+        return;
+
+      // TODO: Do we have a nice utility to match all reduction intrinsics?
+      // We need to confirm reassociation is allowed, because the lanes within
+      // the vector will be in reversed order.
+      IntrinsicInst *II = cast<IntrinsicInst>(Reduce);
+      if (!isa<FPMathOperator>(II) || !II->hasAllowReassoc() ||
+          II->getIntrinsicID() != Intrinsic::vector_reduce_fadd)
+        return;
+
+      // The Phi must be for the current loop, and the reduction must be outside
+      // it.
+      if (!L.contains(Phi) || L.contains(Reduce))
+        return;
+
+      I = Reduce;
+    }
+  }
+
+  Reverses.push_back(Reverse);
+}
+
 /// Given a map of deinterleaves to zext or uitofp casts, remove the operations
 /// and replace them with tbl shuffles.
-static void optimizeSVEDeinterleavedExtends(DeinterleaveMap Deinterleaves) {
+static void
+optimizeSVEDeinterleavedExtends(DeinterleaveMap Deinterleaves,
+                                ArrayRef<IntrinsicInst *> Reverses) {
   for (auto &[Deinterleave, Extends] : Deinterleaves) {
----------------
gbossu wrote:

Is there an assumption that each `(Deinterleave, Extends)` pair will generate a single legal load and `Extends.size()` tbl?

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