[Vectorize] Add StridedLoopUnroll + Versioning for 2-D strided loop nests

llvm/include/llvm/Transforms/Utils/LoopVersioning.h

@@ -46,7 +46,8 @@ class LoopVersioning {
   /// object having no checks and we expect the user to add them.
   LoopVersioning(const LoopAccessInfo &LAI,
                  ArrayRef<RuntimePointerCheck> Checks, Loop *L, LoopInfo *LI,
-                 DominatorTree *DT, ScalarEvolution *SE);
+                 DominatorTree *DT, ScalarEvolution *SE,
+                 bool HoistRuntimeChecks = false);
 
   /// Performs the CFG manipulation part of versioning the loop including
   /// the DominatorTree and LoopInfo updates.
@@ -147,6 +148,7 @@ class LoopVersioning {
   LoopInfo *LI;
   DominatorTree *DT;
   ScalarEvolution *SE;
+  bool HoistRuntimeChecks;
 };
 
 /// Expose LoopVersioning as a pass.  Currently this is only used for

llvm/include/llvm/Transforms/Vectorize/StridedLoopUnroll.h

@@ -0,0 +1,37 @@
+//===----------StridedLoopUnroll.h -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TRANSFORMS_VECTORIZE_LOOPSTRIDEIDIOMVECTORIZE_H
+#define LLVM_LIB_TRANSFORMS_VECTORIZE_LOOPSTRIDEIDIOMVECTORIZE_H
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Scalar/LoopPassManager.h"
+
+namespace llvm {
+
+class ScalarEvolution;
+class LoopInfo;
+
+class StridedLoopUnrollPass : public PassInfoMixin<StridedLoopUnrollPass> {
+public:
+  StridedLoopUnrollPass() = default;
+
+  PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,
+                        LoopStandardAnalysisResults &AR, LPMUpdater &U);
+};
+
+class StridedLoopUnrollVersioningPass
+    : public PassInfoMixin<StridedLoopUnrollVersioningPass> {
+public:
+  StridedLoopUnrollVersioningPass() {}
+
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+} // namespace llvm
+#endif // LLVM_LIB_TRANSFORMS_VECTORIZE_LOOPIDIOMVECTORIZE_H

llvm/lib/Passes/PassBuilder.cpp

@@ -378,6 +378,7 @@
 #include "llvm/Transforms/Vectorize/LoopVectorize.h"
 #include "llvm/Transforms/Vectorize/SLPVectorizer.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h"
+#include "llvm/Transforms/Vectorize/StridedLoopUnroll.h"
 #include "llvm/Transforms/Vectorize/VectorCombine.h"
 #include <optional>
 

llvm/lib/Passes/PassRegistry.def

@@ -538,6 +538,7 @@ FUNCTION_PASS("sjlj-eh-prepare", SjLjEHPreparePass(TM))
 FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("slsr", StraightLineStrengthReducePass())
 FUNCTION_PASS("stack-protector", StackProtectorPass(TM))
+FUNCTION_PASS("strided-loop-unroll-versioning", StridedLoopUnrollVersioningPass())
 FUNCTION_PASS("strip-gc-relocates", StripGCRelocates())
 FUNCTION_PASS("tailcallelim", TailCallElimPass())
 FUNCTION_PASS("transform-warning", WarnMissedTransformationsPass())
@@ -764,6 +765,7 @@ LOOP_PASS("loop-bound-split", LoopBoundSplitPass())
 LOOP_PASS("loop-deletion", LoopDeletionPass())
 LOOP_PASS("loop-idiom", LoopIdiomRecognizePass())
 LOOP_PASS("loop-idiom-vectorize", LoopIdiomVectorizePass())
+LOOP_PASS("strided-loop-unroll", StridedLoopUnrollPass())
 LOOP_PASS("loop-instsimplify", LoopInstSimplifyPass())
 LOOP_PASS("loop-predication", LoopPredicationPass())
 LOOP_PASS("loop-reduce", LoopStrengthReducePass())

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

@@ -38,7 +38,9 @@
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/LoopUnrollPass.h"
 #include "llvm/Transforms/Vectorize/LoopIdiomVectorize.h"
+#include "llvm/Transforms/Vectorize/StridedLoopUnroll.h"
 #include <optional>
 using namespace llvm;
 
@@ -655,6 +657,17 @@ void RISCVTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
     if (Level != OptimizationLevel::O0)
       LPM.addPass(LoopIdiomVectorizePass(LoopIdiomVectorizeStyle::Predicated));
   });
+
+  PB.registerScalarOptimizerLateEPCallback(
+      [=](FunctionPassManager &LPM, OptimizationLevel) {
+        LPM.addPass(StridedLoopUnrollVersioningPass());
+      });
+  PB.registerOptimizerLastEPCallback([=](ModulePassManager &MPM,
+                                         OptimizationLevel Level,
+                                         llvm::ThinOrFullLTOPhase) {
+    MPM.addPass(createModuleToFunctionPassAdaptor(
+        createFunctionToLoopPassAdaptor(StridedLoopUnrollPass())));
+  });
 }
 
 yaml::MachineFunctionInfo *

llvm/lib/Transforms/Utils/LoopVersioning.cpp

@@ -42,9 +42,10 @@ static cl::opt<bool>
 LoopVersioning::LoopVersioning(const LoopAccessInfo &LAI,
                                ArrayRef<RuntimePointerCheck> Checks, Loop *L,
                                LoopInfo *LI, DominatorTree *DT,
-                               ScalarEvolution *SE)
+                               ScalarEvolution *SE, bool HoistRuntimeChecks)
     : VersionedLoop(L), AliasChecks(Checks), Preds(LAI.getPSE().getPredicate()),
-      LAI(LAI), LI(LI), DT(DT), SE(SE) {}
+      LAI(LAI), LI(LI), DT(DT), SE(SE), HoistRuntimeChecks(HoistRuntimeChecks) {
+}
 
 void LoopVersioning::versionLoop(
     const SmallVectorImpl<Instruction *> &DefsUsedOutside) {
@@ -63,8 +64,9 @@ void LoopVersioning::versionLoop(
   SCEVExpander Exp2(*RtPtrChecking.getSE(),
                     VersionedLoop->getHeader()->getDataLayout(),
                     "induction");
-  MemRuntimeCheck = addRuntimeChecks(RuntimeCheckBB->getTerminator(),
-                                     VersionedLoop, AliasChecks, Exp2);
+  MemRuntimeCheck =
+      addRuntimeChecks(RuntimeCheckBB->getTerminator(), VersionedLoop,
+                       AliasChecks, Exp2, HoistRuntimeChecks);
 
   SCEVExpander Exp(*SE, RuntimeCheckBB->getDataLayout(),
                    "scev.check");

llvm/lib/Transforms/Vectorize/CMakeLists.txt

@@ -2,6 +2,7 @@ add_llvm_component_library(LLVMVectorize
   EVLIndVarSimplify.cpp
   LoadStoreVectorizer.cpp
   LoopIdiomVectorize.cpp
+  StridedLoopUnroll.cpp
   LoopVectorizationLegality.cpp
   LoopVectorize.cpp
   SandboxVectorizer/DependencyGraph.cpp