[llvm] [VPlan] Cache the vector loop region lookup (PR #199437)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: Madhur Amilkanthwar (madhur13490)

<details>
<summary>Changes</summary>

Resolves the TODO in VPlan::getVectorLoopRegion() with a mutable cache on VPlan, shared by both overloads.

Measured on an O3 build of the LLVM test suite (~12k TUs): 3,556,267 hits / 126,196 misses (96.57% hit rate).

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


3 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VPlan.cpp (+7-7) 
- (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+10) 
- (modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+2) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index d2a1161984d9a..850b89f203e26 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -1096,18 +1096,18 @@ InstructionCost VPlan::cost(ElementCount VF, VPCostContext &Ctx) {
 }
 
 VPRegionBlock *VPlan::getVectorLoopRegion() {
-  // TODO: Cache if possible.
+  if (VectorLoopRegion)
+    return VectorLoopRegion;
   for (VPBlockBase *B : vp_depth_first_shallow(getEntry()))
-    if (auto *R = dyn_cast<VPRegionBlock>(B))
-      return R->isReplicator() ? nullptr : R;
+    if (auto *R = dyn_cast<VPRegionBlock>(B)) {
+      VectorLoopRegion = R->isReplicator() ? nullptr : R;
+      return VectorLoopRegion;
+    }
   return nullptr;
 }
 
 const VPRegionBlock *VPlan::getVectorLoopRegion() const {
-  for (const VPBlockBase *B : vp_depth_first_shallow(getEntry()))
-    if (auto *R = dyn_cast<VPRegionBlock>(B))
-      return R->isReplicator() ? nullptr : R;
-  return nullptr;
+  return const_cast<VPlan *>(this)->getVectorLoopRegion();
 }
 
 bool VPlan::isOuterLoop() const {
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 7025f39decaf5..040ac9e4e8da1 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -4675,6 +4675,11 @@ class VPlan {
   /// VPlan is destroyed.
   SmallVector<VPBlockBase *> CreatedBlocks;
 
+  /// Cached pointer to the (non-replicator) vector loop region, populated
+  /// lazily by getVectorLoopRegion(). Reset by setEntry() and by
+  /// VPlanTransforms::dissolveLoopRegions() once the loop region is detached.
+  mutable VPRegionBlock *VectorLoopRegion = nullptr;
+
   /// Construct a VPlan with \p Entry to the plan and with \p ScalarHeader
   /// wrapping the original header of the scalar loop. The vector loop will have
   /// index type \p IdxTy.
@@ -4706,8 +4711,13 @@ class VPlan {
   void setEntry(VPBasicBlock *VPBB) {
     Entry = VPBB;
     VPBB->setPlan(this);
+    invalidateVectorLoopRegionCache();
   }
 
+  /// Reset the cached vector loop region pointer. Must be called by
+  /// transforms that change the top-level region structure.
+  void invalidateVectorLoopRegionCache() { VectorLoopRegion = nullptr; }
+
   /// Generate the IR code for this VPlan.
   void execute(VPTransformState *State);
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index e5b7041a002a1..3a1613b5f45d5 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -3936,6 +3936,8 @@ void VPlanTransforms::dissolveLoopRegions(VPlan &Plan) {
   }
   for (VPRegionBlock *R : LoopRegions)
     R->dissolveToCFGLoop();
+  // Drop any cached pointer to the (now-detached) vector loop region.
+  Plan.invalidateVectorLoopRegionCache();
 }
 
 void VPlanTransforms::expandBranchOnTwoConds(VPlan &Plan) {

``````````

</details>


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