From f067bc3c0ad6b8800f63066d654ad15c3dd33fd1 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Tue, 27 Oct 2020 08:59:18 +0000
Subject: [PATCH] [LoopRotation] Allow loop header duplication if vectorization
 is forced.

-Oz normally does not allow loop header duplication so this loop wouldn't be
vectorized.  However the vectorization pragma should override this and allow
for loop rotation.

rdar://problem/49281061

Original patch by Adam Nemet.

Reviewed By: Meinersbur

Differential Revision: https://reviews.llvm.org/D59832
---
 llvm/lib/Transforms/Scalar/LoopRotation.cpp   | 17 ++++++++-
 .../AArch64/Oz-and-forced-vectorize.ll        | 38 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopVectorize/AArch64/Oz-and-forced-vectorize.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopRotation.cpp b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
index f92566ba77ce..73287c7441e9 100644
--- a/llvm/lib/Transforms/Scalar/LoopRotation.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
@@ -39,7 +39,13 @@ LoopRotatePass::LoopRotatePass(bool EnableHeaderDuplication)
 PreservedAnalyses LoopRotatePass::run(Loop &L, LoopAnalysisManager &AM,
                                       LoopStandardAnalysisResults &AR,
                                       LPMUpdater &) {
-  int Threshold = EnableHeaderDuplication ? DefaultRotationThreshold : 0;
+  // Vectorization requires loop-rotation. Use default threshold for loops the
+  // user explicitly marked for vectorization, even when header duplication is
+  // disabled.
+  int Threshold = EnableHeaderDuplication ||
+                          hasVectorizeTransformation(&L) == TM_ForcedByUser
+                      ? DefaultRotationThreshold
+                      : 0;
   const DataLayout &DL = L.getHeader()->getModule()->getDataLayout();
   const SimplifyQuery SQ = getBestSimplifyQuery(AR, DL);
 
@@ -105,9 +111,16 @@ public:
       if (MSSAA)
         MSSAU = MemorySSAUpdater(&MSSAA->getMSSA());
     }
+    // Vectorization requires loop-rotation. Use default threshold for loops the
+    // user explicitly marked for vectorization, even when header duplication is
+    // disabled.
+    int Threshold = hasVectorizeTransformation(L) == TM_ForcedByUser
+                        ? DefaultRotationThreshold
+                        : MaxHeaderSize;
+
     return LoopRotation(L, LI, TTI, AC, &DT, &SE,
                         MSSAU.hasValue() ? MSSAU.getPointer() : nullptr, SQ,
-                        false, MaxHeaderSize, false);
+                        false, Threshold, false);
   }
 };
 } // end namespace
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/Oz-and-forced-vectorize.ll b/llvm/test/Transforms/LoopVectorize/AArch64/Oz-and-forced-vectorize.ll
new file mode 100644
index 000000000000..9dff3dffff2b
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/AArch64/Oz-and-forced-vectorize.ll
@@ -0,0 +1,38 @@
+; RUN: opt -Oz -S -enable-new-pm=0  < %s | FileCheck %s
+; RUN: opt -passes='default<Oz>' -S < %s | FileCheck %s
+
+; Forcing vectorization should allow for more aggressive loop-rotation with
+; -Oz, because LV requires rotated loops. Make sure the loop in @foo is
+; vectorized with -Oz.
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-ios5.0.0"
+
+define void @foo(float* noalias nocapture %ptrA, float* noalias nocapture readonly %ptrB, i64 %size) {
+; CHECK-LABEL: @foo(
+; CHECK: fmul <4 x float>
+;
+entry:
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+  %exitcond = icmp eq i64 %indvars.iv, %size
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+
+for.body:                                         ; preds = %for.cond
+  %arrayidx = getelementptr inbounds float, float* %ptrB, i64 %indvars.iv
+  %0 = load float, float* %arrayidx, align 4
+  %arrayidx2 = getelementptr inbounds float, float* %ptrA, i64 %indvars.iv
+  %1 = load float, float* %arrayidx2, align 4
+  %mul3 = fmul float %0, %1
+  store float %mul3, float* %arrayidx2, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  br label %for.cond, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %for.cond
+  ret void
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.vectorize.enable", i1 true}
-- 
2.34.1