[NFC] clean up memtag-stack code

[fmayer]

we would replace the alloca with tagp for debug instructions, then replace it back with the original alloca. it's easier to just skip the replacement.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Florian Mayer (fmayer)

Changes

we would replace the alloca with tagp for debug instructions, then replace it back with the original alloca. it's easier to just skip the replacement.

---

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

1 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64StackTagging.cpp (+3-9)

diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 961dded132343..15d81c2f44a5b 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -21,7 +21,6 @@
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
-#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/LiveRegUnits.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -520,7 +519,6 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
     assert(Info.AI && SIB.isInterestingAlloca(*Info.AI));
-    TrackingVH<Instruction> OldAI = Info.AI;
     memtag::alignAndPadAlloca(Info, kTagGranuleSize);
     AllocaInst *AI = Info.AI;
     int Tag = NextTag;
@@ -534,7 +532,9 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
                               ConstantInt::get(IRB.getInt64Ty(), Tag)});
     if (Info.AI->hasName())
       TagPCall->setName(Info.AI->getName() + ".tag");
-    Info.AI->replaceAllUsesWith(TagPCall);
+    // Does not replace metadata, so we don't have to handle DPValues.
+    Info.AI->replaceUsesWithIf(
+        TagPCall, [&](Use &U) { return !isa<DbgVariableIntrinsic>(U); });
     TagPCall->setOperand(0, Info.AI);
 
     // Calls to functions that may return twice (e.g. setjmp) confuse the
@@ -574,12 +574,6 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
       for (auto *II : Info.LifetimeEnd)
         II->eraseFromParent();
     }
-
-    // Fixup debug intrinsics to point to the new alloca.
-    for (auto *DVI : Info.DbgVariableIntrinsics)
-      DVI->replaceVariableLocationOp(OldAI, Info.AI);
-    for (auto *DPV : Info.DbgVariableRecords)
-      DPV->replaceVariableLocationOp(OldAI, Info.AI);
   }
 
   // If we have instrumented at least one alloca, all unrecognized lifetime

[fmayer]

Tested this with -DLLVM_EXPERIMENTAL_DEBUGINFO_ITERATORS=On to ensure the comment about DPValues is actually correct

[jmorse]

I can't comment on the general operation of memory tagging; however IMO you should use Value::replaceNonMetadataUsesWith instead of Value::replaceUsesWithIf. Metadata uses get deliberately hidden to avoid affecting codegen decisions, and don't appear on the normal use list, so the isa<DbgVariableIntrinsic> test will never be true. I think replaceNonMetadataUsesWith better expresses what you're aiming for and achieves the same thing.

[fmayer]

I can't comment on the general operation of memory tagging; however IMO you should use Value::replaceNonMetadataUsesWith instead of Value::replaceUsesWithIf. Metadata uses get deliberately hidden to avoid affecting codegen decisions, and don't appear on the normal use list, so the isa<DbgVariableIntrinsic> test will never be true. I think replaceNonMetadataUsesWith better expresses what you're aiming for and achieves the same thing.

Thanks, done.

[jmorse]

LGTM -- I'm not quite sure of the path taken whereby the debuginfo records point at the right alloca, but it's already covered by existing tests.

[fmayer]

I'm not quite sure of the path taken whereby the debuginfo records point at the right alloca,

I don't fully understand, what exactly do you mean?

[jmorse]

I don't fully understand, what exactly do you mean?

Essentially the path taken in the old code: the debug records get RAUW'd to TagPCall, but then later they get replaced back to the "original" Info.AI?

Either way, it's more confusion about what the old code was doing rather than the new code, which is easier to understand.

EDIT: (every time I paste into github it submits an empty comment for me)

[fmayer]

Essentially the path taken in the old code: the debug records get RAUW'd to TagPCall, but then later they get replaced back to the "original" Info.AI?

Not the very original one because it was TrackingVH, so it was the Info.AI after the align and pad call. AFAICT the code was like this by historical accident.

[jmorse]

LGTM then, ship it!