[VPlan] Add -vplan-print-after=<> option

[eas]

Builds on top of #178547.
Next step would be to add UpdateTestChecks support for
VPlan-output-based tests, to be done in a separate PR.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Andrei Elovikov (eas)

Changes

Builds on top of #178547.
Next step would be to add UpdateTestChecks support for
VPlan-output-based tests, to be done in a separate PR.

---

Full diff: https://github.com/llvm/llvm-project/pull/178700.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+7)
• (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.h (+12-2)
• (added) llvm/test/Transforms/LoopVectorize/vplan-print-after.ll (+29)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 4e056ab30e2a6..7c37ee419f79e 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -366,10 +366,17 @@ cl::opt<bool>
                           cl::Hidden,
                           cl::desc("Verfiy VPlans after VPlan transforms."));
 
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 cl::opt<bool> llvm::PrintAfterEachVPlanPass(
     "vplan-print-after-all", cl::init(false), cl::Hidden,
     cl::desc("Print after each VPlanTransforms::runPass."));
 
+cl::list<std::string> llvm::PrintAfterVPlanPasses(
+    "vplan-print-after", cl::Hidden,
+    cl::desc("Print after specified VPlan transformations (substring match "
+             "suffices)."));
+#endif
+
 // This flag enables the stress testing of the VPlan H-CFG construction in the
 // VPlan-native vectorization path. It must be used in conjuction with
 // -enable-vplan-native-path. -vplan-verify-hcfg can also be used to enable the
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index 451597afa73da..617d4f14a54aa 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -34,9 +34,13 @@ class VPRecipeBuilder;
 struct VFRange;
 
 LLVM_ABI_FOR_TEST extern cl::opt<bool> VerifyEachVPlan;
-LLVM_ABI_FOR_TEST extern cl::opt<bool> PrintAfterEachVPlanPass;
 LLVM_ABI_FOR_TEST extern cl::opt<bool> EnableWideActiveLaneMask;
 
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+LLVM_ABI_FOR_TEST extern cl::opt<bool> PrintAfterEachVPlanPass;
+LLVM_ABI_FOR_TEST extern cl::list<std::string> PrintAfterVPlanPasses;
+#endif
+
 struct VPlanTransforms {
   /// Helper to run a VPlan pass \p Pass on \p VPlan, forwarding extra arguments
   /// to the pass. Performs verification/printing after each VPlan pass if
@@ -45,12 +49,18 @@ struct VPlanTransforms {
   static decltype(auto) runPass(StringRef PassName, PassTy &&Pass, VPlan &Plan,
                                 ArgsTy &&...Args) {
     scope_exit PostTransformActions{[&]() {
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
       // Make sure to print before verification, so that output is more useful
       // in case of failures:
-      if (PrintAfterEachVPlanPass) {
+      if (PrintAfterEachVPlanPass ||
+          (PrintAfterVPlanPasses.getNumOccurrences() > 0 &&
+           any_of(PrintAfterVPlanPasses, [PassName](StringRef Entry) {
+             return PassName.contains(Entry);
+           }))) {
         dbgs() << "VPlan after " << PassName << '\n';
         dbgs() << Plan << '\n';
       }
+#endif
       if (VerifyEachVPlan && EnableVerify)
         verifyVPlanIsValid(Plan);
     }};
diff --git a/llvm/test/Transforms/LoopVectorize/vplan-print-after.ll b/llvm/test/Transforms/LoopVectorize/vplan-print-after.ll
new file mode 100644
index 0000000000000..8014a4e480d65
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/vplan-print-after.ll
@@ -0,0 +1,29 @@
+; RUN: opt -passes=loop-vectorize -disable-output  -force-vector-width=4  < %s \
+; RUN:   -vplan-print-after=simplify -vplan-print-after=printFinalVPlan \
+; RUN:   2>&1 | FileCheck %s --implicit-check-not "VPlan after"
+; REQUIRES: asserts
+
+; CHECK:      VPlan after simplifyRecipes
+; CHECK-NEXT: VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK:      VPlan after simplifyBlends
+; CHECK-NEXT: VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK:      VPlan after simplifyRecipes
+; CHECK-NEXT: VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK:      VPlan after printFinalVPlan
+; CHECK-NEXT: VPlan 'Final VPlan for VF={4},UF={1}' {
+
+define void @foo(ptr %ptr, i64 %n) {
+entry:
+  br label %header
+
+header:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %header ]
+  %gep = getelementptr i64, ptr %ptr, i64 %iv
+  store i64 %iv, ptr %gep
+  %iv.next = add nsw i64 %iv, 1
+  %exitcond = icmp slt i64 %iv.next, %n
+  br i1 %exitcond, label %header, label %exit
+
+exit:
+  ret void
+}

[eas]

@fhahn , @lukel97 , technically it includes #178547 in its entirety, but I've based this on the trunk commit before the revert to minimize the diff here, so that review could start productively before #178547 lands.