[llvm] [RemoveDIs][DebugInfo] Hoist DPValues in SpeculativeExecution (PR #80886)

[Stephen Tozer via llvm-commits]

https://github.com/SLTozer created https://github.com/llvm/llvm-project/pull/80886

This patch modifies `SpeculativeExecutionPass::considerHoistingFromTo` to treat DPValues the same way that it treats debug intrinsics, which is to hoist them iff all of their instruction operands within the FromBlock are also being hoisted. This is probably not the ideal behaviour, which would be to not hoist debug info at all in this case, but this patch simply ensures that DPValue behaviour is consistent with debug intrinsic behaviour rather than trying to create new functional changes.

>From d10d73eaa689cc6e370624da0aea1e51537aa3d0 Mon Sep 17 00:00:00 2001
From: Stephen Tozer <Stephen.Tozer at Sony.com>
Date: Tue, 6 Feb 2024 17:16:39 +0000
Subject: [PATCH] Hoist DPValues in SpeculativeExecution

---
 .../Scalar/SpeculativeExecution.cpp           | 91 +++++++++++--------
 .../SpeculativeExecution/PR46267.ll           |  1 +
 2 files changed, 53 insertions(+), 39 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SpeculativeExecution.cpp b/llvm/lib/Transforms/Scalar/SpeculativeExecution.cpp
index 7a5318d4404ca..eebee87c2f514 100644
--- a/llvm/lib/Transforms/Scalar/SpeculativeExecution.cpp
+++ b/llvm/lib/Transforms/Scalar/SpeculativeExecution.cpp
@@ -263,60 +263,73 @@ static InstructionCost ComputeSpeculationCost(const Instruction *I,
 bool SpeculativeExecutionPass::considerHoistingFromTo(
     BasicBlock &FromBlock, BasicBlock &ToBlock) {
   SmallPtrSet<const Instruction *, 8> NotHoisted;
-  const auto AllPrecedingUsesFromBlockHoisted = [&NotHoisted](const User *U) {
-    // Debug variable has special operand to check it's not hoisted.
-    if (const auto *DVI = dyn_cast<DbgVariableIntrinsic>(U)) {
-      return all_of(DVI->location_ops(), [&NotHoisted](Value *V) {
-        if (const auto *I = dyn_cast_or_null<Instruction>(V)) {
-          if (!NotHoisted.contains(I))
-            return true;
-        }
-        return false;
-      });
-    }
-
-    // Usially debug label intrinsic corresponds to label in LLVM IR. In these
-    // cases we should not move it here.
-    // TODO: Possible special processing needed to detect it is related to a
-    // hoisted instruction.
-    if (isa<DbgLabelInst>(U))
-      return false;
-
-    for (const Value *V : U->operand_values()) {
-      if (const Instruction *I = dyn_cast<Instruction>(V)) {
-        if (NotHoisted.contains(I))
-          return false;
+  SmallDenseMap<const Instruction *, SmallVector<DPValue *>> DPValuesToHoist;
+  auto HasNoUnhoistedInstr = [&NotHoisted](auto Values) {
+    for (const Value *V : Values) {
+      if (const auto *I = dyn_cast_or_null<Instruction>(V)) {
+        if (!NotHoisted.contains(I))
+          return true;
       }
     }
-    return true;
+    return false;
   };
+  auto AllPrecedingUsesFromBlockHoisted =
+      [&HasNoUnhoistedInstr](const User *U) {
+        // Debug variable has special operand to check it's not hoisted.
+        if (const auto *DVI = dyn_cast<DbgVariableIntrinsic>(U))
+          return HasNoUnhoistedInstr(DVI->location_ops());
+
+        // Usially debug label intrinsic corresponds to label in LLVM IR. In
+        // these cases we should not move it here.
+        // TODO: Possible special processing needed to detect it is related to a
+        // hoisted instruction.
+        if (isa<DbgLabelInst>(U))
+          return false;
+
+        return HasNoUnhoistedInstr(U->operand_values());
+      };
 
   InstructionCost TotalSpeculationCost = 0;
   unsigned NotHoistedInstCount = 0;
   for (const auto &I : FromBlock) {
-    const InstructionCost Cost = ComputeSpeculationCost(&I, *TTI);
-    if (Cost.isValid() && isSafeToSpeculativelyExecute(&I) &&
-        AllPrecedingUsesFromBlockHoisted(&I)) {
-      TotalSpeculationCost += Cost;
-      if (TotalSpeculationCost > SpecExecMaxSpeculationCost)
-        return false;  // too much to hoist
-    } else {
-      // Debug info intrinsics should not be counted for threshold.
-      if (!isa<DbgInfoIntrinsic>(I))
-        NotHoistedInstCount++;
-      if (NotHoistedInstCount > SpecExecMaxNotHoisted)
-        return false; // too much left behind
-      NotHoisted.insert(&I);
+    // Make note of any DPValues that need hoisting.
+    for (DPValue &DPV : I.getDbgValueRange()) {
+      if (HasNoUnhoistedInstr(DPV.location_ops()))
+        DPValuesToHoist[DPV.getInstruction()].push_back(&DPV);
+
+      const InstructionCost Cost = ComputeSpeculationCost(&I, *TTI);
+      if (Cost.isValid() && isSafeToSpeculativelyExecute(&I) &&
+          AllPrecedingUsesFromBlockHoisted(&I)) {
+        TotalSpeculationCost += Cost;
+        if (TotalSpeculationCost > SpecExecMaxSpeculationCost)
+          return false; // too much to hoist
+      } else {
+        // Debug info intrinsics should not be counted for threshold.
+        if (!isa<DbgInfoIntrinsic>(I))
+          NotHoistedInstCount++;
+        if (NotHoistedInstCount > SpecExecMaxNotHoisted)
+          return false; // too much left behind
+        NotHoisted.insert(&I);
+      }
     }
-  }
 
   for (auto I = FromBlock.begin(); I != FromBlock.end();) {
+    // If any DPValues attached to this instruction should be hoisted, hoist
+    // them now - they will end up attached to either the next hoisted
+    // instruction or the ToBlock terminator.
+    if (DPValuesToHoist.contains(&*I)) {
+      for (auto *DPV : DPValuesToHoist[&*I]) {
+        DPV->removeFromParent();
+        ToBlock.insertDPValueBefore(DPV,
+                                    ToBlock.getTerminator()->getIterator());
+      }
+    }
     // We have to increment I before moving Current as moving Current
     // changes the list that I is iterating through.
     auto Current = I;
     ++I;
     if (!NotHoisted.count(&*Current)) {
-      Current->moveBeforePreserving(ToBlock.getTerminator());
+      Current->moveBefore(ToBlock.getTerminator());
     }
   }
   return true;
diff --git a/llvm/test/Transforms/SpeculativeExecution/PR46267.ll b/llvm/test/Transforms/SpeculativeExecution/PR46267.ll
index 00f80d3e33a6c..c27b492b4b876 100644
--- a/llvm/test/Transforms/SpeculativeExecution/PR46267.ll
+++ b/llvm/test/Transforms/SpeculativeExecution/PR46267.ll
@@ -1,4 +1,5 @@
 ; RUN: opt < %s -S -passes='speculative-execution' | FileCheck %s
+; RUN: opt --try-experimental-debuginfo-iterators < %s -S -passes='speculative-execution' | FileCheck %s
 
 %class.B = type { ptr }