[llvm] [VPlan] Speed up VPSlotTracker by using ModuleSlotTracker (PR #139881)

[Igor Kirillov via llvm-commits]

https://github.com/igogo-x86 updated https://github.com/llvm/llvm-project/pull/139881

>From e0825e4c622cb13d521e48598ccf8aa2aa6906bb Mon Sep 17 00:00:00 2001
From: Igor Kirillov <igor.kirillov at arm.com>
Date: Wed, 14 May 2025 10:57:58 +0000
Subject: [PATCH 1/2] [VPlan] Speed up VPSlotTracker by using ModuleSlotTracker
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Currently, when VPSlotTracker is initialized with a VPlan, its
assignName method calls printAsOperand on each underlying instruction.
Each such call recomputes slot numbers for the entire function, leading
to O(N × M) complexity, where M is the number of instructions in the
loop and N is the number of instructions in the function.

This results in slow debug output for large loops. For example, printing
costs of all instructions becomes O(M² × N), which is especially painful
when enabling verbose dumps.

This patch improves debugging performance by caching slot numbers using
ModuleSlotTracker. It avoids redundant recomputation and makes debug
output significantly faster.
---
 llvm/lib/Transforms/Vectorize/VPlan.cpp      |  5 ++++-
 llvm/lib/Transforms/Vectorize/VPlanHelpers.h | 14 +++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index 167aff737d3fd..a1567fab4eff2 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -1441,7 +1441,10 @@ void VPSlotTracker::assignName(const VPValue *V) {
   std::string Name;
   if (UV) {
     raw_string_ostream S(Name);
-    UV->printAsOperand(S, false);
+    if (MST)
+      UV->printAsOperand(S, false, *MST);
+    else
+      UV->printAsOperand(S, false);
   } else
     Name = VPI->getName();
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
index 1d42c8f5f3737..0aa0133f4fe1c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
@@ -23,6 +23,7 @@
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DebugLoc.h"
+#include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/Support/InstructionCost.h"
 
 namespace llvm {
@@ -382,14 +383,25 @@ class VPSlotTracker {
   /// Number to assign to the next VPValue without underlying value.
   unsigned NextSlot = 0;
 
+  /// Cache slot indexes to avoid recomputing them on each printAsOperand call.
+  std::unique_ptr<ModuleSlotTracker> MST;
+
   void assignName(const VPValue *V);
   void assignNames(const VPlan &Plan);
   void assignNames(const VPBasicBlock *VPBB);
 
 public:
   VPSlotTracker(const VPlan *Plan = nullptr) {
-    if (Plan)
+    if (Plan) {
+      // This check is required to support unit tests with incomplete IR.
+      if (Function *F =
+              Plan->getScalarHeader()->getIRBasicBlock()->getParent()) {
+        Module *M = F->getParent();
+        MST = std::make_unique<ModuleSlotTracker>(M);
+        MST->incorporateFunction(*F);
+      }
       assignNames(*Plan);
+    }
   }
 
   /// Returns the name assigned to \p V, if there is one, otherwise try to

>From d8fe7f8682944244bbdc33d5cb0e40251d7a0310 Mon Sep 17 00:00:00 2001
From: Igor Kirillov <igor.kirillov at arm.com>
Date: Fri, 16 May 2025 10:59:06 +0000
Subject: [PATCH 2/2] Initialize MST lazily

---
 llvm/lib/Transforms/Vectorize/VPlan.cpp      | 17 +++++++++++++++--
 llvm/lib/Transforms/Vectorize/VPlanHelpers.h | 13 +++----------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index a1567fab4eff2..566e265128397 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -1441,10 +1441,23 @@ void VPSlotTracker::assignName(const VPValue *V) {
   std::string Name;
   if (UV) {
     raw_string_ostream S(Name);
-    if (MST)
+    if (MST) {
       UV->printAsOperand(S, false, *MST);
-    else
+    } else if (isa<Instruction>(UV) && !UV->hasName()) {
+      // Lazily create the ModuleSlotTracker when we first hit an unnamed
+      // instruction
+      auto *IUV = cast<Instruction>(UV);
+      // This check is required to support unit tests with incomplete IR.
+      if (IUV->getParent()) {
+        MST = std::make_unique<ModuleSlotTracker>(IUV->getModule());
+        MST->incorporateFunction(*IUV->getFunction());
+      } else {
+        MST = std::make_unique<ModuleSlotTracker>(nullptr);
+      }
+      UV->printAsOperand(S, false, *MST);
+    } else {
       UV->printAsOperand(S, false);
+    }
   } else
     Name = VPI->getName();
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
index 0aa0133f4fe1c..075b25cdff45c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanHelpers.h
@@ -383,7 +383,8 @@ class VPSlotTracker {
   /// Number to assign to the next VPValue without underlying value.
   unsigned NextSlot = 0;
 
-  /// Cache slot indexes to avoid recomputing them on each printAsOperand call.
+  /// Lazily created ModuleSlotTracker, used only when unnamed IR instructions
+  /// require slot tracking.
   std::unique_ptr<ModuleSlotTracker> MST;
 
   void assignName(const VPValue *V);
@@ -392,16 +393,8 @@ class VPSlotTracker {
 
 public:
   VPSlotTracker(const VPlan *Plan = nullptr) {
-    if (Plan) {
-      // This check is required to support unit tests with incomplete IR.
-      if (Function *F =
-              Plan->getScalarHeader()->getIRBasicBlock()->getParent()) {
-        Module *M = F->getParent();
-        MST = std::make_unique<ModuleSlotTracker>(M);
-        MST->incorporateFunction(*F);
-      }
+    if (Plan)
       assignNames(*Plan);
-    }
   }
 
   /// Returns the name assigned to \p V, if there is one, otherwise try to